fix: Improve JWKClient and fix get_session of session_functions

Author: KShivendu

supertokens_python/recipe/session/access_token.py

@@ -22,7 +22,7 @@
 from supertokens_python.utils import get_timestamp_ms
 
 from .exceptions import raise_try_refresh_token_exception
-from .jwks import JWKClient
+from .jwks import JWKClient, JWKSRequestError
 from .jwt import ParsedJWTInfo
 
 
@@ -55,28 +55,36 @@ def get_info_from_access_token(
 
         if jwt_info.version < 3:
             # It won't have kid. So we'll have to try the token against all the keys from all the jwk_clients
+            # If any of them work, we'll use that payload
             for client in jwk_clients:
-                keys = client.get_latest_keys()
+                try:
+                    keys = client.get_latest_keys()
 
-                for k in keys:
-                    try:
-                        payload = jwt.decode(jwt_info.raw_token_string, str(k.key), algorithms=["RS256"])  # type: ignore
-                        break
-                    except DecodeError:
-                        pass
+                    for k in keys:
+                        try:
+                            payload = jwt.decode(jwt_info.raw_token_string, str(k.key), algorithms=["RS256"])  # type: ignore
+                            break
+                        except DecodeError:
+                            pass
+                except JWKSRequestError:
+                    continue
+
+                if payload is not None:
+                    break
 
             if payload is None:
                 raise PyJWKClientError("No key found")
 
-        # Came here means token is v3 or above
-        for client in jwk_clients:
-            matching_key = client.get_matching_key_from_jwt(jwt_info.raw_token_string)
-            payload = jwt.decode(  # type: ignore
-                jwt_info.raw_token_string,
-                matching_key,
-                algorithms=["RS256"],
-                options={"verify_signature": True, "verify_exp": True},
-            )
+        elif jwt_info.version > 3:
+            for client in jwk_clients:
+                matching_key = client.get_matching_key_from_jwt(jwt_info.raw_token_string)
+                payload = jwt.decode(  # type: ignore
+                    jwt_info.raw_token_string,
+                    matching_key,
+                    algorithms=["RS256"],
+                    options={"verify_signature": True, "verify_exp": True},
+                )
+                break
 
         assert payload is not None
 

supertokens_python/recipe/session/constants.py

@@ -12,6 +12,7 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 from __future__ import annotations
+
 from typing import TYPE_CHECKING, List
 
 if TYPE_CHECKING:
@@ -32,4 +33,5 @@
 
 available_token_transfer_methods: List[TokenTransferMethod] = ["cookie", "header"]
 
-JWKCacheMaxAgeInMs = 60000
+JWKCacheMaxAgeInMs = 60 * 1000
+JWKRequestCooldownInMs = 500 * 1000

supertokens_python/recipe/session/jwks.py

@@ -8,22 +8,22 @@
 
 from supertokens_python.utils import get_timestamp_ms
 
-from .constants import JWKCacheMaxAgeInMs
+from .constants import JWKCacheMaxAgeInMs, JWKRequestCooldownInMs
 
 
 class JWKClient:
     def __init__(
         self,
         uri: str,
-        cooldown_duration: int = 500,
+        cooldown_duration: int = JWKRequestCooldownInMs,
         cache_max_age: int = JWKCacheMaxAgeInMs,
     ):
         """A client for retrieving JSON Web Key Sets (JWKS) from a given URI.
 
         Args:
             uri (str): The URI of the JWKS.
-            cooldown_duration (int, optional): The cooldown duration in ms. Defaults to 500.
-            cache_max_age (int, optional): The cache max age in ms. Defaults to 300.
+            cooldown_duration (int, optional): The cooldown duration in ms. Defaults to 500 seconds.
+            cache_max_age (int, optional): The cache max age in ms. Defaults to 5 minutes.
 
         Note: The JSON Web Key Set is fetched when no key matches the selection
         process but only as frequently as the `self.cooldown_duration` option
@@ -63,7 +63,8 @@ def get_latest_keys(self) -> List[PyJWK]:
         if self.jwk_set is None or not self.is_fresh():
             self.reload()
 
-        assert self.jwk_set is not None
+        if self.jwk_set is None:
+            raise JWKSRequestError("Failed to fetch the latest keys")
 
         all_keys: List[PyJWK] = self.jwk_set.keys  # type: ignore
 

supertokens_python/recipe/session/recipe_implementation.py

@@ -56,7 +56,7 @@
     from supertokens_python import AppInfo
     from supertokens_python.querier import Querier
 
-from .constants import JWKCacheMaxAgeInMs
+from .constants import JWKCacheMaxAgeInMs, JWKRequestCooldownInMs
 from .interfaces import SessionContainer
 
 protected_props = [
@@ -80,7 +80,7 @@ def __init__(self, querier: Querier, config: SessionConfig, app_info: AppInfo):
     @property
     def JWK_clients(self) -> List[JWKClient]:
         return [
-            JWKClient(uri, cooldown_duration=500, cache_max_age=JWKCacheMaxAgeInMs)
+            JWKClient(uri, cooldown_duration=JWKRequestCooldownInMs, cache_max_age=JWKCacheMaxAgeInMs)
             for uri in self.querier.get_all_core_urls_for_path(".well-known/jwks.json")
         ]
 

supertokens_python/recipe/session/session_functions.py

@@ -145,32 +145,33 @@ async def get_session(
 
     # If we get here we either have a V2 token that doesn't pass verification or a valid V3> token
     # anti-csrf check if accesstokenInfo is not undefined which means token verification was successful
-
-    if config.anti_csrf == "VIA_TOKEN" and do_anti_csrf_check:
-        if access_token_info is not None:
-            if (
-                anti_csrf_token is None
-                or anti_csrf_token != access_token_info["antiCsrfToken"]
-            ):
-                if anti_csrf_token is None:
-                    log_debug_message(
-                        "getSession: Returning TRY_REFRESH_TOKEN because antiCsrfToken is missing from request"
-                    )
-                    raise_try_refresh_token_exception(
-                        "Provided antiCsrfToken is undefined. If you do not want anti-csrf check for this API, please set doAntiCsrfCheck to false for this API"
-                    )
-                else:
-                    log_debug_message(
-                        "getSession: Returning TRY_REFRESH_TOKEN because the passed antiCsrfToken is not the same as in the access token"
-                    )
-                    raise_try_refresh_token_exception("anti-csrf check failed")
-
-    elif config.anti_csrf == "VIA_CUSTOM_HEADER":
-        # The function should never be called by this (we check this outside the function as well)
-        # There we can add a bit more information to the error, so that's the primary check, this is just making sure.
-        raise Exception(
-            "Please either use VIA_TOKEN, NONE or call with doAntiCsrfCheck false"
-        )
+    
+    if do_anti_csrf_check:
+        if config.anti_csrf == "VIA_TOKEN":
+            if access_token_info is not None:
+                if (
+                    anti_csrf_token is None
+                    or anti_csrf_token != access_token_info["antiCsrfToken"]
+                ):
+                    if anti_csrf_token is None:
+                        log_debug_message(
+                            "getSession: Returning TRY_REFRESH_TOKEN because antiCsrfToken is missing from request"
+                        )
+                        raise_try_refresh_token_exception(
+                            "Provided antiCsrfToken is undefined. If you do not want anti-csrf check for this API, please set doAntiCsrfCheck to false for this API"
+                        )
+                    else:
+                        log_debug_message(
+                            "getSession: Returning TRY_REFRESH_TOKEN because the passed antiCsrfToken is not the same as in the access token"
+                        )
+                        raise_try_refresh_token_exception("anti-csrf check failed")
+
+        elif config.anti_csrf == "VIA_CUSTOM_HEADER":
+            # The function should never be called by this (we check this outside the function as well)
+            # There we can add a bit more information to the error, so that's the primary check, this is just making sure.
+            raise Exception(
+                "Please either use VIA_TOKEN, NONE or call with doAntiCsrfCheck false"
+            )
 
     if (
         access_token_info is not None