scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(mcp): preserve browser OAuth for unrelated bearer tokens

Browse source
This commit is contained in:
Joshua Valluru 2026-09-15 23:08:14 -07:00
parent e035682ed1
commit ee676d59f2
3 changed files with 207 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1,6 +1,8 @@
"""Bridge token flow: litellm identity resolution and the DCR-bridge oauth_delegate mint/refresh pipeline."""
import math
import os
import secrets
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Final, Literal
@ -12,6 +14,9 @@ from typing_extensions import assert_never
from litellm._logging import verbose_logger
from litellm.proxy._experimental.mcp_server.oauth_utils import TOKEN_NO_CACHE_HEADERS
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
_V2_GCM_PREFIX, # pyright: ignore[reportPrivateUsage] # reuse the encrypted credential's format discriminator
)
from litellm.types.mcp_server.mcp_server_manager import MCPServer
if TYPE_CHECKING:
@ -49,6 +54,64 @@ def _litellm_key_from_request(request: Request) -> str | None:
return None
async def oauth_authorization_uses_gateway_credential(request: Request) -> bool:
"""Classify credentials for browser authorize; candidates still require full authorization."""
from litellm.proxy.auth.handle_jwt import JWTHandler # noqa: PLC0415 # proxy import cycle
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the active auth configuration
jwt_handler,
master_key,
user_custom_auth,
)
if "x-litellm-api-key" in request.headers:
return True
token: Final = _litellm_key_from_request(request)
if token is None:
return "authorization" in request.headers
if token.startswith("sk-") or (master_key and secrets.compare_digest(token.encode(), master_key.encode())):
return True
if user_custom_auth is not None or jwt_handler.litellm_jwtauth.oidc_userinfo_enabled:
return True
if not JWTHandler.is_jwt(token):
return await _opaque_bearer_is_gateway_credential(token)
claims: Final = JWTHandler.get_unverified_claims(token)
issuer: Final = claims.get("iss") if claims is not None else None
global_issuer: Final = os.getenv("JWT_ISSUER")
# An unscoped global validator can accept issuers absent from the configured issuer list.
if not isinstance(issuer, str) or not issuer or not global_issuer:
return True
return issuer == global_issuer or any(
issuer == configured.issuer for configured in jwt_handler.litellm_jwtauth.issuers or ()
)
async def _opaque_bearer_is_gateway_credential(token: str) -> bool:
from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import (
is_envelope, # noqa: PLC0415 # envelope imports bridge types
is_refresh_envelope,
)
from litellm.proxy._types import hash_token # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.resolvers.exceptions import KeyNotFoundError # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.resolvers.store import IdentityStore # noqa: PLC0415 # proxy import cycle
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the identity store dependencies
prisma_client,
user_api_key_cache,
)
if is_envelope(token) or is_refresh_envelope(token) or token.startswith(_V2_GCM_PREFIX):
return True
try:
if ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(token) is not None:
return True
await IdentityStore(prisma_client, user_api_key_cache).resolve(hashed_token=hash_token(token))
except KeyNotFoundError:
return False
except Exception as exc: # noqa: BLE001 # an identity lookup fault must not permit cookie fallback
verbose_logger.debug("OAuth bearer ownership could not be checked (%s)", type(exc).__name__)
return True
def _key_is_active(key_obj: "UserAPIKeyAuth") -> bool:
"""``True`` when the presented key is neither blocked nor past its expiry.

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

@ -29,12 +29,12 @@ 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,
authorize_oauth_credential_request,
can_store_oauth_credential,
oauth_authorization_uses_gateway_credential,
)
from litellm.proxy._experimental.mcp_server.faults import (
CallerRejected,
@ -851,10 +851,11 @@ async def _resolve_oauth_authorization_user(
_user_id_from_session_cookie,
)
use_gateway_credential: Final = enforce_binding and await oauth_authorization_uses_gateway_credential(request)
request_user_id: Final = (
await authorize_oauth_credential_request(request, mcp_server.server_id) if enforce_binding else None
await authorize_oauth_credential_request(request, mcp_server.server_id) if use_gateway_credential else None
)
if enforce_binding and request_user_id is None and _litellm_key_from_request(request):
if use_gateway_credential and request_user_id is None:
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:

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

@ -11170,7 +11170,7 @@ async def test_identity_bound_authorization_carries_nonce_and_caller_through_cal
),
)
request = Request({"type": "http", "scheme": "https", "server": ("proxy.example.com", 443),
"path": "/authorize", "query_string": b"", "headers": []})
"path": "/authorize", "query_string": b"", "headers": [(b"authorization", b"Bearer sk-alice")]})
with (
patch( # test-quality-ok: isolate authenticated request resolution from the real encrypted OAuth round trip
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.authorize_oauth_credential_request",
@ -12059,18 +12059,52 @@ async def test_signed_oauth_callback_honors_credential_write_policy(
@pytest.mark.asyncio
@pytest.mark.parametrize("allowed", [False, True])
@pytest.mark.parametrize("credential", [
"jwt", "key", "expired_jwt", "wrong_audience", "bad_signature", "malformed_jwt", "missing_issuer",
"foreign_explicit", "blank_explicit", "unknown_key", "blocked_key", "expired_key", "opaque_record",
"opaque_outage", "opaque_oidc", "opaque_custom", "foreign_unscoped", "foreign_configured", "encrypted", "invalid_encrypted", "envelope", "master",
])
async def test_identity_bound_authorize_preserves_presented_jwt_permissions(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
monkeypatch: pytest.MonkeyPatch,
allowed: bool,
credential: str,
) -> None:
import jwt
from datetime import datetime, timedelta, timezone
from urllib.parse import parse_qs, urlparse
from litellm.proxy._experimental.mcp_server import byok_oauth_endpoints, discoverable_endpoints, mcp_server_manager
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy import proxy_server
from litellm.proxy._types import JWTIssuerConfig, UserAPIKeyAuth, hash_token
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
from litellm.proxy._experimental.mcp_server import 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
handler, signing_key = jwt_oauth_identity
master: Final = "browser-session-test-signing-key-123456789"
monkeypatch.setattr(proxy_server, "master_key", master)
monkeypatch.setattr(proxy_server, "user_custom_auth", (lambda: None) if credential == "opaque_custom" else None)
handler.litellm_jwtauth.oidc_userinfo_enabled = credential == "opaque_oidc"
if credential == "foreign_unscoped":
monkeypatch.delenv("JWT_ISSUER")
if credential == "foreign_configured":
handler.litellm_jwtauth.issuers = [JWTIssuerConfig(
issuer="https://unrelated.example.test", jwks_url="https://idp.example.test/jwks",
audience="litellm-proxy", user_id_jwt_field="identity.user_id",
)]
proxy_server.prisma_client.get_data = AsyncMock(
return_value=None, side_effect=RuntimeError("database unavailable") if credential == "opaque_outage" else None,
)
handler.user_api_key_cache.set_cache("cookie-owner", LiteLLM_UserTable(user_id="cookie-owner"))
key: Final = "opaque-record" if credential == "opaque_record" else "sk-browser-gateway-key"
if credential in ("key", "blocked_key", "expired_key", "opaque_record"):
handler.user_api_key_cache.set_cache(hash_token(key), UserAPIKeyAuth(
token=hash_token(key), user_id="jwt-owner", blocked=credential in ("blocked_key", "opaque_record"),
expires=datetime.now(timezone.utc) - timedelta(seconds=60) if credential == "expired_key" else None,
))
monkeypatch.setenv("LITELLM_SALT_KEY", "authorize-policy-test-salt")
server: Final = MCPServer(
server_id="bound-server", name="bound-server", transport=MCPTransport.http,
@ -12086,21 +12120,120 @@ async def test_identity_bound_authorize_preserves_presented_jwt_permissions(
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",
bearer: Final = (
key if credential in ("key", "blocked_key", "expired_key", "opaque_record", "unknown_key")
else "opaque-bearer" if credential in ("opaque_outage", "opaque_oidc", "opaque_custom")
else "not.a.jwt" if credential == "malformed_jwt"
else "llm_env_invalid" if credential == "envelope"
else "v2:gcm:invalid" if credential == "invalid_encrypted"
else master if credential == "master"
else ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(
LiteLLM_UserTable(user_id="jwt-owner", user_role="internal_user"),
) if credential == "encrypted"
else jwt.encode({"iss": "https://idp.example.test"}, "wrong-signing-key-at-least-32-bytes", algorithm="HS256")
if credential == "bad_signature"
else jwt.encode({"sub": "jwt-owner"}, signing_key, algorithm="RS256") if credential == "missing_issuer"
else _oauth_identity_jwt(
signing_key,
expires_in=-60 if credential == "expired_jwt" else 300,
audience="another-service" if credential == "wrong_audience" else "litellm-proxy",
issuer="https://unrelated.example.test" if credential.startswith("foreign_") or credential == "blank_explicit" else "https://idp.example.test",
)
)
cookie: Final = jwt.encode(
{"user_id": "cookie-owner", "login_method": "sso", "exp": int(time.time()) + 300}, master, algorithm="HS256",
)
response: Final = await discoverable_endpoints.authorize_with_server(
request=_token_request({"Authorization": f"Bearer {_oauth_identity_jwt(signing_key)}"}),
request=_token_request({
"Authorization": f"Bearer {bearer}", "Cookie": f"token={cookie}",
**({"x-litellm-api-key": bearer} if credential == "foreign_explicit" else {}),
**({"x-litellm-api-key": ""} if credential == "blank_explicit" else {}),
}),
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:
if allowed and credential in ("jwt", "key", "foreign_unscoped", "foreign_configured"):
assert redirect.hostname == "upstream.example.test"
assert query["nonce"] and response.headers.get("set-cookie")
assert all(call.args[0].user_id == "jwt-owner" for call in manager.get_allowed_mcp_servers.await_args_list)
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
proxy_server.prisma_client.db.litellm_mcpusercredentials.upsert.assert_not_called()
proxy_server.prisma_client.db.litellm_usertable.create.assert_not_called()
proxy_server.prisma_client.db.litellm_teamtable.create.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize("credential", ["none", "opaque", "foreign_jwt"])
@pytest.mark.parametrize("cookie_state", ["allowed", "server_denied", "expired", "missing"])
async def test_identity_bound_authorize_unrelated_bearer_uses_browser_session(
jwt_oauth_identity: tuple["JWTHandler", "RSAPrivateKey"],
monkeypatch: pytest.MonkeyPatch,
credential: str,
cookie_state: str,
) -> None:
import jwt
from urllib.parse import parse_qs, urlparse
from litellm.models.user import LiteLLM_UserTable
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.mcp_server.mcp_server_manager import MCPOAuthIdentityBinding, MCPServer
handler, signing_key = jwt_oauth_identity
master: Final = "browser-session-test-signing-key-123456789"
monkeypatch.setattr(proxy_server, "master_key", master)
monkeypatch.setattr(proxy_server, "user_custom_auth", None)
monkeypatch.setenv("LITELLM_SALT_KEY", "authorize-policy-test-salt")
handler.user_api_key_cache.set_cache("cookie-owner", LiteLLM_UserTable(user_id="cookie-owner"))
proxy_server.prisma_client.get_data = AsyncMock(return_value=None)
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()
manager.get_allowed_mcp_servers = AsyncMock(return_value=[] if cookie_state == "server_denied" else [server.server_id])
monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager)
bearer: Final = (
_oauth_identity_jwt(signing_key, issuer="https://unrelated.example.test")
if credential == "foreign_jwt" else "unrelated-upstream-bearer"
)
cookie: Final = jwt.encode(
{"user_id": "cookie-owner", "login_method": "sso", "exp": int(time.time()) + (-60 if cookie_state == "expired" else 300)},
master, algorithm="HS256",
)
response: Final = await discoverable_endpoints.authorize_with_server(
request=_token_request({
**({"Authorization": f"Bearer {bearer}"} if credential != "none" else {}),
**({"Cookie": f"token={cookie}"} if cookie_state != "missing" else {}),
}),
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 cookie_state == "allowed":
assert redirect.hostname == "upstream.example.test"
assert query["nonce"] and response.headers.get("set-cookie")
manager.get_allowed_mcp_servers.assert_awaited_once()
assert manager.get_allowed_mcp_servers.call_args.args[0].user_id == "cookie-owner"
elif cookie_state == "server_denied":
assert query["error"] == ["access_denied"]
assert query["state"] == ["client-state"]
else:
assert redirect.path == "/sso/key/generate"
manager.get_allowed_mcp_servers.assert_not_awaited()
proxy_server.prisma_client.db.litellm_mcpusercredentials.upsert.assert_not_called()
proxy_server.prisma_client.db.litellm_usertable.create.assert_not_called()
proxy_server.prisma_client.db.litellm_teamtable.create.assert_not_called()