scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

feat(mcp): seal the authorizing key hash in the dcr_bridge envelope

Browse source 
The mint bound only user_id/server_id into the envelope, which gave admission
no way to reload the caller's key and enforce its current restrictions. Seal the
hashed authorizing key instead (a one-way digest, not a usable credential), so
admission reloads the live UserAPIKeyAuth by it and the key's team/org/tool
permissions and revocation apply per request.

Extract the token endpoint's key resolution into a shared _resolve_active_litellm_key
so the per-user token store (user_id) and the bridge mint (key hash) derive from one
active-key-gated path, and fail the mint closed with invalid_request when no active
key accompanies the request.
This commit is contained in:
Tin Chi Lo 2026-07-10 17:22:59 -07:00
parent 34c6cce705
commit 85255c96fb
2 changed files with 127 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

83
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -351,44 +351,73 @@ def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> Optional[str]:
return key_obj.user_id
async def _extract_user_id_from_request(request: Request) -> Optional[str]:
"""Resolve the LiteLLM ``user_id`` at the OAuth token endpoint so a per-user token is stored
under the same identity the egress later reads it by (``user_api_key_auth.user_id``).
async def _resolve_active_litellm_key(request: Request) -> Optional[Tuple[str, "UserAPIKeyAuth"]]:
"""Resolve the presented litellm key to ``(its hash, the live active key record)``, or ``None``
when the key is absent, unresolvable, or blocked/expired.
Resolves authoritatively via ``get_key_object`` (cache first, then DB) instead of a raw cache
peek. On a multi-replica gateway the token-exchange request can land on a worker whose in-memory
cache never saw the key, and a cross-replica Redis hit deserializes to a plain ``dict`` rather
than a ``UserAPIKeyAuth``; the previous code read only ``Authorization`` and did
``getattr(cached, "user_id")`` with no ``model_type`` rehydration and no DB fallback, so it
silently returned ``None`` and the token was never persisted, which makes the egress 401 on every
reconnect. The resolved key is validated (``_active_key_user_id``) before its identity is trusted,
so a blocked or expired key cannot write. Returns ``None`` when no key is present, the key cannot
be resolved, or it is blocked/expired.
Single resolution path the OAuth token endpoint reuses. Resolves authoritatively via
``get_key_object`` (cache first, then DB) instead of a raw cache peek. On a multi-replica gateway
the token-exchange request can land on a worker whose in-memory cache never saw the key, and a
cross-replica Redis hit deserializes to a plain ``dict`` rather than a ``UserAPIKeyAuth``; the
previous code read only ``Authorization`` and did ``getattr(cached, "user_id")`` with no
``model_type`` rehydration and no DB fallback, so it silently returned ``None``. The resolved key
is validated (``_active_key_user_id``) before it is trusted, so a blocked or expired key resolves
to ``None``. The returned hash is the value ``get_key_object`` and the cache/DB layer key the
record by. Callers derive the ``user_id`` (per-user token store) or seal the hash (dcr_bridge
envelope) from the result.
"""
token = _litellm_key_from_request(request)
if not token:
return None
try:
from litellm.proxy._types import hash_token # noqa: PLC0415
from litellm.proxy.auth.auth_checks import get_key_object # noqa: PLC0415
from litellm.proxy.proxy_server import ( # noqa: PLC0415
from litellm.proxy._types import ( # noqa: PLC0415 # inline import avoids a module-load circular import
hash_token,
)
from litellm.proxy.auth.auth_checks import ( # noqa: PLC0415 # inline import avoids a module-load circular import
get_key_object,
)
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import
prisma_client,
user_api_key_cache,
)
key_hash = hash_token(token)
key_obj = await get_key_object(
hashed_token=hash_token(token),
hashed_token=key_hash,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
return _active_key_user_id(key_obj)
except Exception as exc:
except Exception as exc: # noqa: BLE001 # fail closed to None on any key-resolution error
verbose_logger.debug(
"_extract_user_id_from_request: could not resolve a LiteLLM user_id for the presented "
"key (%s); per-user token will not be stored server-side.",
"_resolve_active_litellm_key: could not resolve the presented key (%s)",
type(exc).__name__,
)
return None
if _active_key_user_id(key_obj) is None:
return None
return key_hash, key_obj
async def _extract_user_id_from_request(request: Request) -> Optional[str]:
"""The LiteLLM ``user_id`` for the token request, so a per-user token is stored under the same
identity the egress later reads it by (``user_api_key_auth.user_id``). ``None`` when no active
key is present. See :func:`_resolve_active_litellm_key` for the resolution and active-key gate.
"""
resolved = await _resolve_active_litellm_key(request)
return _active_key_user_id(resolved[1]) if resolved else None
async def _extract_active_key_hash_from_request(request: Request) -> Optional[str]:
"""The hash of the litellm key that authorized the token request, when it maps to an active key.
A DCR-bridge envelope seals this hash so admission can reload the live ``UserAPIKeyAuth`` record
and enforce the key's current team/org/tool restrictions and revocation, rather than trusting a
frozen identity. The hash is a one-way digest, not a usable credential (the edge rejects a bare
hash presented as a bearer). ``None`` when no active key is present, so no envelope is minted for
a missing, unresolvable, or revoked key.
"""
resolved = await _resolve_active_litellm_key(request)
return resolved[0] if resolved else None
async def _store_per_user_token_server_side(
@ -688,10 +717,12 @@ async def _mint_bridge_delegate_token_response(
) -> JSONResponse:
"""Return the client-held envelope bearer for a DCR-bridge ``oauth_delegate`` token exchange.
The envelope binds the caller's litellm identity (resolved from the token request) to the
The envelope binds the authorizing litellm key (its hash, resolved from the token request) to the
upstream grant, so the client holds one bearer that later admits it and forwards the upstream
token, with nothing stored server-side. Fails closed with an OAuth ``invalid_request`` when no
litellm identity accompanies the token request rather than minting an identity-less credential.
token, with nothing stored server-side. Admission reloads the live key by that hash, so the key's
current restrictions and revocation gate the request. Fails closed with an OAuth
``invalid_request`` when no active litellm key accompanies the token request rather than minting
an unbound credential.
"""
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import
build_bridge_token_response,
@ -708,8 +739,8 @@ async def _mint_bridge_delegate_token_response(
if not master_key:
raise HTTPException(status_code=500, detail="Server misconfigured: master_key is not set")
user_id = await _extract_user_id_from_request(request)
if not user_id:
key_hash = await _extract_active_key_hash_from_request(request)
if not key_hash:
raise HTTPException(
status_code=400,
detail={
@ -727,7 +758,7 @@ async def _mint_bridge_delegate_token_response(
now = datetime.now(timezone.utc)
keys = envelope_keys_from_master_key(master_key)
identity = EnvelopeIdentity(user_id=user_id, server_id=mcp_server.server_id)
identity = EnvelopeIdentity(server_id=mcp_server.server_id, key_hash=key_hash)
sealed = build_bridge_token_response(identity, grant, keys, now)
if not isinstance(sealed, SealedEnvelope):
raise HTTPException(status_code=500, detail="Failed to mint the gateway-bound credential")

78
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -4365,7 +4365,7 @@ async def test_register_bridge_relay_never_persists():
_BRIDGE_MASTER_KEY = "sk-bridge-producer-master-key-0123456789abcdef"
async def _exchange_for_bridge_server(server, upstream_body, user_id):
async def _exchange_for_bridge_server(server, upstream_body, key_hash):
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
exchange_token_with_server,
)
@ -4382,8 +4382,8 @@ async def _exchange_for_bridge_server(server, upstream_body, user_id):
return_value=fake_http_client,
),
patch(
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request",
new=AsyncMock(return_value=user_id),
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_active_key_hash_from_request",
new=AsyncMock(return_value=key_hash),
),
patch("litellm.proxy.proxy_server.master_key", _BRIDGE_MASTER_KEY),
):
@ -4416,7 +4416,7 @@ async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_token
server = _bridge_server(auth_type=MCPAuth.oauth_delegate)
upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600}
response = await _exchange_for_bridge_server(server, upstream, user_id="user-77")
response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77")
body = json.loads(response.body)
token = body["access_token"]
@ -4429,7 +4429,7 @@ async def test_oauth_delegate_bridge_token_exchange_mints_envelope_not_raw_token
keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY)
opened = resolve_bridge_envelope(token, keys, datetime.now(timezone.utc), server.server_id)
assert isinstance(opened, BridgeEnvelopeAdmitted)
assert opened.identity.user_id == "user-77"
assert opened.identity.key_hash == "hashed-litellm-key-77"
assert opened.upstream_authorization.get_secret_value() == "Bearer UPSTREAM-SECRET-TOKEN"
@ -4443,7 +4443,7 @@ async def test_oauth_delegate_bridge_token_exchange_fails_closed_without_litellm
upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600}
with pytest.raises(HTTPException) as exc:
await _exchange_for_bridge_server(server, upstream, user_id=None)
await _exchange_for_bridge_server(server, upstream, key_hash=None)
assert exc.value.status_code == 400
assert exc.value.detail["error"] == "invalid_request"
@ -4457,7 +4457,7 @@ async def test_true_passthrough_bridge_token_exchange_returns_raw_upstream_token
server = _bridge_server(auth_type=MCPAuth.true_passthrough)
upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600}
response = await _exchange_for_bridge_server(server, upstream, user_id="user-77")
response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77")
body = json.loads(response.body)
assert body["access_token"] == "UPSTREAM-SECRET-TOKEN"
@ -4472,7 +4472,7 @@ async def test_non_bridge_oauth_delegate_token_exchange_returns_raw_upstream_tok
server = _bridge_server(auth_type=MCPAuth.oauth_delegate, dcr_bridge=None)
upstream = {"access_token": "UPSTREAM-SECRET-TOKEN", "token_type": "Bearer", "expires_in": 3600}
response = await _exchange_for_bridge_server(server, upstream, user_id="user-77")
response = await _exchange_for_bridge_server(server, upstream, key_hash="hashed-litellm-key-77")
body = json.loads(response.body)
assert body["access_token"] == "UPSTREAM-SECRET-TOKEN"
@ -4822,6 +4822,68 @@ async def test_extract_user_id_rejects_expired_key(proxy_globals):
assert await _extract_user_id_from_request(request) is None
@pytest.mark.asyncio
async def test_extract_active_key_hash_returns_hash_for_active_key(proxy_globals):
"""The dcr_bridge mint seals the hash of the authorizing key so admission can reload the live
record. For an active key the resolver returns exactly hash_token(key), the same value
get_key_object and the whole cache/DB layer key the record by, so the sealed reference resolves
back to this key at admission."""
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_extract_active_key_hash_from_request,
)
from litellm.proxy._types import UserAPIKeyAuth, hash_token
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
key = "sk-alice-key"
cache = UserApiKeyCache()
await cache.async_set_cache(
hash_token(key),
UserAPIKeyAuth(token=hash_token(key), user_id="alice"),
model_type=UserAPIKeyAuth,
)
proxy_globals.user_api_key_cache = cache
proxy_globals.prisma_client = object()
request = _token_request({"x-litellm-api-key": f"Bearer {key}"})
assert await _extract_active_key_hash_from_request(request) == hash_token(key)
@pytest.mark.asyncio
async def test_extract_active_key_hash_rejects_blocked_key(proxy_globals):
"""A blocked key must not yield a hash, so no gateway-bound envelope is minted for a revoked key;
the mint fails closed with invalid_request instead."""
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_extract_active_key_hash_from_request,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
class _FakePrisma:
async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None):
return UserAPIKeyAuth(token=token, user_id="blocked-user", blocked=True)
proxy_globals.user_api_key_cache = UserApiKeyCache()
proxy_globals.prisma_client = _FakePrisma()
request = _token_request({"x-litellm-api-key": "sk-blocked-key"})
assert await _extract_active_key_hash_from_request(request) is None
@pytest.mark.asyncio
async def test_extract_active_key_hash_none_without_litellm_key(proxy_globals):
"""No LiteLLM key on the request yields no hash without consulting the resolver."""
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
_extract_active_key_hash_from_request,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
proxy_globals.user_api_key_cache = UserApiKeyCache()
proxy_globals.prisma_client = object()
request = _token_request({"content-type": "application/json"})
assert await _extract_active_key_hash_from_request(request) is None
@pytest.mark.asyncio
async def test_token_endpoint_uses_client_secret_basic_when_configured():
"""LIT-4091: a server with token_endpoint_auth_method=client_secret_basic must send the