From 48574cc53cb0c164f6ea7a04e6cf3c02bf018673 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 24 Sep 2026 08:10:17 +0000 Subject: [PATCH] feat(mcp): derive bridge and session keys with HKDF-SHA256 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/auth/user_api_key_auth_mcp.py | 17 ++- .../mcp_server/bridge_token_flow.py | 9 +- .../mcp_server/gateway_dcr_flow.py | 36 ++++- .../bridge_credentials.py | 86 ++++++------ .../outbound_credentials/key_derivation.py | 58 ++++++++ .../session_credentials.py | 86 ++++++++---- .../test_bridge_credentials.py | 126 +++++++++++++++++- .../test_session_credentials.py | 119 +++++++++++++++++ 8 files changed, 457 insertions(+), 80 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py 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 43d62ba8b22..4ff913aa89d 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 @@ -25,6 +25,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 ( @@ -869,7 +870,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 @@ -1005,6 +1012,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 @@ -1018,7 +1026,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 e5d11271c67..ed41ee5b286 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 e66504af47a..ab922a67b2d 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -67,6 +67,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent SessionRefreshOpened, SessionSigningConfigError, active_session_signing_keys, + legacy_session_keys_from_master_key, open_session_refresh_bearer, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( @@ -75,8 +76,11 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i MintedSessionToken, OpenedSessionToken, SessionAudience, + SessionExpired, + SessionKeys, SessionPrincipal, SessionSigningKeys, + SessionTokenOpenError, is_session_refresh_token, is_session_token, mint_session_refresh_token, @@ -1236,6 +1240,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( @@ -1395,10 +1400,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): @@ -1459,7 +1467,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 @@ -1469,6 +1483,19 @@ async def revoke_refresh_token(token: str, client_id: str, master_key: str | Non return Response(content="{}", media_type="application/json", headers=TOKEN_NO_CACHE_HEADERS) +def _open_with_legacy( + open_token: Callable[[str, SessionSigningKeys, datetime], OpenedSessionToken | SessionTokenOpenError], + token: str, + keys: SessionSigningKeys, + legacy_keys: SessionKeys | None, + now: datetime, +) -> OpenedSessionToken | SessionTokenOpenError: + opened: Final = open_token(token, keys, now) + if isinstance(opened, OpenedSessionToken) or isinstance(opened, SessionExpired) or legacy_keys is None: + return opened + return open_token(token, legacy_keys, now) + + def _inactive_introspection_response() -> Response: """RFC 7662 section 2.2: any token the gateway cannot vouch for, whatever the reason (wrong family, bad signature, expired, revoked, or a deactivated user), answers 200 @@ -1527,10 +1554,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_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_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..749813df4d0 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 @@ -32,56 +33,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, @@ -136,6 +136,7 @@ def open_bridge_refresh_envelope( 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 +151,9 @@ 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 = open_refresh_envelope(candidate, keys, now) + if not isinstance(opened, OpenedRefreshEnvelope) and legacy_keys is not None: + opened = open_refresh_envelope(candidate, legacy_keys, now) if not isinstance(opened, OpenedRefreshEnvelope): return BridgeRefreshInvalid() if opened.identity.server_id != expected_server_id: @@ -207,6 +210,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 +237,9 @@ def resolve_bridge_envelope( return BridgeEnvelopeInvalid() if not is_envelope(candidate): return NotBridgeEnvelope() - opened: Final = open_envelope(candidate, keys, now) + opened = open_envelope(candidate, keys, now) + if not isinstance(opened, OpenedEnvelope) and legacy_keys is not None: + opened = open_envelope(candidate, 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..98c2218f2eb --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py @@ -0,0 +1,58 @@ +"""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 + +# scrypt work factors (RFC 7914). n=2**15 with r=8/p=1 costs ~50ms and ~32MB per derivation. +_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 + + +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..02b231fed27 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, @@ -38,39 +44,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 +230,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 +249,13 @@ def resolve_session_bearer( return SessionBearerInvalid() if not is_session_token(candidate): return NotSessionBearer() - opened: Final = open_session_token(candidate, keys, now) + opened = open_session_token(candidate, keys, now) + if ( + not isinstance(opened, OpenedSessionToken) + and not isinstance(opened, SessionExpired) + and legacy_keys is not None + ): + opened = open_session_token(candidate, legacy_keys, now) if isinstance(opened, OpenedSessionToken): return SessionBearerAdmitted(principal=opened.principal) return SessionBearerInvalid(expired=isinstance(opened, SessionExpired)) @@ -265,6 +288,7 @@ def open_session_refresh_bearer( 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 +303,13 @@ 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 = open_session_refresh_token(candidate, keys, now) + if ( + not isinstance(opened, OpenedSessionToken) + and not isinstance(opened, SessionExpired) + and legacy_keys is not None + ): + opened = open_session_refresh_token(candidate, legacy_keys, now) if not isinstance(opened, OpenedSessionToken): return SessionRefreshInvalid() if opened.principal.client_id != expected_client_id: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py index f8fb22469f1..6718285ef0a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -8,6 +8,8 @@ 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 datetime import datetime, timedelta, timezone import pytest @@ -23,9 +25,17 @@ 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, ) +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 from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( ENVELOPE_PREFIX, EnvelopeIdentity, @@ -146,11 +156,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 +327,110 @@ 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 = hmac.new(b"\x00" * 32, ikm, hashlib.sha256).digest() + return hmac.new(prk, info + b"\x01", hashlib.sha256).digest()[:32] + + +def _environ(values: dict[str, str]): + 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 = 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 = 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 = envelope_keys_from_master_key(_MASTER_KEY) + session = session_keys_from_master_key(_MASTER_KEY) + derived = { + 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 = envelope_keys_from_master_key(_MASTER_KEY) + result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeAdmitted) + + +def test_legacy_envelope_opens_during_the_grace_window(): + legacy = _legacy_keys(_MASTER_KEY) + token = _sealed_token(legacy) + grace = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + assert grace == legacy + result = 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 = _sealed_token(_legacy_keys(_MASTER_KEY)) + assert legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({})) is None + result = 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 = legacy_envelope_keys_from_master_key( + _MASTER_KEY, + environ=_environ({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 = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: value})) + expected = value.strip().lower() == "true" + assert (result is not None) is expected + + +def test_legacy_refresh_envelope_opens_during_grace(): + legacy = _legacy_keys(_MASTER_KEY) + refresh = _sealed_refresh(legacy) + grace = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + result = 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 = _sealed_refresh(_legacy_keys(_MASTER_KEY)) + result = open_bridge_refresh_envelope(refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID) + assert isinstance(result, BridgeRefreshInvalid) 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 992b1e8632b..396e7b965e1 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 @@ -1,5 +1,6 @@ """Tests for the session-token KDF and the edge/token-endpoint resolvers.""" +import hashlib from datetime import datetime, timedelta, timezone import pytest @@ -11,6 +12,7 @@ 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 import session_credentials from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( NotSessionBearer, SessionBearerAdmitted, @@ -19,11 +21,15 @@ 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, session_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_token import ( SESSION_TTL_SECONDS, AsymmetricSessionKeys, @@ -34,6 +40,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 +237,115 @@ 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: dict[str, str]): + return values.get + + +def _legacy_session_keys(master_key: str) -> SessionKeys: + signing = 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 = resolve_session_bearer(f"Bearer {_access_token()}", KEYS, NOW) + assert isinstance(result, SessionBearerAdmitted) + + +def test_legacy_session_token_opens_during_the_grace_window(): + legacy = _legacy_session_keys(MASTER_KEY) + minted = mint_session_token(PRINCIPAL, legacy, NOW) + assert isinstance(minted, MintedSessionToken) + grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + assert grace == legacy + result = 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 = 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({})) is None + result = 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 = AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem()), kid="k") + grace = legacy_session_keys_from_master_key( + MASTER_KEY, active, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"}) + ) + assert grace is None + + +def test_grace_is_disabled_under_fips(): + grace = legacy_session_keys_from_master_key( + MASTER_KEY, + KEYS, + environ=_environ({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 = _legacy_session_keys(MASTER_KEY) + minted = mint_session_refresh_token(PRINCIPAL, legacy, NOW) + assert isinstance(minted, MintedSessionToken) + grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + result = 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 = mint_session_refresh_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) + assert isinstance(minted, MintedSessionToken) + result = 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 = _legacy_session_keys(MASTER_KEY) + minted = mint_session_token(PRINCIPAL, legacy, NOW) + assert isinstance(minted, MintedSessionToken) + grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + later = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) + result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, later, legacy_keys=grace) + assert isinstance(result, SessionBearerInvalid) + assert result.expired is True + + +def test_fallback_tries_the_active_key_before_legacy(monkeypatch): + calls: list[object] = [] + real = session_credentials.open_session_token + + def spy(candidate, keys, now): + calls.append(keys) + return real(candidate, keys, now) + + monkeypatch.setattr(session_credentials, "open_session_token", spy) + grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + minted = mint_session_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) + assert isinstance(minted, MintedSessionToken) + result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW, legacy_keys=grace) + assert isinstance(result, SessionBearerAdmitted) + assert calls[0] is KEYS