diff --git a/litellm/proxy/auth/v2/authenticators.py b/litellm/proxy/auth/v2/authenticators.py index 484bc09b2c2..c536f012af1 100644 --- a/litellm/proxy/auth/v2/authenticators.py +++ b/litellm/proxy/auth/v2/authenticators.py @@ -78,11 +78,81 @@ class VirtualKeyAuthenticator: ) -# Master key is matched first (exact compare), then virtual keys. authlib-backed -# JWT / OAuth2 nodes slot in next, implementing the same interface. +def _load_jwt_settings() -> Any: + import os + + from litellm.proxy.proxy_server import general_settings + + from .jwt_claims import JWTSettings + + cfg = (general_settings or {}).get("auth_v2_jwt") or {} + jwks_uri = cfg.get("jwks_uri") or os.getenv("AUTH_V2_JWKS_URI") + if not jwks_uri: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="auth_v2: JWT auth received a token but no jwks_uri is configured", + ) + return JWTSettings( + jwks_uri=jwks_uri, + issuer=cfg.get("issuer") or os.getenv("AUTH_V2_JWT_ISSUER"), + audience=cfg.get("audience") or os.getenv("AUTH_V2_JWT_AUDIENCE"), + user_id_claim=cfg.get("user_id_claim", "sub"), + team_claim=cfg.get("team_claim"), + role_claim=cfg.get("role_claim"), + role_map=cfg.get("role_map") or {}, + ) + + +class JWTAuthenticator: + """Verifies a bearer JWT with authlib and maps its claims to an identity.""" + + def can_handle(self, api_key: Optional[str]) -> bool: + return ( + isinstance(api_key, str) + and not api_key.startswith("sk-") + and api_key.count(".") == 2 + ) + + async def authenticate(self, api_key: str, ctx: AuthContext) -> Any: + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + from .jwt_claims import extract_identity + from .jwt_verifier import JWKSProvider, JWTVerificationError, verify + + settings = _load_jwt_settings() + key_set = await JWKSProvider(settings.jwks_uri).get_key_set() + try: + claims = verify( + api_key, key_set, settings.issuer, settings.audience + ) + except JWTVerificationError as e: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=f"auth_v2: {e}", + ) + + identity = extract_identity(claims, settings) + user_role = None + if identity.role is not None: + try: + user_role = LitellmUserRoles(identity.role) + except ValueError: + user_role = None + + return UserAPIKeyAuth( + user_id=identity.user_id, + team_id=identity.team_id, + user_role=user_role, + jwt_claims=claims, + ) + + +# Master key is matched first (exact compare), then virtual keys, then JWTs. +# Dispatch is by credential shape, so the chain is deterministic, not "ask everyone". AUTHENTICATORS: List[Authenticator] = [ MasterKeyAuthenticator(), VirtualKeyAuthenticator(), + JWTAuthenticator(), ] diff --git a/litellm/proxy/auth/v2/jwt_claims.py b/litellm/proxy/auth/v2/jwt_claims.py new file mode 100644 index 00000000000..f612eb6af53 --- /dev/null +++ b/litellm/proxy/auth/v2/jwt_claims.py @@ -0,0 +1,53 @@ +from dataclasses import dataclass, field +from typing import Any, Dict, Optional + + +class JWTClaimError(Exception): + """Raised when verified claims lack the data needed to form an identity.""" + + +@dataclass +class JWTSettings: + jwks_uri: str + issuer: Optional[str] = None + audience: Optional[str] = None + user_id_claim: str = "sub" + team_claim: Optional[str] = None + role_claim: Optional[str] = None + # Maps an upstream role/group value to a litellm role name. Unmapped values + # yield no role rather than trusting an arbitrary IdP string as a role. + role_map: Dict[str, str] = field(default_factory=dict) + + +@dataclass +class JWTIdentity: + user_id: str + team_id: Optional[str] = None + role: Optional[str] = None + + +def extract_identity(claims: Dict[str, Any], settings: JWTSettings) -> JWTIdentity: + """Map verified claims to an identity using the configured claim names. + + The signature is already trusted at this point; this is the LiteLLM-semantic + layer authlib does not own. + """ + user_id = claims.get(settings.user_id_claim) + if not user_id: + raise JWTClaimError( + f"token missing user id claim '{settings.user_id_claim}'" + ) + + team_id = claims.get(settings.team_claim) if settings.team_claim else None + + role: Optional[str] = None + if settings.role_claim: + raw_role = claims.get(settings.role_claim) + if raw_role is not None: + role = settings.role_map.get(raw_role) + + return JWTIdentity( + user_id=str(user_id), + team_id=str(team_id) if team_id is not None else None, + role=role, + ) diff --git a/litellm/proxy/auth/v2/jwt_verifier.py b/litellm/proxy/auth/v2/jwt_verifier.py new file mode 100644 index 00000000000..c5cbb73bae4 --- /dev/null +++ b/litellm/proxy/auth/v2/jwt_verifier.py @@ -0,0 +1,67 @@ +import time +from typing import Any, Dict, Optional + +from authlib.jose import JsonWebKey, jwt +from authlib.jose.errors import JoseError + + +class JWTVerificationError(Exception): + """Raised when a token fails signature or standard-claim validation.""" + + +def build_claims_options(issuer: Optional[str], audience: Optional[str]) -> Dict[str, Any]: + options: Dict[str, Any] = {"exp": {"essential": True}} + if issuer: + options["iss"] = {"essential": True, "value": issuer} + if audience: + options["aud"] = {"essential": True, "value": audience} + return options + + +def verify( + token: str, + key_set: Any, + issuer: Optional[str] = None, + audience: Optional[str] = None, +) -> Dict[str, Any]: + """Verify ``token`` against ``key_set`` and validate exp/iss/aud. + + authlib owns the crypto: signature verification, key selection by ``kid``, + and standard-claim checks. ``key_set`` is an imported JWKS (injected so this + is testable without network). Raises :class:`JWTVerificationError` on any + failure so callers never branch on authlib's internal exception types. + """ + try: + claims = jwt.decode( + token, key_set, claims_options=build_claims_options(issuer, audience) + ) + claims.validate(now=int(time.time())) + return dict(claims) + except JoseError as e: + raise JWTVerificationError(str(e)) from e + except Exception as e: # malformed token, bad header, etc. + raise JWTVerificationError(f"invalid token: {e}") from e + + +class JWKSProvider: + """Fetches and TTL-caches a JWKS document, returning an imported key set.""" + + def __init__(self, jwks_uri: str, ttl_seconds: float = 600.0): + self.jwks_uri = jwks_uri + self.ttl_seconds = ttl_seconds + self._key_set: Any = None + self._fetched_at: float = 0.0 + + async def get_key_set(self) -> Any: + now = time.monotonic() + if self._key_set is not None and (now - self._fetched_at) < self.ttl_seconds: + return self._key_set + + import httpx + + async with httpx.AsyncClient() as client: + response = await client.get(self.jwks_uri) + response.raise_for_status() + self._key_set = JsonWebKey.import_key_set(response.json()) + self._fetched_at = now + return self._key_set diff --git a/pyproject.toml b/pyproject.toml index 100ac8e7805..679cd1871c3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,7 @@ proxy = [ "apscheduler>=3.11.2,<4.0", "fastapi-sso>=0.19.0,<1.0", "casbin>=1.43.0,<2.0", + "authlib>=1.6.5,<2.0", "PyJWT>=2.12.0,<3.0", "python-multipart>=0.0.27,<1.0", "cryptography>=46.0.7,<47.0", diff --git a/tests/test_litellm/proxy/auth/v2/test_jwt_claims.py b/tests/test_litellm/proxy/auth/v2/test_jwt_claims.py new file mode 100644 index 00000000000..5223f04cb51 --- /dev/null +++ b/tests/test_litellm/proxy/auth/v2/test_jwt_claims.py @@ -0,0 +1,51 @@ +import pytest + +from litellm.proxy.auth.v2.jwt_claims import ( + JWTClaimError, + JWTSettings, + extract_identity, +) + +BASE = JWTSettings(jwks_uri="https://idp/jwks") + + +def test_extracts_user_id_from_default_sub_claim(): + ident = extract_identity({"sub": "u1"}, BASE) + assert ident.user_id == "u1" + assert ident.team_id is None + assert ident.role is None + + +def test_custom_user_id_claim(): + settings = JWTSettings(jwks_uri="x", user_id_claim="oid") + assert extract_identity({"oid": "abc"}, settings).user_id == "abc" + + +def test_team_and_role_are_mapped(): + settings = JWTSettings( + jwks_uri="x", + team_claim="team", + role_claim="groups", + role_map={"admins": "proxy_admin"}, + ) + ident = extract_identity({"sub": "u1", "team": "eng", "groups": "admins"}, settings) + assert ident.team_id == "eng" + assert ident.role == "proxy_admin" + + +def test_unmapped_role_value_yields_no_role(): + # An arbitrary IdP role string must not be trusted as a litellm role. + settings = JWTSettings(jwks_uri="x", role_claim="role", role_map={"a": "proxy_admin"}) + assert extract_identity({"sub": "u1", "role": "totally-unknown"}, settings).role is None + + +def test_missing_user_id_raises(): + with pytest.raises(JWTClaimError): + extract_identity({"email": "x@y.z"}, BASE) + + +def test_non_string_ids_are_coerced(): + settings = JWTSettings(jwks_uri="x", team_claim="team") + ident = extract_identity({"sub": 12345, "team": 67}, settings) + assert ident.user_id == "12345" + assert ident.team_id == "67" diff --git a/tests/test_litellm/proxy/auth/v2/test_jwt_verifier.py b/tests/test_litellm/proxy/auth/v2/test_jwt_verifier.py new file mode 100644 index 00000000000..a083b9875d4 --- /dev/null +++ b/tests/test_litellm/proxy/auth/v2/test_jwt_verifier.py @@ -0,0 +1,84 @@ +import time + +import pytest +from authlib.jose import JsonWebKey, jwt + +from litellm.proxy.auth.v2.jwt_verifier import JWTVerificationError, verify + +ISSUER = "https://idp.example" +AUDIENCE = "litellm" + + +@pytest.fixture(scope="module") +def signing(): + key = JsonWebKey.generate_key("RSA", 2048, is_private=True) + kid = key.thumbprint() + public = JsonWebKey.import_key(key.as_pem(is_private=False)).as_dict() + public["kid"] = kid + key_set = JsonWebKey.import_key_set({"keys": [public]}) + return key, kid, key_set + + +def _token(signing, **overrides): + key, kid, _ = signing + claims = { + "sub": "user-42", + "iss": ISSUER, + "aud": AUDIENCE, + "exp": int(time.time()) + 3600, + } + claims.update(overrides) + # jwt.encode returns bytes; real tokens arrive as strings off the header. + return jwt.encode( + {"alg": "RS256", "kid": kid}, claims, key.as_pem(is_private=True) + ).decode("utf-8") + + +def test_valid_token_returns_claims(signing): + _, _, key_set = signing + claims = verify(_token(signing, role="admin"), key_set, ISSUER, AUDIENCE) + assert claims["sub"] == "user-42" + assert claims["role"] == "admin" + + +def test_tampered_signature_is_rejected(signing): + _, _, key_set = signing + bad = _token(signing)[:-4] + "AAAA" + with pytest.raises(JWTVerificationError): + verify(bad, key_set, ISSUER, AUDIENCE) + + +def test_expired_token_is_rejected(signing): + _, _, key_set = signing + with pytest.raises(JWTVerificationError): + verify(_token(signing, exp=int(time.time()) - 10), key_set, ISSUER, AUDIENCE) + + +def test_wrong_audience_is_rejected(signing): + _, _, key_set = signing + with pytest.raises(JWTVerificationError): + verify(_token(signing, aud="someone-else"), key_set, ISSUER, AUDIENCE) + + +def test_wrong_issuer_is_rejected(signing): + _, _, key_set = signing + with pytest.raises(JWTVerificationError): + verify(_token(signing, iss="https://evil.example"), key_set, ISSUER, AUDIENCE) + + +def test_token_signed_by_unknown_key_is_rejected(signing): + _, _, key_set = signing + other = JsonWebKey.generate_key("RSA", 2048, is_private=True) + forged = jwt.encode( + {"alg": "RS256", "kid": "unknown"}, + {"sub": "x", "iss": ISSUER, "aud": AUDIENCE, "exp": int(time.time()) + 60}, + other.as_pem(is_private=True), + ).decode("utf-8") + with pytest.raises(JWTVerificationError): + verify(forged, key_set, ISSUER, AUDIENCE) + + +def test_garbage_is_rejected(signing): + _, _, key_set = signing + with pytest.raises(JWTVerificationError): + verify("not.a.jwt", key_set, ISSUER, AUDIENCE)