diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index d5461f01ed8..425f82794e6 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -903,8 +903,9 @@ class MCPRequestHandler: NotSessionBearer, SessionBearerAdmitted, SessionBearerInvalid, + SessionSigningConfigError, + active_session_signing_keys, resolve_session_bearer, - session_keys_from_master_key, ) from litellm.proxy.proxy_server import master_key @@ -913,7 +914,10 @@ class MCPRequestHandler: await MCPRequestHandler._run_pre_db_read_auth_checks(request=request, route=route) - keys: Final = session_keys_from_master_key(master_key) + keys: Final = active_session_signing_keys(master_key) + if isinstance(keys, SessionSigningConfigError): + verbose_logger.error("mcp gateway session admission rejected: %s", keys.detail) + raise HTTPException(status_code=500, detail="Server misconfigured: mcp_session_token_signing is invalid") result: Final = resolve_session_bearer(authorization_value, keys, datetime.now(timezone.utc)) match result: case SessionBearerAdmitted(): diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 314c80adbc4..853a07972c1 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -65,15 +65,16 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( SessionRefreshOpened, + SessionSigningConfigError, + active_session_signing_keys, open_session_refresh_bearer, - session_keys_from_master_key, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( SESSION_REFRESH_TTL_SECONDS, MintedSessionToken, SessionAudience, - SessionKeys, SessionPrincipal, + SessionSigningKeys, mint_session_refresh_token, mint_session_token, ) @@ -885,7 +886,7 @@ class _SingleUseGuard: return "first" if count == 1 else "replayed" -def _session_token_pair(principal: SessionPrincipal, keys: SessionKeys, now: datetime) -> Response: +def _session_token_pair(principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime) -> Response: access: Final = mint_session_token(principal, keys, now) refresh: Final = mint_session_refresh_token(principal, keys, now) if not isinstance(access, MintedSessionToken) or not isinstance(refresh, MintedSessionToken): @@ -912,7 +913,7 @@ class _ProxyCredentialTokenResponse(TypedDict): def _proxy_credential_response( - minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionKeys, now: datetime + minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime ) -> Response: """The proxy-API token response: the access token is the very credential ``lite login`` stores (accepted on every proxy route with user and team attribution), and @@ -998,7 +999,10 @@ async def aggregate_token( if master_key is None: verbose_logger.error("mcp_gateway_dcr token grant rejected: no master_key configured") return _oauth_error(500, "server_error", "the gateway has no master key configured") - keys: Final = session_keys_from_master_key(master_key) + keys: Final = active_session_signing_keys(master_key) + if isinstance(keys, SessionSigningConfigError): + verbose_logger.error("mcp_gateway_dcr token grant rejected: %s", keys.detail) + return _oauth_error(500, "server_error", "the gateway session signing configuration is invalid") now: Final = datetime.now(timezone.utc) issue: Final = _GrantIssuer( request=request, @@ -1043,7 +1047,7 @@ class _GrantIssuer: self, request: Request, resource: str | None, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, reload_user: ReloadUser, mint_proxy_credential: MintProxyCredential, @@ -1146,7 +1150,7 @@ async def _refresh_token_grant( refresh_token: str | None, client_id: str, resource: str | None, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, issue: _GrantIssuer, ) -> Response: @@ -1182,7 +1186,10 @@ async def revoke_refresh_token(token: str, client_id: str, master_key: str | Non if master_key is None: verbose_logger.error("mcp_gateway_dcr revoke rejected: no master_key configured") return _oauth_error(500, "server_error", "the gateway has no master key configured") - keys: Final = session_keys_from_master_key(master_key) + keys: Final = active_session_signing_keys(master_key) + if isinstance(keys, SessionSigningConfigError): + verbose_logger.error("mcp_gateway_dcr revoke rejected: %s", keys.detail) + return _oauth_error(500, "server_error", "the gateway session signing configuration is invalid") now: Final = datetime.now(timezone.utc) opened: Final = open_session_refresh_bearer(token, keys, now, expected_client_id=client_id) if isinstance(opened, SessionRefreshOpened): diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py index 70a04ac290a..df2bbdba345 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py @@ -20,13 +20,16 @@ from datetime import datetime from functools import lru_cache from typing import Final, Literal, TypeAlias -from pydantic import BaseModel, ConfigDict, SecretStr +from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + AsymmetricSessionKeys, OpenedSessionToken, SessionExpired, SessionKeys, SessionPrincipal, + SessionRotatedPublicKey, + SessionSigningKeys, is_session_refresh_token, is_session_token, open_session_refresh_token, @@ -68,6 +71,99 @@ def session_keys_from_master_key(master_key: str) -> SessionKeys: return SessionKeys(signing_key=SecretStr(signing)) +class SessionSigningPreviousKey(BaseModel): + """One retired key in ``mcp_session_token_signing.previous_public_keys``: its ``kid`` + and the PEM public half (inline or an ``os.environ/`` reference).""" + + model_config = ConfigDict(frozen=True, extra="forbid") + kid: str = Field(min_length=1) + public_key: str = Field(min_length=1) + + +class MCPSessionTokenSigningSettings(BaseModel): + """The ``general_settings.mcp_session_token_signing`` block: opt-in asymmetric signing + for the gateway session tokens. Absent, the gateway keeps the backward-compatible + HS256 key derived from ``master_key``. ``private_key`` and each ``public_key`` accept + a PEM string inline or an ``os.environ/`` (or secret manager) reference.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + algorithm: Literal["RS256"] + kid: str = Field(min_length=1) + private_key: str = Field(min_length=1) + previous_public_keys: tuple[SessionSigningPreviousKey, ...] = () + + +class SessionSigningConfigError(BaseModel): + """``mcp_session_token_signing`` is present but unusable (bad shape, unresolvable + secret reference, or a key that is not a loadable RSA PEM); the caller fails closed + with a server error instead of silently falling back to HS256.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["session_signing_config_error"] = "session_signing_config_error" + detail: str + + +def _resolve_key_material(value: str) -> str | None: + if not value.startswith("os.environ/"): + return value + from litellm.secret_managers.main import get_secret_str # noqa: PLC0415 # heavy import kept off the pure path + + return get_secret_str(value) + + +def resolve_session_signing_keys( + master_key: str, + raw_settings: object | None, +) -> SessionSigningKeys | SessionSigningConfigError: + """Turn the operator's ``mcp_session_token_signing`` setting into signing key material. + + ``None`` (the setting absent) keeps the backward-compatible HS256 key derived from + ``master_key``. A present setting must fully validate into RS256 material; any defect + is a ``SessionSigningConfigError`` value so token issuance and admission fail closed + rather than minting under a key the operator did not intend. + """ + if raw_settings is None: + return session_keys_from_master_key(master_key) + try: + settings: Final = MCPSessionTokenSigningSettings.model_validate(raw_settings) + except ValidationError as exc: + return SessionSigningConfigError(detail=f"mcp_session_token_signing is malformed: {exc}") + private_pem: Final = _resolve_key_material(settings.private_key) + if private_pem is None: + return SessionSigningConfigError(detail="mcp_session_token_signing.private_key reference did not resolve") + resolved_previous: Final = tuple( + (previous.kid, _resolve_key_material(previous.public_key)) for previous in settings.previous_public_keys + ) + unresolved: Final = tuple(kid for kid, pem in resolved_previous if pem is None) + if unresolved: + return SessionSigningConfigError( + detail=f"mcp_session_token_signing.previous_public_keys reference did not resolve for kid(s): {', '.join(unresolved)}" + ) + try: + return AsymmetricSessionKeys( + private_key_pem=SecretStr(private_pem), + kid=settings.kid, + previous_public_keys=tuple( + SessionRotatedPublicKey(kid=kid, public_key_pem=pem) + for kid, pem in resolved_previous + if pem is not None + ), + ) + except ValidationError as exc: + return SessionSigningConfigError( + detail=f"mcp_session_token_signing keys are not usable RSA PEM material: {exc}" + ) + + +def active_session_signing_keys(master_key: str) -> SessionSigningKeys | SessionSigningConfigError: + """Wiring helper for the token endpoint and the admission edge: resolve the signing + keys from the live ``general_settings.mcp_session_token_signing`` block, or derive the + default HS256 key from ``master_key`` when the block is absent.""" + from litellm.proxy.proxy_server import general_settings # noqa: PLC0415 # circular import at module load + + return resolve_session_signing_keys(master_key, general_settings.get("mcp_session_token_signing")) + + class NotSessionBearer(BaseModel): """The bearer is not session-shaped; admission continues on its normal path.""" @@ -116,7 +212,7 @@ def is_session_bearer_shaped(authorization_value: str) -> bool: def resolve_session_bearer( authorization_value: str, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> SessionBearerResult: """Classify an ``Authorization`` value presented at the aggregate MCP edge. @@ -166,7 +262,7 @@ SessionRefreshResult: TypeAlias = SessionRefreshOpened | SessionRefreshInvalid def open_session_refresh_bearer( refresh_value: str, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, expected_client_id: str, ) -> SessionRefreshResult: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py index 2c7b970ca0e..6824f96f927 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py @@ -8,8 +8,11 @@ is therefore a stable REFERENCE, not an authorization: admission reloads the liv record and policy on every request, so deactivating the user (or their team) kills outstanding sessions immediately without a revocation store. -Wire shape: ``llm_session_`` (access) / ``llm_srefresh_`` (refresh) + an HS256 JWT, -the same signing approach as :mod:`.envelope`. Claims are ``iss``/``iat``/``exp`` +Wire shape: ``llm_session_`` (access) / ``llm_srefresh_`` (refresh) + a JWT signed with +the injected key material: HS256 under the default master-key-derived secret (the same +signing approach as :mod:`.envelope`), or RS256 under an operator-provided RSA private +key (:class:`AsymmetricSessionKeys`) so downstream validators hold only the public half. +Claims are ``iss``/``iat``/``exp`` plus ``jti`` (per-mint uniqueness, so two tokens minted in the same second never collide and a future revocation list has a stable handle), ``kind``, ``user_id``, and ``client_id``; ``client_id`` binds the refresh token @@ -31,11 +34,16 @@ injected ``now``); the strict pydantic claims model is the sole, total type gate from __future__ import annotations import secrets +from collections import Counter from datetime import datetime, timedelta +from functools import lru_cache from typing import Final, Literal, TypeAlias import jwt -from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError +from cryptography.exceptions import UnsupportedAlgorithm +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError, field_validator, model_validator SESSION_TOKEN_PREFIX: Final = "llm_session_" """Marker prefix on every serialized session ACCESS token so the admission edge can cheaply @@ -71,6 +79,11 @@ limits while bounding hostile input before JWT parsing.""" _SESSION_JWT_ALGORITHM: Final = "HS256" +_SESSION_RSA_ALGORITHM: Final = "RS256" + +_MIN_RSA_KEY_BITS: Final = 2048 +"""RFC 7518 section 3.3: RS256 requires a key of at least 2048 bits.""" + SessionTokenKind = Literal["session", "session_refresh"] """Which credential a session token is. Stamped into the signed claims and required to match on open, so a signature-valid token of one kind cannot be replayed as the other even if its @@ -120,6 +133,85 @@ class SessionKeys(BaseModel): signing_key: SecretStr = Field(min_length=32) +class SessionRotatedPublicKey(BaseModel): + """The public half of a retired signing key, kept verifiable under its ``kid`` during a + rotation window so tokens minted before the rotation stay valid until they expire.""" + + model_config = ConfigDict(frozen=True) + kid: str = Field(min_length=1) + public_key_pem: str = Field(min_length=1) + + @field_validator("public_key_pem") + @classmethod + def _pem_is_an_rsa_public_key(cls, value: str) -> str: + try: + loaded: Final = serialization.load_pem_public_key(value.encode()) + except (ValueError, TypeError, UnsupportedAlgorithm) as exc: + raise ValueError(f"public_key_pem is not a loadable PEM public key: {exc}") from exc + if not isinstance(loaded, rsa.RSAPublicKey): + raise ValueError("public_key_pem must be an RSA public key in PEM format") # noqa: TRY004 # pydantic validators must raise ValueError + if loaded.key_size < _MIN_RSA_KEY_BITS: + raise ValueError(f"public_key_pem must be an RSA key of at least {_MIN_RSA_KEY_BITS} bits") + return value + + +class AsymmetricSessionKeys(BaseModel): + """Injected RS256 key material: the issuer-held RSA private key and the stable ``kid`` + stamped into every minted token's JOSE header, plus the public halves of previously + rotated keys that verification still accepts while their tokens age out. Downstream + validators never need the private key: :func:`session_public_key_pem` yields the + public half to distribute.""" + + model_config = ConfigDict(frozen=True) + private_key_pem: SecretStr + kid: str = Field(min_length=1) + previous_public_keys: tuple[SessionRotatedPublicKey, ...] = () + + @field_validator("private_key_pem") + @classmethod + def _pem_is_a_strong_rsa_private_key(cls, value: SecretStr) -> SecretStr: + try: + loaded: Final = serialization.load_pem_private_key(value.get_secret_value().encode(), password=None) + except (ValueError, TypeError, UnsupportedAlgorithm) as exc: + raise ValueError(f"private_key_pem is not a loadable unencrypted PEM private key: {exc}") from exc + if not isinstance(loaded, rsa.RSAPrivateKey): + raise ValueError("private_key_pem must be an unencrypted RSA private key in PEM format") # noqa: TRY004 # pydantic validators must raise ValueError + if loaded.key_size < _MIN_RSA_KEY_BITS: + raise ValueError(f"private_key_pem must be an RSA key of at least {_MIN_RSA_KEY_BITS} bits") + return value + + @model_validator(mode="after") + def _kids_are_unique(self) -> AsymmetricSessionKeys: + kids: Final = (self.kid, *(previous.kid for previous in self.previous_public_keys)) + duplicates: Final = tuple(kid for kid, count in Counter(kids).items() if count > 1) + if duplicates: + raise ValueError( + f"every kid must be unique across the current and previous keys; duplicated: {', '.join(duplicates)}" + ) + return self + + +SessionSigningKeys: TypeAlias = SessionKeys | AsymmetricSessionKeys +"""Every key material shape the mints and openers accept: the default master-key-derived +HS256 secret, or operator-configured RS256 RSA keys.""" + + +@lru_cache(maxsize=8) +def _public_key_pem_from_private(private_key_pem: str) -> str: + loaded: Final = serialization.load_pem_private_key(private_key_pem.encode(), password=None) + return ( + loaded.public_key() + .public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo) + .decode() + ) + + +def session_public_key_pem(keys: AsymmetricSessionKeys) -> str: + """The PEM public half of the current RS256 signing key: the only material a downstream + validator (an external gateway verifying ``kid``-matched tokens) ever needs.""" + return _public_key_pem_from_private(keys.private_key_pem.get_secret_value()) + + class MintedSessionToken(BaseModel): """A minted session token: the client-held bearer value and when it expires.""" @@ -221,7 +313,7 @@ def is_session_refresh_token(candidate: str) -> bool: def mint_session_token( principal: SessionPrincipal, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> MintedSessionToken | SessionTokenMintError: """Mint the short-lived session ACCESS token for ``principal``. @@ -241,7 +333,7 @@ def mint_session_token( def mint_session_refresh_token( principal: SessionPrincipal, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> MintedSessionToken | SessionTokenMintError: """Mint the long-lived session REFRESH token for ``principal``. @@ -262,7 +354,7 @@ def mint_session_refresh_token( def open_session_token( candidate: str, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> OpenedSessionToken | SessionTokenOpenError: """Validate a session ACCESS ``candidate`` and recover the principal. @@ -275,7 +367,7 @@ def open_session_token( def open_session_refresh_token( candidate: str, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> OpenedSessionToken | SessionTokenOpenError: """Validate a session REFRESH ``candidate`` and recover the principal. @@ -292,7 +384,7 @@ def _mint( prefix: str, principal: SessionPrincipal, expires_at: datetime, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> MintedSessionToken | SessionTokenTooLarge: """Sign the claims for either token kind and enforce the size cap. Shared by both mints @@ -309,20 +401,33 @@ def _mint( audience=principal.audience, team_id=principal.team_id, ) - token: Final = prefix + jwt.encode( - claims.model_dump(exclude_none=True), keys.signing_key.get_secret_value(), algorithm=_SESSION_JWT_ALGORITHM - ) + token: Final = prefix + _sign_claims(claims, keys) size_bytes: Final = len(token.encode("utf-8")) if size_bytes > MAX_SESSION_TOKEN_BYTES: return SessionTokenTooLarge(size_bytes=size_bytes, max_bytes=MAX_SESSION_TOKEN_BYTES) return MintedSessionToken(token=SecretStr(token), expires_at=expires_at) +def _sign_claims(claims: _SessionClaims, keys: SessionSigningKeys) -> str: + """Sign the claim set under whichever key material was injected: RS256 with the ``kid`` + in the JOSE header (so a validator can pick the right public key), or the default + HS256 secret with no header extras (byte-compatible with every pre-RS256 token).""" + payload: Final = claims.model_dump(exclude_none=True) + if isinstance(keys, AsymmetricSessionKeys): + return jwt.encode( + payload, + keys.private_key_pem.get_secret_value(), + algorithm=_SESSION_RSA_ALGORITHM, + headers={"kid": keys.kid}, + ) + return jwt.encode(payload, keys.signing_key.get_secret_value(), algorithm=_SESSION_JWT_ALGORITHM) + + def _open( candidate: str, prefix: str, expected_kind: SessionTokenKind, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> OpenedSessionToken | SessionTokenOpenError: """Prefix-route, size-bound, signature-verify, kind-check, and expiry-check an @@ -337,7 +442,7 @@ def _open( return SessionMalformed() if len(candidate.encode("utf-8", "surrogatepass")) > MAX_SESSION_TOKEN_BYTES: return SessionMalformed() - claims: Final = _decode_claims(candidate.removeprefix(prefix), keys.signing_key) + claims: Final = _decode_claims(candidate.removeprefix(prefix), keys) if not isinstance(claims, _SessionClaims): return claims if claims.kind != expected_kind: @@ -356,14 +461,51 @@ def _open( ) +class _VerificationMaterial(BaseModel): + model_config = ConfigDict(frozen=True) + key: SecretStr + algorithm: Literal["HS256", "RS256"] + + +def _verification_material( + compact: str, + keys: SessionSigningKeys, +) -> _VerificationMaterial | SessionBadSignature | SessionMalformed: + """Pick the single key and algorithm the candidate is allowed to verify under. + + HS256 mode has exactly one secret. RS256 mode routes by the JOSE header ``kid``: the + current key's derived public half, or a retired key's stored public half during a + rotation window. An unknown or missing ``kid`` is ``SessionBadSignature`` (a foreign + key), and an undecodable header is ``SessionMalformed``. The algorithm is pinned per + key shape, never read from the header, so an HS256 token can never be verified + against a public key or vice versa. + """ + if isinstance(keys, SessionKeys): + return _VerificationMaterial(key=keys.signing_key, algorithm=_SESSION_JWT_ALGORITHM) + try: + header: Final = jwt.get_unverified_header(compact) + except jwt.InvalidTokenError: + return SessionMalformed() + kid: Final = header.get("kid") + if kid == keys.kid: + return _VerificationMaterial(key=SecretStr(session_public_key_pem(keys)), algorithm=_SESSION_RSA_ALGORITHM) + for previous in keys.previous_public_keys: + if previous.kid == kid: + return _VerificationMaterial(key=SecretStr(previous.public_key_pem), algorithm=_SESSION_RSA_ALGORITHM) + return SessionBadSignature() + + def _decode_claims( compact: str, - signing_key: SecretStr, + keys: SessionSigningKeys, ) -> _SessionClaims | SessionBadSignature | SessionMalformed: - """Verify the HS256 signature and shape of an attacker-controlled compact JWT. + """Verify the signature and shape of an attacker-controlled compact JWT. ``compact`` is fully hostile and bounded to ``MAX_SESSION_TOKEN_BYTES`` by the caller. - PyJWT's ``iat``/``nbf``/``exp`` validators are disabled: they raise on hostile claim + The accepted algorithm is pinned by :func:`_verification_material` from the injected + key shape, so ``alg`` confusion (``none``, or HS256 signed with a public key as the + secret) fails before or at signature verification. PyJWT's ``iat``/``nbf``/``exp`` + validators are disabled: they raise on hostile claim types and, for ``iat``/``nbf``, compare against the wall clock rather than the injected ``now`` (``exp`` is checked by the caller against ``now``). Apart from a signature mismatch, every decode failure is ``SessionMalformed``: a non-UTF-8 candidate surfaces @@ -371,11 +513,14 @@ def _decode_claims( ``TypeError`` from PyJWT's claim validators, and a wrong issuer or structurally invalid token as an ``InvalidTokenError``. ``_SessionClaims`` is the total type gate. """ + material: Final = _verification_material(compact, keys) + if not isinstance(material, _VerificationMaterial): + return material try: payload: Final = jwt.decode( compact, - signing_key.get_secret_value(), - algorithms=[_SESSION_JWT_ALGORITHM], + material.key.get_secret_value(), + algorithms=[material.algorithm], issuer=SESSION_ISSUER, options={ "verify_exp": False, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py index 8fa7c15d2d3..00ff06ea082 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py @@ -3,6 +3,9 @@ from datetime import datetime, timedelta, timezone import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( envelope_keys_from_master_key, @@ -13,17 +16,22 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent SessionBearerInvalid, SessionRefreshInvalid, SessionRefreshOpened, + SessionSigningConfigError, is_session_bearer_shaped, open_session_refresh_bearer, resolve_session_bearer, + resolve_session_signing_keys, session_keys_from_master_key, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( SESSION_TTL_SECONDS, + AsymmetricSessionKeys, MintedSessionToken, + SessionKeys, SessionPrincipal, mint_session_refresh_token, mint_session_token, + session_public_key_pem, ) NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc) @@ -133,3 +141,86 @@ def test_refresh_grant_rejects_a_different_client(): def test_refresh_grant_rejects_access_token_presented_as_refresh(): result = open_session_refresh_bearer(_access_token(), KEYS, NOW, expected_client_id="llm_client_abc") assert isinstance(result, SessionRefreshInvalid) + + +def _rsa_private_pem() -> str: + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + return key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ).decode() + + +def test_absent_signing_setting_keeps_the_master_key_hs256_default(): + resolved = resolve_session_signing_keys(MASTER_KEY, None) + assert isinstance(resolved, SessionKeys) + assert resolved.signing_key.get_secret_value() == KEYS.signing_key.get_secret_value() + + +def test_rs256_signing_setting_resolves_inline_pem_material(): + pem = _rsa_private_pem() + resolved = resolve_session_signing_keys( + MASTER_KEY, + {"algorithm": "RS256", "kid": "2026-01", "private_key": pem}, + ) + assert isinstance(resolved, AsymmetricSessionKeys) + assert resolved.kid == "2026-01" + minted = mint_session_token(PRINCIPAL, resolved, NOW) + assert isinstance(minted, MintedSessionToken) + admitted = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", resolved, NOW) + assert isinstance(admitted, SessionBearerAdmitted) + + +def test_rs256_signing_setting_resolves_env_reference(monkeypatch): + monkeypatch.setenv("MCP_SESSION_PRIVATE_KEY", _rsa_private_pem()) + resolved = resolve_session_signing_keys( + MASTER_KEY, + {"algorithm": "RS256", "kid": "2026-01", "private_key": "os.environ/MCP_SESSION_PRIVATE_KEY"}, + ) + assert isinstance(resolved, AsymmetricSessionKeys) + + +def test_rs256_signing_setting_resolves_previous_public_keys(): + old_pem = _rsa_private_pem() + old_keys = AsymmetricSessionKeys(private_key_pem=SecretStr(old_pem), kid="2025-06") + resolved = resolve_session_signing_keys( + MASTER_KEY, + { + "algorithm": "RS256", + "kid": "2026-01", + "private_key": _rsa_private_pem(), + "previous_public_keys": [{"kid": "2025-06", "public_key": session_public_key_pem(old_keys)}], + }, + ) + assert isinstance(resolved, AsymmetricSessionKeys) + minted = mint_session_token(PRINCIPAL, old_keys, NOW) + assert isinstance(minted, MintedSessionToken) + admitted = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", resolved, NOW) + assert isinstance(admitted, SessionBearerAdmitted) + + +@pytest.mark.parametrize( + "raw", + [ + {"algorithm": "HS512", "kid": "k", "private_key": "irrelevant"}, + {"algorithm": "RS256", "kid": "k"}, + {"algorithm": "RS256", "kid": "k", "private_key": "not a pem"}, + {"algorithm": "RS256", "kid": "k", "private_key": "os.environ/UNSET_MCP_SESSION_KEY_VAR"}, + {"algorithm": "RS256", "kid": "k", "private_key": "x", "unexpected": True}, + "not-a-mapping", + ], +) +def test_defective_signing_setting_fails_closed_never_falls_back_to_hs256(raw): + resolved = resolve_session_signing_keys(MASTER_KEY, raw) + assert isinstance(resolved, SessionSigningConfigError) + + +def test_signing_config_error_detail_never_leaks_key_material(): + pem = _rsa_private_pem() + resolved = resolve_session_signing_keys( + MASTER_KEY, + {"algorithm": "RS256", "kid": "k", "private_key": pem, "unexpected": True}, + ) + assert isinstance(resolved, SessionSigningConfigError) + assert pem.splitlines()[1] not in resolved.detail diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py index 36280530eac..2a59e6c1baa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py @@ -4,6 +4,8 @@ from datetime import datetime, timedelta, timezone import jwt import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa from pydantic import SecretStr, ValidationError from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( @@ -13,6 +15,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i SESSION_REFRESH_TTL_SECONDS, SESSION_TOKEN_PREFIX, SESSION_TTL_SECONDS, + AsymmetricSessionKeys, MintedSessionToken, NotASessionToken, OpenedSessionToken, @@ -21,6 +24,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i SessionKeys, SessionMalformed, SessionPrincipal, + SessionRotatedPublicKey, SessionTokenTooLarge, is_session_refresh_token, is_session_token, @@ -28,8 +32,22 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i mint_session_token, open_session_refresh_token, open_session_token, + session_public_key_pem, ) + +def _rsa_private_pem(bits: int = 2048) -> str: + key = rsa.generate_private_key(public_exponent=65537, key_size=bits) + return key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ).decode() + + +_RSA_PEM_A = _rsa_private_pem() +_RSA_PEM_B = _rsa_private_pem() + NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc) KEYS = SessionKeys(signing_key=SecretStr("k" * 32)) OTHER_KEYS = SessionKeys(signing_key=SecretStr("x" * 32)) @@ -264,3 +282,172 @@ def test_signed_claims_with_a_non_string_team_are_rejected(): def test_principal_rejects_an_unknown_audience_at_construction(): with pytest.raises(ValidationError): SessionPrincipal(user_id="user-123", client_id="llm_client_abc", audience="mcp") + + +RSA_KEYS = AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_A), kid="2026-01") +OTHER_RSA_KEYS = AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_B), kid="2025-06") + + +def test_rs256_access_round_trip_with_kid_and_alg_pinned_in_header(): + minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + token = minted.token.get_secret_value() + header = jwt.get_unverified_header(token.removeprefix(SESSION_TOKEN_PREFIX)) + assert header["alg"] == "RS256" + assert header["kid"] == "2026-01" + opened = open_session_token(token, RSA_KEYS, NOW) + assert isinstance(opened, OpenedSessionToken) + assert opened.principal == PRINCIPAL + + +def test_rs256_refresh_round_trip(): + minted = mint_session_refresh_token(PRINCIPAL, RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + token = minted.token.get_secret_value() + opened = open_session_refresh_token(token, RSA_KEYS, NOW) + assert isinstance(opened, OpenedSessionToken) + assert opened.principal == PRINCIPAL + + +def test_rs256_token_verifies_with_public_key_only(): + minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + public_pem = session_public_key_pem(RSA_KEYS) + assert "PUBLIC KEY" in public_pem + assert "PRIVATE" not in public_pem + claims = jwt.decode( + minted.token.get_secret_value().removeprefix(SESSION_TOKEN_PREFIX), + public_pem, + algorithms=["RS256"], + issuer=SESSION_ISSUER, + options={"verify_exp": False}, + ) + assert claims["user_id"] == "user-123" + + +def test_rs256_tampered_signature_is_bad_signature(): + minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + token = minted.token.get_secret_value() + tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb") + assert isinstance(open_session_token(tampered, RSA_KEYS, NOW), SessionBadSignature) + + +def test_rs256_expired_token_is_expired(): + minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) + assert isinstance(open_session_token(minted.token.get_secret_value(), RSA_KEYS, after), SessionExpired) + + +def test_hs256_token_is_rejected_in_rs256_mode(): + assert isinstance(open_session_token(_mint_access(), RSA_KEYS, NOW), SessionBadSignature) + + +def test_hs256_token_claiming_the_current_kid_is_rejected_by_alg_pinning(): + token = SESSION_TOKEN_PREFIX + jwt.encode( + _valid_claims(), + KEYS.signing_key.get_secret_value(), + algorithm="HS256", + headers={"kid": RSA_KEYS.kid}, + ) + assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionMalformed) + + +def test_rs256_token_is_rejected_in_hs256_mode(): + minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + assert isinstance(open_session_token(minted.token.get_secret_value(), KEYS, NOW), SessionMalformed) + + +def test_rs256_token_from_an_unknown_kid_is_bad_signature(): + minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + assert isinstance(open_session_token(minted.token.get_secret_value(), RSA_KEYS, NOW), SessionBadSignature) + + +def test_rs256_token_signed_by_a_foreign_key_claiming_the_current_kid_is_bad_signature(): + token = SESSION_TOKEN_PREFIX + jwt.encode( + _valid_claims(), + _RSA_PEM_B, + algorithm="RS256", + headers={"kid": RSA_KEYS.kid}, + ) + assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionBadSignature) + + +def test_alg_none_token_with_the_current_kid_is_rejected_in_rs256_mode(): + unsigned = jwt.api_jws.encode( + b'{"iss":"litellm-mcp-gateway"}', key=None, algorithm="none", headers={"kid": RSA_KEYS.kid} + ) + assert isinstance(open_session_token(SESSION_TOKEN_PREFIX + unsigned, RSA_KEYS, NOW), SessionMalformed) + + +def test_rotation_previous_public_key_still_verifies_until_removed(): + minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + token = minted.token.get_secret_value() + rotated = AsymmetricSessionKeys( + private_key_pem=SecretStr(_RSA_PEM_A), + kid="2026-01", + previous_public_keys=( + SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS)), + ), + ) + opened = open_session_token(token, rotated, NOW) + assert isinstance(opened, OpenedSessionToken) + assert opened.principal == PRINCIPAL + assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionBadSignature) + + +def test_rotation_window_still_enforces_expiry_and_tamper_on_the_previous_key(): + minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + token = minted.token.get_secret_value() + rotated = AsymmetricSessionKeys( + private_key_pem=SecretStr(_RSA_PEM_A), + kid="2026-01", + previous_public_keys=( + SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS)), + ), + ) + after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) + assert isinstance(open_session_token(token, rotated, after), SessionExpired) + tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb") + assert isinstance(open_session_token(tampered, rotated, NOW), SessionBadSignature) + + +def test_weak_or_garbage_private_key_pem_rejected_at_construction(): + with pytest.raises(ValidationError): + AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem(bits=1024)), kid="weak") + with pytest.raises(ValidationError): + AsymmetricSessionKeys(private_key_pem=SecretStr("not a pem"), kid="junk") + with pytest.raises(ValidationError): + SessionRotatedPublicKey(kid="junk", public_key_pem="not a pem") + with pytest.raises(ValidationError): + SessionRotatedPublicKey(kid="private-half", public_key_pem=_RSA_PEM_A) + + +def test_weak_rotated_public_key_rejected_at_construction(): + weak_public = ( + serialization.load_pem_private_key(_rsa_private_pem(bits=1024).encode(), password=None) + .public_key() + .public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo) + .decode() + ) + with pytest.raises(ValidationError): + SessionRotatedPublicKey(kid="2024-01", public_key_pem=weak_public) + + +def test_duplicate_kids_rejected_at_construction(): + previous = SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS)) + with pytest.raises(ValidationError): + AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_A), kid="2025-06", previous_public_keys=(previous,)) + with pytest.raises(ValidationError): + AsymmetricSessionKeys( + private_key_pem=SecretStr(_RSA_PEM_A), kid="2026-01", previous_public_keys=(previous, previous) + ) + + +def test_asymmetric_keys_repr_never_leaks_the_private_key(): + assert _RSA_PEM_A not in repr(RSA_KEYS)