fix(mcp): vouch for delegate callers with OAuth2 introspection and proxy headers

On delegate servers, a bearer that isn't LiteLLM's own was only tried as an identity through custom auth or JWT auth, so proxies using enable_oauth2_auth or enable_oauth2_proxy_auth admitted those callers anonymously and lost their budgets and attribution

The identity step now picks the check LiteLLM's own auth runs for that token and route, in its order (OAuth2 introspection, then reverse-proxy headers, then JWT), and follows the JWT rules. A rejected identity is anonymous, a refusal after the caller is vouched for fails, and a granted user is admitted by name with the token still forwarded. The route rule for OAuth2 introspection moves into oauth2_introspection_applies so both paths share it
This commit is contained in:
shrey kharbanda 2026-10-28 15:52:39 -07:00
parent a9e049d33b
commit 6f8a42d517
3 changed files with 162 additions and 34 deletions

View file

@ -1102,6 +1102,24 @@ def _matches_routing_override(token_claims: dict, override: "JWTRoutingOverride"
)
def oauth2_introspection_applies(
api_key: str, route: str, *, enable_oauth2_auth: bool, enable_jwt_auth: bool, jwt_handler: JWTHandler
) -> bool:
"""Whether LiteLLM validates ``api_key`` on ``route`` by OAuth2 token introspection, on LLM API, info and
lens routes only: a non-JWT token when OAuth2 auth is on globally, or a JWT matched by a routing override."""
is_jwt: Final = enable_jwt_auth and jwt_handler.is_jwt(token=api_key)
# Routing uses unverified JWT claims only to choose auth path.
# Final authentication is enforced by the selected validator.
route_jwt_to_oauth2: Final = is_jwt and _should_route_jwt_to_oauth2_override(token=api_key, jwt_handler=jwt_handler)
covered_route: Final = (
RouteChecks.is_llm_api_route(route=route)
or RouteChecks.is_info_route(route=route)
or route == "/lens"
or route.startswith("/lens/")
)
return covered_route and ((enable_oauth2_auth and not is_jwt) or route_jwt_to_oauth2)
def _should_route_jwt_to_oauth2_override(token: str, jwt_handler: JWTHandler) -> bool:
routing_overrides: Final = jwt_handler.litellm_jwtauth.routing_overrides
if not routing_overrides:
@ -2216,23 +2234,13 @@ async def user_api_key_auth_builder(
enable_oauth2_auth: Final = general_settings.get("enable_oauth2_auth", False) is True
enable_jwt_auth: Final = general_settings.get("enable_jwt_auth", False) is True
is_jwt = jwt_handler.is_jwt(token=api_key) if enable_jwt_auth else False
# Routing uses unverified JWT claims only to choose auth path.
# Final authentication is enforced by the selected validator.
route_jwt_to_oauth2 = is_jwt and _should_route_jwt_to_oauth2_override(token=api_key, jwt_handler=jwt_handler)
# OAuth2 applies for:
# 1) when global OAuth2 auth is enabled on LLM + info routes
# 2) JWT tokens that explicitly match routing_overrides on LLM + info routes
route_is_lens: Final = route == "/lens" or route.startswith("/lens/")
should_apply_override_oauth2: Final = route_jwt_to_oauth2 and (
RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_info_route(route=route) or route_is_lens
)
should_apply_global_oauth2: Final = enable_oauth2_auth and (
RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_info_route(route=route) or route_is_lens
)
if (should_apply_global_oauth2 and not is_jwt) or should_apply_override_oauth2:
if oauth2_introspection_applies(
api_key,
route,
enable_oauth2_auth=enable_oauth2_auth,
enable_jwt_auth=enable_jwt_auth,
jwt_handler=jwt_handler,
):
from litellm.proxy.proxy_server import premium_user
if premium_user is not True:

View file

@ -8,7 +8,7 @@ from dataclasses import dataclass
from datetime import datetime, timezone
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from fastapi import HTTPException
from pydantic import SecretStr, TypeAdapter
@ -48,6 +48,7 @@ from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports
_get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth
_run_centralized_common_checks, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
oauth2_introspection_applies,
run_centralized_common_checks,
user_api_key_auth,
)
@ -102,6 +103,7 @@ _EMPTY_TOOLSET_GRANTS: Final[Mapping[str, Sequence[str]]] = MappingProxyType({})
OPTIONAL_STRING_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None)
_PROXY_SETTINGS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_CAUSE_CHAIN_MAX_DEPTH: Final = 16
_IdentityCheck: TypeAlias = Literal["oauth2", "oauth2_proxy", "jwt"]
_CACHED_VALUE: Final = TypeAdapter(object)
@ -266,38 +268,78 @@ async def _bearer_is_litellm_credential(token: str) -> bool:
return token.startswith(LITELLM_VIRTUAL_KEY_PREFIX) and await virtual_key_may_exist(token)
async def _upstream_bearer_identity(token: str, litellm_api_key: str, request: Request) -> UserAPIKeyAuth | None:
"""The LiteLLM user an upstream bearer also proves, or ``None`` when it proves none.
Custom auth decides alone, so any rejection is ``None``. A JWT is tried only when JWT auth is on, its
issuer may be LiteLLM's, and its signature verifies. Past that, only a user LiteLLM does not know is
``None``; every other refusal (budget, blocked team, disallowed route) propagates."""
def _identity_check(token: str, route: str, settings: Mapping[str, object]) -> _IdentityCheck | None:
"""The identity check LiteLLM's own auth would run on ``token`` for ``route``, in its order: OAuth2
introspection, then trusted reverse-proxy identity headers, then a JWT whose issuer may be LiteLLM's."""
from litellm.proxy.auth.handle_jwt import JWTHandler # noqa: PLC0415 # proxy import cycle
from litellm.proxy.mcp_server.bridge_token_flow import ( # noqa: PLC0415 # bridge flow imports this module
custom_auth_configured,
jwt_issuer_may_be_litellm,
)
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the active auth configuration
general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict
jwt_handler,
premium_user,
)
enable_jwt_auth: Final = settings.get("enable_jwt_auth") is True
if oauth2_introspection_applies(
token,
route,
enable_oauth2_auth=settings.get("enable_oauth2_auth") is True,
enable_jwt_auth=enable_jwt_auth,
jwt_handler=jwt_handler,
):
return "oauth2" if premium_user is True else None
if settings.get("enable_oauth2_proxy_auth") is True:
return "oauth2_proxy"
if enable_jwt_auth and premium_user is True and JWTHandler.is_jwt(token=token) and jwt_issuer_may_be_litellm(token):
return "jwt"
return None
async def _verify_identity(check: _IdentityCheck, token: str, request: Request) -> None:
from litellm.proxy.auth.oauth2_check import Oauth2Handler # noqa: PLC0415 # proxy import cycle
from litellm.proxy.auth.oauth2_proxy_hook import ( # noqa: PLC0415 # proxy import cycle
handle_oauth2_proxy_request,
)
from litellm.proxy.proxy_server import jwt_handler # noqa: PLC0415 # startup owns the active auth configuration
match check:
case "oauth2":
await Oauth2Handler.check_oauth2_token(token=token)
case "oauth2_proxy":
await handle_oauth2_proxy_request(request=request)
case "jwt":
await jwt_handler.auth_jwt(token=token) # pyright: ignore[reportUnknownMemberType] # returns a bare dict; only whether it raises matters
case _:
assert_never(check)
async def _upstream_bearer_identity(token: str, litellm_api_key: str, request: Request) -> UserAPIKeyAuth | None:
"""The LiteLLM user an upstream bearer also proves, or ``None`` when it proves none.
Custom auth decides alone, so any rejection is ``None``. Otherwise the identity check LiteLLM would run
(OAuth2 introspection, reverse-proxy headers, or a JWT) must pass first. Past that, only a user LiteLLM
does not know is ``None``; every other refusal (budget, blocked team, disallowed route) propagates."""
from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415 # proxy import cycle
from litellm.proxy.mcp_server.bridge_token_flow import ( # noqa: PLC0415 # bridge flow imports this module
custom_auth_configured,
)
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # startup owns the active auth configuration
general_settings, # pyright: ignore[reportUnknownVariableType] # legacy global is declared as a bare dict
)
if custom_auth_configured():
try:
return await user_api_key_auth(api_key=litellm_api_key, request=request)
except (HTTPException, ProxyException):
return None
settings: Final = _PROXY_SETTINGS_ADAPTER.validate_python(general_settings)
if (
settings.get("enable_jwt_auth") is not True
or premium_user is not True
or not JWTHandler.is_jwt(token=token)
or not jwt_issuer_may_be_litellm(token)
):
check: Final = _identity_check(
token, get_request_route(request), _PROXY_SETTINGS_ADAPTER.validate_python(general_settings)
)
if check is None:
return None
try:
await jwt_handler.auth_jwt(token=token) # pyright: ignore[reportUnknownMemberType] # returns a bare dict; only whether it raises matters
await _verify_identity(check, token, request)
except Exception: # noqa: BLE001 # any verification failure means LiteLLM does not vouch for this bearer
return None
try:

