From 08963b744fd59ed4065fb924844bee79f74e1231 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 10 Jul 2026 13:41:43 -0700 Subject: [PATCH] feat(mcp): bind bridge envelope to its server, key derivation via HMAC, add producer + shape helpers --- .../bridge_credentials.py | 77 +++++++++++--- .../test_bridge_credentials.py | 100 ++++++++++++++---- 2 files changed, 144 insertions(+), 33 deletions(-) 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 9b1e8722c54..ef137c68107 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py @@ -1,15 +1,16 @@ -"""Consumer-side helpers for the DCR-bridge ``oauth_delegate`` envelope. +"""Producer and consumer helpers for the DCR-bridge ``oauth_delegate`` envelope. A DCR-bridge ``oauth_delegate`` client presents ONE bearer that is a litellm-signed envelope (see :mod:`.envelope`) carrying both a litellm identity and the upstream OAuth -token. At the MCP admission edge the gateway derives the envelope keys from the proxy -``master_key``, opens the envelope, admits the request under the recovered identity, and -forwards the inner upstream token to the upstream MCP server. This module is the pure -consumer surface; the admission wiring lives in ``user_api_key_auth_mcp.py`` and the -producer (mint) side lands with the token-endpoint flow. +token. The gateway token endpoint mints it (producer) at OAuth issuance, and at the MCP +admission edge the gateway derives the envelope keys from the proxy ``master_key``, opens +it, admits the request under the recovered identity, and forwards the inner upstream token +to the upstream MCP server (consumer). This module is the pure surface for both sides; the +token-endpoint and admission wiring live in their respective call sites. """ import hashlib +import hmac from datetime import datetime from typing import Literal, TypeAlias @@ -18,8 +19,12 @@ from pydantic import BaseModel, ConfigDict, SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( EnvelopeIdentity, EnvelopeKeys, + EnvelopeMintError, OpenedEnvelope, + SealedEnvelope, + UpstreamTokenGrant, is_envelope, + mint_envelope, open_envelope, ) @@ -31,17 +36,40 @@ _BEARER_PREFIX = "bearer " def envelope_keys_from_master_key(master_key: str) -> EnvelopeKeys: """Derive the envelope signing and encryption keys from the proxy master key. - Domain-separated SHA-256 yields two distinct 64-char (256-bit) hex keys from the one - secret, so the producer (mint) and consumer (open) agree on keys without persisting - any, and even a short master key still produces a >= 32-byte signing key (HS256's - requirement). The derivation is deterministic; rotating ``master_key`` invalidates - every outstanding envelope, which is the intended behavior for a signing-key change. + Keyed HMAC-SHA256 over two distinct domain labels yields two independent 64-char + (256-bit) subkeys from the one secret, so the producer (mint) and consumer (open) agree + on keys without persisting any, and even a short master key still produces a >= 32-byte + signing key (HS256's requirement). HMAC keys the derivation on ``master_key`` (the + standard subkey-from-key construction) rather than hashing a concatenation. The + derivation is deterministic; rotating ``master_key`` invalidates every outstanding + envelope, which is the intended behavior for a signing-key change. + + ``master_key`` is the proxy's root signing secret and must be high-entropy: a captured + envelope is an offline oracle for it, exactly as the existing master-key-signed session + tokens already are, so a low-entropy master key compromises the proxy regardless of this + path. A password-style slow KDF is deliberately not used here; it is the wrong tradeoff + for a per-request high-entropy secret. """ - signing = hashlib.sha256((_SIGNING_KEY_DOMAIN + master_key).encode()).hexdigest() - encryption = hashlib.sha256((_ENCRYPTION_KEY_DOMAIN + master_key).encode()).hexdigest() + signing = hmac.new(master_key.encode(), _SIGNING_KEY_DOMAIN.encode(), hashlib.sha256).hexdigest() + encryption = hmac.new(master_key.encode(), _ENCRYPTION_KEY_DOMAIN.encode(), hashlib.sha256).hexdigest() return EnvelopeKeys(signing_key=SecretStr(signing), encryption_key=SecretStr(encryption)) +def build_bridge_token_response( + identity: EnvelopeIdentity, + grant: UpstreamTokenGrant, + keys: EnvelopeKeys, + now: datetime, +) -> SealedEnvelope | EnvelopeMintError: + """Seal ``grant`` for ``identity`` into the client-held bearer the token endpoint returns. + + The producer mirror of :func:`resolve_bridge_envelope`: a thin, pure wrapper over + :func:`mint_envelope` that returns the sealed envelope, or the mint error as a value + (an oversized grant) for the caller to map onto an OAuth error response. + """ + return mint_envelope(identity, grant, keys, now) + + class NotBridgeEnvelope(BaseModel): """The bearer is not an envelope; admission continues on its normal path.""" @@ -76,7 +104,19 @@ def _strip_bearer(value: str) -> str: return value -def resolve_bridge_envelope(authorization_value: str, keys: EnvelopeKeys, now: datetime) -> BridgeEnvelopeResult: +def is_bridge_envelope_shaped(authorization_value: str) -> bool: + """Cheap, keyless test that an ``Authorization`` value carries an envelope (optional + ``Bearer`` scheme stripped). The admission edge engages the bridge arm only for an + envelope, so a plain upstream bearer falls through to normal oauth2 admission.""" + return is_envelope(_strip_bearer(authorization_value)) + + +def resolve_bridge_envelope( + authorization_value: str, + keys: EnvelopeKeys, + now: datetime, + expected_server_id: str, +) -> BridgeEnvelopeResult: """Classify an ``Authorization`` value presented to a bridge ``oauth_delegate`` server. Strips an optional ``Bearer`` scheme, then returns ``NotBridgeEnvelope`` for a @@ -84,6 +124,13 @@ def resolve_bridge_envelope(authorization_value: str, keys: EnvelopeKeys, now: d recovered identity and the upstream ``Authorization`` value to forward for a valid envelope, and ``BridgeEnvelopeInvalid`` for an envelope-shaped bearer that will not open. Never raises: it is total over hostile input via :func:`open_envelope`. + + ``expected_server_id`` is the ``server_id`` of the MCP server the request targets; an + opened envelope whose sealed ``server_id`` does not match is rejected as + ``BridgeEnvelopeInvalid``. Binding here (rather than leaving it to the caller) prevents + replaying an envelope minted for one server against another, which would forward the + first server's upstream credential across a server boundary. The comparison is constant + time so a mismatch does not leak the expected id through timing. """ candidate = _strip_bearer(authorization_value) if not is_envelope(candidate): @@ -91,6 +138,8 @@ def resolve_bridge_envelope(authorization_value: str, keys: EnvelopeKeys, now: d opened = open_envelope(candidate, keys, now) if not isinstance(opened, OpenedEnvelope): return BridgeEnvelopeInvalid() + if not hmac.compare_digest(opened.identity.server_id, expected_server_id): + return BridgeEnvelopeInvalid() grant = opened.grant upstream_authorization = f"{grant.token_type} {grant.access_token.get_secret_value()}" return BridgeEnvelopeAdmitted(identity=opened.identity, upstream_authorization=SecretStr(upstream_authorization)) 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 df6a5de07c3..1aa32d02f60 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 @@ -1,10 +1,11 @@ -"""Spec tests for the DCR-bridge envelope consumer helpers. +"""Spec tests for the DCR-bridge envelope producer and consumer helpers. -These pin the two consumer contracts the admission edge depends on: the master-key key -derivation is deterministic, domain-separated, and always yields a >= 32-byte signing key -(so mint and open agree without persisting keys), and the ``Authorization`` classifier is -total over the three cases admission branches on (non-envelope, valid envelope, and -envelope-shaped-but-unopenable), never leaking the recovered upstream token in a repr. +These pin the contracts the token endpoint and admission edge depend on: the master-key +key derivation is deterministic, domain-separated (keyed HMAC), and always yields a +>= 32-byte signing key; the ``Authorization`` classifier is total over the cases admission +branches on (non-envelope, valid envelope bound to this server, envelope-shaped-but- +unopenable, and envelope minted for a different server); the producer helper round-trips +through the consumer; and no path leaks the upstream token in a repr. """ from datetime import datetime, timedelta, timezone @@ -15,13 +16,16 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti BridgeEnvelopeAdmitted, BridgeEnvelopeInvalid, NotBridgeEnvelope, + build_bridge_token_response, envelope_keys_from_master_key, + is_bridge_envelope_shaped, resolve_bridge_envelope, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( ENVELOPE_PREFIX, EnvelopeIdentity, EnvelopeKeys, + EnvelopeTooLarge, SealedEnvelope, UpstreamTokenGrant, mint_envelope, @@ -31,14 +35,15 @@ _NOW = datetime(2026, 7, 9, 12, 0, 0, tzinfo=timezone.utc) _MASTER_KEY = "sk-master-key-for-derivation-tests-0123456789" _ACCESS_TOKEN = "upstream-access-token-do-not-leak-8f14e45fceea" _IDENTITY = EnvelopeIdentity(user_id="user-123", server_id="srv-456") +_SERVER_ID = _IDENTITY.server_id def _grant() -> UpstreamTokenGrant: return UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN), token_type="Bearer", expires_in=600) -def _sealed_token(keys: EnvelopeKeys, now: datetime = _NOW) -> str: - sealed = mint_envelope(_IDENTITY, _grant(), keys, now) +def _sealed_token(keys: EnvelopeKeys, now: datetime = _NOW, identity: EnvelopeIdentity = _IDENTITY) -> str: + sealed = mint_envelope(identity, _grant(), keys, now) assert isinstance(sealed, SealedEnvelope) return sealed.token.get_secret_value() @@ -66,20 +71,20 @@ def test_key_derivation_signing_key_meets_hs256_floor_for_short_master_key(): def test_derived_keys_round_trip_mint_and_open(): keys = envelope_keys_from_master_key(_MASTER_KEY) - result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW) + result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID) assert isinstance(result, BridgeEnvelopeAdmitted) assert result.identity == _IDENTITY def test_resolve_non_envelope_is_not_bridge_envelope(): keys = envelope_keys_from_master_key(_MASTER_KEY) - assert isinstance(resolve_bridge_envelope("Bearer sk-some-litellm-key", keys, _NOW), NotBridgeEnvelope) - assert isinstance(resolve_bridge_envelope("plain-token", keys, _NOW), NotBridgeEnvelope) + assert isinstance(resolve_bridge_envelope("Bearer sk-some-litellm-key", keys, _NOW, _SERVER_ID), NotBridgeEnvelope) + assert isinstance(resolve_bridge_envelope("plain-token", keys, _NOW, _SERVER_ID), NotBridgeEnvelope) def test_resolve_valid_envelope_returns_identity_and_upstream_authorization(): keys = envelope_keys_from_master_key(_MASTER_KEY) - result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW) + result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID) assert isinstance(result, BridgeEnvelopeAdmitted) assert result.identity == _IDENTITY assert result.upstream_authorization.get_secret_value() == f"Bearer {_ACCESS_TOKEN}" @@ -89,9 +94,9 @@ def test_resolve_strips_optional_bearer_scheme_before_detection(): keys = envelope_keys_from_master_key(_MASTER_KEY) token = _sealed_token(keys) assert token.startswith(ENVELOPE_PREFIX) - bare = resolve_bridge_envelope(token, keys, _NOW) - prefixed = resolve_bridge_envelope(f"Bearer {token}", keys, _NOW) - lower = resolve_bridge_envelope(f"bearer {token}", keys, _NOW) + bare = resolve_bridge_envelope(token, keys, _NOW, _SERVER_ID) + prefixed = resolve_bridge_envelope(f"Bearer {token}", keys, _NOW, _SERVER_ID) + lower = resolve_bridge_envelope(f"bearer {token}", keys, _NOW, _SERVER_ID) assert isinstance(bare, BridgeEnvelopeAdmitted) assert isinstance(prefixed, BridgeEnvelopeAdmitted) assert isinstance(lower, BridgeEnvelopeAdmitted) @@ -102,26 +107,83 @@ def test_resolve_expired_envelope_is_invalid_not_admitted(): keys = envelope_keys_from_master_key(_MASTER_KEY) token = _sealed_token(keys, now=_NOW) later = _NOW + timedelta(seconds=601) - assert isinstance(resolve_bridge_envelope(token, keys, later), BridgeEnvelopeInvalid) + assert isinstance(resolve_bridge_envelope(token, keys, later, _SERVER_ID), BridgeEnvelopeInvalid) def test_resolve_envelope_minted_under_a_different_master_key_is_invalid(): minted = envelope_keys_from_master_key(_MASTER_KEY) other = envelope_keys_from_master_key("a-completely-different-master-key") - assert isinstance(resolve_bridge_envelope(_sealed_token(minted), other, _NOW), BridgeEnvelopeInvalid) + assert isinstance(resolve_bridge_envelope(_sealed_token(minted), other, _NOW, _SERVER_ID), BridgeEnvelopeInvalid) def test_resolve_tampered_envelope_is_invalid(): keys = envelope_keys_from_master_key(_MASTER_KEY) token = _sealed_token(keys) tampered = token[:-4] + ("aaaa" if token[-4:] != "aaaa" else "bbbb") - result = resolve_bridge_envelope(tampered, keys, _NOW) + result = resolve_bridge_envelope(tampered, keys, _NOW, _SERVER_ID) assert isinstance(result, BridgeEnvelopeInvalid) +def test_resolve_envelope_minted_for_another_server_is_invalid(): + """An envelope sealed for server A must be rejected when presented to server B, so a + captured or misrouted envelope cannot forward one server's upstream credential to + another. The valid access token stays sealed; the mismatch alone fails the resolve.""" + keys = envelope_keys_from_master_key(_MASTER_KEY) + other_server_identity = EnvelopeIdentity(user_id=_IDENTITY.user_id, server_id="srv-OTHER") + token = _sealed_token(keys, identity=other_server_identity) + result = resolve_bridge_envelope(token, keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeInvalid) + + +def test_resolve_matching_server_binding_is_admitted(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, "srv-456") + assert isinstance(result, BridgeEnvelopeAdmitted) + + def test_admitted_result_repr_never_leaks_upstream_token(): keys = envelope_keys_from_master_key(_MASTER_KEY) - result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW) + result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID) assert isinstance(result, BridgeEnvelopeAdmitted) assert _ACCESS_TOKEN not in repr(result) assert _ACCESS_TOKEN not in str(result) + + +def test_build_bridge_token_response_round_trips_through_the_consumer(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + sealed = build_bridge_token_response(_IDENTITY, _grant(), keys, _NOW) + assert isinstance(sealed, SealedEnvelope) + assert sealed.token.get_secret_value().startswith(ENVELOPE_PREFIX) + result = resolve_bridge_envelope(sealed.token.get_secret_value(), keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeAdmitted) + assert result.identity == _IDENTITY + assert result.upstream_authorization.get_secret_value() == f"Bearer {_ACCESS_TOKEN}" + + +def test_build_bridge_token_response_oversized_grant_returns_error_value(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + huge = UpstreamTokenGrant(access_token=SecretStr("x" * 20000), token_type="Bearer") + result = build_bridge_token_response(_IDENTITY, huge, keys, _NOW) + assert isinstance(result, EnvelopeTooLarge) + + +def test_build_bridge_token_response_repr_never_leaks_upstream_token(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + sealed = build_bridge_token_response(_IDENTITY, _grant(), keys, _NOW) + assert isinstance(sealed, SealedEnvelope) + assert _ACCESS_TOKEN not in repr(sealed) + assert _ACCESS_TOKEN not in str(sealed) + + +def test_is_bridge_envelope_shaped_detects_envelope_with_and_without_bearer(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + token = _sealed_token(keys) + assert is_bridge_envelope_shaped(token) is True + assert is_bridge_envelope_shaped(f"Bearer {token}") is True + assert is_bridge_envelope_shaped(f"bearer {token}") is True + + +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