scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-11 03:38:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83

Merge bc65ccbba8 into cad87a900f

Browse source
This commit is contained in:
devin-ai-integration[bot] 2026-10-05 20:37:34 +00:00 • committed by GitHub
commit 5666503566
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 968 additions and 80 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -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:

9
litellm/proxy/_experimental/mcp_server/bridge_token_flow.py
View file

@ -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)

22
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
View file

@ -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):

111
litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py
View file

@ -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:

56
litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py Normal file
View file

@ -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"

92
litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py
View file

@ -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:

494
tests/integration/mcp/test_mcp_kdf_rotation.py Normal file
View file

@ -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

137
tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py
View file

@ -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)

110
tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py
View file

@ -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