scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 75

fix(mcp): enforce OAuth write policy across signed callbacks

Browse source
This commit is contained in:
Joshua Valluru 2026-09-15 19:15:34 -07:00
parent 97211bc356
commit c7e4160ee6
4 changed files with 181 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

54
litellm/proxy/_experimental/mcp_server/bridge_token_flow.py
View file

@ -306,18 +306,8 @@ async def _revalidate_active_subject(identity: "EnvelopeIdentity") -> "_KeyResol
async def _extract_user_id_from_request(request: Request, server_id: str | None = None) -> str | None:
"""Resolve identity for binding, or authorize the credential-write action for a target server."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # registry imports auth helpers
global_mcp_server_manager,
)
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
can_access_mcp_server, # noqa: PLC0415 # proxy import cycle
)
from litellm.proxy._types import UserAPIKeyAuth # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.handle_jwt import JWTHandler # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.route_checks import RouteChecks # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.user_api_key_auth import ( # noqa: PLC0415 # proxy import cycle
_run_centralized_common_checks, # pyright: ignore[reportPrivateUsage] # reuse admission policy for the credential-write action
)
token: Final = _litellm_key_from_request(request)
# The OAuth relay is public; the optional server-side write is the same protected action
@ -331,23 +321,39 @@ async def _extract_user_id_from_request(request: Request, server_id: str | None
auth: Final = resolved.key if isinstance(resolved, _ResolvedKey) else resolved
if not isinstance(auth, UserAPIKeyAuth) or not _active_key_user_id(auth):
return None
if write_route is not None and server_id is not None:
try:
RouteChecks.is_virtual_key_allowed_to_call_route(route=write_route, valid_token=auth, request=request)
await _run_centralized_common_checks(
user_api_key_auth_obj=auth,
request=request,
request_data={},
route=write_route,
)
if not await can_access_mcp_server(auth, server_id, global_mcp_server_manager.get_allowed_mcp_servers):
return None
except Exception as exc: # noqa: BLE001 # authorization failure must never write credentials
verbose_logger.debug("OAuth credential write not authorized (%s)", type(exc).__name__)
return None
if server_id is not None and not await can_store_oauth_credential(request, auth, server_id):
return None
return auth.user_id
async def can_store_oauth_credential(request: Request, auth: "UserAPIKeyAuth", server_id: str) -> bool:
"""Apply the same write policy to request credentials and verified signed-callback users."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # registry imports auth helpers
global_mcp_server_manager,
)
from litellm.proxy._experimental.mcp_server.ui_session_utils import (
can_access_mcp_server, # noqa: PLC0415 # proxy import cycle
)
from litellm.proxy.auth.route_checks import RouteChecks # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.user_api_key_auth import ( # noqa: PLC0415 # proxy import cycle
_run_centralized_common_checks, # pyright: ignore[reportPrivateUsage] # reuse admission policy for the credential-write action
)
write_route: Final = f"/v1/mcp/server/{server_id}/oauth-user-credential"
try:
RouteChecks.is_virtual_key_allowed_to_call_route(route=write_route, valid_token=auth, request=request)
await _run_centralized_common_checks(
user_api_key_auth_obj=auth,
request=request,
request_data={},
route=write_route,
)
return await can_access_mcp_server(auth, server_id, global_mcp_server_manager.get_allowed_mcp_servers)
except Exception as exc: # noqa: BLE001 # authorization failure must never write credentials
verbose_logger.debug("OAuth credential write not authorized (%s)", type(exc).__name__)
return False
async def _resolve_jwt_auth(
request: Request,
token: str,

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

@ -29,9 +29,11 @@ from litellm.proxy._experimental.mcp_server.bridge_token_flow import (
_BridgeRefreshReady,
_extract_user_id_from_request,
_finish_bridge_mint,
_litellm_key_from_request, # pyright: ignore[reportPrivateUsage] # shared credential precedence for authorization issuance
_prepare_bridge_mint,
_prepare_bridge_refresh,
_reload_active_user_by_id,
can_store_oauth_credential,
)
from litellm.proxy._experimental.mcp_server.faults import (
CallerRejected,
@ -836,16 +838,29 @@ async def _user_can_reach_mcp_server(user_id: str, server_id: str) -> bool:
return server_id in await global_mcp_server_manager.get_allowed_mcp_servers(admitted)
async def _bridge_authorize_access_denial(
litellm_user_id: str,
async def _resolve_oauth_authorization_user(
request: Request,
mcp_server: MCPServer,
redirect_uri: str,
state: str,
) -> RedirectResponse | None:
"""The denial redirect for a signed-in user who cannot reach the target server, or None to proceed."""
if await _user_can_reach_mcp_server(litellm_user_id, mcp_server.server_id):
return None
return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
enforce_binding: bool,
) -> str | RedirectResponse:
"""Resolve the authorization subject without replacing denied credentials with cookie grants."""
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # proxy import cycle
_user_id_from_session_cookie,
)
request_user_id: Final = (
await _extract_user_id_from_request(request, mcp_server.server_id) if enforce_binding else None
)
if enforce_binding and request_user_id is None and _litellm_key_from_request(request):
return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
user_id: Final = request_user_id or _user_id_from_session_cookie(request)
if user_id is None:
return _redirect_to_litellm_login(request)
if not await _user_can_reach_mcp_server(user_id, mcp_server.server_id):
return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
return user_id
async def authorize_with_server(
@ -911,23 +926,12 @@ async def authorize_with_server(
# Seal the authenticated caller into state so the token exchange cannot select another credential owner.
litellm_user_id: str | None = None
if enforce_binding or (resolved_server.is_dcr_bridge and resolved_server.is_oauth_delegate):
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # inline import avoids a module-load circular import
_user_id_from_session_cookie,
subject: Final = await _resolve_oauth_authorization_user(
request, resolved_server, redirect_uri, state, enforce_binding
)
litellm_user_id = (
await _extract_user_id_from_request(request) if enforce_binding else None
) or _user_id_from_session_cookie(request)
if litellm_user_id is None:
return _redirect_to_litellm_login(request)
denial: Final = await _bridge_authorize_access_denial(
litellm_user_id=litellm_user_id,
mcp_server=resolved_server,
redirect_uri=redirect_uri,
state=state,
)
if denial is not None:
return denial
if isinstance(subject, RedirectResponse):
return subject
litellm_user_id = subject
oauth_nonce: Final = secrets.token_urlsafe(32) if enforce_binding else None
encoded_state: Final = encode_state_with_base_url(
@ -1220,8 +1224,14 @@ async def exchange_token_with_server(
try:
# Identity binding above must retain the verified caller even when a write is
# denied. Authorize persistence separately, immediately before its side effect.
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler
# A sealed code delegates a verified user for this authorized server. Raw
# request credentials retain their own JWT/key restrictions during resolution.
can_store: Final = (
await _user_can_reach_mcp_server(user_id, resolved_server.server_id)
await can_store_oauth_credential(
request, await MCPRequestHandler.reload_admitted_user(user_id), resolved_server.server_id
)
if bridge_identity is not None
else await _extract_user_id_from_request(request, resolved_server.server_id) == user_id
)

1
tests/proxy_unit_tests/test_user_api_key_auth.py
View file

@ -1069,6 +1069,7 @@ async def test_jwt_non_admin_team_route_access(monkeypatch):
mock_jwt_response = {
"is_proxy_admin": False,
"jwt_claims": {},
"team_id": None,
"team_object": None,
"user_id": None,

118
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -11171,8 +11171,8 @@ async def test_identity_bound_authorization_carries_nonce_and_caller_through_cal
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._extract_user_id_from_request",
new=AsyncMock(return_value="alice")),
patch( # test-quality-ok: isolate user access lookup while testing nonce and caller preservation
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._bridge_authorize_access_denial",
new=AsyncMock(return_value=None)),
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._user_can_reach_mcp_server",
new=AsyncMock(return_value=True)),
):
authorized = await authorize_with_server(
request, server, "client", "http://127.0.0.1:6274/callback", state="client-state",
@ -11972,3 +11972,117 @@ async def test_oauth_write_denial_does_not_erase_identity_binding(
assert denied.value.status_code == 403
assert denied.value.detail == {"error": "oauth_principal_mismatch"}
manager.get_allowed_mcp_servers.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("admin_only", [False, True])
async def test_signed_oauth_callback_honors_credential_write_policy(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
monkeypatch: pytest.MonkeyPatch,
admin_only: bool,
) -> None:
import httpx
import litellm
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server import discoverable_endpoints, mcp_server_manager
from litellm.proxy._types import MCPTransport
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server: Final = MCPServer(
server_id="signed-server", name="signed-server", transport=MCPTransport.http,
auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="client",
token_url="https://upstream.example.test/token",
)
monkeypatch.setattr(proxy_server, "general_settings", {
"enable_jwt_auth": True,
"admin_only_routes": [f"/v1/mcp/server/{server.server_id}/oauth-user-credential"] if admin_only else [],
})
monkeypatch.setenv("LITELLM_SALT_KEY", "signed-oauth-test-salt")
manager: Final = MagicMock()
manager.get_allowed_mcp_servers = AsyncMock(return_value=[server.server_id])
manager.invalidate_user_oauth_token_cache = AsyncMock()
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
table: Final = proxy_server.prisma_client.db.litellm_mcpusercredentials
table.find_unique = AsyncMock(return_value=None)
table.upsert = AsyncMock()
clients: Final = LLMClientCache()
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", clients)
def upstream_response(outbound: httpx.Request) -> httpx.Response:
assert outbound.url == server.token_url
assert b"code=upstream-code" in outbound.content
return httpx.Response(200, json={"access_token": "upstream-token", "token_type": "Bearer"})
async with httpx.AsyncClient(transport=httpx.MockTransport(upstream_response)) as transport:
upstream: Final = AsyncHTTPHandler()
await upstream.client.aclose()
upstream.client = transport
clients.set_cache("async_httpx_client" + httpxSpecialProvider.Oauth2Check, upstream)
response: Final = await discoverable_endpoints.exchange_token_with_server(
request=_token_request({}, path="/signed-server/token"), mcp_server=server,
grant_type="authorization_code",
code=discoverable_endpoints.seal_bridge_authorization_code("upstream-code", "jwt-owner", server.server_id),
redirect_uri="http://localhost/callback", client_id="client", client_secret=None, code_verifier=None,
)
assert response.status_code == 200
assert json.loads(response.body)["access_token"] == "upstream-token"
if admin_only:
table.upsert.assert_not_awaited()
else:
table.upsert.assert_awaited_once()
assert table.upsert.call_args.kwargs["where"]["user_id_server_id"] == {
"user_id": "jwt-owner", "server_id": server.server_id,
}
@pytest.mark.asyncio
@pytest.mark.parametrize("allowed", [False, True])
async def test_identity_bound_authorize_preserves_presented_jwt_permissions(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
monkeypatch: pytest.MonkeyPatch,
allowed: bool,
) -> None:
from urllib.parse import parse_qs, urlparse
from litellm.proxy._experimental.mcp_server import byok_oauth_endpoints, discoverable_endpoints, mcp_server_manager
from litellm.proxy._types import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
_, signing_key = jwt_oauth_identity
monkeypatch.setenv("LITELLM_SALT_KEY", "authorize-policy-test-salt")
server: Final = MCPServer(
server_id="bound-server", name="bound-server", transport=MCPTransport.http,
auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", client_id="client",
authorization_url="https://upstream.example.test/authorize", token_url="https://upstream.example.test/token",
oauth_identity_binding=MCPOAuthIdentityBinding(
mode="enforce", issuer="https://upstream.example.test", audiences=["client"],
),
)
manager: Final = MagicMock()
# The full user roster permits the server; the presented JWT may have narrower access.
manager.get_allowed_mcp_servers = AsyncMock(
side_effect=lambda auth: [server.server_id] if allowed or auth.mcp_admitted_user_subject else [],
)
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
monkeypatch.setattr( # test-quality-ok: session-cookie decoder is the separate authentication boundary; a valid cookie must not override a denied explicit credential
byok_oauth_endpoints, "_user_id_from_session_cookie", lambda request: "jwt-owner",
)
response: Final = await discoverable_endpoints.authorize_with_server(
request=_token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key)}"}),
mcp_server=server, client_id="client", redirect_uri="http://127.0.0.1:6274/callback",
state="client-state", code_challenge="pkce-challenge", code_challenge_method="S256",
)
redirect: Final = urlparse(response.headers["location"])
query: Final = parse_qs(redirect.query)
if allowed:
assert redirect.hostname == "upstream.example.test"
assert query["nonce"] and response.headers.get("set-cookie")
else:
assert redirect.hostname == "127.0.0.1"
assert query["error"] == ["access_denied"]
assert query["state"] == ["client-state"]
assert "set-cookie" not in response.headers