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 f575265a5e7..6efe299780c 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 @@ -27,6 +27,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti NotBridgeEnvelope, envelope_keys_from_master_key, is_bridge_envelope_shaped, + legacy_envelope_keys_from_master_key, resolve_bridge_envelope, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( @@ -868,7 +869,13 @@ class MCPRequestHandler: await MCPRequestHandler._run_pre_db_read_auth_checks(request=request, route=route) keys: Final = envelope_keys_from_master_key(master_key) - result: Final = resolve_bridge_envelope(authorization_value, keys, datetime.now(timezone.utc), server.server_id) + result: Final = resolve_bridge_envelope( + authorization_value, + keys, + datetime.now(timezone.utc), + server.server_id, + legacy_keys=legacy_envelope_keys_from_master_key(master_key), + ) match result: case BridgeEnvelopeAdmitted(): return result @@ -998,6 +1005,7 @@ class MCPRequestHandler: SessionBearerInvalid, SessionSigningConfigError, active_session_signing_keys, + legacy_session_keys_from_master_key, resolve_session_bearer, ) from litellm.proxy.proxy_server import master_key @@ -1011,7 +1019,12 @@ class MCPRequestHandler: 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)) + result: Final = resolve_session_bearer( + authorization_value, + keys, + datetime.now(timezone.utc), + legacy_keys=legacy_session_keys_from_master_key(master_key, keys), + ) match result: case SessionBearerAdmitted(): try: diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index 77bdbd26b35..d2b4ce6040f 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -867,6 +867,7 @@ async def _prepare_bridge_refresh( from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import BridgeRefreshOpened, envelope_keys_from_master_key, + legacy_envelope_keys_from_master_key, open_bridge_refresh_envelope, ) from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import @@ -878,7 +879,13 @@ async def _prepare_bridge_refresh( if not refresh_value: return "invalid_refresh" keys: Final = envelope_keys_from_master_key(master_key) - opened: Final = open_bridge_refresh_envelope(refresh_value, keys, datetime.now(timezone.utc), mcp_server.server_id) + opened: Final = open_bridge_refresh_envelope( + refresh_value, + keys, + datetime.now(timezone.utc), + mcp_server.server_id, + legacy_keys=legacy_envelope_keys_from_master_key(master_key), + ) if not isinstance(opened, BridgeRefreshOpened): return "invalid_refresh" failure: Final = await _revalidate_active_subject(opened.identity) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 1136410fd18..85654bcb6a9 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -68,6 +68,8 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent SessionRefreshOpened, SessionSigningConfigError, active_session_signing_keys, + legacy_session_keys_from_master_key, + open_session_credential_with_legacy, open_session_refresh_bearer, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( @@ -76,6 +78,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i MintedSessionToken, OpenedSessionToken, SessionAudience, + SessionKeys, SessionPrincipal, SessionSigningKeys, is_session_refresh_token, @@ -1241,6 +1244,7 @@ async def aggregate_token( keys=keys, now=now, issue=issue, + legacy_keys=legacy_session_keys_from_master_key(master_key, keys), ) if grant_type == TOKEN_EXCHANGE_GRANT_TYPE: return await _token_exchange_grant( @@ -1400,10 +1404,13 @@ async def _refresh_token_grant( keys: SessionSigningKeys, now: datetime, issue: _GrantIssuer, + legacy_keys: SessionKeys | None = None, ) -> Response: if not refresh_token: return _oauth_error(400, "invalid_request", "refresh_token is required") - opened: Final = open_session_refresh_bearer(refresh_token, keys, now, expected_client_id=client_id) + opened: Final = open_session_refresh_bearer( + refresh_token, keys, now, expected_client_id=client_id, legacy_keys=legacy_keys + ) if not isinstance(opened, SessionRefreshOpened): return _oauth_error(400, "invalid_grant", "the refresh token is invalid for this client") if _resource_conflicts_with_scope(request, resource, opened.principal.resource_server_id): @@ -1464,7 +1471,13 @@ async def revoke_refresh_token(token: str, client_id: str, master_key: str | Non 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) + opened: Final = open_session_refresh_bearer( + token, + keys, + now, + expected_client_id=client_id, + legacy_keys=legacy_session_keys_from_master_key(master_key, keys), + ) if isinstance(opened, SessionRefreshOpened): burned: Final = await _SingleUseGuard(cache).claim( f"{_USED_REFRESH_CACHE_PREFIX}{opened.jti}", SESSION_REFRESH_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS @@ -1532,10 +1545,11 @@ async def introspect_gateway_token( verbose_logger.error("mcp_gateway_dcr introspect rejected: %s", keys.detail) return _oauth_error(500, "server_error", keys.detail) now: Final = datetime.now(timezone.utc) + legacy_keys: Final = legacy_session_keys_from_master_key(master_key, keys) if is_session_token(token): - opened = open_session_token(token, keys, now) + opened = open_session_credential_with_legacy(open_session_token, token, keys, legacy_keys, now) elif is_session_refresh_token(token): - opened = open_session_refresh_token(token, keys, now) + opened = open_session_credential_with_legacy(open_session_refresh_token, token, keys, legacy_keys, now) else: return _inactive_introspection_response() if not isinstance(opened, OpenedSessionToken): diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py index ab5fa65480e..fb5a2e38236 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py @@ -9,7 +9,8 @@ to the upstream MCP server (consumer). This module is the pure surface for both token-endpoint and admission wiring live in their respective call sites. """ -import hashlib +import os +from collections.abc import Callable from datetime import datetime from functools import lru_cache from typing import Final, Literal, TypeAlias @@ -20,6 +21,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import EnvelopeIdentity, EnvelopeKeys, EnvelopeMintError, + EnvelopeOpenError, OpenedEnvelope, OpenedRefreshEnvelope, RefreshCredential, @@ -32,56 +34,55 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import open_envelope, open_refresh_envelope, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( + hkdf_sha256, + legacy_kdf_grace_enabled, + legacy_scrypt, +) _SIGNING_KEY_DOMAIN: Final = b"litellm-mcp-bridge:envelope-signing:" _ENCRYPTION_KEY_DOMAIN: Final = b"litellm-mcp-bridge:envelope-encryption:" -# scrypt work factors (RFC 7914). n=2**15 with r=8/p=1 costs ~50ms and ~32MB per derivation, which -# makes offline guessing of a candidate master key memory-hard rather than a bare hash comparison. -_SCRYPT_N: Final = 2**15 -_SCRYPT_R: Final = 8 -_SCRYPT_P: Final = 1 -# scrypt's working-set is ~128 * N * r * p bytes; cap at twice that so the maxmem ceiling scales -# with every work factor and a future p or r bump does not trip "memory limit exceeded". -_SCRYPT_MAXMEM: Final = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2 -_DERIVED_KEY_BYTES: Final = 32 - @lru_cache(maxsize=8) def envelope_keys_from_master_key(master_key: str) -> EnvelopeKeys: """Derive the envelope signing and encryption keys from the proxy master key. - A memory-hard scrypt KDF (RFC 7914) over two distinct domain-label salts yields two - independent 256-bit subkeys from the one secret, so the producer (mint) and consumer - (open) agree on keys without persisting any. scrypt is used rather than a bare hash or - HMAC so that a captured envelope is not a cheap offline oracle for the master key: each - candidate guess costs a full memory-hard derivation, which is what protects a deployment - whose master key is weaker than it should be. The result is cached (the master key is - fixed for a process), so the KDF runs once per key and adds nothing to the per-request - admission path. The derivation is deterministic; rotating ``master_key`` invalidates - every outstanding envelope, which is the intended behavior for a signing-key change. + HKDF-SHA256 (RFC 5869) over the master key, once per domain label as ``info``, yields + two independent 256-bit subkeys from the one secret, so the producer (mint) and consumer + (open) agree on keys without persisting any. No memory-hard guessing cost is needed: + the master key is a high-entropy secret, so a captured envelope is not a cheap offline + oracle for it either way. The result is cached (the master key is fixed for a process), + so the KDF runs once per key and adds nothing to the per-request admission path. The + derivation is deterministic; rotating ``master_key`` invalidates every outstanding + envelope, which is the intended behavior for a signing-key change. """ - signing: Final = hashlib.scrypt( - master_key.encode(), - salt=_SIGNING_KEY_DOMAIN, - n=_SCRYPT_N, - r=_SCRYPT_R, - p=_SCRYPT_P, - maxmem=_SCRYPT_MAXMEM, - dklen=_DERIVED_KEY_BYTES, - ).hex() - encryption: Final = hashlib.scrypt( - master_key.encode(), - salt=_ENCRYPTION_KEY_DOMAIN, - n=_SCRYPT_N, - r=_SCRYPT_R, - p=_SCRYPT_P, - maxmem=_SCRYPT_MAXMEM, - dklen=_DERIVED_KEY_BYTES, - ).hex() + return EnvelopeKeys( + signing_key=SecretStr(hkdf_sha256(master_key, _SIGNING_KEY_DOMAIN)), + encryption_key=SecretStr(hkdf_sha256(master_key, _ENCRYPTION_KEY_DOMAIN)), + ) + + +@lru_cache(maxsize=8) +def _legacy_envelope_keys(master_key: str) -> EnvelopeKeys | None: + """The pre-rotation scrypt subkeys, kept so envelopes minted before the KDF rotation + can still be opened during the opt-in grace window. ``None`` when scrypt is unavailable.""" + signing: Final = legacy_scrypt(master_key, _SIGNING_KEY_DOMAIN) + encryption: Final = legacy_scrypt(master_key, _ENCRYPTION_KEY_DOMAIN) + if signing is None or encryption is None: + return None return EnvelopeKeys(signing_key=SecretStr(signing), encryption_key=SecretStr(encryption)) +def legacy_envelope_keys_from_master_key( + master_key: str, environ: Callable[[str], str | None] = os.environ.get +) -> EnvelopeKeys | None: + """The legacy envelope keys when the operator enabled the KDF grace window, else ``None``.""" + if not legacy_kdf_grace_enabled(environ): + return None + return _legacy_envelope_keys(master_key) + + def build_bridge_token_response( identity: EnvelopeIdentity, grant: UpstreamTokenGrant, @@ -131,11 +132,40 @@ class BridgeRefreshInvalid(BaseModel): BridgeRefreshResult: TypeAlias = BridgeRefreshOpened | BridgeRefreshInvalid +def _open_envelope_with_fallback( + candidate: str, + keys: EnvelopeKeys, + legacy_keys: EnvelopeKeys | None, + now: datetime, +) -> OpenedEnvelope | EnvelopeOpenError: + """Open an access envelope under the active keys, then under the legacy keys when + grace supplied them.""" + primary: Final = open_envelope(candidate, keys, now) + if isinstance(primary, OpenedEnvelope) or legacy_keys is None: + return primary + return open_envelope(candidate, legacy_keys, now) + + +def _open_refresh_envelope_with_fallback( + candidate: str, + keys: EnvelopeKeys, + legacy_keys: EnvelopeKeys | None, + now: datetime, +) -> OpenedRefreshEnvelope | EnvelopeOpenError: + """Open a refresh envelope under the active keys, then under the legacy keys when + grace supplied them.""" + primary: Final = open_refresh_envelope(candidate, keys, now) + if isinstance(primary, OpenedRefreshEnvelope) or legacy_keys is None: + return primary + return open_refresh_envelope(candidate, legacy_keys, now) + + def open_bridge_refresh_envelope( refresh_value: str, keys: EnvelopeKeys, now: datetime, expected_server_id: str, + legacy_keys: EnvelopeKeys | None = None, ) -> BridgeRefreshResult: """Open a refresh envelope a bridge ``oauth_delegate`` client presented on a refresh_token grant. @@ -150,7 +180,7 @@ def open_bridge_refresh_envelope( candidate: Final = _strip_bearer(refresh_value) if not is_refresh_envelope(candidate): return BridgeRefreshInvalid() - opened: Final = open_refresh_envelope(candidate, keys, now) + opened: Final = _open_refresh_envelope_with_fallback(candidate, keys, legacy_keys, now) if not isinstance(opened, OpenedRefreshEnvelope): return BridgeRefreshInvalid() if opened.identity.server_id != expected_server_id: @@ -207,6 +237,7 @@ def resolve_bridge_envelope( keys: EnvelopeKeys, now: datetime, expected_server_id: str, + legacy_keys: EnvelopeKeys | None = None, ) -> BridgeEnvelopeResult: """Classify an ``Authorization`` value presented to a bridge ``oauth_delegate`` server. @@ -233,7 +264,7 @@ def resolve_bridge_envelope( return BridgeEnvelopeInvalid() if not is_envelope(candidate): return NotBridgeEnvelope() - opened: Final = open_envelope(candidate, keys, now) + opened: Final = _open_envelope_with_fallback(candidate, keys, legacy_keys, now) if not isinstance(opened, OpenedEnvelope): return BridgeEnvelopeInvalid() if opened.identity.server_id != expected_server_id: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py new file mode 100644 index 00000000000..9ea90386bc8 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py @@ -0,0 +1,56 @@ +"""Key derivation for the MCP outbound-credential token families. + +Bridge envelopes and gateway session tokens are keyed by HKDF-SHA256 (RFC 5869) over the +high-entropy proxy ``master_key`` with a per-family domain label as ``info``. The legacy +scrypt derivation is kept for one release as an opt-in read fallback so tokens minted +before the rotation still open during the grace window; it is never used under FIPS. +""" + +import hashlib +import os +from collections.abc import Callable +from typing import Final + +from cryptography.hazmat.primitives.hashes import SHA256 +from cryptography.hazmat.primitives.kdf.hkdf import HKDF + +from litellm.proxy.common_utils.fips import is_fips_mode + +LEGACY_KDF_GRACE_ENV_VAR: Final = "LITELLM_MCP_LEGACY_KDF_GRACE" +DERIVED_KEY_BYTES: Final = 32 + +# RFC 7914 work factors of the pre-rotation derivation; maxmem is twice the working set. +_SCRYPT_N: Final = 2**15 +_SCRYPT_R: Final = 8 +_SCRYPT_P: Final = 1 +_SCRYPT_MAXMEM: Final = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2 + + +def hkdf_sha256(master_key: str, info: bytes) -> str: + """Derive a 256-bit subkey from ``master_key`` bound to the domain label ``info``.""" + return HKDF(algorithm=SHA256(), length=DERIVED_KEY_BYTES, salt=None, info=info).derive(master_key.encode()).hex() + + +def legacy_scrypt(master_key: str, salt: bytes) -> str | None: + """The pre-rotation derivation. ``None`` when hashlib.scrypt is unavailable (some + OpenSSL builds, including FIPS providers, do not expose it).""" + try: + return hashlib.scrypt( + master_key.encode(), + salt=salt, + n=_SCRYPT_N, + r=_SCRYPT_R, + p=_SCRYPT_P, + maxmem=_SCRYPT_MAXMEM, + dklen=DERIVED_KEY_BYTES, + ).hex() + except (AttributeError, ValueError): + return None + + +def legacy_kdf_grace_enabled(environ: Callable[[str], str | None] = os.environ.get) -> bool: + """True only when the operator opted the legacy read fallback in and the deployment is + not in FIPS mode (scrypt is not a FIPS-approved KDF).""" + if is_fips_mode(environ): + return False + return (environ(LEGACY_KDF_GRACE_ENV_VAR) or "").strip().lower() == "true" 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 df2bbdba345..396a821b0fa 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py @@ -8,20 +8,26 @@ recovered litellm user (consumer), reloading the live user record and policy bef anything runs. This module is the pure surface for both sides; the token-endpoint and admission wiring live in their respective call sites. -The signing key is derived with the same memory-hard scrypt construction as +The signing key is derived with the same HKDF-SHA256 construction as :func:`~.bridge_credentials.envelope_keys_from_master_key` but under a distinct domain label, so session tokens and bridge envelopes never share key material: a token of one family is unverifiable in the other by key separation, on top of the distinct issuers, prefixes, and claim shapes. """ -import hashlib +import os +from collections.abc import Callable from datetime import datetime from functools import lru_cache from typing import Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError +from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( + hkdf_sha256, + legacy_kdf_grace_enabled, + legacy_scrypt, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( AsymmetricSessionKeys, OpenedSessionToken, @@ -30,6 +36,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i SessionPrincipal, SessionRotatedPublicKey, SessionSigningKeys, + SessionTokenOpenError, is_session_refresh_token, is_session_token, open_session_refresh_token, @@ -38,39 +45,49 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i _SESSION_SIGNING_KEY_DOMAIN: Final = b"litellm-mcp-gateway:session-signing:" -# scrypt work factors (RFC 7914), identical to the envelope KDF: memory-hard so a captured -# session token is not a cheap offline oracle for the master key. -_SCRYPT_N: Final = 2**15 -_SCRYPT_R: Final = 8 -_SCRYPT_P: Final = 1 -_SCRYPT_MAXMEM: Final = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2 -_DERIVED_KEY_BYTES: Final = 32 - @lru_cache(maxsize=8) def session_keys_from_master_key(master_key: str) -> SessionKeys: """Derive the session signing key from the proxy master key. - A memory-hard scrypt KDF (RFC 7914) over a session-specific domain-label salt yields a - 256-bit subkey from the one secret, so the producer (mint) and consumer (open) agree on - the key without persisting any. The domain label differs from both envelope labels in - :mod:`.bridge_credentials`, so compromise or misuse of one token family never crosses - into the other. The result is cached (the master key is fixed for a process); rotating - ``master_key`` invalidates every outstanding session, which is the intended behavior - for a signing-key change. + HKDF-SHA256 (RFC 5869) over the master key with a session-specific domain label as + ``info`` yields a 256-bit subkey from the one secret, so the producer (mint) and + consumer (open) agree on the key without persisting any. The label differs from both + envelope labels in :mod:`.bridge_credentials`, so compromise or misuse of one token + family never crosses into the other. The result is cached (the master key is fixed for + a process); rotating ``master_key`` invalidates every outstanding session, which is + the intended behavior for a signing-key change. """ - signing: Final = hashlib.scrypt( - master_key.encode(), - salt=_SESSION_SIGNING_KEY_DOMAIN, - n=_SCRYPT_N, - r=_SCRYPT_R, - p=_SCRYPT_P, - maxmem=_SCRYPT_MAXMEM, - dklen=_DERIVED_KEY_BYTES, - ).hex() + return SessionKeys(signing_key=SecretStr(hkdf_sha256(master_key, _SESSION_SIGNING_KEY_DOMAIN))) + + +@lru_cache(maxsize=8) +def _legacy_session_keys(master_key: str) -> SessionKeys | None: + """The pre-rotation scrypt signing key, kept so session tokens minted before the KDF + rotation can still be opened during the opt-in grace window. ``None`` when scrypt is + unavailable.""" + signing: Final = legacy_scrypt(master_key, _SESSION_SIGNING_KEY_DOMAIN) + if signing is None: + return None return SessionKeys(signing_key=SecretStr(signing)) +def legacy_session_keys_from_master_key( + master_key: str, + active: SessionSigningKeys, + environ: Callable[[str], str | None] = os.environ.get, +) -> SessionKeys | None: + """The legacy session key when the operator enabled the KDF grace window, else ``None``. + + The fallback only applies to the default master-key HS256 path: a configured + ``mcp_session_token_signing`` RS256 key was never scrypt-derived, so an + :class:`AsymmetricSessionKeys` ``active`` gets no legacy key even with grace on. + """ + if not isinstance(active, SessionKeys) or not legacy_kdf_grace_enabled(environ): + return None + return _legacy_session_keys(master_key) + + 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).""" @@ -214,6 +231,7 @@ def resolve_session_bearer( authorization_value: str, keys: SessionSigningKeys, now: datetime, + legacy_keys: SessionKeys | None = None, ) -> SessionBearerResult: """Classify an ``Authorization`` value presented at the aggregate MCP edge. @@ -232,7 +250,7 @@ def resolve_session_bearer( return SessionBearerInvalid() if not is_session_token(candidate): return NotSessionBearer() - opened: Final = open_session_token(candidate, keys, now) + opened: Final = open_session_credential_with_legacy(open_session_token, candidate, keys, legacy_keys, now) if isinstance(opened, OpenedSessionToken): return SessionBearerAdmitted(principal=opened.principal) return SessionBearerInvalid(expired=isinstance(opened, SessionExpired)) @@ -260,11 +278,29 @@ class SessionRefreshInvalid(BaseModel): SessionRefreshResult: TypeAlias = SessionRefreshOpened | SessionRefreshInvalid +def open_session_credential_with_legacy( + open_token: Callable[[str, SessionSigningKeys, datetime], OpenedSessionToken | SessionTokenOpenError], + token: str, + keys: SessionSigningKeys, + legacy_keys: SessionKeys | None, + now: datetime, +) -> OpenedSessionToken | SessionTokenOpenError: + """Open a session credential under the active keys, then under the legacy keys when + grace supplied them. An :class:`OpenedSessionToken` or :class:`SessionExpired` under + the active keys is final: expiry must still report as expiry, so only a signature or + format failure falls through to the legacy derivation.""" + primary: Final = open_token(token, keys, now) + if isinstance(primary, (OpenedSessionToken, SessionExpired)) or legacy_keys is None: + return primary + return open_token(token, legacy_keys, now) + + def open_session_refresh_bearer( refresh_value: str, keys: SessionSigningKeys, now: datetime, expected_client_id: str, + legacy_keys: SessionKeys | None = None, ) -> SessionRefreshResult: """Open a session refresh token presented on a ``refresh_token`` grant. @@ -279,7 +315,7 @@ def open_session_refresh_bearer( candidate: Final = _strip_bearer(refresh_value) if not is_session_refresh_token(candidate): return SessionRefreshInvalid() - opened: Final = open_session_refresh_token(candidate, keys, now) + opened: Final = open_session_credential_with_legacy(open_session_refresh_token, candidate, keys, legacy_keys, now) if not isinstance(opened, OpenedSessionToken): return SessionRefreshInvalid() if opened.principal.client_id != expected_client_id: diff --git a/tests/integration/mcp/test_mcp_kdf_rotation.py b/tests/integration/mcp/test_mcp_kdf_rotation.py new file mode 100644 index 00000000000..9fdce17d7d8 --- /dev/null +++ b/tests/integration/mcp/test_mcp_kdf_rotation.py @@ -0,0 +1,494 @@ +"""Bridge envelopes and gateway session tokens are keyed by HKDF-SHA256 over the master key. + +Every cell mints a bearer in the test under keys the test derives itself (HKDF-SHA256 for the current +construction, scrypt for the legacy one) and presents it to a real proxy. The proxy is the only party +that derives keys from the master key, so admission proves which KDF it runs and whether the legacy +grace window is honoured. +""" + +import hashlib +import uuid +from datetime import datetime, timezone +from pathlib import Path +from typing import Final + +import httpx +import pytest +from cryptography.hazmat.primitives.hashes import SHA256 +from cryptography.hazmat.primitives.kdf.hkdf import HKDF +from integration._support.client import Gateway +from integration._support.mcp import McpPeer, mcp_peer, register_mcp +from integration._support.oauth_server import AuthorizationServer, oauth_server +from integration._support.process import owned_proxy +from pydantic import SecretStr +from typing_extensions import ReadOnly, TypedDict + +from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + EnvelopeKeys, + OpenedEnvelope, + OpenedRefreshEnvelope, + RefreshCredential, + SealedEnvelope, + UpstreamTokenGrant, + key_hash_identity, + mint_envelope, + mint_refresh_envelope, + open_envelope, + open_refresh_envelope, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + MintedSessionToken, + OpenedSessionToken, + SessionKeys, + SessionPrincipal, + mint_session_refresh_token, + mint_session_token, + open_session_refresh_token, +) +from litellm.proxy.utils import hash_token + +ACCEPT: Final = {"Accept": "application/json, text/event-stream"} +TOOLS_LIST: Final = {"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {}} +ENVELOPE_SIGNING: Final = b"litellm-mcp-bridge:envelope-signing:" +ENVELOPE_ENCRYPTION: Final = b"litellm-mcp-bridge:envelope-encryption:" +SESSION_SIGNING: Final = b"litellm-mcp-gateway:session-signing:" +GRACE: Final = {"LITELLM_MCP_LEGACY_KDF_GRACE": "true"} + + +class _RefreshGrantForm(TypedDict): + grant_type: ReadOnly[str] + refresh_token: ReadOnly[str] + client_id: ReadOnly[str] + + +class _RevokeForm(TypedDict): + token: ReadOnly[str] + client_id: ReadOnly[str] + + +class _IntrospectForm(TypedDict): + token: ReadOnly[str] + + +class _ClientRegistration(TypedDict): + redirect_uris: ReadOnly[tuple[str, ...]] + client_name: ReadOnly[str] + + +class _ObjectPermission(TypedDict): + mcp_servers: ReadOnly[tuple[str, ...]] + + +def _hkdf(master_key: str, info: bytes) -> str: + return HKDF(algorithm=SHA256(), length=32, salt=None, info=info).derive(master_key.encode()).hex() + + +def _scrypt(master_key: str, salt: bytes) -> str: + return hashlib.scrypt(master_key.encode(), salt=salt, n=2**15, r=8, p=1, maxmem=2**27, dklen=32).hex() + + +def _hkdf_envelope_keys(master_key: str) -> EnvelopeKeys: + return EnvelopeKeys( + signing_key=SecretStr(_hkdf(master_key, ENVELOPE_SIGNING)), + encryption_key=SecretStr(_hkdf(master_key, ENVELOPE_ENCRYPTION)), + ) + + +def _scrypt_envelope_keys(master_key: str) -> EnvelopeKeys: + return EnvelopeKeys( + signing_key=SecretStr(_scrypt(master_key, ENVELOPE_SIGNING)), + encryption_key=SecretStr(_scrypt(master_key, ENVELOPE_ENCRYPTION)), + ) + + +def _hkdf_session_keys(master_key: str) -> SessionKeys: + return SessionKeys(signing_key=SecretStr(_hkdf(master_key, SESSION_SIGNING))) + + +def _scrypt_session_keys(master_key: str) -> SessionKeys: + return SessionKeys(signing_key=SecretStr(_scrypt(master_key, SESSION_SIGNING))) + + +def _envelope(identity_server: str, key: str, keys: EnvelopeKeys, upstream_token: str) -> str: + sealed: Final = mint_envelope( + key_hash_identity(server_id=identity_server, key_hash=hash_token(key)), + UpstreamTokenGrant(access_token=SecretStr(upstream_token), token_type="Bearer", expires_in=600), + keys, + datetime.now(timezone.utc), + ) + assert isinstance(sealed, SealedEnvelope), sealed + return sealed.token.get_secret_value() + + +def _session(user_id: str, keys: SessionKeys, client_id: str = "dcr-integration") -> str: + minted: Final = mint_session_token( + SessionPrincipal(user_id=user_id, client_id=client_id), keys, datetime.now(timezone.utc) + ) + assert isinstance(minted, MintedSessionToken), minted + return minted.token.get_secret_value() + + +def _refresh_envelope(identity_server: str, key: str, keys: EnvelopeKeys, upstream_refresh: str) -> str: + sealed: Final = mint_refresh_envelope( + key_hash_identity(server_id=identity_server, key_hash=hash_token(key)), + RefreshCredential(refresh_token=SecretStr(upstream_refresh), scope="tools.read tools.call", expires_in=600), + keys, + datetime.now(timezone.utc), + ) + assert isinstance(sealed, SealedEnvelope), sealed + return sealed.token.get_secret_value() + + +def _session_refresh(user_id: str, client_id: str, keys: SessionKeys) -> str: + minted: Final = mint_session_refresh_token( + SessionPrincipal(user_id=user_id, client_id=client_id), keys, datetime.now(timezone.utc) + ) + assert isinstance(minted, MintedSessionToken), minted + return minted.token.get_secret_value() + + +def _upstream_refresh(auth: AuthorizationServer) -> str: + issued: Final = auth.issue("authorization_code", "kdf-client", "integration-user", "tools.read tools.call") + token: Final = issued["refresh_token"] + assert isinstance(token, str) and token, issued + return token + + +def _register_gateway_client(client: httpx.Client) -> str: + registered: Final = client.post( + "/register", + json=_ClientRegistration(redirect_uris=("http://127.0.0.1/callback",), client_name="kdf-rotation"), + ) + assert registered.status_code == 201, registered.text + client_id: Final = registered.json().get("client_id") + assert isinstance(client_id, str) and client_id, registered.text + return client_id + + +def _register_bridge(scenario, peer: McpPeer, alias: str, issuer: str) -> str: + return register_mcp( + scenario, + peer, + alias, + auth_type="oauth_delegate", + dcr_bridge=True, + issuer=issuer, + authorization_url=issuer + "/authorize", + token_url=issuer + "/token", + registration_url=issuer + "/register", + ) + + +def _tools_list(client: httpx.Client, path: str, bearer: str) -> httpx.Response: + return client.post(path, headers={**ACCEPT, "Authorization": f"Bearer {bearer}"}, json=TOOLS_LIST) + + +def _peer_authorizations(peer: McpPeer) -> tuple[bytes | None, ...]: + return tuple( + value if isinstance(value := call["headers"].get(b"authorization"), bytes) else None + for call in peer.drain() + if isinstance(call["headers"], dict) + ) + + +def _assert_admitted(response: httpx.Response, peer: McpPeer, upstream_token: str) -> None: + assert response.status_code == 200, response.text + assert "tools" in response.text, response.text + forwarded: Final = _peer_authorizations(peer) + assert forwarded and set(forwarded) == {f"Bearer {upstream_token}".encode()}, forwarded + + +def _assert_rejected(response: httpx.Response, peer: McpPeer) -> None: + assert response.status_code == 401, response.text + assert "invalid_token" in response.headers.get("www-authenticate", ""), response.headers + assert peer.drain() == () + + +def test_envelope_minted_under_hkdf_sha256_keys_is_admitted(gateway: Gateway) -> None: + with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario: + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) + upstream_token: Final = "up-" + uuid.uuid4().hex + peer.drain() + envelope: Final = _envelope(identity, key, _hkdf_envelope_keys(gateway.key), upstream_token) + _assert_admitted(_tools_list(gateway.client, f"/{alias}/mcp", envelope), peer, upstream_token) + + +def test_envelope_minted_under_legacy_scrypt_keys_is_rejected_without_grace(gateway: Gateway) -> None: + with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario: + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) + peer.drain() + envelope: Final = _envelope(identity, key, _scrypt_envelope_keys(gateway.key), "up-" + uuid.uuid4().hex) + _assert_rejected(_tools_list(gateway.client, f"/{alias}/mcp", envelope), peer) + + +@pytest.mark.parametrize("kdf", ("hkdf", "scrypt")) +def test_envelope_of_either_kdf_is_admitted_during_the_legacy_grace_window( + gateway: Gateway, tmp_path: Path, kdf: str +) -> None: + with ( + owned_proxy(gateway, tmp_path, GRACE) as graced, + mcp_peer() as peer, + oauth_server() as auth, + graced.scenario() as scenario, + ): + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) + upstream_token: Final = "up-" + uuid.uuid4().hex + peer.drain() + keys: Final = _hkdf_envelope_keys(graced.key) if kdf == "hkdf" else _scrypt_envelope_keys(graced.key) + envelope: Final = _envelope(identity, key, keys, upstream_token) + _assert_admitted(_tools_list(graced.client, f"/{alias}/mcp", envelope), peer, upstream_token) + + +def test_envelope_under_a_foreign_master_key_is_rejected_during_grace(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy(gateway, tmp_path, GRACE) as graced, + mcp_peer() as peer, + oauth_server() as auth, + graced.scenario() as scenario, + ): + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) + peer.drain() + foreign: Final = "sk-foreign-" + uuid.uuid4().hex + for keys in (_hkdf_envelope_keys(foreign), _scrypt_envelope_keys(foreign)): + _assert_rejected(_tools_list(graced.client, f"/{alias}/mcp", _envelope(identity, key, keys, "up")), peer) + + +def test_grace_variable_set_to_anything_but_true_keeps_legacy_envelopes_rejected( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + owned_proxy(gateway, tmp_path, {"LITELLM_MCP_LEGACY_KDF_GRACE": "yes"}) as proxy, + mcp_peer() as peer, + oauth_server() as auth, + proxy.scenario() as scenario, + ): + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) + peer.drain() + envelope: Final = _envelope(identity, key, _scrypt_envelope_keys(proxy.key), "up") + _assert_rejected(_tools_list(proxy.client, f"/{alias}/mcp", envelope), peer) + + +def test_session_token_minted_under_hkdf_sha256_key_is_admitted(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + token: Final = _session(scenario.user(), _hkdf_session_keys(gateway.key)) + response: Final = _tools_list(gateway.client, "/mcp", token) + assert response.status_code == 200, response.text + assert "tools" in response.text, response.text + + +def test_session_token_minted_under_legacy_scrypt_key_is_rejected_without_grace(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + token: Final = _session(scenario.user(), _scrypt_session_keys(gateway.key)) + response: Final = _tools_list(gateway.client, "/mcp", token) + assert response.status_code == 401, response.text + assert "invalid_token" in response.headers.get("www-authenticate", ""), response.headers + + +@pytest.mark.parametrize("kdf", ("hkdf", "scrypt")) +def test_session_token_of_either_kdf_is_admitted_during_the_legacy_grace_window( + gateway: Gateway, tmp_path: Path, kdf: str +) -> None: + with owned_proxy(gateway, tmp_path, GRACE) as graced, graced.scenario() as scenario: + keys: Final = _hkdf_session_keys(graced.key) if kdf == "hkdf" else _scrypt_session_keys(graced.key) + response: Final = _tools_list(graced.client, "/mcp", _session(scenario.user(), keys)) + assert response.status_code == 200, response.text + assert "tools" in response.text, response.text + + +def test_tampered_session_token_is_rejected_during_grace(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy(gateway, tmp_path, GRACE) as graced, graced.scenario() as scenario: + token: Final = _session(scenario.user(), _scrypt_session_keys(graced.key)) + tampered: Final = token[:-2] + ("AA" if token[-2:] != "AA" else "BB") + response: Final = _tools_list(graced.client, "/mcp", tampered) + assert response.status_code == 401, response.text + assert "invalid_token" in response.headers.get("www-authenticate", ""), response.headers + + +def test_bridge_refresh_envelope_minted_under_hkdf_sha256_keys_is_renewed(gateway: Gateway) -> None: + with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario: + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) + envelope: Final = _refresh_envelope(identity, key, _hkdf_envelope_keys(gateway.key), _upstream_refresh(auth)) + response: Final = gateway.client.post( + f"/{alias}/token", + headers=(("x-litellm-api-key", key),), + data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=envelope, client_id="kdf-client"), + ) + assert response.status_code == 200, response.text + renewed: Final = response.json() + assert renewed["access_token"] != envelope, response.text + opened: Final = open_refresh_envelope( + renewed["refresh_token"], _hkdf_envelope_keys(gateway.key), datetime.now(timezone.utc) + ) + assert isinstance(opened, OpenedRefreshEnvelope), opened + + +def test_bridge_refresh_envelope_minted_under_legacy_scrypt_keys_is_rejected_without_grace( + gateway: Gateway, +) -> None: + with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario: + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) + envelope: Final = _refresh_envelope(identity, key, _scrypt_envelope_keys(gateway.key), _upstream_refresh(auth)) + response: Final = gateway.client.post( + f"/{alias}/token", + headers=(("x-litellm-api-key", key),), + data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=envelope, client_id="kdf-client"), + ) + assert response.status_code == 400, response.text + assert response.json()["error"] == "invalid_grant", response.text + assert auth.token_requests() == (), "a rejected envelope must never reach the upstream token endpoint" + + +@pytest.mark.parametrize("kdf", ("hkdf", "scrypt")) +def test_bridge_refresh_envelope_of_either_kdf_is_renewed_during_the_legacy_grace_window( + gateway: Gateway, tmp_path: Path, kdf: str +) -> None: + with ( + owned_proxy(gateway, tmp_path, GRACE) as graced, + mcp_peer() as peer, + oauth_server() as auth, + graced.scenario() as scenario, + ): + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) + peer.drain() + keys: Final = _hkdf_envelope_keys(graced.key) if kdf == "hkdf" else _scrypt_envelope_keys(graced.key) + envelope: Final = _refresh_envelope(identity, key, keys, _upstream_refresh(auth)) + response: Final = graced.client.post( + f"/{alias}/token", + headers=(("x-litellm-api-key", key),), + data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=envelope, client_id="kdf-client"), + ) + assert response.status_code == 200, response.text + renewed: Final = response.json() + opened: Final = open_refresh_envelope( + renewed["refresh_token"], _hkdf_envelope_keys(graced.key), datetime.now(timezone.utc) + ) + assert isinstance(opened, OpenedRefreshEnvelope), opened + access: Final = open_envelope( + renewed["access_token"], _hkdf_envelope_keys(graced.key), datetime.now(timezone.utc) + ) + assert isinstance(access, OpenedEnvelope), access + _assert_admitted( + _tools_list(graced.client, f"/{alias}/mcp", renewed["access_token"]), + peer, + access.grant.access_token.get_secret_value(), + ) + + +def test_session_refresh_token_minted_under_hkdf_sha256_key_is_renewed(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + client_id: Final = _register_gateway_client(gateway.client) + refresh: Final = _session_refresh(scenario.user(), client_id, _hkdf_session_keys(gateway.key)) + response: Final = gateway.client.post( + "/token", data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=refresh, client_id=client_id) + ) + assert response.status_code == 200, response.text + renewed: Final = response.json() + assert renewed["refresh_token"] != refresh, response.text + opened: Final = open_session_refresh_token( + renewed["refresh_token"], _hkdf_session_keys(gateway.key), datetime.now(timezone.utc) + ) + assert isinstance(opened, OpenedSessionToken), opened + + +def test_session_refresh_token_minted_under_legacy_scrypt_key_is_rejected_without_grace(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + client_id: Final = _register_gateway_client(gateway.client) + refresh: Final = _session_refresh(scenario.user(), client_id, _scrypt_session_keys(gateway.key)) + response: Final = gateway.client.post( + "/token", data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=refresh, client_id=client_id) + ) + assert response.status_code == 400, response.text + assert response.json()["error"] == "invalid_grant", response.text + + +def test_session_refresh_token_minted_under_legacy_scrypt_key_is_renewed_during_the_legacy_grace_window( + gateway: Gateway, tmp_path: Path +) -> None: + with owned_proxy(gateway, tmp_path, GRACE) as graced, graced.scenario() as scenario: + client_id: Final = _register_gateway_client(graced.client) + refresh: Final = _session_refresh(scenario.user(), client_id, _scrypt_session_keys(graced.key)) + response: Final = graced.client.post( + "/token", data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=refresh, client_id=client_id) + ) + assert response.status_code == 200, response.text + renewed: Final = response.json() + opened: Final = open_session_refresh_token( + renewed["refresh_token"], _hkdf_session_keys(graced.key), datetime.now(timezone.utc) + ) + assert isinstance(opened, OpenedSessionToken), opened + admitted: Final = _tools_list(graced.client, "/mcp", renewed["access_token"]) + assert admitted.status_code == 200, admitted.text + + +def test_session_refresh_token_minted_under_legacy_scrypt_key_is_revoked_during_grace( + gateway: Gateway, tmp_path: Path +) -> None: + with owned_proxy(gateway, tmp_path, GRACE) as graced, graced.scenario() as scenario: + client_id: Final = _register_gateway_client(graced.client) + refresh: Final = _session_refresh(scenario.user(), client_id, _scrypt_session_keys(graced.key)) + revoked: Final = graced.client.post("/revoke", data=_RevokeForm(token=refresh, client_id=client_id)) + assert revoked.status_code == 200, revoked.text + response: Final = graced.client.post( + "/token", data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=refresh, client_id=client_id) + ) + assert response.status_code == 400, response.text + assert "invalid_grant" in response.text, response.text + + +def test_session_refresh_token_minted_under_legacy_scrypt_key_revokes_cleanly_without_grace( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + client_id: Final = _register_gateway_client(gateway.client) + refresh: Final = _session_refresh(scenario.user(), client_id, _scrypt_session_keys(gateway.key)) + revoked: Final = gateway.client.post("/revoke", data=_RevokeForm(token=refresh, client_id=client_id)) + assert revoked.status_code == 200, revoked.text + response: Final = gateway.client.post( + "/token", data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=refresh, client_id=client_id) + ) + assert response.status_code == 400, response.text + assert "invalid_grant" in response.text, response.text + + +@pytest.mark.parametrize(("kdf", "active"), (("hkdf", True), ("scrypt", False))) +def test_session_access_token_introspects_by_its_kdf_without_grace(gateway: Gateway, kdf: str, active: bool) -> None: + with gateway.scenario() as scenario: + keys: Final = _hkdf_session_keys(gateway.key) if kdf == "hkdf" else _scrypt_session_keys(gateway.key) + key: Final = scenario.key(user_id=scenario.user()) + response: Final = gateway.client.post( + "/introspect", + headers=(("x-litellm-api-key", key),), + data=_IntrospectForm(token=_session(scenario.user(), keys)), + ) + assert response.status_code == 200, response.text + assert response.json()["active"] is active, response.text + + +def test_session_access_token_minted_under_legacy_scrypt_key_introspects_active_during_grace( + gateway: Gateway, tmp_path: Path +) -> None: + with owned_proxy(gateway, tmp_path, GRACE) as graced, graced.scenario() as scenario: + key: Final = scenario.key(user_id=scenario.user()) + response: Final = graced.client.post( + "/introspect", + headers=(("x-litellm-api-key", key),), + data=_IntrospectForm(token=_session(scenario.user(), _scrypt_session_keys(graced.key))), + ) + assert response.status_code == 200, response.text + assert response.json()["active"] is True, response.text diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py index f8fb22469f1..673c418041b 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -8,7 +8,12 @@ unopenable, and envelope minted for a different server); the producer helper rou through the consumer; and no path leaks the upstream token in a repr. """ +import hashlib +import hmac +from collections.abc import Callable, Mapping from datetime import datetime, timedelta, timezone +from types import MappingProxyType +from typing import Final import pytest from pydantic import SecretStr @@ -23,6 +28,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti build_bridge_token_response, envelope_keys_from_master_key, is_bridge_envelope_shaped, + legacy_envelope_keys_from_master_key, open_bridge_refresh_envelope, resolve_bridge_envelope, ) @@ -37,6 +43,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity, mint_envelope, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( + LEGACY_KDF_GRACE_ENV_VAR, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + session_keys_from_master_key, +) +from litellm.proxy.common_utils.fips import FIPS_MODE_ENV_VAR _NOW = datetime(2026, 7, 9, 12, 0, 0, tzinfo=timezone.utc) _MASTER_KEY = "sk-master-key-for-derivation-tests-0123456789" @@ -146,11 +159,10 @@ def test_key_derivation_signing_key_meets_hs256_floor_for_short_master_key(): assert len(keys.signing_key.get_secret_value()) >= 32 -def test_key_derivation_is_cached_so_the_memory_hard_kdf_runs_once_per_key(): - """The scrypt KDF is intentionally expensive to resist offline guessing, so it must be cached: - repeated calls for the same master key return the identical object rather than re-deriving, - keeping the per-request admission path free. Returning a distinct object each call would mean - the cache was dropped and every open would pay the memory-hard cost.""" +def test_key_derivation_is_cached_so_the_hkdf_runs_once_per_key(): + """Repeated calls for the same master key return the identical object rather than + re-deriving, keeping the per-request admission path free. Returning a distinct object + each call would mean the cache was dropped and every open would pay the KDF cost.""" first = envelope_keys_from_master_key("sk-cache-probe-key-9988776655") assert envelope_keys_from_master_key("sk-cache-probe-key-9988776655") is first @@ -318,3 +330,118 @@ def test_is_bridge_envelope_shaped_rejects_non_envelope_bearer(): assert is_bridge_envelope_shaped("Bearer sk-some-litellm-key") is False assert is_bridge_envelope_shaped("plain-upstream-token") is False assert is_bridge_envelope_shaped("") is False + + +def _rfc5869(ikm: bytes, info: bytes) -> bytes: + """Independent RFC 5869 HKDF-SHA256, hand-written on hmac/hashlib so a wrong KDF + construction in the product cannot agree with it.""" + prk: Final = hmac.new(b"\x00" * 32, ikm, hashlib.sha256).digest() + return hmac.new(prk, info + b"\x01", hashlib.sha256).digest()[:32] + + +def _environ(values: Mapping[str, str]) -> Callable[[str], str | None]: + return values.get + + +def _legacy_keys(master_key: str) -> EnvelopeKeys: + def derive(salt: bytes) -> str: + return hashlib.scrypt( + master_key.encode(), salt=salt, n=2**15, r=8, p=1, maxmem=128 * 2**15 * 8 * 2, dklen=32 + ).hex() + + return EnvelopeKeys( + signing_key=SecretStr(derive(b"litellm-mcp-bridge:envelope-signing:")), + encryption_key=SecretStr(derive(b"litellm-mcp-bridge:envelope-encryption:")), + ) + + +def test_envelope_key_derivation_matches_an_independent_rfc5869_vector(): + keys: Final = envelope_keys_from_master_key("sk-1234") + assert keys.signing_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-bridge:envelope-signing:").hex() + assert ( + keys.encryption_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-bridge:envelope-encryption:").hex() + ) + # literal pin computed from the same hand implementation (RFC 5869, salt = 32 zero bytes) + assert keys.signing_key.get_secret_value() == "331032bfe2b8d86bd00586ea69c0d854140019ca0e0adbe4dd0cc9e06cb36605" + + +def test_session_key_derivation_matches_an_independent_rfc5869_vector(): + keys: Final = session_keys_from_master_key("sk-1234") + assert keys.signing_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-gateway:session-signing:").hex() + assert keys.signing_key.get_secret_value() == "511f07b16ceeb5fba1660033785d4b9a4a6e80c749a51d639a688c287671a0e6" + + +def test_the_three_derived_keys_are_pairwise_distinct(): + envelope: Final = envelope_keys_from_master_key(_MASTER_KEY) + session: Final = session_keys_from_master_key(_MASTER_KEY) + derived: Final = frozenset( + ( + envelope.signing_key.get_secret_value(), + envelope.encryption_key.get_secret_value(), + session.signing_key.get_secret_value(), + ) + ) + assert len(derived) == 3 + + +def test_new_envelopes_open_under_the_hkdf_keys_with_no_legacy_fallback(): + keys: Final = envelope_keys_from_master_key(_MASTER_KEY) + result: Final = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeAdmitted) + + +def test_legacy_envelope_opens_during_the_grace_window(): + legacy: Final = _legacy_keys(_MASTER_KEY) + token: Final = _sealed_token(legacy) + grace: Final = legacy_envelope_keys_from_master_key( + _MASTER_KEY, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) + assert grace == legacy + result: Final = resolve_bridge_envelope( + token, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID, legacy_keys=grace + ) + assert isinstance(result, BridgeEnvelopeAdmitted) + assert result.identity == _IDENTITY + + +def test_legacy_envelope_is_invalid_without_grace(): + token: Final = _sealed_token(_legacy_keys(_MASTER_KEY)) + assert legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ(MappingProxyType({}))) is None + result: Final = resolve_bridge_envelope(token, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeInvalid) + + +def test_grace_is_disabled_under_fips(): + grace: Final = legacy_envelope_keys_from_master_key( + _MASTER_KEY, + environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"})), + ) + assert grace is None + + +@pytest.mark.parametrize("value", ("yes", "1", "on", "TRUE ", "True")) +def test_grace_is_enabled_only_by_case_insensitive_true(value: str): + result: Final = legacy_envelope_keys_from_master_key( + _MASTER_KEY, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: value})) + ) + expected: Final = value.strip().lower() == "true" + assert (result is not None) is expected + + +def test_legacy_refresh_envelope_opens_during_grace(): + legacy: Final = _legacy_keys(_MASTER_KEY) + refresh: Final = _sealed_refresh(legacy) + grace: Final = legacy_envelope_keys_from_master_key( + _MASTER_KEY, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) + result: Final = open_bridge_refresh_envelope( + refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID, legacy_keys=grace + ) + assert isinstance(result, BridgeRefreshOpened) + assert result.refresh.refresh_token.get_secret_value() == _UPSTREAM_REFRESH + + +def test_legacy_refresh_envelope_is_invalid_without_grace(): + refresh: Final = _sealed_refresh(_legacy_keys(_MASTER_KEY)) + result: Final = open_bridge_refresh_envelope(refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID) + assert isinstance(result, BridgeRefreshInvalid) diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py index 992b1e8632b..13886982b37 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py @@ -1,6 +1,10 @@ """Tests for the session-token KDF and the edge/token-endpoint resolvers.""" +import hashlib +from collections.abc import Callable, Mapping from datetime import datetime, timedelta, timezone +from types import MappingProxyType +from typing import Final import pytest from cryptography.hazmat.primitives import serialization @@ -11,6 +15,9 @@ from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( envelope_keys_from_master_key, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( + LEGACY_KDF_GRACE_ENV_VAR, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( NotSessionBearer, SessionBearerAdmitted, @@ -19,6 +26,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent SessionRefreshOpened, SessionSigningConfigError, is_session_bearer_shaped, + legacy_session_keys_from_master_key, open_session_refresh_bearer, resolve_session_bearer, resolve_session_signing_keys, @@ -34,6 +42,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i mint_session_token, session_public_key_pem, ) +from litellm.proxy.common_utils.fips import FIPS_MODE_ENV_VAR NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc) MASTER_KEY = "sk-master-key-for-tests" @@ -230,3 +239,104 @@ def test_signing_config_error_detail_never_leaks_key_material(): ) assert isinstance(resolved, SessionSigningConfigError) assert pem.splitlines()[1] not in resolved.detail + + +def _environ(values: Mapping[str, str]) -> Callable[[str], str | None]: + return values.get + + +def _legacy_session_keys(master_key: str) -> SessionKeys: + signing: Final = hashlib.scrypt( + master_key.encode(), + salt=b"litellm-mcp-gateway:session-signing:", + n=2**15, + r=8, + p=1, + maxmem=128 * 2**15 * 8 * 2, + dklen=32, + ).hex() + return SessionKeys(signing_key=SecretStr(signing)) + + +def test_new_session_tokens_open_under_the_hkdf_key_with_no_legacy_fallback(): + result: Final = resolve_session_bearer(f"Bearer {_access_token()}", KEYS, NOW) + assert isinstance(result, SessionBearerAdmitted) + + +def test_legacy_session_token_opens_during_the_grace_window(): + legacy: Final = _legacy_session_keys(MASTER_KEY) + minted: Final = mint_session_token(PRINCIPAL, legacy, NOW) + assert isinstance(minted, MintedSessionToken) + grace: Final = legacy_session_keys_from_master_key( + MASTER_KEY, KEYS, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) + assert grace == legacy + result: Final = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW, legacy_keys=grace) + assert isinstance(result, SessionBearerAdmitted) + assert result.principal == PRINCIPAL + + +def test_legacy_session_token_is_invalid_without_grace(): + minted: Final = mint_session_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) + assert isinstance(minted, MintedSessionToken) + assert legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ(MappingProxyType({}))) is None + result: Final = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW) + assert isinstance(result, SessionBearerInvalid) + assert result.expired is False + + +def test_legacy_session_key_is_never_derived_for_configured_rs256_keys(): + active: Final = AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem()), kid="k") + grace: Final = legacy_session_keys_from_master_key( + MASTER_KEY, active, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) + assert grace is None + + +def test_grace_is_disabled_under_fips(): + grace: Final = legacy_session_keys_from_master_key( + MASTER_KEY, + KEYS, + environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"})), + ) + assert grace is None + + +def test_legacy_session_refresh_token_opens_during_grace(): + legacy: Final = _legacy_session_keys(MASTER_KEY) + minted: Final = mint_session_refresh_token(PRINCIPAL, legacy, NOW) + assert isinstance(minted, MintedSessionToken) + grace: Final = legacy_session_keys_from_master_key( + MASTER_KEY, KEYS, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) + result: Final = open_session_refresh_bearer( + minted.token.get_secret_value(), + KEYS, + NOW, + expected_client_id="llm_client_abc", + legacy_keys=grace, + ) + assert isinstance(result, SessionRefreshOpened) + assert result.principal == PRINCIPAL + + +def test_legacy_session_refresh_token_is_invalid_without_grace(): + minted: Final = mint_session_refresh_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) + assert isinstance(minted, MintedSessionToken) + result: Final = open_session_refresh_bearer( + minted.token.get_secret_value(), KEYS, NOW, expected_client_id="llm_client_abc" + ) + assert isinstance(result, SessionRefreshInvalid) + + +def test_expired_token_under_the_primary_key_still_reports_expired_during_grace(): + legacy: Final = _legacy_session_keys(MASTER_KEY) + minted: Final = mint_session_token(PRINCIPAL, legacy, NOW) + assert isinstance(minted, MintedSessionToken) + grace: Final = legacy_session_keys_from_master_key( + MASTER_KEY, KEYS, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) + later: Final = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) + result: Final = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, later, legacy_keys=grace) + assert isinstance(result, SessionBearerInvalid) + assert result.expired is True