scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

fix(mcp): require audiences and prove refresh-token ownership for identity binding

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-02 15:52:47 +00:00
parent dae4259b00
commit 50874f9fc9
5 changed files with 211 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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

93
litellm/proxy/_experimental/mcp_server/oauth_identity_binding.py
View file

@ -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

4
litellm/types/mcp_server/mcp_server_manager.py
View file

@ -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

117
tests/test_litellm/proxy/_experimental/mcp_server/test_oauth_identity_binding.py
View file

@ -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"

6
tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
View file

@ -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)