fix(invalid_token): Retry cognito signed token issues (#57)

* fix(invalid_token): Retry cognito signed token issues

Author: Levence

staxapp/auth.py

@@ -13,13 +13,14 @@
 
 
 class StaxAuth:
-    def __init__(self, config_branch):
+    def __init__(self, config_branch, max_retries: int = 3):
         config = StaxConfig.api_config
 
         self.identity_pool = config.get(config_branch).get("identityPoolId")
         self.user_pool = config.get(config_branch).get("userPoolId")
         self.client_id = config.get(config_branch).get("userPoolWebClientId")
         self.aws_region = config.get(config_branch).get("region")
+        self.max_retries = max_retries
 
     def requests_auth(self, username, password, **kwargs):
         if username is None:
@@ -83,23 +84,35 @@ def sts_from_cognito_identity_pool(self, token, cognito_client=None, **kwargs):
                 region_name=self.aws_region,
                 config=BotoConfig(signature_version=UNSIGNED),
             )
-        try:
-            id = cognito_client.get_id(
-                IdentityPoolId=self.identity_pool,
-                Logins={
-                    f"cognito-idp.{self.aws_region}.amazonaws.com/{self.user_pool}": token
-                },
-            )
-            id_creds = cognito_client.get_credentials_for_identity(
-                IdentityId=id["IdentityId"],
-                Logins={
-                    f"cognito-idp.{self.aws_region}.amazonaws.com/{self.user_pool}": token
-                },
-            )
-        except ClientError as e:
+
+        for i in range(self.max_retries):
+            try:
+                id = cognito_client.get_id(
+                    IdentityPoolId=self.identity_pool,
+                    Logins={
+                        f"cognito-idp.{self.aws_region}.amazonaws.com/{self.user_pool}": token
+                    },
+                )
+                id_creds = cognito_client.get_credentials_for_identity(
+                    IdentityId=id["IdentityId"],
+                    Logins={
+                        f"cognito-idp.{self.aws_region}.amazonaws.com/{self.user_pool}": token
+                    },
+                )
+                break
+            except ClientError as e:
+                # AWS eventual consistency, attempt to retry up to 3 times
+                if "Couldn't verify signed token" in str(e):
+                    continue
+                else:
+                    raise InvalidCredentialsException(
+                        f"Unexpected Client Error. Error details: {e}"
+                    )
+        else:
             raise InvalidCredentialsException(
-                f"Unexpected Client Error. Error details: {e}"
+                "Retries Exceeded: Unexpected Client Error"
             )
+
         return id_creds
 
     def sigv4_signed_auth_headers(self, id_creds):

tests/test_auth.py

@@ -134,15 +134,39 @@ def testCredsClient(self):
         Test the cognito client is invoked and throws an error
         """
         sa = StaxAuth("ApiAuth")
+
+        # Test Invalid Credentials
         token = jwt.encode({"sub": "unittest"}, "secret", algorithm="HS256")
         jwt_token = jwt.decode(token, verify=False)
         with self.assertRaises(InvalidCredentialsException):
             sa.sts_from_cognito_identity_pool(jwt_token.get("sub"))
 
+        # Test "Couldn't verify signed token" retry
+        expected_parameters = {
+            "IdentityPoolId": sa.identity_pool,
+            "Logins": {
+                f"cognito-idp.{sa.aws_region}.amazonaws.com/{sa.user_pool}": "unittest"
+            }
+        }
+        for i in range(sa.max_retries):
+            self.cognito_stub.add_client_error(
+                "get_id",
+                service_error_code="NotAuthorizedException",
+                service_message="Invalid login token. Couldn't verify signed token.",
+                expected_params=expected_parameters,
+            )
+        self.cognito_stub.activate()
+        
+        with self.assertRaises(InvalidCredentialsException) as e:
+            sa.sts_from_cognito_identity_pool(jwt_token.get("sub"), cognito_client=self.cognito_client)
+
+        self.assertEqual(str(e.exception), "InvalidCredentialsException: Retries Exceeded: Unexpected Client Error")
+        self.assertEqual(len(self.cognito_stub._queue), 0)
+
     def testAuthErrors(self):
         """
-      Test that errors are thrown when keys are invalid
-      """
+        Test that errors are thrown when keys are invalid
+        """
         sa = StaxAuth("ApiAuth")
         # Test with no username
         with self.assertRaises(InvalidCredentialsException):