View file

@ -2826,6 +2826,84 @@ class TestMCPDelegateAuthToUpstream:
assert exc_info.value is refusal
@staticmethod
@contextlib.contextmanager
def _idp_auth_enabled(mechanism: str, verify: AsyncMock, premium: bool = True):
from litellm.proxy import proxy_server
from litellm.proxy.auth.oauth2_check import Oauth2Handler
settings = {"oauth2": {"enable_oauth2_auth": True}, "oauth2_proxy": {"enable_oauth2_proxy_auth": True}}
verifier = (
patch.object(Oauth2Handler, "check_oauth2_token", verify)
if mechanism == "oauth2"
else patch("litellm.proxy.auth.oauth2_proxy_hook.handle_oauth2_proxy_request", verify)
)
with (
patch.dict(proxy_server.general_settings, settings[mechanism]),
patch.object(proxy_server, "premium_user", premium),
patch.object(proxy_server, "user_custom_auth", None),
verifier,
):
yield
@pytest.mark.parametrize("mechanism", ["oauth2", "oauth2_proxy"])
@pytest.mark.parametrize(
"verify_effect, allowed_server_ids, expected_user_id",
[
(None, ["test-server-id"], "idp-user"),
(None, [], None),
(ValueError("Token is not active"), ["test-server-id"], None),
],
ids=["granted-user", "user-not-granted", "identity-rejected"],
)
async def test_delegate_idp_identity_follows_the_jwt_rules(
self, mechanism, verify_effect, allowed_server_ids, expected_user_id
):
"""OAuth2 introspection and trusted reverse-proxy headers admit a granted user by name, and a user
who is not granted or not vouched for anonymously. The bearer is forwarded either way."""
admission = AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user"))
with (
TestMCPDelegateAuthToUpstream._idp_auth_enabled(mechanism, AsyncMock(side_effect=verify_effect)),
TestMCPDelegateAuthToUpstream._delegate_admission(admission, allowed_server_ids),
):
auth_result, _, _, _, oauth2_headers, _ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token")
)
assert auth_result.user_id == expected_user_id
assert auth_result.mcp_forwards_caller_authorization is True
assert oauth2_headers == {"Authorization": "Bearer opaque-idp-token"}
@pytest.mark.parametrize("mechanism", ["oauth2", "oauth2_proxy"])
async def test_delegate_idp_user_refusal_propagates(self, mechanism):
"""Once the provider vouches for the caller, LiteLLM's budget and permission refusals apply."""
refusal = ProxyException(message="Budget has been exceeded", type="budget_exceeded", param=None, code=400)
with (
TestMCPDelegateAuthToUpstream._idp_auth_enabled(mechanism, AsyncMock()),
TestMCPDelegateAuthToUpstream._delegate_admission(AsyncMock(side_effect=refusal)),
):
with pytest.raises(ProxyException) as exc_info:
await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token")
)
assert exc_info.value is refusal
async def test_delegate_oauth2_without_premium_never_calls_the_provider(self):
"""OAuth2 introspection is enterprise-only, so without a license the bearer is the upstream's alone."""
introspect = AsyncMock()
with (
TestMCPDelegateAuthToUpstream._idp_auth_enabled("oauth2", introspect, premium=False),
TestMCPDelegateAuthToUpstream._delegate_admission(AsyncMock(return_value=UserAPIKeyAuth(user_id="idp-user"))),
):
auth_result, *_ = await MCPRequestHandler.process_mcp_request(
TestMCPDelegateAuthToUpstream._bearer_scope("opaque-idp-token")
)
assert auth_result.user_id is None
assert auth_result.mcp_forwards_caller_authorization is True
introspect.assert_not_awaited()
async def test_delegate_jwt_from_unconfigured_issuer_is_forwarded_without_jwt_auth(self):
"""A JWT from an issuer LiteLLM's JWT auth is not configured for never reaches JWT auth."""
auth_jwt = AsyncMock(return_value={"sub": "jwt-user"})