feat(mcp): bind bridge envelope to its server, key derivation via HMAC, add producer + shape helpers

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

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