scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(mcp): enforce team block and alias-priority token injection on bridge admission

Browse source 
Two follow-ups on the envelope admission arm flagged in review.

Team revocation bypass: _reload_admitted_key checked only the key's own
blocked/expires, so blocking a key's team left every envelope minted under it
live until expiry. Reload the team and reject a blocked team, mirroring
common_checks, so a team block revokes its envelopes immediately.

Caller-overridable upstream token: egress resolves the per-server auth header
alias-first, but injection keyed under server_name, so for a server with a
distinct alias a caller-forwarded x-mcp-{alias}-authorization sat at the
higher-priority slot and paired the admitted identity with an attacker's
upstream credential. Inject under alias-first so the sealed token owns the slot
egress resolves.
This commit is contained in:
Tin Chi Lo 2026-07-10 17:58:43 -07:00
parent 50cc2c01cf
commit c59f16f42e
2 changed files with 89 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -508,7 +508,13 @@ class MCPRequestHandler:
``PassthroughConfig`` override; the envelope ``Authorization`` the leak-defense
strips never reaches the upstream. A new headers dict is returned rather than
mutating the input. Fails closed with a 401 on an invalid or expired envelope, or
when the referenced key is missing, blocked, or expired.
when the referenced key is missing, blocked, or expired, or its team is blocked.
The sealed token is keyed alias-first, matching the order egress resolves
(``lookup_mcp_server_auth_in_headers`` tries ``alias`` before ``server_name``). Keying
under ``server_name`` would leave a caller-supplied ``x-mcp-{alias}-authorization`` at the
higher-priority alias slot, pairing the admitted identity with an attacker's upstream
credential; the alias-keyed injection overwrites any such caller value.
"""
from litellm.proxy.proxy_server import master_key
@ -519,7 +525,7 @@ class MCPRequestHandler:
result = resolve_bridge_envelope(authorization_value, keys, datetime.now(timezone.utc), server.server_id)
match result:
case BridgeEnvelopeAdmitted():
header_key = server.server_name or server.alias
header_key = server.alias or server.server_name
if header_key is None:
raise HTTPException(status_code=500, detail="Server misconfigured: MCP server has no routable name")
admitted = await MCPRequestHandler._reload_admitted_key(result.identity.key_hash)
@ -533,7 +539,7 @@ class MCPRequestHandler:
@staticmethod
async def _reload_admitted_key(key_hash: str) -> UserAPIKeyAuth:
"""Reload the live key record an admitted envelope references.
"""Reload the live key record an admitted envelope references and re-check live policy.
Resolving the current ``UserAPIKeyAuth`` (cache first, then DB) is what stops the
envelope from carrying frozen authority: the key's present team/org/object-permission
@ -542,7 +548,9 @@ class MCPRequestHandler:
unrestricted identity. ``get_key_object`` raises for a hash with no key row; a
blocked or expired row is rejected explicitly because ``get_key_object`` resolves a
row without applying those checks (the main ``user_api_key_auth`` pipeline enforces
them downstream, which this admission path bypasses).
them downstream, which this admission path bypasses). The key's team is reloaded and
rejected when blocked, mirroring ``common_checks``, so blocking a team revokes every
envelope minted under its keys rather than leaving them live until expiry.
"""
from litellm.proxy.auth.auth_checks import get_key_object
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
@ -559,8 +567,33 @@ class MCPRequestHandler:
raise HTTPException(status_code=401, detail="Invalid or expired credential") from None
if not MCPRequestHandler._admitted_key_is_active(key_object):
raise HTTPException(status_code=401, detail="Invalid or expired credential")
await MCPRequestHandler._reject_if_admitted_team_blocked(key_object)
return key_object
@staticmethod
async def _reject_if_admitted_team_blocked(key_object: UserAPIKeyAuth) -> None:
"""Reload the key's team and fail closed with a 401 when it is blocked or no longer
resolves. ``get_key_object`` returns the key row without any team validation, so this
admission path applies the same live team-block gate ``common_checks`` runs on the
standard auth pipeline; without it, blocking a team would not revoke envelopes already
minted under its keys until they expired."""
team_id = key_object.team_id
if not team_id:
return
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
try:
team_object = await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
except (ProxyException, HTTPException):
raise HTTPException(status_code=401, detail="Invalid or expired credential") from None
if team_object.blocked is True:
raise HTTPException(status_code=401, detail="Invalid or expired credential")
@staticmethod
def _admitted_key_is_active(key_object: UserAPIKeyAuth) -> bool:
"""False when the referenced key is blocked or past its expiry, so a revoked key

56
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

@ -4962,14 +4962,17 @@ class TestMCPDcrBridgeDelegateAdmission:
@staticmethod
@contextlib.contextmanager
def _patch_key_reload(*, return_value=None, side_effect=None):
def _patch_key_reload(*, return_value=None, side_effect=None, team_blocked=False):
"""Patch the live-key reload dependencies used by ``_reload_admitted_key``: the
``get_key_object`` lookup plus the ``prisma_client`` / ``user_api_key_cache`` globals it
reads. Yields the ``get_key_object`` mock so callers can assert the sealed ``key_hash`` was
the reload key."""
``get_key_object`` lookup, the ``get_team_object`` lookup its team-block gate runs, and the
``prisma_client`` / ``user_api_key_cache`` globals they read. The team resolves unblocked by
default; ``team_blocked=True`` simulates an admin blocking the key's team. Yields the
``get_key_object`` mock so callers can assert the sealed ``key_hash`` was the reload key."""
get_key_object = AsyncMock(return_value=return_value, side_effect=side_effect)
get_team_object = AsyncMock(return_value=MagicMock(blocked=team_blocked))
with (
patch("litellm.proxy.auth.auth_checks.get_key_object", get_key_object),
patch("litellm.proxy.auth.auth_checks.get_team_object", get_team_object),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
):
@ -5102,6 +5105,29 @@ class TestMCPDcrBridgeDelegateAdmission:
assert exc_info.value.status_code == 401
async def test_blocked_team_envelope_fails_closed_401(self):
"""Blocking the key's TEAM must revoke its envelopes immediately: the reloaded key is active
but its team is blocked, so admission 401s. Without the live team re-check, a caller could
keep executing tools after an admin blocked the team, until the envelope expired."""
envelope = self._mint_bridge_envelope(key_hash=self._KEY_HASH)
scope = {
"type": "http",
"method": "POST",
"path": "/mcp/bridge_delegate_server",
"headers": [(b"authorization", f"Bearer {envelope}".encode("latin-1"))],
}
with (
patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_mgr,
patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY),
self._patch_key_reload(return_value=self._reloaded_key(), team_blocked=True),
):
mock_mgr.get_mcp_server_by_name.return_value = self._bridge_delegate_server()
with pytest.raises(HTTPException) as exc_info:
await MCPRequestHandler.process_mcp_request(scope)
assert exc_info.value.status_code == 401
async def test_alias_only_server_injects_under_alias_egress_can_resolve(self):
"""When server_name is None, the inner token must be keyed under the alias (which egress
resolves), never under server.name (which egress never looks up), so the forwarded token is
@ -5129,6 +5155,28 @@ class TestMCPDcrBridgeDelegateAdmission:
assert mcp_server_auth_headers == {"bridge_alias": {"Authorization": "Bearer inner-upstream-access-token"}}
async def test_sealed_token_wins_over_caller_forwarded_alias_header(self):
"""When a bridge server has both a server_name and a distinct alias, the sealed inner token
must occupy the alias slot, the identifier egress resolves first. Otherwise a caller who
forwards x-mcp-{alias}-authorization keeps that entry at the higher-priority slot and pairs
the admitted identity with an attacker-chosen upstream credential."""
envelope = self._mint_bridge_envelope()
attacker_forwarded = {"bridge_alias": {"Authorization": "Bearer ATTACKER-UPSTREAM-TOKEN"}}
with (
patch("litellm.proxy.proxy_server.master_key", self._MASTER_KEY),
self._patch_key_reload(return_value=self._reloaded_key()),
):
_auth, new_headers = await MCPRequestHandler._admit_dcr_bridge_delegate(
server=self._bridge_delegate_server(server_name="bridge_name", alias="bridge_alias"),
authorization_value=f"Bearer {envelope}",
mcp_server_auth_headers=attacker_forwarded,
)
# The sealed token owns the alias slot, overwriting the caller's value; the attacker token
# survives nowhere egress would resolve.
assert new_headers == {"bridge_alias": {"Authorization": "Bearer inner-upstream-access-token"}}
async def test_server_with_no_alias_or_server_name_is_not_admitted_via_bridge_arm(self):
"""A bridge server egress cannot route to (no alias and no server_name) must not take the
envelope arm; it fails closed to normal oauth2 admission rather than admitting and dropping