mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(mcp): preserve identity checks across cached OAuth credentials
This commit is contained in:
parent
7c1d060b7a
commit
55ab5ee53d
14 changed files with 311 additions and 54 deletions
|
|
@ -1 +1,3 @@
|
|||
Read @CLAUDE.md for coding guidelines
|
||||
|
||||
Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review; never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or lint budgets to satisfy a check
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
|
|||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast
|
||||
|
||||
from fastapi import HTTPException
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -1692,13 +1693,18 @@ async def refresh_user_oauth_token(
|
|||
)
|
||||
return None
|
||||
|
||||
binding_proof: Final = await enforce_oauth_identity_binding(
|
||||
server=server,
|
||||
token_response=body,
|
||||
litellm_user_id=user_id,
|
||||
grant_type="refresh_token",
|
||||
refresh_ownership=RefreshTokenPresented(refresh_token),
|
||||
)
|
||||
try:
|
||||
binding_proof: Final = await enforce_oauth_identity_binding(
|
||||
server=server,
|
||||
token_response=body,
|
||||
litellm_user_id=user_id,
|
||||
grant_type="refresh_token",
|
||||
refresh_ownership=RefreshTokenPresented(refresh_token),
|
||||
)
|
||||
except HTTPException as exc:
|
||||
if exc.status_code != 403:
|
||||
raise
|
||||
return None
|
||||
|
||||
access_token: Final[str | None] = body.get("access_token")
|
||||
if not access_token:
|
||||
|
|
@ -1812,8 +1818,14 @@ async def resolve_user_oauth_access_token(
|
|||
|
||||
binding: Final = server.oauth_identity_binding
|
||||
enforce_binding: Final = binding is not None and binding.mode == "enforce"
|
||||
if enforce_binding:
|
||||
await mcp_per_user_token_cache.delete(user_id, server_id)
|
||||
if prefetched_creds is None and enforce_binding and binding is not None:
|
||||
bound_token: Final = await mcp_per_user_token_cache.get_token(user_id, server_id)
|
||||
if bound_token is not None:
|
||||
if await credential_binding_matches(
|
||||
binding, user_id, server_id, {"identity_binding_proof": bound_token.identity_binding_proof}
|
||||
):
|
||||
return bound_token.access_token
|
||||
await mcp_per_user_token_cache.delete(user_id, server_id)
|
||||
if prefetched_creds is None and not enforce_binding:
|
||||
cached_token: Final = await mcp_per_user_token_cache.get(user_id, server_id)
|
||||
if cached_token is not None:
|
||||
|
|
@ -1848,7 +1860,9 @@ async def resolve_user_oauth_access_token(
|
|||
access_token: Final[str] = cred["access_token"]
|
||||
if prefetched_creds is None:
|
||||
ttl: Final = _compute_per_user_token_ttl(server, _remaining_token_seconds(cred.get("expires_at")))
|
||||
await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl)
|
||||
await mcp_per_user_token_cache.set(
|
||||
user_id, server_id, access_token, ttl, identity_binding_proof=cred.get("identity_binding_proof")
|
||||
)
|
||||
return access_token
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
|
|
|
|||
|
|
@ -28,6 +28,8 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import (
|
|||
build_upstream_oauth2_token_request,
|
||||
resolve_upstream_resource,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import OAuthToken
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_codec import OAuthTokenCacheCodec
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
|
|
@ -233,8 +235,17 @@ class MCPPerUserTokenCache:
|
|||
def _cache_key(self, user_id: str, server_id: str) -> str:
|
||||
return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}"
|
||||
|
||||
def _codec(self) -> OAuthTokenCacheCodec:
|
||||
return OAuthTokenCacheCodec(
|
||||
encrypt_value_helper,
|
||||
lambda blob: decrypt_value_helper(blob, key="mcp_per_user_token", exception_type="debug"),
|
||||
)
|
||||
|
||||
async def get(self, user_id: str, server_id: str) -> str | None:
|
||||
"""Return the plaintext access_token, or None on miss/error."""
|
||||
token: Final = await self.get_token(user_id, server_id)
|
||||
return token.access_token if token is not None else None
|
||||
|
||||
async def get_token(self, user_id: str, server_id: str) -> OAuthToken | None:
|
||||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
|
||||
|
||||
|
|
@ -242,12 +253,7 @@ class MCPPerUserTokenCache:
|
|||
encrypted: Final = await user_api_key_cache.async_get_cache(key)
|
||||
if encrypted is None:
|
||||
return None
|
||||
plaintext: Final = decrypt_value_helper(
|
||||
encrypted,
|
||||
key="mcp_per_user_token",
|
||||
exception_type="debug",
|
||||
)
|
||||
return plaintext or None
|
||||
return self._codec().decode(encrypted)
|
||||
except Exception as exc:
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.get failed for user=%s server=%s: %s",
|
||||
|
|
@ -263,13 +269,16 @@ class MCPPerUserTokenCache:
|
|||
server_id: str,
|
||||
access_token: str,
|
||||
ttl: int,
|
||||
identity_binding_proof: str | None = None,
|
||||
) -> None:
|
||||
"""Store NaCl-encrypted access_token in Redis with the given TTL."""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
|
||||
|
||||
key: Final = self._cache_key(user_id, server_id)
|
||||
encrypted: Final = encrypt_value_helper(access_token)
|
||||
encrypted: Final = self._codec().encode(
|
||||
OAuthToken(access_token=access_token, identity_binding_proof=identity_binding_proof)
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(key, encrypted, ttl=ttl)
|
||||
verbose_logger.debug(
|
||||
"MCPPerUserTokenCache.set: cached token for user=%s server=%s ttl=%ds",
|
||||
|
|
|
|||
|
|
@ -342,7 +342,7 @@ async def _evaluate_binding(
|
|||
claims: Final = _decode_id_token(id_token, binding, signing_key)
|
||||
if isinstance(claims, _BindingRejection):
|
||||
return claims
|
||||
if grant_type == "authorization_code":
|
||||
if grant_type == "authorization_code" and (binding.mode == "enforce" or expected_nonce is not None):
|
||||
nonce: Final = claims.get("nonce")
|
||||
if not expected_nonce or not isinstance(nonce, str) or not hmac.compare_digest(nonce, expected_nonce):
|
||||
return _BindingRejection(
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ import time
|
|||
from collections.abc import Awaitable, Callable
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import (
|
||||
TokenEndpointAuthConfigError,
|
||||
|
|
@ -93,6 +95,14 @@ class AuthorizationCodeRefresher:
|
|||
self._identity_validator = identity_validator
|
||||
|
||||
async def refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None:
|
||||
try:
|
||||
return await self._refresh(user_id, server_id, token)
|
||||
except HTTPException as exc:
|
||||
if exc.status_code != 403:
|
||||
raise
|
||||
return None
|
||||
|
||||
async def _refresh(self, user_id: str, server_id: str, token: OAuthToken) -> OAuthToken | None:
|
||||
if token.refresh_token is None:
|
||||
return None
|
||||
server: Final = self._server_lookup(server_id)
|
||||
|
|
@ -162,4 +172,5 @@ class AuthorizationCodeRefresher:
|
|||
expires_at=self._clock() + expires_in if expires_in is not None else None,
|
||||
refresh_token=new_refresh,
|
||||
scopes=scopes,
|
||||
identity_binding_proof=binding_proof,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -41,6 +41,7 @@ class OAuthToken:
|
|||
expires_at: float | None = None
|
||||
refresh_token: str | None = None
|
||||
scopes: tuple[str, ...] = ()
|
||||
identity_binding_proof: str | None = None
|
||||
|
||||
def __repr__(self) -> str:
|
||||
has_refresh: Final = self.refresh_token is not None
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
from collections.abc import Callable, Mapping
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -39,7 +40,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.token_cache_cod
|
|||
OAuthTokenCacheCodec,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.v2_token_store import (
|
||||
CredentialReader,
|
||||
V2PerUserTokenStore,
|
||||
)
|
||||
|
||||
|
|
@ -146,12 +146,26 @@ def _runtime_backend_and_coordinator() -> tuple[TokenCacheBackend | None, Refres
|
|||
return backend, coordinator, True
|
||||
|
||||
|
||||
async def _read_bound_credential(
|
||||
server_lookup: ServerLookup, user_id: str, server_id: str
|
||||
) -> Mapping[str, object] | None:
|
||||
credential: Final = await _read_credential(user_id, server_id)
|
||||
server: Final = server_lookup(server_id)
|
||||
binding: Final = server.oauth_identity_binding if server else None
|
||||
if credential is not None and binding is not None and binding.mode == "enforce":
|
||||
if not await credential_binding_matches(binding, user_id, server_id, credential):
|
||||
return None
|
||||
return credential
|
||||
|
||||
|
||||
def _build_per_user_oauth_token_store(
|
||||
server_lookup: ServerLookup,
|
||||
) -> tuple[CachedOAuthTokenStore, bool]:
|
||||
backend, coordinator, uses_redis = _runtime_backend_and_coordinator()
|
||||
refresher: Final = AuthorizationCodeRefresher(server_lookup, _post_token_endpoint, _persist_credential)
|
||||
refreshing: Final = RefreshingTokenStore(V2PerUserTokenStore(_read_credential), refresher, coordinator=coordinator)
|
||||
refreshing: Final = RefreshingTokenStore(
|
||||
V2PerUserTokenStore(partial(_read_bound_credential, server_lookup)), refresher, coordinator=coordinator
|
||||
)
|
||||
return CachedOAuthTokenStore(refreshing, default_ttl_seconds=_DEFAULT_TTL_SECONDS, backend=backend), uses_redis
|
||||
|
||||
|
||||
|
|
@ -176,10 +190,8 @@ class LazyPerUserOAuthTokenStore:
|
|||
*,
|
||||
store_builder: StoreBuilder = _build_per_user_oauth_token_store,
|
||||
redis_available: Callable[[], bool] = _redis_cache_is_available,
|
||||
credential_reader: CredentialReader = _read_credential,
|
||||
) -> None:
|
||||
self._server_lookup = server_lookup
|
||||
self._credential_reader = credential_reader
|
||||
self._store_builder = store_builder
|
||||
self._redis_available = redis_available
|
||||
self._store: InvalidatableOAuthTokenStore | None = None
|
||||
|
|
@ -188,13 +200,18 @@ class LazyPerUserOAuthTokenStore:
|
|||
self._local_fetches = 0
|
||||
|
||||
async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None:
|
||||
token: Final = await self._fetch_token(user_id, server_id)
|
||||
server: Final = self._server_lookup(server_id)
|
||||
binding: Final = server.oauth_identity_binding if server else None
|
||||
if binding is not None and binding.mode == "enforce":
|
||||
credential: Final = await self._credential_reader(user_id, server_id)
|
||||
await self.invalidate(user_id, server_id)
|
||||
if credential is None or not await credential_binding_matches(binding, user_id, server_id, credential):
|
||||
if token is not None and binding is not None and binding.mode == "enforce":
|
||||
if not await credential_binding_matches(
|
||||
binding, user_id, server_id, {"identity_binding_proof": token.identity_binding_proof}
|
||||
):
|
||||
await self.invalidate(user_id, server_id)
|
||||
return None
|
||||
return token
|
||||
|
||||
async def _fetch_token(self, user_id: str, server_id: str) -> OAuthToken | None:
|
||||
if self._uses_redis:
|
||||
store = self._store
|
||||
if store is not None:
|
||||
|
|
|
|||
|
|
@ -1,24 +1,26 @@
|
|||
"""Serialize + encrypt boundary for caching an OAuth token in a shared (Redis) cache.
|
||||
|
||||
A cross-replica cache must serialize the token, and a plaintext bearer in Redis is a leak, so this
|
||||
encrypts the value (NaCl in production via the injected ``encrypt``, identity in tests). It caches
|
||||
**only** the ``access_token``: the hot path needs just the bearer, expiry is carried by the cache
|
||||
entry's TTL (set from the token's ``expires_at`` by the cache), and the long-lived refresh_token stays
|
||||
in the DB - the refresh path is always a cache miss that re-reads it - so it never reaches Redis. A
|
||||
decoded token therefore carries only the bearer (``expires_at`` and ``refresh_token`` both None); the
|
||||
TTL, not the value, bounds its life. An empty/undecryptable blob (e.g. master-key rotation) is a miss.
|
||||
Shared cache values contain an encrypted access token and optional identity-binding proof.
|
||||
Refresh tokens remain in the database; cache TTL bounds the access token's lifetime.
|
||||
Legacy bearer-only entries decode without proof and cannot satisfy identity enforcement.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
|
||||
OAuthToken,
|
||||
)
|
||||
|
||||
_BOUND_PREFIX: Final = "litellm-bound-oauth-v1:"
|
||||
_BOUND_PAYLOAD: Final = TypeAdapter(dict[str, str])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class OAuthTokenCacheCodec:
|
||||
|
|
@ -26,10 +28,30 @@ class OAuthTokenCacheCodec:
|
|||
decrypt: Callable[[str], str | None]
|
||||
|
||||
def encode(self, token: OAuthToken) -> str:
|
||||
if token.identity_binding_proof is not None:
|
||||
return self.encrypt(
|
||||
_BOUND_PREFIX
|
||||
+ json.dumps(
|
||||
{
|
||||
"access_token": token.access_token,
|
||||
"identity_binding_proof": token.identity_binding_proof,
|
||||
}
|
||||
)
|
||||
)
|
||||
return self.encrypt(token.access_token)
|
||||
|
||||
def decode(self, blob: str) -> OAuthToken | None:
|
||||
access_token: Final = self.decrypt(blob)
|
||||
if not access_token:
|
||||
return None
|
||||
if access_token.startswith(_BOUND_PREFIX):
|
||||
try:
|
||||
payload: Final = _BOUND_PAYLOAD.validate_json(access_token[len(_BOUND_PREFIX) :])
|
||||
except ValidationError:
|
||||
return None
|
||||
bearer: Final = payload.get("access_token")
|
||||
proof: Final = payload.get("identity_binding_proof")
|
||||
if not bearer or not proof:
|
||||
return None
|
||||
return OAuthToken(access_token=bearer, identity_binding_proof=proof)
|
||||
return OAuthToken(access_token=access_token, refresh_token=None)
|
||||
|
|
|
|||
|
|
@ -46,11 +46,13 @@ def _to_oauth_token(payload: Mapping[str, object]) -> OAuthToken | None:
|
|||
return None
|
||||
refresh_token: Final = payload.get("refresh_token")
|
||||
expires_at: Final = payload.get("expires_at")
|
||||
binding_proof: Final = payload.get("identity_binding_proof")
|
||||
return OAuthToken(
|
||||
access_token=access_token,
|
||||
expires_at=_iso_to_epoch(expires_at) if isinstance(expires_at, str) else None,
|
||||
refresh_token=refresh_token if isinstance(refresh_token, str) else None,
|
||||
scopes=_to_scopes(payload.get("scopes")),
|
||||
identity_binding_proof=binding_proof if isinstance(binding_proof, str) else None,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -293,17 +293,18 @@ async def test_refresh_uses_admin_entered_token_url_when_issuer_yield_empties_re
|
|||
@pytest.mark.asyncio
|
||||
async def test_identity_rejection_never_persists_or_returns_refreshed_token():
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
validator = AsyncMock(side_effect=HTTPException(status_code=403, detail="oauth_principal_mismatch"))
|
||||
persist = AsyncMock()
|
||||
refresher = AuthorizationCodeRefresher(
|
||||
_lookup(_Server()), _endpoint({"access_token": "foreign-token", "id_token": "foreign-identity"}),
|
||||
persist, identity_validator=validator,
|
||||
_lookup(_Server()),
|
||||
_endpoint({"access_token": "foreign-token", "id_token": "foreign-identity"}),
|
||||
persist,
|
||||
identity_validator=validator,
|
||||
)
|
||||
with pytest.raises(HTTPException) as error:
|
||||
await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt"))
|
||||
assert error.value.status_code == 403
|
||||
assert await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")) is None
|
||||
persist.assert_not_awaited()
|
||||
|
||||
|
||||
|
|
@ -314,10 +315,13 @@ async def test_verified_refresh_preserves_binding_proof_in_storage():
|
|||
validator = AsyncMock(return_value="verified-binding")
|
||||
persist = AsyncMock()
|
||||
refresher = AuthorizationCodeRefresher(
|
||||
_lookup(_Server()), _endpoint({"access_token": "new", "refresh_token": "rotated"}),
|
||||
persist, identity_validator=validator,
|
||||
_lookup(_Server()),
|
||||
_endpoint({"access_token": "new", "refresh_token": "rotated"}),
|
||||
persist,
|
||||
identity_validator=validator,
|
||||
)
|
||||
token = await refresher.refresh("alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt"))
|
||||
assert token.access_token == "new"
|
||||
assert token.refresh_token == "rotated"
|
||||
assert token.identity_binding_proof == "verified-binding"
|
||||
assert persist.await_args.kwargs["identity_binding_proof"] == "verified-binding"
|
||||
|
|
|
|||
|
|
@ -254,20 +254,86 @@ async def test_enforcement_invalidates_cached_legacy_credentials_before_use():
|
|||
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
|
||||
|
||||
server = MCPServer(
|
||||
server_id="srv", name="srv", transport=MCPTransport.http, auth_type=MCPAuth.oauth2,
|
||||
server_id="srv",
|
||||
name="srv",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
oauth_identity_binding=MCPOAuthIdentityBinding(
|
||||
mode="enforce", issuer="https://idp.example.com", audiences=["client"],
|
||||
mode="enforce",
|
||||
issuer="https://idp.example.com",
|
||||
audiences=["client"],
|
||||
),
|
||||
)
|
||||
cached = _RecordingStore("belongs-to-bob")
|
||||
|
||||
async def read_legacy(user_id, server_id):
|
||||
return {"access_token": "belongs-to-bob", "refresh_token": "bobs-refresh"}
|
||||
|
||||
store = LazyPerUserOAuthTokenStore(
|
||||
lambda server_id: server, store_builder=lambda lookup: (cached, False),
|
||||
redis_available=lambda: False, credential_reader=read_legacy,
|
||||
lambda server_id: server,
|
||||
store_builder=lambda lookup: (cached, False),
|
||||
redis_available=lambda: False,
|
||||
)
|
||||
assert await store.fetch("alice", "srv") is None
|
||||
assert cached.calls == []
|
||||
assert cached.calls == [("alice", "srv")]
|
||||
assert cached.invalidations == [("alice", "srv")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforced_cache_hit_avoids_credential_read_and_rejects_changed_policy(monkeypatch):
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.oauth_identity_binding import current_binding_proof
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials import per_user_oauth_store as module
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
server = MCPServer(
|
||||
server_id="srv",
|
||||
name="srv",
|
||||
transport="http",
|
||||
auth_type="oauth2",
|
||||
oauth_identity_binding={
|
||||
"mode": "enforce",
|
||||
"issuer": "https://idp.example",
|
||||
"audiences": ["client"],
|
||||
"caller_field": "user_id",
|
||||
"principal_claim": "sub",
|
||||
},
|
||||
)
|
||||
proof = await current_binding_proof(server.oauth_identity_binding, "alice", "srv")
|
||||
read = AsyncMock(return_value={"access_token": "alice-token", "identity_binding_proof": proof})
|
||||
monkeypatch.setattr(module, "_read_credential", read)
|
||||
monkeypatch.setattr(module, "_runtime_backend_and_coordinator", lambda: (None, None, False))
|
||||
store = LazyPerUserOAuthTokenStore(lambda _: server, redis_available=lambda: False)
|
||||
assert (await store.fetch("alice", "srv")).access_token == "alice-token"
|
||||
assert (await store.fetch("alice", "srv")).access_token == "alice-token"
|
||||
read.assert_awaited_once_with("alice", "srv")
|
||||
server.oauth_identity_binding = server.oauth_identity_binding.model_copy(update={"audiences": ["changed"]})
|
||||
assert await store.fetch("alice", "srv") is None
|
||||
read.assert_awaited_once()
|
||||
assert await store.fetch("alice", "srv") is None
|
||||
assert read.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_expired_unverified_credential_never_reaches_refresh(monkeypatch):
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials import per_user_oauth_store as module
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
server = MCPServer(
|
||||
server_id="srv",
|
||||
name="srv",
|
||||
transport="http",
|
||||
auth_type="oauth2",
|
||||
token_url="https://idp.example/token",
|
||||
oauth_identity_binding={"mode": "enforce", "issuer": "https://idp.example", "audiences": ["client"]},
|
||||
)
|
||||
read = AsyncMock(
|
||||
return_value={"access_token": "bob", "refresh_token": "bob-refresh", "expires_at": "2000-01-01T00:00:00Z"}
|
||||
)
|
||||
post = AsyncMock()
|
||||
monkeypatch.setattr(module, "_read_credential", read)
|
||||
monkeypatch.setattr(module, "_post_token_endpoint", post)
|
||||
monkeypatch.setattr(module, "_runtime_backend_and_coordinator", lambda: (None, None, False))
|
||||
store = LazyPerUserOAuthTokenStore(lambda _: server, redis_available=lambda: False)
|
||||
assert await store.fetch("alice", "srv") is None
|
||||
post.assert_not_awaited()
|
||||
|
|
|
|||
|
|
@ -52,3 +52,17 @@ def test_undecryptable_blob_is_a_miss():
|
|||
def test_empty_plaintext_is_a_miss():
|
||||
codec = OAuthTokenCacheCodec(encrypt=lambda s: s, decrypt=lambda b: b)
|
||||
assert codec.decode("") is None
|
||||
|
||||
|
||||
def test_bound_token_round_trip_preserves_proof_without_refresh_secret():
|
||||
codec = _wrapping_codec()
|
||||
blob = codec.encode(OAuthToken("alice-token", refresh_token="private-refresh", identity_binding_proof="proof"))
|
||||
assert "private-refresh" not in blob
|
||||
decoded = codec.decode(blob)
|
||||
assert decoded == OAuthToken("alice-token", identity_binding_proof="proof")
|
||||
|
||||
|
||||
def test_malformed_bound_entries_fail_closed():
|
||||
codec = _wrapping_codec()
|
||||
for payload in ("not-json", "{}", '{"access_token":"at"}', '{"access_token":1,"identity_binding_proof":"p"}'):
|
||||
assert codec.decode("enc:litellm-bound-oauth-v1:" + payload) is None
|
||||
|
|
|
|||
|
|
@ -1629,3 +1629,78 @@ async def test_enforcement_rejects_preexisting_unverified_credential():
|
|||
cred={"access_token": "belongs-to-bob", "refresh_token": "bobs-refresh-token"},
|
||||
)
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_identity_rejection_returns_reauthentication_without_persisting(monkeypatch):
|
||||
from fastapi import HTTPException
|
||||
from litellm.proxy._experimental.mcp_server import db as module
|
||||
|
||||
validator = AsyncMock(side_effect=HTTPException(status_code=403, detail="oauth_principal_mismatch"))
|
||||
monkeypatch.setattr(module, "enforce_oauth_identity_binding", validator)
|
||||
result, captured = await _run_refresh(
|
||||
monkeypatch, _refresh_server(), {"access_token": "bob", "refresh_token": "rotated"}
|
||||
)
|
||||
assert result is None
|
||||
assert captured["data"]["grant_type"] == "refresh_token"
|
||||
module.store_user_oauth_credential.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_verified_legacy_cache_reads_avoid_database_and_reject_policy_changes(monkeypatch):
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._experimental.mcp_server import db as module
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import mcp_per_user_token_cache
|
||||
from litellm.proxy._experimental.mcp_server.oauth_identity_binding import current_binding_proof
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
server = MCPServer(
|
||||
server_id="srv",
|
||||
name="srv",
|
||||
transport="http",
|
||||
auth_type="oauth2",
|
||||
oauth_identity_binding={
|
||||
"mode": "enforce",
|
||||
"issuer": "https://idp.example",
|
||||
"audiences": ["client"],
|
||||
"caller_field": "user_id",
|
||||
"principal_claim": "sub",
|
||||
},
|
||||
)
|
||||
proof = await current_binding_proof(server.oauth_identity_binding, "alice", "srv")
|
||||
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
|
||||
read = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(module, "get_user_oauth_credential", read)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
|
||||
await mcp_per_user_token_cache.set("alice", "srv", "alice-token", 60, identity_binding_proof=proof)
|
||||
assert await module.resolve_user_oauth_access_token("alice", server) == "alice-token"
|
||||
assert await module.resolve_user_oauth_access_token("alice", server) == "alice-token"
|
||||
read.assert_not_awaited()
|
||||
server.oauth_identity_binding = server.oauth_identity_binding.model_copy(update={"audiences": ["changed"]})
|
||||
assert await module.resolve_user_oauth_access_token("alice", server) is None
|
||||
read.assert_awaited_once()
|
||||
assert await mcp_per_user_token_cache.get_token("alice", "srv") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unverified_legacy_cache_cannot_bypass_enforcement(monkeypatch):
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._experimental.mcp_server import db as module
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import mcp_per_user_token_cache
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
server = MCPServer(
|
||||
server_id="srv",
|
||||
name="srv",
|
||||
transport="http",
|
||||
auth_type="oauth2",
|
||||
oauth_identity_binding={"mode": "enforce", "issuer": "https://idp.example", "audiences": ["client"]},
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
|
||||
monkeypatch.setattr(module, "get_user_oauth_credential", AsyncMock(return_value={"access_token": "bob"}))
|
||||
await mcp_per_user_token_cache.set("alice", "srv", "bob", 60)
|
||||
assert await module.resolve_user_oauth_access_token("alice", server) is None
|
||||
assert await mcp_per_user_token_cache.get("alice", "srv") is None
|
||||
|
|
|
|||
|
|
@ -13,14 +13,14 @@ from pydantic import ValidationError
|
|||
|
||||
from litellm.proxy._experimental.mcp_server.oauth_identity_binding import (
|
||||
RefreshOwnershipProven,
|
||||
VerifiedRefreshToken,
|
||||
RefreshTokenPresented,
|
||||
current_binding_proof,
|
||||
VerifiedRefreshToken,
|
||||
_discover_jwks_url,
|
||||
_fetch_issuer_jwks,
|
||||
_load_caller_principal,
|
||||
_load_stored_refresh_token,
|
||||
_select_signing_key,
|
||||
current_binding_proof,
|
||||
enforce_oauth_identity_binding,
|
||||
)
|
||||
from litellm.types.mcp import MCPAuth, MCPTransport
|
||||
|
|
@ -471,19 +471,20 @@ async def test_refresh_with_mismatched_id_token_rejected():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_mode_logs_but_does_not_reject():
|
||||
async def test_audit_mode_logs_but_does_not_reject(caplog):
|
||||
token: Final = _sign_id_token({"email": "mallory@example.com", "email_verified": True})
|
||||
result: Final = await enforce_oauth_identity_binding(
|
||||
server=_server(mode="audit"),
|
||||
token_response={"access_token": "at", "id_token": token},
|
||||
litellm_user_id="user-a",
|
||||
grant_type="authorization_code",
|
||||
expected_nonce="test-nonce",
|
||||
refresh_ownership=None,
|
||||
jwks_fetcher=_jwks_fetcher,
|
||||
caller_principal_loader=_caller_loader("alice@example.com"),
|
||||
)
|
||||
assert result is None
|
||||
assert "oauth_principal_mismatch" in caplog.text
|
||||
assert "nonce" not in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -642,6 +643,25 @@ async def test_binding_proof_rejects_changed_user_or_policy():
|
|||
def test_identity_binding_rejects_modes_without_gateway_credential_custody(auth_type):
|
||||
with pytest.raises(ValidationError, match="gateway-managed per-user"):
|
||||
MCPServer(
|
||||
server_id="srv", name="srv", transport=MCPTransport.http, auth_type=auth_type,
|
||||
server_id="srv",
|
||||
name="srv",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=auth_type,
|
||||
oauth_identity_binding=_server().oauth_identity_binding,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_matching_login_without_nonce_does_not_report_failure(caplog):
|
||||
token: Final = _sign_id_token({"email": "alice@example.com", "email_verified": True})
|
||||
result: Final = await enforce_oauth_identity_binding(
|
||||
server=_server(mode="audit"),
|
||||
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 result is None
|
||||
assert "oauth_identity_binding audit" not in caplog.text
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue