Improve JWKS automatic rotation

Author: RobertoPrevato

.flake8

@@ -6,3 +6,4 @@ max-complexity = 18
 select = B,C,E,F,W,T4,B9
 per-file-ignores =
   guardpost/__init__.py:F401
+  tests/test_jwks.py:E501

CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.0.1] - 2023-03-20 :sun_with_face:
+- Improves the automatic rotation of `JWKS`: when validating `JWTs`, `JWKS` are
+  refreshed automatically if an unknown `kid` is encountered, and `JWKS` were
+  last fetched more than `refresh_time` seconds ago (by default 120 seconds).
+- Corrects an inconsistency in how `claims` are read in the `User` class.
+
 ## [1.0.0] - 2023-01-07 :star:
 - Adds built-in support for dependency injection, using the new `ContainerProtocol`
   in `rodi` v2.

guardpost/__about__.py

@@ -1 +1 @@
-__version__ = "1.0.0"
+__version__ = "1.0.1"

guardpost/authentication.py

@@ -26,11 +26,14 @@ def __init__(
 
     @property
     def sub(self) -> Optional[str]:
-        return self["sub"]
+        return self.get("sub")
 
     def is_authenticated(self) -> bool:
         return bool(self.authentication_mode)
 
+    def get(self, key: str):
+        return self.claims.get(key)
+
     def __getitem__(self, item):
         return self.claims[item]
 
@@ -44,15 +47,15 @@ def has_claim_value(self, name: str, value: str) -> bool:
 class User(Identity):
     @property
     def id(self) -> Optional[str]:
-        return self["id"] or self.sub
+        return self.get("id") or self.sub
 
     @property
     def name(self) -> Optional[str]:
-        return self["name"]
+        return self.get("name")
 
     @property
     def email(self) -> Optional[str]:
-        return self["email"]
+        return self.get("email")
 
 
 class AuthenticationHandler(ABC):

guardpost/authorization.py

@@ -86,7 +86,6 @@ def _get_message(forced_failure, failed_requirements):
 
 
 class AuthorizationContext:
-
     __slots__ = ("identity", "requirements", "_succeeded", "_failed_forced")
 
     def __init__(self, identity: Identity, requirements: Sequence[Requirement]):
@@ -222,7 +221,6 @@ async def _handle_with_policy(self, policy: Policy, identity: Identity, scope: A
         with AuthorizationContext(
             identity, list(self._get_requirements(policy, scope))
         ) as context:
-
             for requirement in context.requirements:
                 if _is_async_handler(type(requirement)):  # type: ignore
                     await requirement.handle(context)

guardpost/jwks/__init__.py

@@ -70,6 +70,9 @@ def from_dict(cls, value) -> "JWK":
 class JWKS:
     keys: List[JWK]
 
+    def update(self, new_set: "JWKS"):
+        self.keys = list({key.kid: key for key in self.keys + new_set.keys}.values())
+
     @classmethod
     def from_dict(cls, value) -> "JWKS":
         if "keys" not in value:

guardpost/jwks/caching.py

@@ -1,19 +1,23 @@
 import time
 from typing import Optional
 
-from . import JWKS, KeysProvider
+from . import JWK, JWKS, KeysProvider
 
 
 class CachingKeysProvider(KeysProvider):
     """
     Kind of KeysProvider that can cache the result of another KeysProvider.
     """
 
-    def __init__(self, keys_provider: KeysProvider, cache_time: float) -> None:
+    def __init__(
+        self, keys_provider: KeysProvider, cache_time: float, refresh_time: float = 120
+    ) -> None:
         """
         Creates a new instance of CachingKeysProvider bound to a given KeysProvider,
         and caching its result up to an optional amount of seconds described by
         cache_time. Expiration is disabled if `cache_time` <= 0.
+        JWKS are refreshed anyway if an unknown `kid` is encountered and the set was
+        fetched more than `refresh_time` seconds ago.
         """
         super().__init__()
 
@@ -22,6 +26,7 @@ def __init__(self, keys_provider: KeysProvider, cache_time: float) -> None:
 
         self._keys: Optional[JWKS] = None
         self._cache_time = cache_time
+        self._refresh_time = refresh_time
         self._last_fetch_time: float = 0
         self._keys_provider = keys_provider
 
@@ -34,6 +39,14 @@ async def _fetch_keys(self) -> JWKS:
         self._last_fetch_time = time.time()
         return self._keys
 
+    async def _refresh_keys(self) -> JWKS:
+        new_set = await self._fetch_keys()
+        if self._keys is None:  # pragma: no cover
+            self._keys = new_set
+        else:
+            self._keys.update(new_set)
+        return self._keys
+
     async def get_keys(self) -> JWKS:
         if self._keys is not None:
             if self._cache_time > 0 and (
@@ -43,3 +56,27 @@ async def get_keys(self) -> JWKS:
             else:
                 return self._keys
         return await self._fetch_keys()
+
+    async def get_key(self, kid: str) -> Optional[JWK]:
+        """
+        Tries to get a JWK by kid. If the JWK is not found and the last time the keys
+        were fetched is older than `refresh_time` (default 120 seconds), it fetches
+        again the JWKS from the source.
+        """
+        jwks = await self.get_keys()
+
+        for jwk in jwks.keys.copy():
+            if jwk.kid is not None and jwk.kid == kid:
+                return jwk
+
+        if (
+            self._refresh_time > 0
+            and time.time() - self._last_fetch_time >= self._refresh_time
+        ):
+            jwks = await self._refresh_keys()
+
+            for jwk in jwks.keys.copy():
+                if jwk.kid is not None and jwk.kid == kid:
+                    return jwk
+
+        return None

guardpost/jwts/__init__.py

@@ -44,7 +44,8 @@ def __init__(
         require_kid: bool = True,
         keys_provider: Optional[KeysProvider] = None,
         keys_url: Optional[str] = None,
-        cache_time: float = 10800
+        cache_time: float = 10800,
+        refresh_time: float = 120,
     ) -> None:
         """
         Creates a new instance of JWTValidator. This class only supports validating
@@ -72,9 +73,15 @@ def __init__(
         keys_url : Optional[str], optional
             If provided, keys are obtained from the given URL through HTTP GET.
             This parameter is ignored if `keys_provider` is given.
-        cache_time : float, optional
-            If >= 0, JWKS are cached in memory and stored for the given amount in
-            seconds. By default 10800 (3 hours).
+        cache_time : float
+            JWKS are cached in memory and stored for the given amount in seconds.
+            By default 10800 (3 hours). Regardless of this parameter, JWKS are refreshed
+            automatically if an unknown kid is met and JWKS were last fetched more than
+            `refresh_time` earlier (in seconds).
+        refresh_time : float
+            JWKS are refreshed automatically if an unknown `kid` is encountered, and
+            JWKS were last fetched more than `refresh_time` seconds ago (by default
+            120 seconds)
         """
         if keys_provider:
             pass
@@ -89,26 +96,24 @@ def __init__(
                 "`authority`, or `keys_provider`."
             )
 
-        if cache_time:
-            keys_provider = CachingKeysProvider(keys_provider, cache_time)
+        keys_provider = CachingKeysProvider(keys_provider, cache_time, refresh_time)
 
         self._valid_issuers = list(valid_issuers)
         self._valid_audiences = list(valid_audiences)
         self._algorithms = list(algorithms)
-        self._keys_provider: KeysProvider = keys_provider
+        self._keys_provider = keys_provider
         self.require_kid = require_kid
         self.logger = get_logger()
 
     async def get_jwks(self) -> JWKS:
         return await self._keys_provider.get_keys()
 
     async def get_jwk(self, kid: str) -> JWK:
-        jwks = await self.get_jwks()
+        key = await self._keys_provider.get_key(kid)
 
-        for jwk in jwks.keys:
-            if jwk.kid is not None and jwk.kid == kid:
-                return jwk
-        raise InvalidAccessToken("kid not recognized")
+        if key is None:
+            raise InvalidAccessToken("kid not recognized")
+        return key
 
     def _validate_jwt_by_key(
         self, access_token: str, jwk: JWK

tests/test_authentication.py

@@ -112,7 +112,6 @@ async def authenticate(self, context: Request):
 
 @pytest.mark.asyncio
 async def test_strategy_throws_for_missing_context():
-
     strategy = AuthenticationStrategy()
 
     with raises(ValueError):

tests/test_common.py

@@ -147,7 +147,6 @@ def test_authorization_strategy_set_default_fluent():
 
 
 def test_unauthorized_error_supports_error_and_description():
-
     error = UnauthorizedError(
         None,
         [],