Throttle repeated PyJWKClient refreshes

Author: jpadilla

jwt/jwks_client.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
 import json
+import math
+import threading
+import time
 import urllib.request
 from functools import lru_cache
 from ssl import SSLContext
@@ -38,6 +41,7 @@ def __init__(
         headers: dict[str, Any] | None = None,
         timeout: float = 30,
         ssl_context: SSLContext | None = None,
+        cooldown_duration: float = 30,
     ):
         """A client for retrieving signing keys from a JWKS endpoint.
 
@@ -54,6 +58,14 @@ def __init__(
         - ``lifespan``: Time in seconds before the cached JWK Set expires.
           Defaults to ``300`` (5 minutes). Must be greater than 0.
 
+        Unknown key IDs can trigger one forced refresh after the cooldown
+        period configured by ``cooldown_duration``. Every successful fetch
+        starts this cooldown, including the initial fetch and cache-expiry
+        fetches. A newly rotated key may therefore wait for the cooldown
+        period before it is fetched; set ``cooldown_duration`` to ``0`` to
+        disable this behavior. The cooldown is bypassed when
+        ``cache_jwk_set`` is ``False``.
+
         **Tier 2 — Signing key cache** (disabled by default):
         Caches individual signing keys (looked up by ``kid``) using an LRU
         cache with **no time-based expiration**. Keys are evicted only when
@@ -80,6 +92,9 @@ def __init__(
         :type timeout: float
         :param ssl_context: Optional SSL context for the request.
         :type ssl_context: ssl.SSLContext or None
+        :param cooldown_duration: Minimum time in seconds between forced
+            refreshes after an unknown key ID. Defaults to ``30``.
+        :type cooldown_duration: float
         """
         if headers is None:
             headers = {}
