diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 92f43e782d0..7639aa37266 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -126,6 +126,19 @@ def decode_state_hash(encrypted_state: str) -> dict: return state_data +def _get_validated_client_redirect_uri( + request: Request, state_data: Dict[str, Any] +) -> str: + """Return a trusted (same-origin, loopback, or ops-allowlisted) + client redirect URI from OAuth state. + """ + redirect_uri = state_data.get("client_redirect_uri") or state_data.get("base_url") + if not redirect_uri or not isinstance(redirect_uri, str): + raise HTTPException(status_code=400, detail="Invalid redirect URI") + validate_trusted_redirect_uri(request, redirect_uri) + return redirect_uri + + def _append_query_params(url: str, params: Dict[str, str]) -> str: parsed = urlparse(url) query_params = parse_qsl(parsed.query, keep_blank_values=True) @@ -677,11 +690,6 @@ async def token_endpoint( async def callback(request: Request, code: str, state: str): try: state_data = decode_state_hash(state) - redirect_uri = state_data.get("client_redirect_uri") or state_data.get( - "base_url" - ) - if not redirect_uri or not isinstance(redirect_uri, str): - raise HTTPException(status_code=400, detail="Invalid redirect URI") original_state = state_data["original_state"] # Re-validate the client redirect URI at the sink. /authorize @@ -690,7 +698,7 @@ async def callback(request: Request, code: str, state: str): # expiry and remain valid indefinitely. Validating here blocks # the open-redirect + code-theft primitive even for pre-fix # states while permitting same-origin / allowlisted clients. - validate_trusted_redirect_uri(request, redirect_uri) + redirect_uri = _get_validated_client_redirect_uri(request, state_data) params = {"code": code, "state": original_state} complete_returned_url = _append_query_params(redirect_uri, params) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index f34bfb5f399..aa86cb2ba7d 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2835,6 +2835,30 @@ if MCP_AVAILABLE: ) return user_api_key_auth.model_copy(update={"object_permission": updated_op}) + def _get_passthrough_resource_metadata_url(scope: Scope, server_name: str) -> str: + request = StarletteRequest(scope) + base_url = get_request_base_url(request) + _path = scope.get("_original_path") or scope.get("path", "") or "" + + if _path.startswith(f"/{server_name}/mcp"): + return f"{base_url}/.well-known/oauth-protected-resource/{server_name}/mcp" + return f"{base_url}/.well-known/oauth-protected-resource/mcp/{server_name}" + + def _get_passthrough_www_authenticate( + scope: Scope, + server_name: str, + invalid_token: bool = False, + ) -> str: + resource_metadata_url = _get_passthrough_resource_metadata_url( + scope=scope, + server_name=server_name, + ) + params = [] + if invalid_token: + params.append('error="invalid_token"') + params.append(f'resource_metadata="{resource_metadata_url}"') + return "Bearer " + ", ".join(params) + async def _raise_preemptive_401_for_unauthenticated_servers( scope: Scope, mcp_servers: Optional[List[str]], @@ -2895,18 +2919,10 @@ if MCP_AVAILABLE: server, oauth2_headers, mcp_server_auth_headers ) ): - request = StarletteRequest(scope) - base_url = get_request_base_url(request) - _path = scope.get("_original_path") or scope.get("path", "") or "" - - # Pick the well-known resource-metadata form that matches the inbound - # route pattern so the metadata's `resource` field round-trips to what - # the client actually hit (RFC 9728 §3.2). - if _path.startswith(f"/{server_name}/mcp"): - resource_metadata_url = f"{base_url}/.well-known/oauth-protected-resource/{server_name}/mcp" - else: - resource_metadata_url = f"{base_url}/.well-known/oauth-protected-resource/mcp/{server_name}" - www_authenticate = f'Bearer resource_metadata="{resource_metadata_url}"' + www_authenticate = _get_passthrough_www_authenticate( + scope=scope, + server_name=server_name, + ) raise HTTPException( status_code=401, detail="Unauthorized", @@ -3045,30 +3061,16 @@ if MCP_AVAILABLE: for srv in passthrough_servers ] ) - request = StarletteRequest(scope) - base_url = get_request_base_url(request) - _path = scope.get("_original_path") or scope.get("path", "") or "" for srv, (probe_status, _) in zip(passthrough_servers, probe_results): if probe_status == 401: - # Token is missing or expired — direct the client at the - # gateway's oauth-protected-resource well-known URL (which - # proxies the upstream IdP's metadata), matching the - # pre-emptive 401 path in - # _raise_preemptive_401_for_unauthenticated_servers. The - # gateway is not the authorization server for pass-through - # servers, so emitting ``authorization_uri=`` would point - # clients at the wrong AS metadata. - if _path.startswith(f"/{srv.name}/mcp"): - resource_metadata_url = ( - f"{base_url}/.well-known/oauth-protected-resource/" - f"{srv.name}/mcp" - ) - else: - resource_metadata_url = ( - f"{base_url}/.well-known/oauth-protected-resource/mcp/" - f"{srv.name}" - ) - www_authenticate = f'Bearer resource_metadata="{resource_metadata_url}"' + # Token is missing or expired: keep pass-through clients on the + # protected-resource discovery flow so they re-authorize against + # the upstream IdP metadata proxied by LiteLLM. + www_authenticate = _get_passthrough_www_authenticate( + scope=scope, + server_name=srv.name, + invalid_token=True, + ) raise HTTPException( status_code=401, detail="Unauthorized", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 23f2867152a..d7ed315efa9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15858,6 +15858,7 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): ) if toolset is not None: scope = dict(request.scope) + scope["_original_path"] = scope.get("path", "") scope["path"] = "/mcp" token = _mcp_active_toolset_id.set(toolset.toolset_id) try: diff --git a/tests/test_litellm/interactions/test_openapi_compliance.py b/tests/test_litellm/interactions/test_openapi_compliance.py index 8bd887c939f..aededaaca77 100644 --- a/tests/test_litellm/interactions/test_openapi_compliance.py +++ b/tests/test_litellm/interactions/test_openapi_compliance.py @@ -166,6 +166,7 @@ class TestResponseCompliance: "created", "updated", "role", + "steps", "usage", ] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 241ff89b7e6..82c1d97faa9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -1,7 +1,7 @@ import json import os import sys -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call as mock_call, patch import pytest from fastapi.testclient import TestClient diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index bb77b9e9a9c..97c67e10ee1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -5,23 +5,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException -TRUSTED_PROXY_IP = "10.0.0.5" -TRUSTED_PROXY_RANGES = ["10.0.0.0/8"] - - -def set_request_from_trusted_proxy(mock_request): - mock_request.client = MagicMock() - mock_request.client.host = TRUSTED_PROXY_IP - - -@pytest.fixture -def trusted_proxy_origin_headers(): - with patch( - "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", - return_value=True, - ): - yield - # Fixture to mock IP address check for all MCP tests # This prevents tests from failing due to IP-based access control @@ -63,7 +46,7 @@ def trust_xff(): ``test_get_request_base_url_xff_trust_gate``. """ with patch( - "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy", return_value=True, ): yield @@ -592,7 +575,6 @@ async def test_authorize_endpoint_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) # Mock the encryption functions with patch( @@ -661,7 +643,6 @@ async def test_token_endpoint_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm-proxy.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) # Mock httpx client response mock_response = MagicMock() @@ -740,7 +721,6 @@ async def test_oauth_protected_resource_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) # Call the endpoint response = await oauth_protected_resource_mcp( @@ -796,7 +776,6 @@ async def test_oauth_authorization_server_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) # Call the endpoint response = await oauth_authorization_server_mcp( @@ -835,7 +814,6 @@ async def test_register_client_respects_x_forwarded_proto(): mock_request = MagicMock(spec=Request) mock_request.base_url = "http://proxy.litellm.example/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy - set_request_from_trusted_proxy(mock_request) with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", @@ -898,7 +876,6 @@ async def test_authorize_endpoint_respects_x_forwarded_host(): "X-Forwarded-Proto": "https", "X-Forwarded-Host": "proxy.example.com", } - set_request_from_trusted_proxy(mock_request) # Mock the encryption functions with patch( @@ -971,7 +948,6 @@ async def test_token_endpoint_respects_x_forwarded_host(): "X-Forwarded-Proto": "https", "X-Forwarded-Host": "proxy.example.com", } - set_request_from_trusted_proxy(mock_request) # Mock httpx client response mock_response = MagicMock() @@ -1129,11 +1105,7 @@ async def test_token_endpoint_respects_x_forwarded_host(): ], ) def test_get_request_base_url_comprehensive( - base_url, - x_forwarded_proto, - x_forwarded_host, - x_forwarded_port, - expected_url, + base_url, x_forwarded_proto, x_forwarded_host, x_forwarded_port, expected_url ): """Comprehensive test for get_request_base_url with various header combinations. @@ -1152,7 +1124,6 @@ def test_get_request_base_url_comprehensive( mock_request = MagicMock(spec=Request) mock_request.base_url = base_url - set_request_from_trusted_proxy(mock_request) headers = {} if x_forwarded_proto: @@ -1168,7 +1139,7 @@ def test_get_request_base_url_comprehensive( mock_request.headers.get = mock_get with patch( - "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy", return_value=True, ): result = get_request_base_url(mock_request) @@ -2053,10 +2024,6 @@ async def test_oauth_callback_redirects_with_state(): except ImportError: pytest.skip("MCP discoverable endpoints not available") - mock_request = MagicMock() - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - # Mock the state decoding mock_state_data = { "base_url": "http://localhost:3000/ui/mcp/oauth/callback", @@ -2103,10 +2070,6 @@ async def test_oauth_callback_preserves_client_redirect_uri_query(): except ImportError: pytest.skip("MCP discoverable endpoints not available") - mock_request = MagicMock() - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" ) as mock_decode: @@ -2144,10 +2107,6 @@ async def test_oauth_callback_handles_invalid_state(): except ImportError: pytest.skip("MCP discoverable endpoints not available") - mock_request = MagicMock() - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - # Mock state decoding to raise an exception with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" @@ -2570,10 +2529,6 @@ async def test_callback_revalidates_loopback_on_decoded_base_url(): callback, ) - mock_request = MagicMock() - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" ) as mock_decode: @@ -2600,10 +2555,6 @@ async def test_callback_revalidates_loopback_on_decoded_client_redirect_uri(): callback, ) - mock_request = MagicMock() - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" ) as mock_decode: @@ -2630,10 +2581,6 @@ async def test_callback_rejects_state_missing_redirect_uri(): callback, ) - mock_request = MagicMock() - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" ) as mock_decode: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 83a5fea8758..1cea264f73d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -109,11 +109,7 @@ async def test_mcp_server_tool_call_body_contains_request_data(): assert body["arguments"] == tool_arguments -def test_prepare_mcp_server_headers_passthrough_omits_authorization(): - """Pass-through servers must not forward the inbound Authorization header - (LiteLLM API key) to the upstream server. The upstream must provide its own - bearer token via a server-specific header. - """ +def test_prepare_mcp_server_headers_case_insensitive_extra_headers(): try: from litellm.proxy._experimental.mcp_server.server import ( _prepare_mcp_server_headers, @@ -122,82 +118,9 @@ def test_prepare_mcp_server_headers_passthrough_omits_authorization(): pytest.skip("MCP server not available") server = MCPServer( - server_id="server-passthrough", + server_id="server-case", name="server", transport=MCPTransport.http, - auth_type=MCPAuth.none, # Explicitly none for pass-through - extra_headers=["Authorization"], - ) - - server_auth_header, extra_headers = _prepare_mcp_server_headers( - server=server, - mcp_server_auth_headers=None, - mcp_auth_header=None, - oauth2_headers=None, - raw_headers={"authorization": "Bearer sk-litellm-key"}, - ) - - assert server_auth_header is None - # Authorization should NOT be forwarded for pass-through servers - assert extra_headers is None - - -def test_prepare_mcp_server_headers_passthrough_forwards_other_headers(): - """Pass-through servers should forward other headers (not Authorization) - from the raw request.""" - try: - from litellm.proxy._experimental.mcp_server.server import ( - _prepare_mcp_server_headers, - ) - except ImportError: - pytest.skip("MCP server not available") - - server = MCPServer( - server_id="server-passthrough-headers", - name="server", - transport=MCPTransport.http, - auth_type=MCPAuth.none, # Pass-through mode - extra_headers=["Authorization", "x-request-id", "x-trace-id"], - ) - - server_auth_header, extra_headers = _prepare_mcp_server_headers( - server=server, - mcp_server_auth_headers=None, - mcp_auth_header=None, - oauth2_headers=None, - raw_headers={ - "authorization": "Bearer sk-litellm-key", - "x-request-id": "req-123", - "x-trace-id": "trace-456", - }, - ) - - assert server_auth_header is None - # Authorization should be omitted, but other headers forwarded - assert extra_headers == {"x-request-id": "req-123", "x-trace-id": "trace-456"} - - -def test_prepare_mcp_server_headers_passthrough_forwards_authorization_with_explicit_admission(): - """Transparent OAuth pass-through: when LiteLLM admission used the explicit - `x-litellm-api-key` header, the inbound `Authorization` header is - unambiguously the upstream OAuth bearer and MUST be forwarded. - - Regression for EAI-506 V5/V6 — a standards-compliant MCP client (e.g. - OpenCode) completes PKCE against the upstream IdP and sends the resulting - token as plain `Authorization: Bearer ` per the MCP spec. - """ - try: - from litellm.proxy._experimental.mcp_server.server import ( - _prepare_mcp_server_headers, - ) - except ImportError: - pytest.skip("MCP server not available") - - server = MCPServer( - server_id="server-passthrough-explicit-admission", - name="server", - transport=MCPTransport.http, - auth_type=MCPAuth.none, extra_headers=["Authorization"], ) @@ -208,20 +131,15 @@ def test_prepare_mcp_server_headers_passthrough_forwards_authorization_with_expl oauth2_headers=None, raw_headers={ "x-litellm-api-key": "Bearer sk-litellm-key", - "authorization": "Bearer upstream-okta-token", + "authorization": "Bearer token", }, ) assert server_auth_header is None - assert extra_headers == {"Authorization": "Bearer upstream-okta-token"} + assert extra_headers == {"Authorization": "Bearer token"} def test_prepare_mcp_server_headers_passthrough_strips_authorization_without_admission_header(): - """Counterpart to the explicit-admission test: without `x-litellm-api-key`, - the inbound `Authorization` may itself be the LiteLLM admission key, so we - strip it to avoid leaking the gateway credential upstream. This preserves - the security guarantee introduced in commit 3753970cc9. - """ try: from litellm.proxy._experimental.mcp_server.server import ( _prepare_mcp_server_headers, diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 90c5d4f4fc5..99a65e0814b 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -2754,3 +2754,564 @@ def test_build_decode_kwargs_no_warning_when_scoped( if "neither JWT_AUDIENCE nor JWT_ISSUER" in r.getMessage() ] assert matching == [] + + +def _base64url_encode_int(value: int) -> str: + import base64 + + value_bytes = value.to_bytes((value.bit_length() + 7) // 8, "big") + return base64.urlsafe_b64encode(value_bytes).decode("utf-8").rstrip("=") + + +def _get_rsa_key_and_jwk(kid: str): + from cryptography.hazmat.primitives.asymmetric import rsa + + private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + public_numbers = private_key.public_key().public_numbers() + jwk = { + "kty": "RSA", + "n": _base64url_encode_int(value=public_numbers.n), + "e": _base64url_encode_int(value=public_numbers.e), + "kid": kid, + "alg": "RS256", + "use": "sig", + } + return private_key, jwk + + +def _encode_rsa_jwt( + private_key, + issuer: str, + audience: str, + kid: str, + extra_claims: Optional[dict] = None, +) -> str: + import time + + import jwt + from cryptography.hazmat.primitives import serialization + + private_key_pem = private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + current_time = int(time.time()) + claims = { + "sub": "test-subject", + "iss": issuer, + "aud": audience, + "iat": current_time, + "exp": current_time + 300, + } + if extra_claims: + claims.update(extra_claims) + + return jwt.encode( + claims, + private_key_pem, + algorithm="RS256", + headers={"kid": kid}, + ) + + +def _get_jwt_handler_with_issuer_keys(issuers: list, keys_by_url: dict) -> JWTHandler: + from litellm.caching.dual_cache import DualCache + + cache = DualCache() + for jwks_url, keys in keys_by_url.items(): + cache.set_cache( + key=f"litellm_jwt_auth_keys_{jwks_url}", + value=keys, + ) + + jwt_handler = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth(issuers=issuers), + ) + return jwt_handler + + +@pytest.mark.asyncio +async def test_get_public_key_fetches_and_caches_jwks_response(): + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.dual_cache import DualCache + + jwt_handler = JWTHandler() + cache = DualCache() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth(public_key_ttl=123), + ) + expected_key_id = "cached-key" + _, jwk = _get_rsa_key_and_jwk(kid=expected_key_id) + mock_response = MagicMock() + mock_response.json.return_value = {"keys": [jwk]} + jwt_handler.http_handler.get = AsyncMock(return_value=mock_response) + + public_key = await jwt_handler._get_public_key_from_jwks_url( + jwks_url="https://issuer.example.com/keys", + kid=expected_key_id, + ) + + assert public_key == jwk + cached_keys = await cache.async_get_cache( + key="litellm_jwt_auth_keys_https://issuer.example.com/keys" + ) + assert cached_keys == [jwk] + + +@pytest.mark.asyncio +async def test_get_public_key_tries_next_jwks_url_when_kid_missing(monkeypatch): + from litellm.caching.dual_cache import DualCache + + first_jwks_url = "https://first.example.com/keys" + second_jwks_url = "https://second.example.com/keys" + monkeypatch.setenv( + "JWT_PUBLIC_KEY_URL", f"{first_jwks_url}, {second_jwks_url},," + ) + _, first_jwk = _get_rsa_key_and_jwk(kid="first-key") + _, second_jwk = _get_rsa_key_and_jwk(kid="second-key") + cache = DualCache() + cache.set_cache(key=f"litellm_jwt_auth_keys_{first_jwks_url}", value=[first_jwk]) + cache.set_cache( + key=f"litellm_jwt_auth_keys_{second_jwks_url}", value=[second_jwk] + ) + jwt_handler = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth(), + ) + + public_key = await jwt_handler.get_public_key(kid="second-key") + + assert public_key == second_jwk + + +def test_get_jwks_url_for_issuer_falls_back_to_discovery_document(): + jwt_handler = JWTHandler() + issuer_config = LiteLLM_JWTAuth( + issuers=[{"issuer": "https://issuer.example.com/tenant/"}] + ).issuers[0] + + jwks_url = jwt_handler._get_jwks_url_for_issuer(issuer_config=issuer_config) + + assert ( + jwks_url + == "https://issuer.example.com/tenant/.well-known/openid-configuration" + ) + + +@pytest.mark.asyncio +async def test_multi_issuer_jwt_validates_selected_issuer_and_maps_claims( + monkeypatch, +): + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False) + + issuer_one = "https://issuer-one.example.com" + issuer_two = "https://issuer-two.example.com" + issuer_one_jwks_url = f"{issuer_one}/keys" + issuer_two_jwks_url = f"{issuer_two}/keys" + shared_kid = "shared-kid" + + _, issuer_one_jwk = _get_rsa_key_and_jwk(kid=shared_kid) + issuer_two_private_key, issuer_two_jwk = _get_rsa_key_and_jwk(kid=shared_kid) + + jwt_handler = _get_jwt_handler_with_issuer_keys( + issuers=[ + { + "issuer": issuer_one, + "jwks_url": issuer_one_jwks_url, + "audience": "audience-one", + "user_id_jwt_field": "email", + "user_email_jwt_field": "email", + }, + { + "issuer": issuer_two, + "jwks_url": issuer_two_jwks_url, + "audience": "audience-two", + "user_id_jwt_field": "repository_owner", + "team_id_jwt_field": "repository", + }, + ], + keys_by_url={ + issuer_one_jwks_url: [issuer_one_jwk], + issuer_two_jwks_url: [issuer_two_jwk], + }, + ) + + token = _encode_rsa_jwt( + private_key=issuer_two_private_key, + issuer=issuer_two, + audience="audience-two", + kid=shared_kid, + extra_claims={ + "repository_owner": "example-org", + "repository": "example-org/litellm-fork", + }, + ) + + claims = await jwt_handler.auth_jwt(token=token) + + assert claims[JWTHandler.LITELLM_JWT_ISSUER_CLAIM] == issuer_two + assert jwt_handler.get_user_id(token=claims, default_value=None) == "example-org" + assert jwt_handler.get_team_id(token=claims, default_value=None) == ( + "example-org/litellm-fork" + ) + + +@pytest.mark.asyncio +async def test_multi_issuer_jwt_maps_kubernetes_namespace_claim(monkeypatch): + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False) + + issuer = "https://oidc.eks.eu-west-1.amazonaws.com/id/test-cluster" + jwks_url = f"{issuer}/keys" + private_key, jwk = _get_rsa_key_and_jwk(kid="k8s-key") + jwt_handler = _get_jwt_handler_with_issuer_keys( + issuers=[ + { + "issuer": issuer, + "jwks_url": jwks_url, + "audience": None, + "disable_audience_validation": True, + "user_id_jwt_field": "kubernetes\\.io.namespace", + } + ], + keys_by_url={jwks_url: [jwk]}, + ) + token = _encode_rsa_jwt( + private_key=private_key, + issuer=issuer, + audience="kubernetes.default.svc", + kid="k8s-key", + extra_claims={"kubernetes.io": {"namespace": "example-namespace"}}, + ) + + claims = await jwt_handler.auth_jwt(token=token) + + assert ( + jwt_handler.get_user_id(token=claims, default_value=None) == "example-namespace" + ) + + +@pytest.mark.asyncio +async def test_multi_issuer_jwt_rejects_unknown_issuer(monkeypatch): + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False) + + configured_issuer = "https://issuer.example.com" + private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key") + jwt_handler = _get_jwt_handler_with_issuer_keys( + issuers=[ + { + "issuer": configured_issuer, + "jwks_url": f"{configured_issuer}/keys", + "audience": "expected-audience", + } + ], + keys_by_url={f"{configured_issuer}/keys": [jwk]}, + ) + token = _encode_rsa_jwt( + private_key=private_key, + issuer="https://unknown-issuer.example.com", + audience="expected-audience", + kid="issuer-key", + ) + + with pytest.raises(Exception) as exc: + await jwt_handler.auth_jwt(token=token) + + assert "Unsupported JWT issuer" in str(exc.value) + + +@pytest.mark.asyncio +async def test_multi_issuer_jwt_rejects_wrong_audience(monkeypatch): + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False) + + issuer = "https://issuer.example.com" + jwks_url = f"{issuer}/keys" + private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key") + jwt_handler = _get_jwt_handler_with_issuer_keys( + issuers=[ + { + "issuer": issuer, + "jwks_url": jwks_url, + "audience": "expected-audience", + } + ], + keys_by_url={jwks_url: [jwk]}, + ) + token = _encode_rsa_jwt( + private_key=private_key, + issuer=issuer, + audience="wrong-audience", + kid="issuer-key", + ) + + with pytest.raises(Exception) as exc: + await jwt_handler.auth_jwt(token=token) + + assert "Validation fails" in str(exc.value) + + +@pytest.mark.asyncio +async def test_multi_issuer_jwt_same_kid_does_not_cross_issuer_keys(monkeypatch): + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False) + + issuer_one = "https://issuer-one.example.com" + issuer_two = "https://issuer-two.example.com" + issuer_one_jwks_url = f"{issuer_one}/keys" + issuer_two_jwks_url = f"{issuer_two}/keys" + shared_kid = "shared-kid" + issuer_one_private_key, issuer_one_jwk = _get_rsa_key_and_jwk(kid=shared_kid) + _, issuer_two_jwk = _get_rsa_key_and_jwk(kid=shared_kid) + jwt_handler = _get_jwt_handler_with_issuer_keys( + issuers=[ + { + "issuer": issuer_one, + "jwks_url": issuer_one_jwks_url, + "audience": "audience-one", + }, + { + "issuer": issuer_two, + "jwks_url": issuer_two_jwks_url, + "audience": "audience-two", + }, + ], + keys_by_url={ + issuer_one_jwks_url: [issuer_one_jwk], + issuer_two_jwks_url: [issuer_two_jwk], + }, + ) + token = _encode_rsa_jwt( + private_key=issuer_one_private_key, + issuer=issuer_two, + audience="audience-two", + kid=shared_kid, + ) + + with pytest.raises(Exception) as exc: + await jwt_handler.auth_jwt(token=token) + + assert "Validation fails" in str(exc.value) + + +@pytest.mark.asyncio +async def test_multi_issuer_jwt_missing_mapped_claim_fails_closed(monkeypatch): + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False) + + issuer = "https://issuer.example.com" + jwks_url = f"{issuer}/keys" + private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key") + jwt_handler = _get_jwt_handler_with_issuer_keys( + issuers=[ + { + "issuer": issuer, + "jwks_url": jwks_url, + "audience": "expected-audience", + "user_id_jwt_field": "email", + } + ], + keys_by_url={jwks_url: [jwk]}, + ) + token = _encode_rsa_jwt( + private_key=private_key, + issuer=issuer, + audience="expected-audience", + kid="issuer-key", + ) + + with pytest.raises(Exception) as exc: + await jwt_handler.auth_jwt(token=token) + + assert "missing required mapped claim: email" in str(exc.value) + assert "Validation fails" not in str(exc.value) + + +@pytest.mark.asyncio +async def test_multi_issuer_jwt_requires_audience_unless_explicitly_disabled( + monkeypatch, +): + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False) + + issuer = "https://issuer.example.com" + jwks_url = f"{issuer}/keys" + private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key") + jwt_handler = _get_jwt_handler_with_issuer_keys( + issuers=[ + { + "issuer": issuer, + "jwks_url": jwks_url, + } + ], + keys_by_url={jwks_url: [jwk]}, + ) + token = _encode_rsa_jwt( + private_key=private_key, + issuer=issuer, + audience="some-other-client", + kid="issuer-key", + ) + + with pytest.raises(Exception) as exc: + await jwt_handler.auth_jwt(token=token) + + assert "must configure audience" in str(exc.value) + + +@pytest.mark.asyncio +async def test_global_jwt_ignores_user_supplied_internal_claims(monkeypatch): + from litellm.caching.dual_cache import DualCache + + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.delenv("JWT_ISSUER", raising=False) + + jwks_url = "https://global-issuer.example.com/keys" + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url) + + private_key, jwk = _get_rsa_key_and_jwk(kid="global-key") + cache = DualCache() + cache.set_cache(key=f"litellm_jwt_auth_keys_{jwks_url}", value=[jwk]) + + jwt_handler = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth( + user_id_jwt_field="email", + user_email_jwt_field="email", + team_id_jwt_field="team.id", + team_ids_jwt_field="teams", + org_id_jwt_field="org.id", + end_user_id_jwt_field="end_user.id", + ), + ) + token = _encode_rsa_jwt( + private_key=private_key, + issuer="https://global-issuer.example.com", + audience="some-other-client", + kid="global-key", + extra_claims={ + "email": "real-user@example.com", + "team": {"id": "real-team"}, + "teams": ["real-team", "secondary-team"], + "org": {"id": "real-org"}, + "end_user": {"id": "real-end-user"}, + JWTHandler.LITELLM_JWT_ISSUER_CLAIM: "https://issuer.example.com", + JWTHandler.LITELLM_USER_ID_CLAIM: "victim-user", + JWTHandler.LITELLM_USER_EMAIL_CLAIM: "victim@example.com", + JWTHandler.LITELLM_TEAM_ID_CLAIM: "victim-team", + JWTHandler.LITELLM_TEAM_IDS_CLAIM: ["victim-team"], + JWTHandler.LITELLM_ORG_ID_CLAIM: "victim-org", + JWTHandler.LITELLM_END_USER_ID_CLAIM: "victim-end-user", + }, + ) + + claims = await jwt_handler.auth_jwt(token=token) + + assert jwt_handler.get_user_id(token=claims, default_value=None) == ( + "real-user@example.com" + ) + assert jwt_handler.get_user_email(token=claims, default_value=None) == ( + "real-user@example.com" + ) + assert jwt_handler.get_team_id(token=claims, default_value=None) == "real-team" + assert jwt_handler.get_team_ids_from_jwt(token=claims) == [ + "real-team", + "secondary-team", + ] + assert jwt_handler.get_org_id(token=claims, default_value=None) == "real-org" + assert jwt_handler.get_end_user_id(token=claims, default_value=None) == ( + "real-end-user" + ) + + +@pytest.mark.asyncio +async def test_multi_issuer_jwt_strips_unmapped_internal_claims(monkeypatch): + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False) + + issuer = "https://issuer.example.com" + jwks_url = f"{issuer}/keys" + private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key") + jwt_handler = _get_jwt_handler_with_issuer_keys( + issuers=[ + { + "issuer": issuer, + "jwks_url": jwks_url, + "audience": "expected-audience", + "user_email_jwt_field": "email", + } + ], + keys_by_url={jwks_url: [jwk]}, + ) + token = _encode_rsa_jwt( + private_key=private_key, + issuer=issuer, + audience="expected-audience", + kid="issuer-key", + extra_claims={ + "email": "real-user@example.com", + JWTHandler.LITELLM_USER_ID_CLAIM: "victim-user", + JWTHandler.LITELLM_TEAM_ID_CLAIM: "victim-team", + }, + ) + + claims = await jwt_handler.auth_jwt(token=token) + + assert JWTHandler.LITELLM_USER_ID_CLAIM not in claims + assert JWTHandler.LITELLM_TEAM_ID_CLAIM not in claims + assert jwt_handler.get_user_id(token=claims, default_value=None) is None + assert jwt_handler.get_team_id(token=claims, default_value=None) is None + assert jwt_handler.get_user_email(token=claims, default_value=None) == ( + "real-user@example.com" + ) + + +@pytest.mark.asyncio +async def test_multi_issuer_jwt_does_not_emit_unscoped_global_warning( + monkeypatch, caplog +): + import logging + + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.delenv("JWT_ISSUER", raising=False) + monkeypatch.delenv("JWT_PUBLIC_KEY_URL", raising=False) + JWTHandler._unscoped_jwt_warning_emitted = False + + issuer = "https://issuer.example.com" + jwks_url = f"{issuer}/keys" + private_key, jwk = _get_rsa_key_and_jwk(kid="issuer-key") + jwt_handler = _get_jwt_handler_with_issuer_keys( + issuers=[ + { + "issuer": issuer, + "jwks_url": jwks_url, + "audience": "expected-audience", + } + ], + keys_by_url={jwks_url: [jwk]}, + ) + token = _encode_rsa_jwt( + private_key=private_key, + issuer=issuer, + audience="expected-audience", + kid="issuer-key", + ) + + with caplog.at_level(logging.WARNING): + await jwt_handler.auth_jwt(token=token) + + assert "Tokens minted by any application" not in caplog.text + assert JWTHandler._unscoped_jwt_warning_emitted is False diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index fcbef9b4294..a5c8320a5cc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -2,7 +2,7 @@ import os import sys import types import json -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta from types import SimpleNamespace from typing import List, Optional from unittest.mock import AsyncMock, MagicMock, patch