Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py: 52%
76 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""Producer and consumer helpers for the DCR-bridge ``oauth_delegate`` envelope.
3A DCR-bridge ``oauth_delegate`` client presents ONE bearer that is a litellm-signed
4envelope (see :mod:`.envelope`) carrying both a litellm identity and the upstream OAuth
5token. The gateway token endpoint mints it (producer) at OAuth issuance, and at the MCP
6admission edge the gateway derives the envelope keys from the proxy ``master_key``, opens
7it, admits the request under the recovered identity, and forwards the inner upstream token
8to the upstream MCP server (consumer). This module is the pure surface for both sides; the
9token-endpoint and admission wiring live in their respective call sites.
10"""
12import hashlib
13from datetime import datetime
14from functools import lru_cache
15from typing import Final, Literal, TypeAlias
17from pydantic import BaseModel, ConfigDict, SecretStr
19from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
20 EnvelopeIdentity,
21 EnvelopeKeys,
22 EnvelopeMintError,
23 OpenedEnvelope,
24 OpenedRefreshEnvelope,
25 RefreshCredential,
26 SealedEnvelope,
27 UpstreamTokenGrant,
28 is_envelope,
29 is_refresh_envelope,
30 mint_envelope,
31 mint_refresh_envelope,
32 open_envelope,
33 open_refresh_envelope,
34)
36_SIGNING_KEY_DOMAIN: Final = b"litellm-mcp-bridge:envelope-signing:"
37_ENCRYPTION_KEY_DOMAIN: Final = b"litellm-mcp-bridge:envelope-encryption:"
39# scrypt work factors (RFC 7914). n=2**15 with r=8/p=1 costs ~50ms and ~32MB per derivation, which
40# makes offline guessing of a candidate master key memory-hard rather than a bare hash comparison.
41_SCRYPT_N: Final = 2**15
42_SCRYPT_R: Final = 8
43_SCRYPT_P: Final = 1
44# scrypt's working-set is ~128 * N * r * p bytes; cap at twice that so the maxmem ceiling scales
45# with every work factor and a future p or r bump does not trip "memory limit exceeded".
46_SCRYPT_MAXMEM: Final = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2
47_DERIVED_KEY_BYTES: Final = 32
50@lru_cache(maxsize=8)
51def envelope_keys_from_master_key(master_key: str) -> EnvelopeKeys:
52 """Derive the envelope signing and encryption keys from the proxy master key.
54 A memory-hard scrypt KDF (RFC 7914) over two distinct domain-label salts yields two
55 independent 256-bit subkeys from the one secret, so the producer (mint) and consumer
56 (open) agree on keys without persisting any. scrypt is used rather than a bare hash or
57 HMAC so that a captured envelope is not a cheap offline oracle for the master key: each
58 candidate guess costs a full memory-hard derivation, which is what protects a deployment
59 whose master key is weaker than it should be. The result is cached (the master key is
60 fixed for a process), so the KDF runs once per key and adds nothing to the per-request
61 admission path. The derivation is deterministic; rotating ``master_key`` invalidates
62 every outstanding envelope, which is the intended behavior for a signing-key change.
63 """
64 signing: Final = hashlib.scrypt(
65 master_key.encode(),
66 salt=_SIGNING_KEY_DOMAIN,
67 n=_SCRYPT_N,
68 r=_SCRYPT_R,
69 p=_SCRYPT_P,
70 maxmem=_SCRYPT_MAXMEM,
71 dklen=_DERIVED_KEY_BYTES,
72 ).hex()
73 encryption: Final = hashlib.scrypt(
74 master_key.encode(),
75 salt=_ENCRYPTION_KEY_DOMAIN,
76 n=_SCRYPT_N,
77 r=_SCRYPT_R,
78 p=_SCRYPT_P,
79 maxmem=_SCRYPT_MAXMEM,
80 dklen=_DERIVED_KEY_BYTES,
81 ).hex()
82 return EnvelopeKeys(signing_key=SecretStr(signing), encryption_key=SecretStr(encryption))
85def build_bridge_token_response(
86 identity: EnvelopeIdentity,
87 grant: UpstreamTokenGrant,
88 keys: EnvelopeKeys,
89 now: datetime,
90) -> SealedEnvelope | EnvelopeMintError:
91 """Seal ``grant`` for ``identity`` into the client-held bearer the token endpoint returns.
93 The producer mirror of :func:`resolve_bridge_envelope`: a thin, pure wrapper over
94 :func:`mint_envelope` that returns the sealed envelope, or the mint error as a value
95 for the caller to map onto an OAuth error response.
96 """
97 return mint_envelope(identity, grant, keys, now)
100def build_bridge_refresh_token_response(
101 identity: EnvelopeIdentity,
102 refresh: RefreshCredential,
103 keys: EnvelopeKeys,
104 now: datetime,
105) -> SealedEnvelope | EnvelopeMintError:
106 """Seal ``refresh`` for ``identity`` into the long-lived refresh envelope the token endpoint returns
107 alongside the access envelope, so the client can renew without re-authenticating. A thin, pure
108 wrapper over :func:`mint_refresh_envelope`; returns the mint error as a value for the caller to map.
109 """
110 return mint_refresh_envelope(identity, refresh, keys, now)
113class BridgeRefreshOpened(BaseModel):
114 """A valid refresh envelope presented to the token endpoint: the identity to re-validate and renew
115 under, and the upstream refresh grant to exchange."""
117 model_config = ConfigDict(frozen=True)
118 tag: Literal["opened"] = "opened"
119 identity: EnvelopeIdentity
120 refresh: RefreshCredential
123class BridgeRefreshInvalid(BaseModel):
124 """The presented refresh grant is not a valid refresh envelope for this server (not refresh-shaped,
125 will not open, or minted for a different server); the token endpoint fails the refresh closed."""
127 model_config = ConfigDict(frozen=True)
128 tag: Literal["invalid"] = "invalid"
131BridgeRefreshResult: TypeAlias = BridgeRefreshOpened | BridgeRefreshInvalid
134def open_bridge_refresh_envelope(
135 refresh_value: str,
136 keys: EnvelopeKeys,
137 now: datetime,
138 expected_server_id: str,
139) -> BridgeRefreshResult:
140 """Open a refresh envelope a bridge ``oauth_delegate`` client presented on a refresh_token grant.
142 The token-endpoint mirror of :func:`resolve_bridge_envelope`: strips an optional ``Bearer`` scheme,
143 then returns ``BridgeRefreshOpened`` with the recovered identity and upstream refresh grant, or
144 ``BridgeRefreshInvalid`` for anything that is not a valid refresh envelope for this server. Never
145 raises; total over hostile input via :func:`open_refresh_envelope`. ``expected_server_id`` binds the
146 envelope to the server the request targets, so a refresh envelope minted for one server cannot renew
147 against another. A raw upstream refresh token (not envelope-shaped) is ``BridgeRefreshInvalid``: this
148 mode never hands the client a bare upstream refresh token, so it must never accept one.
149 """
150 candidate: Final = _strip_bearer(refresh_value)
151 if not is_refresh_envelope(candidate):
152 return BridgeRefreshInvalid()
153 opened: Final = open_refresh_envelope(candidate, keys, now)
154 if not isinstance(opened, OpenedRefreshEnvelope):
155 return BridgeRefreshInvalid()
156 if opened.identity.server_id != expected_server_id:
157 return BridgeRefreshInvalid()
158 return BridgeRefreshOpened(identity=opened.identity, refresh=opened.refresh)
161class NotBridgeEnvelope(BaseModel):
162 """The bearer is not an envelope; admission continues on its normal path."""
164 model_config = ConfigDict(frozen=True)
165 tag: Literal["not_bridge_envelope"] = "not_bridge_envelope"
168class BridgeEnvelopeAdmitted(BaseModel):
169 """A valid envelope: the identity to admit under and the full upstream ``Authorization``
170 value (``token_type access_token``) to forward to the upstream MCP server."""
172 model_config = ConfigDict(frozen=True)
173 tag: Literal["admitted"] = "admitted"
174 identity: EnvelopeIdentity
175 upstream_authorization: SecretStr
178class BridgeEnvelopeInvalid(BaseModel):
179 """The bearer is envelope-shaped but did not open (expired, tampered, wrong key);
180 admission must fail closed rather than fall through to normal validation."""
182 model_config = ConfigDict(frozen=True)
183 tag: Literal["invalid"] = "invalid"
186BridgeEnvelopeResult: TypeAlias = NotBridgeEnvelope | BridgeEnvelopeAdmitted | BridgeEnvelopeInvalid
189def _strip_bearer(value: str) -> str:
190 parts: Final = value.split(None, 1)
191 if len(parts) == 2 and parts[0].lower() == "bearer": 191 ↛ 192line 191 didn't jump to line 192 because the condition on line 191 was never true
192 return parts[1]
193 return value
196def is_bridge_envelope_shaped(authorization_value: str) -> bool:
197 """Cheap, keyless test that an ``Authorization`` value carries an envelope of either kind (optional
198 ``Bearer`` scheme stripped). The admission edge engages the bridge arm for an access envelope (to
199 admit) and for a refresh envelope (to reject it explicitly, since a refresh credential is never
200 usable at the tool-call edge); a plain upstream bearer falls through to normal oauth2 admission."""
201 candidate: Final = _strip_bearer(authorization_value)
202 return is_envelope(candidate) or is_refresh_envelope(candidate)
205def resolve_bridge_envelope(
206 authorization_value: str,
207 keys: EnvelopeKeys,
208 now: datetime,
209 expected_server_id: str,
210) -> BridgeEnvelopeResult:
211 """Classify an ``Authorization`` value presented to a bridge ``oauth_delegate`` server.
213 Strips an optional ``Bearer`` scheme, then returns ``NotBridgeEnvelope`` for a
214 non-envelope bearer (normal admission continues), ``BridgeEnvelopeAdmitted`` with the
215 recovered identity and the upstream ``Authorization`` value to forward for a valid
216 envelope, and ``BridgeEnvelopeInvalid`` for an envelope-shaped bearer that will not
217 open. Never raises: it is total over hostile input via :func:`open_envelope`.
219 A refresh envelope is ``BridgeEnvelopeInvalid`` here: it is a valid gateway credential but only ever
220 presented back to the token endpoint, never usable to authenticate a tool call, so admission must
221 fail it closed rather than let it fall through to another arm.
223 ``expected_server_id`` is the ``server_id`` of the MCP server the request targets; an
224 opened envelope whose sealed ``server_id`` does not match is rejected as
225 ``BridgeEnvelopeInvalid``. Binding here (rather than leaving it to the caller) prevents
226 replaying an envelope minted for one server against another, which would forward the
227 first server's upstream credential across a server boundary. ``server_id`` is not a
228 secret (the caller targets that server), so a plain equality check is sufficient and,
229 unlike ``hmac.compare_digest`` on ``str``, does not raise on a non-ASCII server_id.
230 """
231 candidate: Final = _strip_bearer(authorization_value)
232 if is_refresh_envelope(candidate):
233 return BridgeEnvelopeInvalid()
234 if not is_envelope(candidate):
235 return NotBridgeEnvelope()
236 opened: Final = open_envelope(candidate, keys, now)
237 if not isinstance(opened, OpenedEnvelope):
238 return BridgeEnvelopeInvalid()
239 if opened.identity.server_id != expected_server_id:
240 return BridgeEnvelopeInvalid()
241 grant: Final = opened.grant
242 authorization_scheme: Final = "Bearer" if grant.token_type.lower() == "bearer" else grant.token_type
243 upstream_authorization: Final = f"{authorization_scheme} {grant.access_token.get_secret_value()}"
244 return BridgeEnvelopeAdmitted(identity=opened.identity, upstream_authorization=SecretStr(upstream_authorization))