@@ -98,6 +113,18 @@ def __init__(
         self.headers = headers
         self.timeout = timeout
         self.ssl_context = ssl_context
+        if cooldown_duration < 0:
+            raise PyJWKClientError(
+                "Cooldown duration must be greater than or equal to 0, "
+                f'the input is "{cooldown_duration}"'
+            )
+        if not math.isfinite(cooldown_duration):
+            raise PyJWKClientError(
+                f'Cooldown duration must be finite, the input is "{cooldown_duration}"'
+            )
+        self.cooldown_duration = cooldown_duration
+        self._last_successful_fetch: float | None = None
+        self._client_lock = threading.RLock()
 
         if cache_jwk_set:
             # Init jwt set cache with default or given lifespan.
@@ -147,6 +174,7 @@ def fetch_data(self) -> Any:
         # wipe that breaks legitimate auth.
         if self.jwk_set_cache is not None:
             self.jwk_set_cache.put(jwk_set)
+        self._last_successful_fetch = time.monotonic()
         return jwk_set
 
     def get_jwk_set(self, refresh: bool = False) -> PyJWKSet:
@@ -186,6 +214,10 @@ def get_signing_keys(self, refresh: bool = False) -> list[PyJWK]:
         :raises PyJWKClientError: If no signing keys are found.
         """
         jwk_set = self.get_jwk_set(refresh)
+        return self._get_signing_keys_from_jwk_set(jwk_set)
+
+    @staticmethod
+    def _get_signing_keys_from_jwk_set(jwk_set: PyJWKSet) -> list[PyJWK]:
         signing_keys = [
             jwk_set_key
             for jwk_set_key in jwk_set.keys
@@ -200,8 +232,9 @@ def get_signing_keys(self, refresh: bool = False) -> list[PyJWK]:
     def get_signing_key(self, kid: str) -> PyJWK:
         """Return the signing key matching the given ``kid``.
 
-        If no match is found in the current JWK Set, the set is
-        refreshed from the endpoint and the lookup is retried once.
+        If no match is found in the current JWK Set, the set is refreshed
+        from the endpoint and the lookup is retried once when the refresh
+        cooldown permits it.
 
         :param kid: The key ID to look up.
         :type kid: str
@@ -210,20 +243,28 @@ def get_signing_key(self, kid: str) -> PyJWK:
         :raises PyJWKClientError: If no matching key is found after
             refreshing.
         """
-        signing_keys = self.get_signing_keys()
-        signing_key = self.match_kid(signing_keys, kid)
-
-        if not signing_key:
-            # If no matching signing key from the jwk set, refresh the jwk set and try again.
-            signing_keys = self.get_signing_keys(refresh=True)
+        with self._client_lock:
+            signing_keys = self.get_signing_keys()
             signing_key = self.match_kid(signing_keys, kid)
 
             if not signing_key:
-                raise PyJWKClientError(
-                    f'Unable to find a signing key that matches: "{kid}"'
+                cooling_down = (
+                    self.jwk_set_cache is not None
+                    and self._last_successful_fetch is not None
+                    and time.monotonic() - self._last_successful_fetch
+                    < self.cooldown_duration
                 )
+                if not cooling_down:
+                    signing_keys = self.get_signing_keys(refresh=True)
+                    self._last_successful_fetch = time.monotonic()
+                    signing_key = self.match_kid(signing_keys, kid)
 
-        return signing_key
+                if not signing_key:
+                    raise PyJWKClientError(
+                        f'Unable to find a signing key that matches: "{kid}"'
+                    )
+
+            return signing_key
 
     def get_signing_key_from_jwt(self, token: str | bytes) -> PyJWK:
         """Return the signing key for a JWT by reading its ``kid`` header.

tests/test_jwks_client.py

@@ -386,7 +386,7 @@ def test_failed_request_should_raise_connection_error(self) -> None:
 
     def test_get_jwt_set_refresh_cache(self) -> None:
         url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
-        jwks_client = PyJWKClient(url)
+        jwks_client = PyJWKClient(url, cooldown_duration=0)
 
         kid = "NEE1QURBOTM4MzI5RkFDNTYxOTU1MDg2ODgwQ0UzMTk1QjYyRkRFQw"
 
@@ -411,6 +411,145 @@ def test_get_jwt_set_no_matching_kid_after_second_attempt(self) -> None:
             ):
                 jwks_client.get_signing_key(kid)
 
+    def test_unknown_kid_refresh_is_cooled_down(self) -> None:
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+        kid = "unknown-kid"
+        jwks_client = PyJWKClient(url, cooldown_duration=30)
+
+        with mock.patch("urllib.request.build_opener") as build_opener_mock:
+            opener = mock.Mock()
+            build_opener_mock.return_value = opener
+            response = mock.Mock()
+            response.__enter__ = mock.Mock(return_value=response)
+            response.__exit__ = mock.Mock()
+            response.read.return_value = json.dumps(RESPONSE_DATA_NO_MATCHING_KID)
+            opener.open.return_value = response
+
+            for _ in range(2):
+                with pytest.raises(PyJWKClientError, match="matches"):
+                    jwks_client.get_signing_key(kid)
+
+        assert opener.open.call_count == 1
+
+    def test_unknown_kid_refresh_runs_again_after_cooldown(self) -> None:
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+        kid = "unknown-kid"
+        jwks_client = PyJWKClient(url, cooldown_duration=0.01)
+
+        with mock.patch("urllib.request.build_opener") as build_opener_mock:
+            opener = mock.Mock()
+            build_opener_mock.return_value = opener
+            response = mock.Mock()
+            response.__enter__ = mock.Mock(return_value=response)
+            response.__exit__ = mock.Mock()
+            response.read.return_value = json.dumps(RESPONSE_DATA_NO_MATCHING_KID)
+            opener.open.return_value = response
+
+            clock = [0.0]
+            with mock.patch(
+                "jwt.jwks_client.time.monotonic", side_effect=lambda: clock[0]
+            ):
+                with pytest.raises(PyJWKClientError, match="matches"):
+                    jwks_client.get_signing_key(kid)
+                clock[0] = 0.02
+                with pytest.raises(PyJWKClientError, match="matches"):
+                    jwks_client.get_signing_key(kid)
+
+        assert opener.open.call_count == 2
+
+    def test_unknown_kid_refresh_serializes_concurrent_misses(self) -> None:
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+        kid = "unknown-kid"
+        jwks_client = PyJWKClient(url, cooldown_duration=30)
+        refresh_started = threading.Event()
+        second_done = threading.Event()
+        release_refresh = threading.Event()
+
+        with mock.patch("urllib.request.build_opener") as build_opener_mock:
+            opener = mock.Mock()
+            build_opener_mock.return_value = opener
+            response = mock.Mock()
+            response.__enter__ = mock.Mock(return_value=response)
+            response.__exit__ = mock.Mock()
+            response.read.return_value = json.dumps(RESPONSE_DATA_NO_MATCHING_KID)
+
+            opener.open.return_value = response
+            jwks_client.get_jwk_set()
+            jwks_client._last_successful_fetch = 0
+
+            def open_response(*args: object, **kwargs: object) -> mock.Mock:
+                if opener.open.call_count == 1:
+                    refresh_started.set()
+                    release_refresh.wait(timeout=5)
+                return response
+
+            opener.open.reset_mock()
+            opener.open.side_effect = open_response
+            errors: list[Exception] = []
+
+            def lookup() -> None:
+                try:
+                    jwks_client.get_signing_key(kid)
+                except PyJWKClientError as error:
+                    errors.append(error)
+
+            first = threading.Thread(target=lookup)
+
+            def second_lookup() -> None:
+                lookup()
+                second_done.set()
+
+            second = threading.Thread(target=second_lookup)
+            with mock.patch("jwt.jwks_client.time.monotonic", return_value=31):
+                first.start()
+                assert refresh_started.wait(timeout=5)
+                second.start()
+                assert not second_done.wait(timeout=0.1)
+                release_refresh.set()
+                first.join(timeout=5)
+                second.join(timeout=5)
+
+        assert len(errors) == 2
+        assert not first.is_alive()
+        assert not second.is_alive()
+        assert opener.open.call_count == 1
+
+    def test_unknown_kid_refresh_ignores_cooldown_when_cache_disabled(self) -> None:
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+        kid = "unknown-kid"
+        jwks_client = PyJWKClient(url, cache_jwk_set=False, cooldown_duration=30)
+
+        with mock.patch("urllib.request.build_opener") as build_opener_mock:
+            opener = mock.Mock()
+            build_opener_mock.return_value = opener
+            response = mock.Mock()
+            response.__enter__ = mock.Mock(return_value=response)
+            response.__exit__ = mock.Mock()
+            response.read.side_effect = [
+                json.dumps(RESPONSE_DATA_NO_MATCHING_KID),
+                json.dumps(RESPONSE_DATA_NO_MATCHING_KID),
+                json.dumps(RESPONSE_DATA_NO_MATCHING_KID),
+                json.dumps(RESPONSE_DATA_NO_MATCHING_KID),
+            ]
+            opener.open.return_value = response
+
+            for _ in range(2):
+                with pytest.raises(PyJWKClientError, match="matches"):
+                    jwks_client.get_signing_key(kid)
+
+        assert opener.open.call_count == 4
+
+    @pytest.mark.parametrize(
+        "cooldown_duration", [float("nan"), float("inf"), float("-inf")]
+    )
+    def test_unknown_kid_refresh_rejects_non_finite_cooldown(
+        self, cooldown_duration: float
+    ) -> None:
+        url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"
+
+        with pytest.raises(PyJWKClientError, match="Cooldown duration"):
+            PyJWKClient(url, cooldown_duration=cooldown_duration)
+
     def test_get_jwt_set_invalid_lifespan(self) -> None:
         url = "https://dev-87evx9ru.auth0.com/.well-known/jwks.json"