From 50874f9fc9cdc9e925d0fc7abda06a178a303980 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 2 Sep 2026 15:52:47 +0000 Subject: [PATCH] fix(mcp): require audiences and prove refresh-token ownership for identity binding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/discoverable_endpoints.py | 9 ++ .../mcp_server/oauth_identity_binding.py | 93 +++++++++++--- .../types/mcp_server/mcp_server_manager.py | 4 +- .../mcp_server/test_oauth_identity_binding.py | 117 +++++++++++++++++- .../test_mcp_management_endpoints.py | 6 +- 5 files changed, 211 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 92128a82856..88f2c47f5fc 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -56,6 +56,8 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( revoke_refresh_token, ) from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( + RefreshOwnershipProven, + RefreshTokenPresented, enforce_oauth_identity_binding, ) from litellm.proxy._experimental.mcp_server.oauth_utils import ( @@ -1046,7 +1048,13 @@ async def exchange_token_with_server( refresh_request_scope = scope or bridge_upstream_scope if refresh_request_scope: token_data["scope"] = refresh_request_scope + refresh_ownership = ( + RefreshOwnershipProven() + if bridge_upstream_refresh is not None + else RefreshTokenPresented(upstream_refresh_token) + ) else: + refresh_ownership = None if not code: raise HTTPException( status_code=400, @@ -1145,6 +1153,7 @@ async def exchange_token_with_server( token_response=token_response, litellm_user_id=resolved_user_id, grant_type=grant_type, + refresh_ownership=refresh_ownership, ) # Store server-side when the server is configured for per-user OAuth and diff --git a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py index f132f59ae99..d4339c15e1a 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py @@ -14,6 +14,7 @@ from typing import Final, Literal import jwt from fastapi import HTTPException +from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache @@ -37,6 +38,7 @@ _jwks_cache: Final = InMemoryCache(default_ttl=_JWKS_CACHE_TTL_SECONDS) JwksFetcher = Callable[[MCPOAuthIdentityBinding], Awaitable[list[Mapping[str, object]]]] CallerPrincipalLoader = Callable[[str, MCPOAuthIdentityBinding], Awaitable[str | None]] +StoredRefreshTokenLoader = Callable[[str, str], Awaitable[str | None]] _RejectionCode = Literal["oauth_principal_mismatch", "oauth_identity_binding_failed"] @@ -47,6 +49,19 @@ class _BindingRejection: description: str +@dataclass(frozen=True, slots=True) +class RefreshOwnershipProven: + """The gateway itself unwrapped the upstream refresh token from a sealed per-user envelope.""" + + +@dataclass(frozen=True, slots=True) +class RefreshTokenPresented: + refresh_token: str + + +RefreshOwnership = RefreshOwnershipProven | RefreshTokenPresented | None + + async def _fetch_issuer_jwks(binding: MCPOAuthIdentityBinding) -> list[Mapping[str, object]]: jwks_url: Final[str] = binding.jwks_url or await _discover_jwks_url(binding.issuer) cached: Final = await _jwks_cache.async_get_cache(jwks_url) @@ -98,11 +113,8 @@ def _decode_id_token( signing_key.key, algorithms=list(_ALLOWED_ID_TOKEN_ALGORITHMS), issuer=binding.issuer, - audience=binding.audiences if binding.audiences else None, - options={ - "require": ["iss", "exp"], - "verify_aud": bool(binding.audiences), - }, + audience=binding.audiences, + options={"require": ["iss", "exp"]}, ) except jwt.InvalidTokenError as exc: return _BindingRejection( @@ -121,7 +133,11 @@ def _upstream_principal( code="oauth_identity_binding_failed", description=f"id_token has no usable '{binding.principal_claim}' claim", ) - if binding.principal_claim == "email" and binding.require_email_verified and claims.get("email_verified") is not True: + if ( + binding.principal_claim == "email" + and binding.require_email_verified + and claims.get("email_verified") is not True + ): return _BindingRejection( code="oauth_identity_binding_failed", description="id_token email is not verified (email_verified is not true)", @@ -142,6 +158,26 @@ async def _load_caller_principal(litellm_user_id: str, binding: MCPOAuthIdentity return loaded.user_email +async def _load_stored_refresh_token(litellm_user_id: str, server_id: str) -> str | None: + try: + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + get_user_oauth_credential, + ) + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 + + prisma_client: Final = get_prisma_client_or_throw( + "Database not connected. Cannot verify OAuth refresh token ownership." + ) + cred: Final = await get_user_oauth_credential( + prisma_client=prisma_client, + user_id=litellm_user_id, + server_id=server_id, + ) + return cred.get("refresh_token") if cred else None + except Exception: # noqa: BLE001 # a credential lookup failure must fail closed + return None + + def _principals_match(upstream: str, caller: str, binding: MCPOAuthIdentityBinding) -> bool: if binding.principal_claim == "email" or binding.caller_field == "user_email": return upstream.strip().casefold() == caller.strip().casefold() @@ -153,17 +189,41 @@ async def _evaluate_binding( token_response: Mapping[str, object], litellm_user_id: str | None, grant_type: str, + server_id: str, + refresh_ownership: RefreshOwnership, jwks_fetcher: JwksFetcher, caller_principal_loader: CallerPrincipalLoader, + stored_refresh_token_loader: StoredRefreshTokenLoader, ) -> _BindingRejection | None: id_token: Final = token_response.get("id_token") if not isinstance(id_token, str) or not id_token: - if grant_type == "refresh_token": - return None - return _BindingRejection( - code="oauth_identity_binding_failed", - description="the upstream token response carries no id_token to bind the credential to a principal", - ) + if grant_type != "refresh_token": + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the upstream token response carries no id_token to bind the credential to a principal", + ) + match refresh_ownership: + case RefreshOwnershipProven(): + return None + case None: + return _BindingRejection( + code="oauth_identity_binding_failed", + description="refresh_token grant without an id_token carries no refresh token to prove ownership of", + ) + case RefreshTokenPresented(refresh_token): + if not litellm_user_id: + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the request carries no resolvable LiteLLM user identity to bind the credential to", + ) + stored: Final = await stored_refresh_token_loader(litellm_user_id, server_id) + if stored is None or stored != refresh_token: + return _BindingRejection( + code="oauth_identity_binding_failed", + description="the presented refresh_token is not the caller's stored credential for this server", + ) + return None + assert_never(refresh_ownership) if not litellm_user_id: return _BindingRejection( code="oauth_identity_binding_failed", @@ -204,15 +264,17 @@ async def enforce_oauth_identity_binding( token_response: Mapping[str, object], litellm_user_id: str | None, grant_type: str, + refresh_ownership: RefreshOwnership, jwks_fetcher: JwksFetcher = _fetch_issuer_jwks, caller_principal_loader: CallerPrincipalLoader = _load_caller_principal, + stored_refresh_token_loader: StoredRefreshTokenLoader = _load_stored_refresh_token, ) -> None: """Validate the exchanged token's upstream principal against the LiteLLM caller. No-op when the server has no binding or it is disabled. In enforce mode a failure raises 403 before the caller returns, stores, or caches the token; in audit mode failures are logged only. - A refresh_token grant without an id_token is allowed in both modes: the stored credential keeps - the binding established at the original authorization_code exchange. + A refresh_token grant without an id_token is allowed only when the presented refresh token matches + the caller's stored credential or a sealed bridge envelope already proved ownership. """ binding: Final = server.oauth_identity_binding if binding is None or binding.mode == "disabled": @@ -222,8 +284,11 @@ async def enforce_oauth_identity_binding( token_response=token_response, litellm_user_id=litellm_user_id, grant_type=grant_type, + server_id=server.server_id, + refresh_ownership=refresh_ownership, jwks_fetcher=jwks_fetcher, caller_principal_loader=caller_principal_loader, + stored_refresh_token_loader=stored_refresh_token_loader, ) if rejection is None: return diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index cc3f6c8e8d8..cebc7decac4 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,7 +1,7 @@ from datetime import datetime from typing import Any, Final, Literal -from pydantic import BaseModel, ConfigDict, field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator from litellm.types.mcp import ( DEFAULT_SUBJECT_TOKEN_TYPE, @@ -52,7 +52,7 @@ class MCPOAuthIdentityBinding(BaseModel): mode: Literal["disabled", "audit", "enforce"] = "disabled" issuer: str jwks_url: str | None = None - audiences: list[str] = [] + audiences: list[str] = Field(min_length=1) principal_claim: str = "email" caller_field: Literal["user_email", "user_id"] = "user_email" require_email_verified: bool = True diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py index 69eaabeda4a..de59020c44c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py @@ -7,8 +7,11 @@ import pytest from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from fastapi import HTTPException +from pydantic import ValidationError from litellm.proxy._experimental.mcp_server.oauth_identity_binding import ( + RefreshOwnershipProven, + RefreshTokenPresented, enforce_oauth_identity_binding, ) from litellm.types.mcp import MCPTransport @@ -54,6 +57,13 @@ def _caller_loader(email: str | None): return load +def _stored_refresh_token_loader(refresh_token: str | None): + async def load(_user_id: str, _server_id: str) -> str | None: + return refresh_token + + return load + + def _server(mode: str = "enforce", **binding_overrides: object) -> MCPServer: return MCPServer( server_id="srv-1", @@ -77,6 +87,7 @@ async def test_matching_principal_passes(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -92,6 +103,7 @@ async def test_mismatched_principal_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -108,6 +120,7 @@ async def test_missing_id_token_rejected_on_authorization_code(): token_response={"access_token": "at"}, litellm_user_id="user-a", grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -116,18 +129,89 @@ async def test_missing_id_token_rejected_on_authorization_code(): @pytest.mark.asyncio -async def test_refresh_without_id_token_allowed(): +async def test_refresh_without_id_token_allowed_when_presented_token_matches_stored_credential(): result: Final = await enforce_oauth_identity_binding( server=_server(), token_response={"access_token": "at"}, litellm_user_id="user-a", grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented("rt-1"), jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"), ) assert result is None +@pytest.mark.asyncio +async def test_refresh_without_id_token_rejects_different_presented_token(): + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented("rt-stolen"), + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_refresh_without_id_token_rejects_missing_stored_token(): + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=RefreshTokenPresented("rt-1"), + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader(None), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + +@pytest.mark.asyncio +async def test_refresh_without_id_token_passes_when_bridge_proves_ownership(): + async def fail_if_called(_user_id: str, _server_id: str) -> str | None: + raise AssertionError("stored refresh token loader should not be called") + + result: Final = await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=RefreshOwnershipProven(), + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=fail_if_called, + ) + assert result is None + + +@pytest.mark.asyncio +async def test_refresh_without_id_token_rejects_without_ownership_proof(): + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at"}, + litellm_user_id="user-a", + grant_type="refresh_token", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + stored_refresh_token_loader=_stored_refresh_token_loader("rt-1"), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" + + @pytest.mark.asyncio async def test_refresh_with_mismatched_id_token_rejected(): token: Final = _sign_id_token({"email": "mallory@example.com", "email_verified": True}) @@ -137,6 +221,7 @@ async def test_refresh_with_mismatched_id_token_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="refresh_token", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -150,6 +235,7 @@ async def test_audit_mode_logs_but_does_not_reject(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -165,6 +251,7 @@ async def test_unverified_email_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -187,6 +274,7 @@ async def test_wrong_issuer_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id="user-a", grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -202,6 +290,7 @@ async def test_no_litellm_identity_rejected(): token_response={"access_token": "at", "id_token": token}, litellm_user_id=None, grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader("alice@example.com"), ) @@ -215,6 +304,7 @@ async def test_disabled_binding_is_noop(): token_response={"access_token": "at"}, litellm_user_id=None, grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader(None), ) @@ -234,7 +324,32 @@ async def test_no_binding_is_noop(): token_response={"access_token": "at"}, litellm_user_id=None, grant_type="authorization_code", + refresh_ownership=None, jwks_fetcher=_jwks_fetcher, caller_principal_loader=_caller_loader(None), ) assert result is None + + +def test_identity_binding_requires_non_empty_audiences(): + with pytest.raises(ValidationError): + MCPOAuthIdentityBinding(mode="enforce", issuer=ISSUER, audiences=[]) + with pytest.raises(ValidationError): + MCPOAuthIdentityBinding(mode="enforce", issuer=ISSUER) + + +@pytest.mark.asyncio +async def test_wrong_audience_rejected(): + token: Final = _sign_id_token({"aud": "other-client", "email": "alice@example.com", "email_verified": True}) + with pytest.raises(HTTPException) as exc_info: + await enforce_oauth_identity_binding( + server=_server(), + token_response={"access_token": "at", "id_token": token}, + litellm_user_id="user-a", + grant_type="authorization_code", + refresh_ownership=None, + jwks_fetcher=_jwks_fetcher, + caller_principal_loader=_caller_loader("alice@example.com"), + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["error"] == "oauth_identity_binding_failed" diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index bfbbf816e69..f3510e7e0a8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -4674,7 +4674,11 @@ async def test_store_mcp_oauth_user_credential_blocked_when_identity_binding_enf name=server_id, url="https://mcp.example.com", transport=MCPTransport.http, - oauth_identity_binding=MCPOAuthIdentityBinding(mode="enforce", issuer="https://idp.example.com"), + oauth_identity_binding=MCPOAuthIdentityBinding( + mode="enforce", + issuer="https://idp.example.com", + audiences=["litellm-client"], + ), ) store_mock = AsyncMock(return_value=None)