From ba4853a75fb9676362da17f67d0f64bd18afd4e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Padilla?= Date: Thu, 10 Sep 2026 14:10:52 -0400 Subject: [PATCH] Throttle repeated PyJWKClient refreshes --- jwt/jwks_client.py | 63 ++++++++++++++--- tests/test_jwks_client.py | 141 +++++++++++++++++++++++++++++++++++++- 2 files changed, 192 insertions(+), 12 deletions(-) Index: pyjwt-2.12.1/jwt/jwks_client.py =================================================================== --- pyjwt-2.12.1.orig/jwt/jwks_client.py +++ pyjwt-2.12.1/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 @@ -25,6 +28,7 @@ class PyJWKClient: 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. @@ -41,6 +45,14 @@ class PyJWKClient: - ``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 @@ -67,6 +79,9 @@ class PyJWKClient: :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 = {} @@ -85,6 +100,18 @@ class PyJWKClient: 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. @@ -132,6 +159,7 @@ class PyJWKClient: # 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: @@ -171,6 +199,10 @@ class PyJWKClient: :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 @@ -185,8 +217,9 @@ class PyJWKClient: 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 @@ -195,20 +228,28 @@ class PyJWKClient: :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) + + if not signing_key: + raise PyJWKClientError( + f'Unable to find a signing key that matches: "{kid}"' + ) - return signing_key + 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. Index: pyjwt-2.12.1/tests/test_jwks_client.py =================================================================== --- pyjwt-2.12.1.orig/tests/test_jwks_client.py +++ pyjwt-2.12.1/tests/test_jwks_client.py @@ -324,7 +324,7 @@ class TestPyJWKClient: 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" @@ -349,6 +349,145 @@ class TestPyJWKClient: ): 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"