refactor: Implement and use custom JWKClient

Author: KShivendu

supertokens_python/recipe/session/access_token.py

@@ -22,6 +22,7 @@
 from supertokens_python.utils import get_timestamp_ms
 
 from .exceptions import raise_try_refresh_token_exception
+from .jwks import JWKClient
 from .jwt import ParsedJWTInfo
 
 
@@ -45,37 +46,37 @@ def sanitize_number(n: Any) -> Union[Union[int, float], None]:
 
 def get_info_from_access_token(
     jwt_info: ParsedJWTInfo,
-    jwk_clients: List[jwt.PyJWKClient],
+    jwk_clients: List[JWKClient],
     do_anti_csrf_check: bool,
 ):
+    # TODO: Add different tests to verify this works as expected
     try:
         payload: Optional[Dict[str, Any]] = None
 
+        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
+            for client in jwk_clients:
+                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
+
+            if payload is None:
+                raise PyJWKClientError("No key found")
+
+        # Came here means token is v3 or above
         for client in jwk_clients:
-            try:
-                # TODO: verify this works as expected
-                signing_key: str = client.get_signing_key_from_jwt(jwt_info.raw_token_string).key  # type: ignore
-                payload = jwt.decode(  # type: ignore
-                    jwt_info.raw_token_string,
-                    signing_key,
-                    algorithms=["RS256"],
-                    options={"verify_signature": True, "verify_exp": True},
-                )
-            except PyJWKClientError as e:
-                # If no kid is present, this error is thrown
-                # So we'll have to try the token against all the keys if it's v2
-                if jwt_info.version == 2:
-                    for client in jwk_clients:
-                        keys = client.get_signing_keys()
-                        for k in keys:
-                            try:
-                                payload = jwt.decode(jwt_info.raw_token_string, str(k.key), algorithms=["RS256"])  # type: ignore
-                            except DecodeError:
-                                pass
-                    if payload is None:
-                        raise e
-            except DecodeError as e:
-                raise e
+            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},
+            )
 
         assert payload is not None
 
@@ -90,7 +91,7 @@ def get_info_from_access_token(
             user_id = sanitize_string(payload.get("sub"))
             expiry_time = sanitize_number(
                 payload.get("exp", 0) * 1000
-            )  # FIXME: Is adding 0 as default okay?
+            )  # FIXME: Is using 0 as default okay?
             time_created = sanitize_number(payload.get("iat", 0) * 1000)
             user_data = payload
 

supertokens_python/recipe/session/constants.py

@@ -31,3 +31,5 @@
 ACCESS_CONTROL_EXPOSE_HEADERS = "Access-Control-Expose-Headers"
 
 available_token_transfer_methods: List[TokenTransferMethod] = ["cookie", "header"]
+
+JWKCacheMaxAgeInMs = 60000

supertokens_python/recipe/session/jwks.py

