mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(jwt): retry JWKS fetches, serve stale keys, and return 503 when the IdP is unreachable (#37690)
A JWKS fetch had no retry, so a single connect timeout to the identity provider failed authentication outright, and once the cached copy expired there was nothing to fall back on. How that surfaced depended on the outage shape: httpx.ConnectTimeout was missing from DB_CONNECTION_ERROR_TYPES so it fell through to the generic auth handler as a 401 with an empty detail, while a read timeout took the database path and reported a healthy database as unreachable. Transport failures are now retried three times with a short backoff, and the last-known-good JWKS stays usable for a bounded window past public_key_ttl. That window is public_key_stale_ttl, a new config field defaulting to 3600s and settable to 0 to fail closed. It is checked on every read against the current setting rather than baked into the cache entry when it is written, so lowering it binds immediately instead of waiting for entries written under the old value to age out, which matters because a shared cache survives the restart an operator performs to make the change take effect. A copy whose write time cannot be established is not servable. Only httpx.TransportError unlocks the stale copy, so an identity provider that answers at all, including with a narrowed key set, revokes on the next refresh. Every stale serve logs the kid it authenticated, how long ago that copy was refreshed, and how long until it stops being trusted. A sustained outage is remembered for 30s per key url, so it costs one fetch per window instead of three timeouts per request serialised behind the refresh lock. Non-200 JWKS responses now raise instead of being cached as the key set, which previously let an error body overwrite the last-known-good copy. An unreachable identity provider with no cached copy left returns 503 auth_provider_unavailable. Resolves LIT-5524 Co-authored-by: Yassin Kortam <yassin@berri.ai>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
b2aff8be0f
commit
52403d7a8d
4 changed files with 841 additions and 41 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 = "<html>not json</html>"
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue