Add stale_if_error option to PyJWKClient

CHANGELOG.rst

@@ -7,6 +7,17 @@ This project adheres to `Semantic Versioning <https://semver.org/>`__.
 `Unreleased <https://github.com/jpadilla/pyjwt/compare/2.15.1...HEAD>`__
 -------------------------------------------------------------------------
 
+Added
+~~~~~
+
+- Add ``stale_if_error`` to ``PyJWKClient`` to keep serving an expired JWK Set
+  for up to that many seconds when refreshing it fails with a connection error,
+  a timeout, or an HTTP 500, 502, 503 or 504 response. Disabled by default.
+  After a fallback, fetches pause for ``cooldown_duration``.
+  ``JWKSetCache.get()`` and ``JWKSetCache.is_expired()`` accept a matching
+  ``max_stale`` argument,
+  in `#1239 <https://github.com/jpadilla/pyjwt/pull/1239>`__.
+
 Changed
 ~~~~~~~
 

jwt/jwk_set_cache.py

@@ -32,16 +32,28 @@ def put(self, jwk_set: Optional[Union[PyJWKSet, dict[str, Any]]]) -> None:
 
         self.jwk_set_with_timestamp = PyJWTSetWithTimestamp(jwk_set)
 
-    def get(self) -> Optional[PyJWKSet]:
-        if self.jwk_set_with_timestamp is None or self.is_expired():
+    def get(self, max_stale: float = 0) -> Optional[PyJWKSet]:
+        """Return the cached ``PyJWKSet`` if it is not expired.
+
+        Can return keys that are up to *max_stale* seconds
+        past expiration. Returns ``None`` if the cache is empty
+        or expired past *max_stale* tolerance.
+        """
+        if self.jwk_set_with_timestamp is None or self.is_expired(max_stale):
             return None
 
         return self.jwk_set_with_timestamp.get_jwk_set()
 
-    def is_expired(self) -> bool:
+    def is_expired(self, max_stale: float = 0) -> bool:
+        """Return whether the currently cached ``PyJWKSet`` is expired.
+
+        Returns ``True`` if the cached ``PyJWKSet`` is more than
+        *max_stale* seconds past its lifetime. Returns ``False``
+        otherwise. Always returns ``False`` for a lifespan of ``-1``.
+        """
         return (
             self.jwk_set_with_timestamp is not None
             and self.lifespan > -1
             and time.monotonic()
-            > self.jwk_set_with_timestamp.get_timestamp() + self.lifespan
+            > self.jwk_set_with_timestamp.get_timestamp() + self.lifespan + max_stale
         )

jwt/jwks_client.py

@@ -43,6 +43,7 @@ def __init__(
         timeout: float = 30,
         ssl_context: SSLContext | None = None,
         cooldown_duration: float = 30,
+        stale_if_error: float = 0,
     ):
         """A client for retrieving signing keys from a JWKS endpoint.
 
@@ -67,6 +68,16 @@ def __init__(
         disable this behavior. The cooldown is bypassed when
         ``cache_jwk_set`` is ``False``.
 
+        If refreshing an expired JWK Set fails, ``stale_if_error`` lets the
+        client keep serving it for up to that many seconds past expiry
+        instead of raising. This covers connection failures, timeouts, and
+        HTTP 500, 502, 503 and 504 responses; other errors still raise, as
+        does ``get_jwk_set(refresh=True)``. Pass ``float("inf")`` to serve it
+        for as long as fetching fails. Stale keys are served silently, and
+        a key removed from the endpoint stays usable until the stale set is
+        discarded. After a fallback, the client serves the stale set without
+        fetching for ``cooldown_duration`` seconds.
+
         **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
@@ -94,8 +105,13 @@ def __init__(
         :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``.
+            refreshes after an unknown key ID, and between fetches after a
+            ``stale_if_error`` fallback. Defaults to ``30``.
         :type cooldown_duration: float
+        :param stale_if_error: Time in seconds past expiry that a cached
+            JWK Set may still be served when refreshing it fails. Defaults
+            to ``0`` (disabled).
+        :type stale_if_error: float
         """
         if headers is None:
             headers = {}
@@ -124,7 +140,26 @@ def __init__(
                 f'Cooldown duration must be finite, the input is "{cooldown_duration}"'
             )
         self.cooldown_duration = cooldown_duration
+
+        if stale_if_error < 0:
+            raise PyJWKClientError(
+                "Stale if error must be greater than or equal to 0, "
+                f'the input is "{stale_if_error}"'
+            )
+
+        if math.isnan(stale_if_error):
+            raise PyJWKClientError(
+                f'Stale if error must not be NaN, the input is "{stale_if_error}"'
+            )
+
+        if stale_if_error > 0 and not cache_jwk_set:
+            raise PyJWKClientError(
+                "Stale if error must be 0 if JWK Set caching is disabled"
+            )
+        self.stale_if_error = stale_if_error
+
         self._last_successful_fetch: float | None = None
+        self._last_failed_fetch: float | None = None
         self._client_lock = threading.RLock()
 
         if cache_jwk_set:
@@ -187,6 +222,19 @@ def fetch_data(self) -> Any:
         self._last_successful_fetch = time.monotonic()
         return jwk_set
 
+    @staticmethod
+    def _should_serve_stale_on(cause: BaseException | None) -> bool:
+        if cause is None or isinstance(cause, http.client.HTTPException):
+            return True
+
+        if isinstance(cause, HTTPError):
+            return cause.code in (500, 502, 503, 504)
+
+        if isinstance(cause, OSError):
+            return True
+
+        return False
+
     def get_jwk_set(self, refresh: bool = False) -> PyJWKSet:
         """Return the JWK Set, using the cache when available.
 
@@ -203,9 +251,35 @@ def get_jwk_set(self, refresh: bool = False) -> PyJWKSet:
         if self.jwk_set_cache is not None and not refresh:
             data = self.jwk_set_cache.get()
 
+        skip_stale_if_error = self.stale_if_error == 0 or refresh
+
+        # After a stale_if_error fallback, keep serving the stale set
+        # without fetching until the cooldown has passed.
+        if data is None and not skip_stale_if_error:
+            if (
+                self.jwk_set_cache is not None
+                and self._last_failed_fetch is not None
+                and time.monotonic() < self._last_failed_fetch + self.cooldown_duration
+            ):
+                data = self.jwk_set_cache.get(max_stale=self.stale_if_error)
+
         if data is None:
-            data = self.fetch_data()
-            fetched = True
+            try:
+                data = self.fetch_data()
+                fetched = True
+                self._last_failed_fetch = None
+            except PyJWKClientConnectionError as e:
+                if skip_stale_if_error or self.jwk_set_cache is None:
+                    raise
+
+                cause = e.__cause__
+                if self._should_serve_stale_on(cause):
+                    data = self.jwk_set_cache.get(max_stale=self.stale_if_error)
+
+                if data is None:
+                    raise
+
+                self._last_failed_fetch = time.monotonic()
 
         # A cache hit is already parsed, so serve it as-is. Only a fresh
         # fetch reaches the check below, which still matters because