@@ -0,0 +1,101 @@
+import json
+import urllib.request
+from typing import List, Optional
+from urllib.error import URLError
+
+from jwt import PyJWK, PyJWKSet
+from jwt.api_jwt import decode_complete as decode_token  # type: ignore
+
+from supertokens_python.utils import get_timestamp_ms
+
+from .constants import JWKCacheMaxAgeInMs
+
+
+class JWKClient:
+    def __init__(
+        self,
+        uri: str,
+        cooldown_duration: int = 500,
+        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.
+
+        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
+        allows to prevent abuse. The `self.cache_max_age` option is used to
+        determine how long the JWKS is cached for.
+
+        Whenever you make a call to `get_signing_key_from_jwt`, the JWKS
+        is fetched if it is older than `self.cache_max_age` ms unless
+        cooldown is active.
+        """
+        self.uri = uri
+        self.cooldown_duration = cooldown_duration
+        self.cache_max_age = cache_max_age
+        self.timeout_sec = 5
+        self.last_fetch_time: int = 0
+        self.jwk_set: Optional[PyJWKSet] = None
+
+    def reload(self):
+        try:
+            with urllib.request.urlopen(self.uri, timeout=self.timeout_sec) as response:
+                self.jwk_set = PyJWKSet.from_dict(json.load(response))  # type: ignore
+                self.last_fetch_time = get_timestamp_ms()
+        except URLError as e:
+            raise JWKSRequestError(f'Failed to fetch data from the url, err: "{e}"')
+
+    def is_cooling_down(self) -> bool:
+        return (self.last_fetch_time > 0) and (
+            get_timestamp_ms() - self.last_fetch_time < self.cooldown_duration
+        )
+
+    def is_fresh(self) -> bool:
+        return (self.last_fetch_time > 0) and (
+            get_timestamp_ms() - self.last_fetch_time < self.cache_max_age
+        )
+
+    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
+
+        all_keys: List[PyJWK] = self.jwk_set.keys  # type: ignore
+
+        return all_keys
+
+    def get_matching_key_from_jwt(self, token: str) -> str:
+        header = decode_token(token, options={"verify_signature": False})["header"]
+        kid: str = header["kid"]  # type: ignore
+
+        if self.jwk_set is None or not self.is_fresh():
+            self.reload()
+
+        assert self.jwk_set is not None
+
+        try:
+            return str(self.jwk_set[kid].key)  # type: ignore
+        except KeyError:
+            if not self.is_cooling_down():
+                # One more attempt to fetch the latest keys
+                # and then try to find the key again.
+                self.reload()
+                try:
+                    return str(self.jwk_set[kid].key)  # type: ignore
+                except KeyError:
+                    pass
+
+        raise JWKSKeyNotFoundError("No key found for the given kid")
+
+
+class JWKSKeyNotFoundError(Exception):
+    pass
+
+
+class JWKSRequestError(Exception):
+    pass

supertokens_python/recipe/session/recipe_implementation.py

@@ -16,8 +16,6 @@
 import json
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
 
-import jwt
-
 from supertokens_python.logger import log_debug_message
 from supertokens_python.normalised_url_path import NormalisedURLPath
 from supertokens_python.utils import resolve
@@ -48,16 +46,17 @@
     SessionInformationResult,
     SessionObj,
 )
+from .jwks import JWKClient
 from .jwt import ParsedJWTInfo, parse_jwt_without_signature_verification
 from .session_class import Session
 from .utils import SessionConfig, validate_claims_in_payload
 
-
 if TYPE_CHECKING:
     from typing import List, Union
     from supertokens_python import AppInfo
     from supertokens_python.querier import Querier
 
+from .constants import JWKCacheMaxAgeInMs
 from .interfaces import SessionContainer
 
 protected_props = [
@@ -79,10 +78,9 @@ def __init__(self, querier: Querier, config: SessionConfig, app_info: AppInfo):
         self.app_info = app_info
 
     @property
-    def JWK_Clients(self) -> List[jwt.PyJWKClient]:
-        # FIXME: Find params OR Implement caching
+    def JWK_clients(self) -> List[JWKClient]:
         return [
-            jwt.PyJWKClient(uri)
+            JWKClient(uri, cooldown_duration=500, 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

@@ -19,6 +19,7 @@
 from supertokens_python.recipe.session.interfaces import SessionInformationResult
 
 from .access_token import get_info_from_access_token
+from .constants import JWKCacheMaxAgeInMs
 from .jwt import ParsedJWTInfo
 
 if TYPE_CHECKING:
@@ -28,9 +29,6 @@
 from supertokens_python.normalised_url_path import NormalisedURLPath
 from supertokens_python.process_state import AllowedProcessStates, ProcessState
 
-JWKCacheMaxAgeInMs = 60000
-
-
 from .exceptions import (
     TryRefreshTokenError,
     raise_token_theft_exception,
@@ -83,7 +81,7 @@ async def get_session(
     try:
         access_token_info = get_info_from_access_token(
             parsed_access_token,
-            recipe_implementation.JWK_Clients,  # FIXME: Use JWKS
+            recipe_implementation.JWK_clients,
             config.anti_csrf == "VIA_TOKEN" and do_anti_csrf_check,
         )