diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6b471ed795b..669924077e3 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3756,6 +3756,11 @@ class ProxyErrorTypes(str, enum.Enum): General authentication error """ + auth_provider_unavailable = "auth_provider_unavailable" + """ + The identity provider needed to authenticate the request (e.g. its JWKS endpoint) is unreachable + """ + internal_server_error = "internal_server_error" """ Internal server error @@ -3846,6 +3851,7 @@ class ProxyErrorTypes(str, enum.Enum): DB_CONNECTION_ERROR_TYPES: Final = ( httpx.ConnectError, + httpx.ConnectTimeout, httpx.ReadError, httpx.ReadTimeout, ) @@ -4524,6 +4530,9 @@ class JWTIssuerConfig(BaseModel): return self +DEFAULT_JWKS_STALE_TTL: Final = 3600 + + class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): """ A class to define the roles and permissions for a LiteLLM Proxy w/ JWT Auth. @@ -4539,6 +4548,8 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): - user_allowed_email_subdomain: If specified, only emails from specified subdomain will be allowed to access proxy. - end_user_id_jwt_field: The field in the JWT token that stores the end-user ID (maps to `LiteLLMEndUserTable`). Turn this off by setting to `None`. Enables end-user cost tracking. Use this for external customers. - public_key_ttl: Default - 600s. TTL for caching public JWT keys. + - public_key_stale_ttl: Default - 3600s. Extra time past `public_key_ttl` that the last-known-good JWKS response + stays usable while the identity provider is unreachable. Set to 0 to fail closed instead. - public_allowed_routes: list of allowed routes for authenticated but unknown litellm role jwt tokens. - enforce_rbac: If true, enforce RBAC for all routes. - custom_validate: A custom function to validates the JWT token. @@ -4589,6 +4600,15 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): user_id_upsert: bool = Field(default=False, description="If user doesn't exist, upsert them into the db.") end_user_id_jwt_field: str | None = None public_key_ttl: float = 600 + public_key_stale_ttl: float = Field( + default=DEFAULT_JWKS_STALE_TTL, + ge=0, + description=( + "Seconds beyond `public_key_ttl` that the last-known-good JWKS response stays usable while the identity " + "provider is unreachable. Bounds how long a signing key the provider has since removed can still be " + "trusted. Set to 0 to fail closed and reject requests as soon as the cached keys expire." + ), + ) public_allowed_routes: list[str] = ["public_routes"] enforce_rbac: bool = False roles_jwt_field: str | None = None # v2 on role mappings diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 1e3265af967..39e6ca9a369 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -8,12 +8,16 @@ JWT token must have 'litellm_proxy_admin' in scope. from __future__ import annotations +import asyncio import fnmatch import hashlib import os import re -from typing import Any, Final, Literal, NoReturn, cast +import time +from collections.abc import Awaitable, Callable +from typing import Any, Final, Literal, NoReturn, TypeVar, cast +import httpx import jwt from cryptography import x509 from cryptography.hazmat.backends import default_backend @@ -25,6 +29,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.llms.custom_httpx.httpx_handler import HTTPHandler from litellm.proxy._types import ( + DEFAULT_JWKS_STALE_TTL, RBAC_ROLES, JWKKeyValue, JWTAuthBuilderResult, @@ -74,6 +79,32 @@ class NoMatchingJWTPublicKeyError(Exception): """Raised when a JWKS endpoint returns no key matching the requested ``kid``.""" +class JWKSUnreachableError(Exception): + """Raised when an IdP's JWKS / OIDC discovery endpoint is unreachable and no cached copy is left to fall back on.""" + + +JWKS_FETCH_ATTEMPTS: Final = 3 +JWKS_FETCH_RETRY_BACKOFF_SECONDS: Final = 0.25 +JWKS_UNREACHABLE_BACKOFF_SECONDS: Final = 30 +STALE_CACHE_KEY_PREFIX: Final = "litellm_stale_" +STALE_WRITTEN_AT_CACHE_KEY_PREFIX: Final = "litellm_stale_written_at_" +UNREACHABLE_CACHE_KEY_PREFIX: Final = "litellm_jwks_unreachable_" + +_CachedValueT = TypeVar("_CachedValueT", bound=JWKKeyValue | str) + + +def jwks_unavailable_exception(error: JWKSUnreachableError) -> ProxyException: + return ProxyException( + message=( + "Service Unavailable, the identity provider's JWKS endpoint is temporarily " + f"unreachable, so the JWT signature could not be verified. Please retry shortly. Error: {error}" + ), + type=ProxyErrorTypes.auth_provider_unavailable, + param="None", + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) + + class JWTHandler: """ - treat the sub id passed in as the user id @@ -121,6 +152,8 @@ class JWTHandler: ) -> None: self.http_handler = HTTPHandler() self.leeway = 0 + # Per-cache-key locks so a TTL lapse triggers one refresh instead of one per in-flight request. + self._refresh_locks: dict[str, asyncio.Lock] = {} # mutable-ok: lock registry, keyed by JWKS url def update_environment( self, @@ -611,13 +644,151 @@ class JWTHandler: if ".well-known/openid-configuration" not in url: return url - cache_key: Final = f"litellm_oidc_discovery_{url}" - cached_jwks_uri: Final = await self.user_api_key_cache.async_get_cache(cache_key) - if cached_jwks_uri is not None: - return cached_jwks_uri + return await self._cached_with_stale_fallback( + cache_key=f"litellm_oidc_discovery_{url}", + ttl=self._get_public_key_cache_ttl(), + refresh=lambda: self._fetch_jwks_uri_from_discovery(url), + log_context="an OIDC discovery lookup", + ) + async def _get_with_transient_retries(self, url: str) -> httpx.Response: + """GET ``url``, retrying transport failures so one IdP blip does not fail the request.""" + for attempt in range(1, JWKS_FETCH_ATTEMPTS): + try: + return await self.http_handler.get(url) + except httpx.TransportError as e: + verbose_proxy_logger.warning( + "JWT Auth: %s fetching %s (attempt %s/%s), retrying: %s", + type(e).__name__, + url, + attempt, + JWKS_FETCH_ATTEMPTS, + e, + ) + await asyncio.sleep(JWKS_FETCH_RETRY_BACKOFF_SECONDS * attempt) + + try: + return await self.http_handler.get(url) + except httpx.TransportError as e: + raise JWKSUnreachableError(f"{type(e).__name__} fetching {url} after {JWKS_FETCH_ATTEMPTS} attempts") from e + + async def _get_cached_value(self, cache_key: str) -> _CachedValueT | None: + cached: Final = await self.user_api_key_cache.async_get_cache(cache_key) + return cast("_CachedValueT | None", cached) # cast-ok: cache reads are untyped + + async def _get_cached_timestamp(self, cache_key: str) -> float | None: + cached: Final = await self.user_api_key_cache.async_get_cache(cache_key) + # A JSON round-trip through Redis hands a whole-number epoch back as an int. + return float(cached) if isinstance(cached, (int, float)) else None + + async def _put_cached_value(self, cache_key: str, value: JWKKeyValue | str | float, ttl: float) -> None: + await self.user_api_key_cache.async_set_cache(key=cache_key, value=value, ttl=ttl) + + async def _cached_with_stale_fallback( + self, + cache_key: str, + ttl: float, + refresh: Callable[[], Awaitable[_CachedValueT]], + log_context: str, + ) -> _CachedValueT: + """Read ``cache_key``, refreshing it through a single-flight lock on a miss.""" + cached: Final[_CachedValueT | None] = await self._get_cached_value(cache_key) + if cached is not None: + return cached + + lock: Final = self._refresh_locks.setdefault(cache_key, asyncio.Lock()) + async with lock: + cached_after_lock: Final[_CachedValueT | None] = await self._get_cached_value(cache_key) + if cached_after_lock is not None: + return cached_after_lock + return await self._refresh_or_serve_stale( + cache_key=cache_key, ttl=ttl, refresh=refresh, log_context=log_context + ) + + async def _refresh_or_serve_stale( + self, + cache_key: str, + ttl: float, + refresh: Callable[[], Awaitable[_CachedValueT]], + log_context: str, + ) -> _CachedValueT: + """Refresh ``cache_key`` from the IdP, falling back to the last-known-good copy when it is unreachable. + + Signing keys rotate rarely, so a last-known-good key beats failing authentication during an IdP blip. + How long a key the IdP has since removed stays trusted is bounded by ``public_key_ttl`` + + ``public_key_stale_ttl`` measured from when the copy was taken, and that bound is enforced here on every + read rather than baked into the cache entry's own expiry. An operator who lowers ``public_key_stale_ttl``, + or sets it to 0 to fail closed, is usually doing it mid-incident, and a copy written under the old longer + setting would otherwise stay servable until it aged out on its own. A copy whose write time cannot be + established is not servable, so the bound cannot be dodged by losing the timestamp. + """ + stale_ttl: Final = self._get_public_key_stale_ttl() + outcome: Final = await self._refresh_or_record_outage( + cache_key=cache_key, ttl=ttl, stale_ttl=stale_ttl, refresh=refresh + ) + if not isinstance(outcome, JWKSUnreachableError): + return outcome + if stale_ttl <= 0: + raise outcome + + stale: Final[_CachedValueT | None] = await self._get_cached_value(f"{STALE_CACHE_KEY_PREFIX}{cache_key}") + age: Final = await self._stale_copy_age(cache_key) + lifetime: Final = ttl + stale_ttl + if stale is None or age is None or age > lifetime: + raise outcome + verbose_proxy_logger.warning( + "JWT Auth: identity provider unreachable, authenticating %s against a stale JWKS copy of %s " + "(last refreshed %.0fs ago, stops being trusted in %.0fs). Refresh failed: %s", + log_context, + cache_key, + age, + max(lifetime - age, 0), + outcome, + ) + return stale + + async def _stale_copy_age(self, cache_key: str) -> float | None: + written_at: Final = await self._get_cached_timestamp(f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{cache_key}") + return None if written_at is None else time.time() - written_at + + async def _refresh_or_record_outage( + self, + cache_key: str, + ttl: float, + stale_ttl: float, + refresh: Callable[[], Awaitable[_CachedValueT]], + ) -> _CachedValueT | JWKSUnreachableError: + """Refresh ``cache_key``, returning the outage as a value rather than raising it. + + A failed refresh is remembered for ``JWKS_UNREACHABLE_BACKOFF_SECONDS`` so a sustained outage costs one + fetch per window instead of one per request serialised behind the refresh lock. + """ + unreachable_cache_key: Final = f"{UNREACHABLE_CACHE_KEY_PREFIX}{cache_key}" + recent_failure: Final[str | None] = await self._get_cached_value(unreachable_cache_key) + if recent_failure is not None: + return JWKSUnreachableError(recent_failure) + + try: + refreshed: Final = await refresh() + except JWKSUnreachableError as e: + await self._put_cached_value( + cache_key=unreachable_cache_key, value=str(e), ttl=JWKS_UNREACHABLE_BACKOFF_SECONDS + ) + return e + + await self._put_cached_value(cache_key=cache_key, value=refreshed, ttl=ttl) + if stale_ttl > 0: + await self._put_cached_value( + cache_key=f"{STALE_CACHE_KEY_PREFIX}{cache_key}", value=refreshed, ttl=ttl + stale_ttl + ) + await self._put_cached_value( + cache_key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{cache_key}", value=time.time(), ttl=ttl + stale_ttl + ) + return refreshed + + async def _fetch_jwks_uri_from_discovery(self, url: str) -> str: verbose_proxy_logger.debug("JWT Auth: Fetching OIDC discovery document from %s", url) - response: Final = await self.http_handler.get(url) + response: Final = await self._get_with_transient_retries(url) if response.status_code != 200: raise Exception( f"JWT Auth: OIDC discovery endpoint {url} returned status {response.status_code}: {response.text}" @@ -632,11 +803,6 @@ class JWTHandler: raise Exception(f"JWT Auth: OIDC discovery document at {url} does not contain a 'jwks_uri' field.") verbose_proxy_logger.debug("JWT Auth: Resolved OIDC discovery %s -> jwks_uri=%s", url, jwks_uri) - await self.user_api_key_cache.async_set_cache( - key=cache_key, - value=jwks_uri, - ttl=self._get_public_key_cache_ttl(), - ) return jwks_uri def _get_public_key_cache_ttl(self) -> float: @@ -645,33 +811,36 @@ class JWTHandler: return 600 return litellm_jwtauth.public_key_ttl + def _get_public_key_stale_ttl(self) -> float: + litellm_jwtauth: Final = getattr(self, "litellm_jwtauth", None) + if litellm_jwtauth is None: + return DEFAULT_JWKS_STALE_TTL + return litellm_jwtauth.public_key_stale_ttl + + async def _fetch_jwks_keys(self, resolved_jwks_url: str) -> JWKKeyValue: + response: Final = await self._get_with_transient_retries(resolved_jwks_url) + if response.status_code != 200: + raise Exception( + f"JWT Auth: JWKS endpoint {resolved_jwks_url} returned status {response.status_code}: {response.text}" + ) + + try: + response_json: Final = response.json() + except Exception as e: + verbose_proxy_logger.error("Error parsing response: %s. Original Response: %s", e, response.text) + raise Exception(f"Error parsing response: {e}. Check server logs for original response.") + + keys: Final = response_json["keys"] if "keys" in response_json else response_json + return cast(JWKKeyValue, keys) # cast-ok: JWTKeyItem declares only `kid`, validating would drop key material + async def _get_public_key_from_jwks_url(self, jwks_url: str, kid: str | None) -> dict: resolved_jwks_url: Final = await self._resolve_jwks_url(jwks_url) - cache_key: Final = f"litellm_jwt_auth_keys_{resolved_jwks_url}" - - cached_keys: Final = await self.user_api_key_cache.async_get_cache(cache_key) - - if cached_keys is None: - response: Final = await self.http_handler.get(resolved_jwks_url) - - try: - response_json: Final = response.json() - except Exception as e: - verbose_proxy_logger.error("Error parsing response: %s. Original Response: %s", e, response.text) - raise Exception(f"Error parsing response: {e}. Check server logs for original response.") - - if "keys" in response_json: - keys: JWKKeyValue = response_json["keys"] - else: - keys = response_json - - await self.user_api_key_cache.async_set_cache( - key=cache_key, - value=keys, - ttl=self._get_public_key_cache_ttl(), - ) - else: - keys = cached_keys + keys: Final = await self._cached_with_stale_fallback( + cache_key=f"litellm_jwt_auth_keys_{resolved_jwks_url}", + ttl=self._get_public_key_cache_ttl(), + refresh=lambda: self._fetch_jwks_keys(resolved_jwks_url), + log_context=f"kid={kid}", + ) public_key: Final = self.parse_keys(keys=keys, kid=kid) if public_key is not None: @@ -692,6 +861,9 @@ class JWTHandler: return await self._get_public_key_from_jwks_url(jwks_url=key_url, kid=kid) except NoMatchingJWTPublicKeyError as e: verbose_proxy_logger.debug("JWT Auth: No matching public key found at %s: %s", key_url, e) + except JWKSUnreachableError as e: + verbose_proxy_logger.error("JWT Auth: JWKS endpoint %s unreachable: %s", key_url, e) + raise jwks_unavailable_exception(e) from e raise NoMatchingJWTPublicKeyError(f"No matching public key found. keys={keys_url_list}, kid={kid}") @@ -969,10 +1141,14 @@ class JWTHandler: ) async def _auth_jwt_with_issuer(self, token: str, issuer_config: JWTIssuerConfig, kid: str | None) -> dict: - public_key: Final = await self._get_public_key_from_jwks_url( - jwks_url=self._get_jwks_url_for_issuer(issuer_config=issuer_config), - kid=kid, - ) + try: + public_key: Final = await self._get_public_key_from_jwks_url( + jwks_url=self._get_jwks_url_for_issuer(issuer_config=issuer_config), + kid=kid, + ) + except JWKSUnreachableError as e: + raise jwks_unavailable_exception(e) from e + try: payload: Final = self._decode_jwt_with_public_key( token=token, diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 3840c90d691..a9e12beb54b 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -1,10 +1,16 @@ +import asyncio +import re +import time +from collections.abc import Mapping, Sequence from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch from fastapi import HTTPException +import httpx import pytest from litellm.proxy._types import ( + DEFAULT_JWKS_STALE_TTL, JWTLiteLLMRoleMap, LiteLLM_JWTAuth, LiteLLM_TeamMembership, @@ -15,7 +21,16 @@ from litellm.proxy._types import ( ProxyErrorTypes, ProxyException, ) -from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler +from litellm.caching.dual_cache import DualCache +from litellm.proxy.auth.handle_jwt import ( + JWKS_FETCH_ATTEMPTS, + STALE_CACHE_KEY_PREFIX, + STALE_WRITTEN_AT_CACHE_KEY_PREFIX, + JWKSUnreachableError, + JWTAuthManager, + JWTHandler, + NoMatchingJWTPublicKeyError, +) @pytest.mark.asyncio @@ -3921,6 +3936,7 @@ async def test_get_public_key_fetches_and_caches_jwks_response(): expected_key_id = "cached-key" _, jwk = _get_rsa_key_and_jwk(kid=expected_key_id) mock_response = MagicMock() + mock_response.status_code = 200 mock_response.json.return_value = {"keys": [jwk]} jwt_handler.http_handler.get = AsyncMock(return_value=mock_response) @@ -3936,6 +3952,560 @@ async def test_get_public_key_fetches_and_caches_jwks_response(): assert cached_keys == [jwk] +class _ScriptedJWKSEndpoint: + """Injected stand-in for ``JWTHandler.http_handler`` with scripted per-call outcomes. + + Each outcome is either an exception to raise or a JSON body to return; the + last outcome repeats for any further calls. + """ + + def __init__( + self, + outcomes: Sequence[Exception | Mapping[str, object] | MagicMock], + delay: float = 0.0, + ) -> None: + self.outcomes = outcomes + self.delay = delay + self.call_count = 0 + + async def get( + self, + url: str, + params: Mapping[str, str] | None = None, + headers: Mapping[str, str] | None = None, + ) -> MagicMock: + self.call_count += 1 + if self.delay: + await asyncio.sleep(self.delay) + outcome = self.outcomes[min(self.call_count - 1, len(self.outcomes) - 1)] + if isinstance(outcome, Exception): + raise outcome + if isinstance(outcome, MagicMock): + return outcome + response = MagicMock() + response.status_code = 200 + response.json.return_value = outcome + return response + + +def _get_jwt_handler_with_scripted_endpoint( + cache: "DualCache", + endpoint: _ScriptedJWKSEndpoint, + public_key_ttl: float = 600, + public_key_stale_ttl: float = DEFAULT_JWKS_STALE_TTL, +) -> JWTHandler: + jwt_handler = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth( + public_key_ttl=public_key_ttl, + public_key_stale_ttl=public_key_stale_ttl, + ), + ) + jwt_handler.http_handler = endpoint + return jwt_handler + + +@pytest.mark.asyncio +async def test_get_public_key_retries_transient_jwks_fetch_failure(): + """A single connect timeout to the IdP must be retried, not surfaced to the caller.""" + from litellm.caching.dual_cache import DualCache + + _, jwk = _get_rsa_key_and_jwk(kid="retried-key") + endpoint = _ScriptedJWKSEndpoint((httpx.ConnectTimeout("connect timed out"), {"keys": [jwk]})) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(DualCache(), endpoint) + + public_key = await jwt_handler._get_public_key_from_jwks_url( + jwks_url="https://issuer.example.com/keys", + kid="retried-key", + ) + + assert public_key == jwk + assert endpoint.call_count == 2 + + +@pytest.mark.asyncio +async def test_get_public_key_serves_stale_keys_when_jwks_refresh_fails(): + """Once the TTL lapses, an unreachable IdP must not invalidate a still-valid signing key.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://issuer.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="stale-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint) + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="stale-key") == jwk + + await cache.async_delete_cache(key=f"litellm_jwt_auth_keys_{jwks_url}") + endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) + + public_key = await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="stale-key") + + assert public_key == jwk + + +@pytest.mark.asyncio +async def test_stale_jwks_window_is_the_configured_grace_past_a_long_public_key_ttl(): + """The stale window is `public_key_stale_ttl` past the active entry, whatever `public_key_ttl` is set to. + + Deriving the window from `public_key_ttl` instead would collapse it to nothing on the long TTLs that + make the fallback worth having. + """ + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://long-ttl-issuer.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="long-ttl-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint( + cache, + endpoint, + public_key_ttl=90000, + public_key_stale_ttl=3600, + ) + + await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="long-ttl-key") + + active_key = f"litellm_jwt_auth_keys_{jwks_url}" + active_deadline = cache.in_memory_cache.ttl_dict[active_key] + stale_deadline = cache.in_memory_cache.ttl_dict[f"{STALE_CACHE_KEY_PREFIX}{active_key}"] + + assert stale_deadline - active_deadline == pytest.approx(3600, abs=1) + + +@pytest.mark.asyncio +async def test_long_public_key_ttl_still_serves_stale_keys_when_the_idp_is_unreachable(): + """A long `public_key_ttl` must not leave the stale fallback inert once that TTL finally lapses.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://long-ttl-fallback.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="long-ttl-fallback-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_ttl=604800) + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="long-ttl-fallback-key") == jwk + + await cache.async_delete_cache(key=f"litellm_jwt_auth_keys_{jwks_url}") + endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="long-ttl-fallback-key") == jwk + + +@pytest.mark.asyncio +async def test_removed_signing_key_stops_being_trusted_once_the_stale_window_expires(monkeypatch): + """The stale fallback is bounded: past its window a key the IdP dropped is no longer served.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://revoking-issuer.example.com/keys" + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url) + + _, jwk = _get_rsa_key_and_jwk(kid="revoked-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint) + + assert await jwt_handler.get_public_key(kid="revoked-key") == jwk + + active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}" + await cache.async_delete_cache(key=active_cache_key) + endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) + + assert await jwt_handler.get_public_key(kid="revoked-key") == jwk + + await cache.async_delete_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_cache_key}") + + with pytest.raises(ProxyException) as exc_info: + await jwt_handler.get_public_key(kid="revoked-key") + + assert exc_info.value.code == "503" + assert exc_info.value.type == ProxyErrorTypes.auth_provider_unavailable + + +@pytest.mark.asyncio +async def test_key_removed_from_a_reachable_jwks_is_rejected_without_consulting_the_stale_copy(): + """A reachable IdP always wins: dropping a key revokes it immediately, stale copy included.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://rotating-issuer.example.com/keys" + _, retired_jwk = _get_rsa_key_and_jwk(kid="retired-key") + _, current_jwk = _get_rsa_key_and_jwk(kid="current-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [retired_jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint) + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="retired-key") == retired_jwk + + active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}" + await cache.async_delete_cache(key=active_cache_key) + endpoint.outcomes = ({"keys": [current_jwk]},) + + with pytest.raises(NoMatchingJWTPublicKeyError): + await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="retired-key") + + assert await cache.async_get_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_cache_key}") == [current_jwk] + + +@pytest.mark.asyncio +async def test_zero_public_key_stale_ttl_fails_closed_instead_of_serving_stale_keys(): + """`public_key_stale_ttl=0` is the escape hatch for deployments that cannot trust an unrefreshed key.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://fail-closed-issuer.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="fail-closed-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_stale_ttl=0) + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="fail-closed-key") == jwk + assert await cache.async_get_cache(key=f"{STALE_CACHE_KEY_PREFIX}litellm_jwt_auth_keys_{jwks_url}") is None + + await cache.async_delete_cache(key=f"litellm_jwt_auth_keys_{jwks_url}") + endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) + + with pytest.raises(JWKSUnreachableError): + await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="fail-closed-key") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("lowered_stale_ttl", [0, 30]) +async def test_lowering_public_key_stale_ttl_stops_serving_a_copy_cached_under_the_old_setting(lowered_stale_ttl): + """Lowering the window has to bite immediately: an operator does this mid-incident, on a shared cache. + + The stale entry keeps whatever expiry it was written with, so enforcing the bound only at write time would + leave a copy taken under the old, longer setting servable until it aged out on its own. + """ + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://relaxed-then-tightened.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="tightened-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + generous = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_stale_ttl=86400) + + assert await generous._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="tightened-key") == jwk + + active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}" + await cache.async_delete_cache(key=active_cache_key) + assert await cache.async_get_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_cache_key}") == [jwk] + + # The operator tightens the window and restarts; the cache, and its long-lived copy, survive. + endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) + tightened = _get_jwt_handler_with_scripted_endpoint( + cache, endpoint, public_key_stale_ttl=lowered_stale_ttl + ) + await cache.async_set_cache( + key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}", + value=time.time() - 7200, + ttl=86400, + ) + + with pytest.raises(JWKSUnreachableError): + await tightened._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="tightened-key") + + +@pytest.mark.asyncio +async def test_zero_public_key_stale_ttl_fails_closed_even_for_a_freshly_written_copy(): + """`0` must fail closed on its own, not merely because the copy happens to be older than `public_key_ttl`. + + The active entry can disappear before it expires, through cache eviction or a flush, which leaves a stale + copy younger than `public_key_ttl`. Bounding only on age would still serve it. + """ + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://evicted-active-entry.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="fresh-copy-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + generous = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_ttl=600, public_key_stale_ttl=3600) + + assert await generous._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="fresh-copy-key") == jwk + + active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}" + await cache.async_delete_cache(key=active_cache_key) + written_at = await cache.async_get_cache(key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}") + assert time.time() - written_at < 600 + + endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) + fail_closed = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_ttl=600, public_key_stale_ttl=0) + + with pytest.raises(JWKSUnreachableError): + await fail_closed._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="fresh-copy-key") + + +@pytest.mark.asyncio +async def test_stale_copy_with_no_recorded_write_time_is_not_served(): + """The bound is enforced from the recorded write time, so losing it must fail closed, never open.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://undated-copy.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="undated-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint) + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="undated-key") == jwk + + active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}" + await cache.async_delete_cache(key=active_cache_key) + await cache.async_delete_cache(key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}") + assert await cache.async_get_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_cache_key}") == [jwk] + endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) + + with pytest.raises(JWKSUnreachableError): + await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="undated-key") + + +@pytest.mark.asyncio +async def test_increasing_public_key_stale_ttl_only_extends_within_the_new_bound(): + """Raising the window re-measures from the copy's refresh time; it does not bless whatever is cached.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://widened-window.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="widened-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + narrow = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_ttl=600, public_key_stale_ttl=60) + + assert await narrow._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="widened-key") == jwk + + active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}" + written_at_key = f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}" + await cache.async_delete_cache(key=active_cache_key) + endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) + widened = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_ttl=600, public_key_stale_ttl=3600) + + # Older than the widened bound of 600 + 3600, so widening must not revive it. + await cache.async_set_cache(key=written_at_key, value=time.time() - 5000, ttl=86400) + with pytest.raises(JWKSUnreachableError): + await widened._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="widened-key") + + # Inside the widened bound, so it is servable again. + await cache.async_set_cache(key=written_at_key, value=time.time() - 1000, ttl=86400) + assert await widened._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="widened-key") == jwk + + +@pytest.mark.asyncio +async def test_stale_copy_written_at_survives_a_whole_number_epoch(): + """A Redis JSON round-trip can return the epoch as an int, and that must not read as a missing timestamp. + + Rejecting it would fail closed on a copy that is well inside the window, in the shared-cache deployment + the stale fallback exists to serve. + """ + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://int-epoch.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="int-epoch-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint) + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="int-epoch-key") == jwk + + active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}" + await cache.async_delete_cache(key=active_cache_key) + await cache.async_set_cache( + key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}", + value=int(time.time()) - 60, + ttl=86400, + ) + endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="int-epoch-key") == jwk + + +@pytest.mark.asyncio +async def test_stale_copy_with_a_malformed_write_time_is_not_served(): + """An unreadable refresh timestamp is indistinguishable from an unbounded one, so it fails closed.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://malformed-timestamp.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="malformed-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint) + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="malformed-key") == jwk + + active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}" + await cache.async_delete_cache(key=active_cache_key) + await cache.async_set_cache( + key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}", + value="whenever", + ttl=86400, + ) + endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) + + with pytest.raises(JWKSUnreachableError): + await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="malformed-key") + + +@pytest.mark.asyncio +async def test_public_key_stale_ttl_defaults_to_one_hour(): + """The default is the exposure bound for a key the IdP revoked mid-outage, so it stays short deliberately.""" + assert LiteLLM_JWTAuth().public_key_stale_ttl == 3600 + + +@pytest.mark.asyncio +async def test_stale_fallback_warns_with_the_kid_and_how_stale_the_jwks_copy_is(caplog): + """Serving an unrefreshed signing key is a security-relevant event, so it must be legible in the logs.""" + import logging + + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://warned-issuer.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="warned-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint, public_key_stale_ttl=1800) + + await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="warned-key") + + active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}" + await cache.async_delete_cache(key=active_cache_key) + await cache.async_set_cache( + key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}", + value=time.time() - 120, + ttl=600, + ) + endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) + + caplog.set_level(logging.WARNING) + await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="warned-key") + + warnings = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + stale_warnings = [m for m in warnings if "stale JWKS copy" in m] + assert len(stale_warnings) == 1 + assert "kid=warned-key" in stale_warnings[0] + assert jwks_url in stale_warnings[0] + + freshness = re.search(r"last refreshed (\d+)s ago, stops being trusted in (\d+)s", stale_warnings[0]) + assert freshness is not None + age, remaining = int(freshness.group(1)), int(freshness.group(2)) + assert age == pytest.approx(120, abs=2) + assert remaining == pytest.approx(600 + 1800 - 120, abs=2) + + +@pytest.mark.asyncio +async def test_unparseable_jwks_response_does_not_fall_back_to_the_stale_copy(): + """Only an unreachable IdP unlocks the stale copy. A reachable one that answers badly must surface the error.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://garbled-issuer.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="garbled-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint) + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="garbled-key") == jwk + + await cache.async_delete_cache(key=f"litellm_jwt_auth_keys_{jwks_url}") + garbled = MagicMock() + garbled.status_code = 200 + garbled.text = "not json" + garbled.json.side_effect = ValueError("Expecting value: line 1 column 1") + endpoint.outcomes = (garbled,) + + with pytest.raises(Exception, match="Error parsing response"): + await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="garbled-key") + + +@pytest.mark.asyncio +async def test_jwks_error_response_is_not_cached_over_the_last_known_good_keys(): + """An IdP error body must never be stored as the key set, least of all as the stale copy.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://erroring-issuer.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="erroring-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint) + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="erroring-key") == jwk + + active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}" + await cache.async_delete_cache(key=active_cache_key) + server_error = MagicMock() + server_error.status_code = 503 + server_error.text = '{"error": "upstream unavailable"}' + server_error.json.return_value = {"error": "upstream unavailable"} + endpoint.outcomes = (server_error,) + + with pytest.raises(Exception, match="returned status 503"): + await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="erroring-key") + + assert await cache.async_get_cache(key=active_cache_key) is None + assert await cache.async_get_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_cache_key}") == [jwk] + + +@pytest.mark.asyncio +async def test_sustained_jwks_outage_refetches_once_per_backoff_window_not_once_per_request(): + """Without a backoff, every request during an outage pays three timeouts serialised behind the refresh lock.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://flooded-issuer.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="flooded-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint) + + await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="flooded-key") + + await cache.async_delete_cache(key=f"litellm_jwt_auth_keys_{jwks_url}") + endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) + calls_before_outage = endpoint.call_count + + public_keys = await asyncio.gather( + *[jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="flooded-key") for _ in range(6)] + ) + + assert public_keys == [jwk] * 6 + assert endpoint.call_count - calls_before_outage == JWKS_FETCH_ATTEMPTS + + +@pytest.mark.asyncio +async def test_get_public_key_raises_503_when_jwks_unreachable_and_no_cached_keys(monkeypatch): + """An unreachable IdP is an infra failure: 503, never a 401 that clients read as bad credentials.""" + from litellm.caching.dual_cache import DualCache + + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "https://issuer.example.com/keys") + endpoint = _ScriptedJWKSEndpoint((httpx.ConnectTimeout("connect timed out"),)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(DualCache(), endpoint) + + with pytest.raises(ProxyException) as exc_info: + await jwt_handler.get_public_key(kid="any-key") + + assert exc_info.value.code == "503" + assert exc_info.value.type == ProxyErrorTypes.auth_provider_unavailable + assert "ConnectTimeout" in exc_info.value.message + assert endpoint.call_count == 3 + + +@pytest.mark.asyncio +async def test_get_public_key_coalesces_concurrent_jwks_refreshes(): + """Concurrent requests in the TTL-expiry window share one JWKS fetch.""" + from litellm.caching.dual_cache import DualCache + + _, jwk = _get_rsa_key_and_jwk(kid="coalesced-key") + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},), delay=0.05) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(DualCache(), endpoint) + + public_keys = await asyncio.gather( + *[ + jwt_handler._get_public_key_from_jwks_url( + jwks_url="https://coalesce.example.com/keys", + kid="coalesced-key", + ) + for _ in range(5) + ] + ) + + assert public_keys == [jwk] * 5 + assert endpoint.call_count == 1 + + @pytest.mark.asyncio async def test_get_public_key_tries_next_jwks_url_when_kid_missing(monkeypatch): from litellm.caching.dual_cache import DualCache @@ -4140,6 +4710,38 @@ async def test_auth_jwt_issuer_path_expired_token_raises_401(monkeypatch): assert "Token Expired" in exc_info.value.message +@pytest.mark.asyncio +async def test_auth_jwt_issuer_path_unreachable_jwks_raises_503(monkeypatch): + """The issuer-scoped path must report an unreachable IdP as 503, not as a credential failure.""" + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False) + + issuer = "https://unreachable-issuer.example.com" + jwks_url = f"{issuer}/keys" + private_key, _ = _get_rsa_key_and_jwk(kid="unreachable-kid") + + jwt_handler = _get_jwt_handler_with_issuer_keys( + issuers=[{"issuer": issuer, "jwks_url": jwks_url, "audience": "my-audience"}], + keys_by_url={}, + ) + endpoint = _ScriptedJWKSEndpoint((httpx.ConnectTimeout("connect timed out"),)) + jwt_handler.http_handler = endpoint + + token = _encode_rsa_jwt( + private_key=private_key, + issuer=issuer, + audience="my-audience", + kid="unreachable-kid", + ) + + with pytest.raises(ProxyException) as exc_info: + await jwt_handler.auth_jwt(token=token) + + assert exc_info.value.code == "503" + assert exc_info.value.type == ProxyErrorTypes.auth_provider_unavailable + assert endpoint.call_count == 3 + + @pytest.mark.asyncio async def test_multi_issuer_jwt_maps_kubernetes_namespace_claim(monkeypatch): monkeypatch.delenv("JWT_AUDIENCE", raising=False) diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 474e571e592..84a9ddfacff 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -117,6 +117,8 @@ def test_is_database_connection_generic_errors(): TimeoutError("timed out"), OSError("network is unreachable"), asyncio.TimeoutError(), + httpx.ConnectError("connection refused"), + httpx.ConnectTimeout("connect timed out"), HTTPClientClosedError(), ClientNotConnectedError(), PrismaError("can't reach database server"),