From c59f16f42ea54e799c8d73c5875165d21c84fe02 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 10 Jul 2026 17:58:43 -0700 Subject: [PATCH] fix(mcp): enforce team block and alias-priority token injection on bridge admission 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. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 41 ++++++++++++-- .../auth/test_user_api_key_auth_mcp.py | 56 +++++++++++++++++-- 2 files changed, 89 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index adb1812059b..bee621b3721 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 9873f3ac7c5..d0d75231b59 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -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