mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
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:
parent
a9e049d33b
commit
6f8a42d517
3 changed files with 162 additions and 34 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"})
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue