From 8ca2f14f81b832ec56961e24a6f8cf212f93d0ba Mon Sep 17 00:00:00 2001 From: gym-cmd <186399764+gym-cmd@users.noreply.github.com> Date: Fri, 15 May 2026 21:05:40 +0100 Subject: [PATCH] fix(mcp): address oauth passthrough review findings --- .../mcp_server/auth/user_api_key_auth_mcp.py | 52 ++- .../mcp_server/discoverable_endpoints.py | 79 ++--- .../_experimental/mcp_server/oauth_utils.py | 28 +- litellm/proxy/auth/ip_address_utils.py | 10 +- .../mcp_management_endpoints.py | 8 +- .../mcp_server/test_discoverable_endpoints.py | 45 +-- tests/proxy_unit_tests/test_jwt.py | 8 +- .../auth/test_user_api_key_auth_mcp.py | 314 +++++++++++++++++- .../mcp_server/test_discoverable_endpoints.py | 121 ++++++- .../mcp_server/test_mcp_oauth_passthrough.py | 131 +++++--- .../proxy/auth/test_mcp_ip_filtering.py | 52 ++- .../test_mcp_management_endpoints.py | 31 ++ 12 files changed, 738 insertions(+), 141 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 49655bbc144..c1bb6edd772 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -14,6 +14,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.auth.ip_address_utils import IPAddressUtils def _parse_mcp_server_names_from_path(path: str) -> Optional[List[str]]: @@ -29,7 +30,7 @@ def _parse_mcp_server_names_from_path(path: str) -> Optional[List[str]]: def _is_mcp_passthrough_cold_start( - scope: Scope, mcp_servers: Optional[List[str]] + scope: Scope, mcp_servers: Optional[List[str]], client_ip: Optional[str] ) -> bool: """True when the request targets a pass-through server with no auth headers — the cold-start OAuth discovery case per RFC 9728 / MCP Authorization spec. @@ -42,12 +43,32 @@ def _is_mcp_passthrough_cold_start( ) for name in mcp_servers: - server = global_mcp_server_manager.get_mcp_server_by_name(name) + server = global_mcp_server_manager.get_mcp_server_by_name( + name, client_ip=client_ip + ) if server is not None and getattr(server, "is_oauth_passthrough", False): return True return False +def _is_litellm_auth_admission_error(exc: Exception) -> bool: + if isinstance(exc, HTTPException): + return exc.status_code == 401 + if isinstance(exc, ProxyException): + try: + return int(exc.code) == 401 + except (TypeError, ValueError): + return False + return False + + +def _has_client_supplied_mcp_auth( + mcp_auth_header: Optional[str], + mcp_server_auth_headers: Dict[str, Dict[str, str]], +) -> bool: + return bool(mcp_auth_header) or bool(mcp_server_auth_headers) + + class MCPRequestHandler: """ Class to handle MCP request processing, including: @@ -188,7 +209,9 @@ class MCPRequestHandler: "401", "403", ) and MCPRequestHandler._target_servers_use_oauth2( - path=request.url.path, mcp_servers=mcp_servers + path=request.url.path, + mcp_servers=mcp_servers, + client_ip=IPAddressUtils.get_mcp_client_ip(request), ): verbose_logger.debug( "MCP OAuth2: target server is OAuth2-mode, treating " @@ -202,15 +225,24 @@ class MCPRequestHandler: validated_user_api_key_auth = await user_api_key_auth( api_key=litellm_api_key, request=request ) - except (HTTPException, ProxyException): + except (HTTPException, ProxyException) as exc: # Cold-start MCP OAuth discovery: RFC 9728 / MCP Authorization spec # require unauthenticated requests to protected resources to receive # 401 + WWW-Authenticate. Defer to _raise_preemptive_401_for_unauthenticated_servers # for pass-through servers instead of surfacing a generic admission error. path = scope.get("path", "") mcp_servers_from_path = _parse_mcp_server_names_from_path(path) - if _is_mcp_passthrough_cold_start( - scope, mcp_servers_from_path or mcp_servers + client_ip = IPAddressUtils.get_mcp_client_ip(request) + if ( + mcp_servers_from_path is not None + and not _has_client_supplied_mcp_auth( + mcp_auth_header, + mcp_server_auth_headers, + ) + and _is_litellm_auth_admission_error(exc) + and _is_mcp_passthrough_cold_start( + scope, mcp_servers_from_path, client_ip=client_ip + ) ): verbose_logger.debug( "MCP pass-through cold start: deferring admission to route 401 emitter" @@ -250,7 +282,9 @@ class MCPRequestHandler: return [] @staticmethod - def _target_servers_use_oauth2(path: str, mcp_servers: Optional[List[str]]) -> bool: + def _target_servers_use_oauth2( + path: str, mcp_servers: Optional[List[str]], client_ip: Optional[str] + ) -> bool: """ True only when EVERY MCP server the request targets is configured for ``auth_type == oauth2``. If any target is non-OAuth2 — or if the target @@ -279,7 +313,9 @@ class MCPRequestHandler: return False for name in target_names: - server = global_mcp_server_manager.get_mcp_server_by_name(name) + server = global_mcp_server_manager.get_mcp_server_by_name( + name, client_ip=client_ip + ) if server is None or server.auth_type != MCPAuth.oauth2: return False return True diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 303c9d2e5fa..fc8930fd81f 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -13,6 +13,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, + get_request_base_url, validate_trusted_redirect_uri, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils @@ -30,55 +31,33 @@ from litellm.types.mcp_server.mcp_server_manager import MCPServer # Keyed by (server_id, resource_url) → (expires_at_epoch, payload). _OAUTH_METADATA_CACHE: Dict[Tuple[str, str], Tuple[float, dict]] = {} _OAUTH_METADATA_CACHE_TTL_SECONDS = 300 +_OAUTH_METADATA_CACHE_MAX_SIZE = 128 router = APIRouter( tags=["mcp"], ) -def get_request_base_url(request: Request) -> str: - """ - Get the base URL for the request, considering X-Forwarded-* headers. +def _prune_oauth_metadata_cache(now: Optional[float] = None) -> None: + now = now if now is not None else time.time() + expired_cache_keys = [ + cache_key + for cache_key, (expires_at, _payload) in _OAUTH_METADATA_CACHE.items() + if expires_at <= now + ] + for cache_key in expired_cache_keys: + _OAUTH_METADATA_CACHE.pop(cache_key, None) - X-Forwarded-Proto / X-Forwarded-Host / X-Forwarded-Port are only honoured - when the request comes from a configured trusted proxy - (``use_x_forwarded_for`` enabled AND caller in ``mcp_trusted_proxy_ranges``). - Otherwise the request's literal ``base_url`` is returned, so an - untrusted caller cannot poison OAuth-discovery / redirect_uri values - by injecting headers. + if len(_OAUTH_METADATA_CACHE) <= _OAUTH_METADATA_CACHE_MAX_SIZE: + return - Args: - request: FastAPI Request object - - Returns: - The reconstructed base URL (e.g., "https://proxy.example.com") - """ - base_url = str(request.base_url).rstrip("/") - parsed = urlparse(base_url) - - if not IPAddressUtils.is_request_from_trusted_proxy(request): - return base_url - - x_forwarded_proto = request.headers.get("X-Forwarded-Proto") - x_forwarded_host = request.headers.get("X-Forwarded-Host") - x_forwarded_port = request.headers.get("X-Forwarded-Port") - - scheme = x_forwarded_proto if x_forwarded_proto else parsed.scheme - - if x_forwarded_host: - # X-Forwarded-Host may already include port (e.g., "example.com:8080") - if ":" in x_forwarded_host and not x_forwarded_host.startswith("["): - netloc = x_forwarded_host - elif x_forwarded_port: - netloc = f"{x_forwarded_host}:{x_forwarded_port}" - else: - netloc = x_forwarded_host - else: - netloc = parsed.netloc - if x_forwarded_port and ":" not in netloc: - netloc = f"{netloc}:{x_forwarded_port}" - - return urlunparse((scheme, netloc, parsed.path, "", "", "")) + overflow = len(_OAUTH_METADATA_CACHE) - _OAUTH_METADATA_CACHE_MAX_SIZE + cache_keys_by_expiry = sorted( + _OAUTH_METADATA_CACHE, + key=lambda cache_key: _OAUTH_METADATA_CACHE[cache_key][0], + ) + for cache_key in cache_keys_by_expiry[:overflow]: + _OAUTH_METADATA_CACHE.pop(cache_key, None) def encode_state_with_base_url( @@ -438,6 +417,11 @@ async def exchange_token_with_server( headers={"Accept": "application/json"}, data=token_data, ) + if response is None: + raise HTTPException( + status_code=502, + detail="MCP upstream token endpoint returned no response", + ) response.raise_for_status() token_response = response.json() @@ -543,6 +527,11 @@ async def register_client_with_server( headers=headers, json=register_data, ) + if response is None: + raise HTTPException( + status_code=502, + detail="MCP upstream registration endpoint returned no response", + ) response.raise_for_status() token_response = response.json() @@ -737,8 +726,10 @@ async def fetch_upstream_oauth_protected_resource( return None cache_key = (mcp_server.server_id, mcp_server.url) + now = time.time() + _prune_oauth_metadata_cache(now) cached = _OAUTH_METADATA_CACHE.get(cache_key) - if cached is not None and cached[0] > time.time(): + if cached is not None and cached[0] > now: return cached[1] host_base = f"{upstream.scheme}://{upstream.netloc}" @@ -769,10 +760,12 @@ async def fetch_upstream_oauth_protected_resource( except Exception: continue if isinstance(payload, dict): + now = time.time() _OAUTH_METADATA_CACHE[cache_key] = ( - time.time() + _OAUTH_METADATA_CACHE_TTL_SECONDS, + now + _OAUTH_METADATA_CACHE_TTL_SECONDS, payload, ) + _prune_oauth_metadata_cache(now) return payload if len(network_errors) == len(candidates): diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 3f99784f2c4..57d444c5c0c 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -8,6 +8,7 @@ from fastapi import HTTPException, Request from litellm._logging import verbose_logger from litellm.proxy.auth.ip_address_utils import IPAddressUtils +from litellm.proxy.utils import get_proxy_base_url # RFC 6749 §5.1 / OAuth 2.1 draft-15 §4.1.3: token-endpoint responses # must not be cached — both success and error bodies may reveal secrets. @@ -96,7 +97,7 @@ def validate_loopback_redirect_uri(redirect_uri: str) -> None: def validate_trusted_redirect_uri(request: Request, redirect_uri: str) -> None: - """Accept same-origin (proxy's own origin) OR loopback ``redirect_uri``. + """Accept trusted same-origin (proxy's own origin) OR loopback ``redirect_uri``. Same-origin is required for the LiteLLM UI's OAuth flow: the UI redirects to ``/ui/mcp/oauth/callback`` which is not loopback @@ -105,6 +106,10 @@ def validate_trusted_redirect_uri(request: Request, redirect_uri: str) -> None: proxy, so the open-redirect / code-theft primitive that motivated :func:`validate_loopback_redirect_uri` does not apply here. + The trusted origin must come from explicit ``PROXY_BASE_URL`` config or + from X-Forwarded-* headers supplied by a configured trusted proxy. The raw + Host / request base URL is intentionally not trusted for this decision. + Loopback continues to be accepted for native MCP clients (per OAuth 2.1 §4.1.2.1 + RFC 8252 §7.3). @@ -122,17 +127,32 @@ def validate_trusted_redirect_uri(request: Request, redirect_uri: str) -> None: if parsed.fragment: raise HTTPException(status_code=400, detail="invalid_request") - try: - proxy_base = urlparse(get_request_base_url(request)) + configured_proxy_base_url = get_proxy_base_url() + if configured_proxy_base_url: + proxy_base = urlparse(configured_proxy_base_url) if ( parsed.netloc + and proxy_base.scheme in ("http", "https") and parsed.scheme == proxy_base.scheme and parsed.netloc.lower() == proxy_base.netloc.lower() ): return + + try: + if IPAddressUtils.is_request_from_trusted_proxy( + request + ) and request.headers.get("X-Forwarded-Host"): + proxy_base = urlparse(get_request_base_url(request)) + if ( + parsed.netloc + and proxy_base.scheme in ("http", "https") + and parsed.scheme == proxy_base.scheme + and parsed.netloc.lower() == proxy_base.netloc.lower() + ): + return except Exception as exc: verbose_logger.warning( - "validate_trusted_redirect_uri: could not determine proxy origin, " + "validate_trusted_redirect_uri: could not determine trusted proxy origin, " "falling back to loopback-only check. error=%s", exc, ) diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index 39d3282942f..16a5bc61d31 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -153,8 +153,9 @@ class IPAddressUtils: verbose_proxy_logger.warning( "use_x_forwarded_for is enabled but mcp_trusted_proxy_ranges " "is not configured. X-Forwarded-* headers will NOT be " - "trusted, so MCP OAuth discovery URLs will use the proxy's " - "literal base URL. Set mcp_trusted_proxy_ranges in " + "trusted, so MCP OAuth discovery URLs and access-control " + "client IPs will use the proxy's literal request values. " + "Set mcp_trusted_proxy_ranges in " "general_settings to your reverse-proxy CIDR(s) to allow " "X-Forwarded-* through." ) @@ -200,6 +201,11 @@ class IPAddressUtils: # If XFF is enabled, validate the request comes from a trusted proxy if use_xff and "x-forwarded-for" in request.headers: trusted_ranges = general_settings.get("mcp_trusted_proxy_ranges") + if not trusted_ranges: + IPAddressUtils.is_request_from_trusted_proxy( + request, general_settings=general_settings + ) + return _get_request_ip_address(request, use_x_forwarded_for=False) if trusted_ranges: # Validate direct connection is from trusted proxy direct_ip = request.client.host if request.client else None diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index de81449fd41..b2e21768842 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1595,12 +1595,12 @@ if MCP_AVAILABLE: @router.get( "/server/oauth/{server_id}/authorize", include_in_schema=False, - dependencies=[Depends(user_api_key_auth)], + dependencies=[Depends(_mcp_oauth_user_api_key_auth)], ) async def mcp_authorize( request: Request, server_id: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: UserAPIKeyAuth = Depends(_mcp_oauth_user_api_key_auth), client_id: Optional[str] = None, redirect_uri: str = Query(...), state: str = "", @@ -1640,12 +1640,12 @@ if MCP_AVAILABLE: @router.post( "/server/oauth/{server_id}/token", include_in_schema=False, - dependencies=[Depends(user_api_key_auth)], + dependencies=[Depends(_mcp_oauth_user_api_key_auth)], ) async def mcp_token( request: Request, server_id: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: UserAPIKeyAuth = Depends(_mcp_oauth_user_api_key_auth), grant_type: str = Form(...), code: Optional[str] = Form(None), redirect_uri: Optional[str] = Form(None), diff --git a/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index f3744468bf5..2a8768df722 100644 --- a/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -15,12 +15,15 @@ def set_request_from_trusted_proxy(mock_request): @pytest.fixture def trusted_proxy_origin_headers(): - with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy", - return_value=True, - ), patch( - "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", - return_value=True, + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy", + return_value=True, + ), + patch( + "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", + return_value=True, + ), ): yield @@ -77,7 +80,7 @@ async def test_authorize_endpoint_includes_response_type(): request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", - redirect_uri="http://127.0.0.1:60108/callback", + redirect_uri="http://127.0.0.1:60108/callback", state="test_state", ) @@ -175,7 +178,6 @@ async def test_token_endpoint_forwards_code_verifier(): from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.proxy._types import MCPTransport from fastapi import Request - import httpx except ImportError: pytest.skip("MCP discoverable endpoints not available") @@ -490,7 +492,7 @@ async def test_authorize_endpoint_respects_x_forwarded_proto( request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", - redirect_uri="http://127.0.0.1:60108/callback", + redirect_uri="http://127.0.0.1:60108/callback", state="test_state", ) @@ -567,7 +569,7 @@ async def test_token_endpoint_respects_x_forwarded_proto( mock_get_client.return_value = mock_async_client # Call token endpoint - response = await token_endpoint( + await token_endpoint( request=mock_request, grant_type="authorization_code", code="test_code", @@ -908,7 +910,7 @@ async def test_authorize_endpoint_respects_x_forwarded_host( request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", - redirect_uri="http://127.0.0.1:60108/callback", + redirect_uri="http://127.0.0.1:60108/callback", state="test_state", ) @@ -989,7 +991,7 @@ async def test_token_endpoint_respects_x_forwarded_host( mock_get_client.return_value = mock_async_client # Call token endpoint - response = await token_endpoint( + await token_endpoint( request=mock_request, grant_type="authorization_code", code="test_code", @@ -1224,14 +1226,17 @@ def test_validate_trusted_redirect_uri_rejects_spoofed_forwarded_host(): mock_request.client = MagicMock() mock_request.client.host = "203.0.113.10" - with patch( - "litellm.proxy.proxy_server.general_settings", - { - "use_x_forwarded_for": True, - "mcp_trusted_proxy_ranges": TRUSTED_PROXY_RANGES, - }, - create=True, - ), pytest.raises(HTTPException): + with ( + patch( + "litellm.proxy.proxy_server.general_settings", + { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": TRUSTED_PROXY_RANGES, + }, + create=True, + ), + pytest.raises(HTTPException), + ): validate_trusted_redirect_uri( mock_request, "https://attacker.example.com/callback", diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/proxy_unit_tests/test_jwt.py index 46438206420..e25a5d86513 100644 --- a/tests/proxy_unit_tests/test_jwt.py +++ b/tests/proxy_unit_tests/test_jwt.py @@ -1715,9 +1715,7 @@ async def test_multi_issuer_jwt_validates_selected_issuer_and_maps_claims( 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_user_id(token=claims, default_value=None) == ("example-org") assert jwt_handler.get_team_id(token=claims, default_value=None) == ( "example-org/litellm-fork" ) @@ -1753,7 +1751,9 @@ async def test_multi_issuer_jwt_maps_kubernetes_namespace_claim(monkeypatch): claims = await jwt_handler.auth_jwt(token=token) - assert jwt_handler.get_user_id(token=claims, default_value=None) == "example-namespace" + assert ( + jwt_handler.get_user_id(token=claims, default_value=None) == "example-namespace" + ) @pytest.mark.asyncio 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 6e0dadcd4d8..c9f2c454bb0 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,12 +1,9 @@ import json import os import sys -from unittest import mock from unittest.mock import AsyncMock, MagicMock, patch -import orjson import pytest -from fastapi import FastAPI, Request from fastapi.testclient import TestClient sys.path.insert( @@ -19,7 +16,6 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) from litellm.proxy._types import SpecialHeaders, UserAPIKeyAuth -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @pytest.mark.asyncio @@ -342,7 +338,7 @@ class TestMCPRequestHandler: with patch( "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", side_effect=mock_user_api_key_auth, - ) as mock_auth: + ): # Call the method ( auth_result, @@ -885,6 +881,278 @@ class TestMCPPublicRouteGuard: assert isinstance(auth_result, UserAPIKeyAuth) +@pytest.mark.asyncio +class TestMCPPassthroughColdStartAdmission: + @staticmethod + def _make_passthrough_server(): + server = MagicMock() + server.is_oauth_passthrough = True + return server + + async def test_cold_start_ignores_header_without_path_target(self): + from fastapi import HTTPException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [(b"x-mcp-servers", b"passthrough_server")], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = ( + TestMCPPassthroughColdStartAdmission._make_passthrough_server() + ) + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + mock_mgr.get_mcp_server_by_name.assert_not_called() + + async def test_cold_start_rejects_server_specific_authorization_header(self): + from fastapi import HTTPException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/passthrough_server", + "headers": [ + ( + b"x-mcp-passthrough_server-authorization", + b"Bearer upstream-token", + ) + ], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = ( + TestMCPPassthroughColdStartAdmission._make_passthrough_server() + ) + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + + async def test_cold_start_rejects_legacy_mcp_auth_header(self): + from fastapi import HTTPException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/passthrough_server", + "headers": [(b"x-mcp-auth", b"Bearer upstream-token")], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = ( + TestMCPPassthroughColdStartAdmission._make_passthrough_server() + ) + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + + async def test_cold_start_fails_closed_when_client_ip_hides_server(self): + from fastapi import HTTPException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/passthrough_server", + "headers": [], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.IPAddressUtils.get_mcp_client_ip", + return_value="203.0.113.10", + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = None + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + mock_mgr.get_mcp_server_by_name.assert_called_once_with( + "passthrough_server", client_ip="203.0.113.10" + ) + + async def test_cold_start_propagates_non_401_http_error(self): + from fastapi import HTTPException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/passthrough_server", + "headers": [], + } + + async def mock_user_api_key_auth_forbidden(api_key, request): + raise HTTPException(status_code=403, detail="Forbidden") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_forbidden, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = ( + TestMCPPassthroughColdStartAdmission._make_passthrough_server() + ) + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 403 + + async def test_cold_start_propagates_non_auth_proxy_exception(self): + from litellm.proxy._types import ProxyException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/passthrough_server", + "headers": [], + } + + async def mock_user_api_key_auth_server_error(api_key, request): + raise ProxyException( + message="Internal error", + type="server_error", + param=None, + code=500, + ) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_server_error, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = ( + TestMCPPassthroughColdStartAdmission._make_passthrough_server() + ) + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request(scope) + + async def test_cold_start_allows_401_for_path_passthrough_target(self): + from fastapi import HTTPException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/passthrough_server", + "headers": [], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = ( + TestMCPPassthroughColdStartAdmission._make_passthrough_server() + ) + (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + + assert isinstance(auth_result, UserAPIKeyAuth) + mock_mgr.get_mcp_server_by_name.assert_called_once_with( + "passthrough_server", client_ip="" + ) + + async def test_cold_start_allows_proxy_exception_401_for_path_target(self): + from litellm.proxy._types import ProxyException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/passthrough_server", + "headers": [], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise ProxyException( + message="Authentication Error", + type="auth_error", + param="api_key", + code=401, + ) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = ( + TestMCPPassthroughColdStartAdmission._make_passthrough_server() + ) + (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + + assert isinstance(auth_result, UserAPIKeyAuth) + mock_mgr.get_mcp_server_by_name.assert_called_once_with( + "passthrough_server", client_ip="" + ) + + @pytest.mark.asyncio class TestMCPOAuth2FallbackTargetGating: """ @@ -1000,6 +1268,41 @@ class TestMCPOAuth2FallbackTargetGating: (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) assert isinstance(auth_result, UserAPIKeyAuth) + async def test_fallback_blocked_when_client_ip_hides_oauth2_target(self): + from fastapi import HTTPException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/hidden_oauth2_server", + "headers": [(b"authorization", b"Bearer upstream-token")], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.IPAddressUtils.get_mcp_client_ip", + return_value="203.0.113.10", + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = None + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + + assert exc_info.value.status_code == 401 + mock_mgr.get_mcp_server_by_name.assert_called_once_with( + "hidden_oauth2_server", client_ip="203.0.113.10" + ) + async def test_fallback_blocked_when_any_target_in_header_is_not_oauth2(self): """ x-mcp-servers can list multiple targets. If ANY of them is non-OAuth2, @@ -1480,7 +1783,6 @@ class TestMCPAccessGroupsE2E: mock_auth.assert_called_once() -@pytest.mark.asyncio def test_mcp_path_based_server_segregation(monkeypatch): # Import the MCP server FastAPI app and context getter from litellm.proxy._experimental.mcp_server.server import app, get_auth_context 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 6cf1e8184f5..dc3449573cd 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 @@ -49,7 +49,7 @@ def trust_xff(): ``test_get_request_base_url_xff_trust_gate``. """ with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy", + "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", return_value=True, ): yield @@ -1154,7 +1154,7 @@ def test_get_request_base_url_comprehensive( mock_request.headers.get = mock_get with patch( - "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy", + "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", return_value=True, ): result = get_request_base_url(mock_request) @@ -1312,20 +1312,121 @@ def test_validate_trusted_redirect_uri_rejects_spoofed_forwarded_host(): mock_request.client = MagicMock() mock_request.client.host = "203.0.113.10" - with patch( - "litellm.proxy.proxy_server.general_settings", - { - "use_x_forwarded_for": True, - "mcp_trusted_proxy_ranges": TRUSTED_PROXY_RANGES, - }, - create=True, - ), pytest.raises(HTTPException): + with ( + patch( + "litellm.proxy.proxy_server.general_settings", + { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": TRUSTED_PROXY_RANGES, + }, + create=True, + ), + pytest.raises(HTTPException), + ): validate_trusted_redirect_uri( mock_request, "https://attacker.example.com/callback", ) +def test_validate_trusted_redirect_uri_rejects_raw_host_same_origin(): + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + except ImportError: + pytest.skip("MCP OAuth utilities not available") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://attacker.example.com/" + mock_request.headers = {} + mock_request.client = MagicMock() + mock_request.client.host = "203.0.113.10" + + with ( + patch( + "litellm.proxy._experimental.mcp_server.oauth_utils.get_proxy_base_url", + return_value=None, + ), + patch( + "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", + return_value=False, + ), + pytest.raises(HTTPException), + ): + validate_trusted_redirect_uri( + mock_request, + "https://attacker.example.com/callback", + ) + + +def test_validate_trusted_redirect_uri_rejects_trusted_proxy_raw_host_fallback(): + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + except ImportError: + pytest.skip("MCP OAuth utilities not available") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://attacker.example.com/" + mock_request.headers = {"X-Forwarded-Proto": "https"} + mock_request.client = MagicMock() + mock_request.client.host = TRUSTED_PROXY_IP + + with ( + patch( + "litellm.proxy._experimental.mcp_server.oauth_utils.get_proxy_base_url", + return_value=None, + ), + patch( + "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", + return_value=True, + ), + pytest.raises(HTTPException), + ): + validate_trusted_redirect_uri( + mock_request, + "https://attacker.example.com/callback", + ) + + +def test_validate_trusted_redirect_uri_allows_configured_proxy_base_url_origin(): + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + except ImportError: + pytest.skip("MCP OAuth utilities not available") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + mock_request.headers = {} + mock_request.client = MagicMock() + mock_request.client.host = "203.0.113.10" + + with ( + patch( + "litellm.proxy._experimental.mcp_server.oauth_utils.get_proxy_base_url", + return_value="https://gateway.example.com", + ), + patch( + "litellm.proxy._experimental.mcp_server.oauth_utils.IPAddressUtils.is_request_from_trusted_proxy", + return_value=False, + ), + ): + validate_trusted_redirect_uri( + mock_request, + "https://gateway.example.com/ui/mcp/oauth/callback", + ) + + def test_validate_trusted_redirect_uri_allows_forwarded_origin_from_trusted_proxy(): try: from fastapi import Request diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough.py index d6e489f3384..ab237d58d2c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough.py @@ -134,9 +134,9 @@ async def test_oauth_protected_resource_passthrough_proxies_upstream_metadata(): global_mcp_server_manager.registry.clear() passthrough_server = MCPServer( server_id="passthrough-1", - name="knowledge_qa", - server_name="knowledge_qa", - alias="knowledge_qa", + name="sample_docs", + server_name="sample_docs", + alias="sample_docs", url="https://upstream.example.com/mcp", transport=MCPTransport.http, auth_type=MCPAuth.none, @@ -163,7 +163,7 @@ async def test_oauth_protected_resource_passthrough_proxies_upstream_metadata(): ): result = await _build_oauth_protected_resource_response( request=_make_request(), - mcp_server_name="knowledge_qa", + mcp_server_name="sample_docs", use_standard_pattern=True, ) @@ -171,7 +171,7 @@ async def test_oauth_protected_resource_passthrough_proxies_upstream_metadata(): "https://okta.example.com/oauth2/default" ] # resource is normalized to the gateway URL so bearers are sent back to us - assert result["resource"].endswith("/mcp/knowledge_qa") + assert result["resource"].endswith("/mcp/sample_docs") assert result["scopes_supported"] == ["openid", "profile"] @@ -184,9 +184,9 @@ async def test_oauth_protected_resource_passthrough_cache_hit(): global_mcp_server_manager.registry.clear() passthrough_server = MCPServer( server_id="passthrough-2", - name="knowledge_qa", - server_name="knowledge_qa", - alias="knowledge_qa", + name="sample_docs", + server_name="sample_docs", + alias="sample_docs", url="https://upstream.example.com/mcp", transport=MCPTransport.http, auth_type=MCPAuth.none, @@ -209,18 +209,74 @@ async def test_oauth_protected_resource_passthrough_cache_hit(): ): await _build_oauth_protected_resource_response( request=_make_request(), - mcp_server_name="knowledge_qa", + mcp_server_name="sample_docs", use_standard_pattern=True, ) await _build_oauth_protected_resource_response( request=_make_request(), - mcp_server_name="knowledge_qa", + mcp_server_name="sample_docs", use_standard_pattern=True, ) assert mock_client.get.await_count == 1 +def test_oauth_metadata_cache_prunes_to_max_size(): + now = 1_000_000.0 + max_size = discoverable_endpoints._OAUTH_METADATA_CACHE_MAX_SIZE + + for index in range(max_size + 10): + _OAUTH_METADATA_CACHE[(f"server-{index}", f"https://upstream/{index}")] = ( + now + index + 1, + {"index": index}, + ) + + discoverable_endpoints._prune_oauth_metadata_cache(now) + + assert len(_OAUTH_METADATA_CACHE) == max_size + assert ("server-0", "https://upstream/0") not in _OAUTH_METADATA_CACHE + assert ( + f"server-{max_size + 9}", + f"https://upstream/{max_size + 9}", + ) in _OAUTH_METADATA_CACHE + + +@pytest.mark.asyncio +async def test_oauth_metadata_cache_expired_entry_is_refetched(): + passthrough_server = MCPServer( + server_id="expired-cache-server", + name="sample_docs", + server_name="sample_docs", + alias="sample_docs", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.none, + extra_headers=["Authorization"], + ) + _OAUTH_METADATA_CACHE[(passthrough_server.server_id, passthrough_server.url)] = ( + 0, + {"authorization_servers": ["https://stale.example.com"]}, + ) + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "authorization_servers": ["https://fresh.example.com"], + } + mock_client = MagicMock() + mock_client.get = AsyncMock(return_value=mock_response) + + with patch.object( + discoverable_endpoints, "get_async_httpx_client", return_value=mock_client + ): + result = await discoverable_endpoints.fetch_upstream_oauth_protected_resource( + passthrough_server + ) + + assert result == {"authorization_servers": ["https://fresh.example.com"]} + assert mock_client.get.await_count == 1 + + @pytest.mark.asyncio async def test_oauth_protected_resource_passthrough_network_error_returns_502(): from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( @@ -230,9 +286,9 @@ async def test_oauth_protected_resource_passthrough_network_error_returns_502(): global_mcp_server_manager.registry.clear() passthrough_server = MCPServer( server_id="passthrough-3", - name="knowledge_qa", - server_name="knowledge_qa", - alias="knowledge_qa", + name="sample_docs", + server_name="sample_docs", + alias="sample_docs", url="https://upstream.example.com/mcp", transport=MCPTransport.http, auth_type=MCPAuth.none, @@ -251,7 +307,7 @@ async def test_oauth_protected_resource_passthrough_network_error_returns_502(): with pytest.raises(HTTPException) as exc_info: await _build_oauth_protected_resource_response( request=_make_request(), - mcp_server_name="knowledge_qa", + mcp_server_name="sample_docs", use_standard_pattern=True, ) @@ -262,9 +318,9 @@ async def test_oauth_protected_resource_passthrough_network_error_returns_502(): async def test_fetch_upstream_metadata_returns_none_when_not_all_candidates_network_fail(): passthrough_server = MCPServer( server_id="passthrough-partial-network", - name="knowledge_qa", - server_name="knowledge_qa", - alias="knowledge_qa", + name="sample_docs", + server_name="sample_docs", + alias="sample_docs", url="https://upstream.example.com/mcp", transport=MCPTransport.http, auth_type=MCPAuth.none, @@ -381,7 +437,7 @@ async def test_fetch_tools_from_passthrough_raises_on_upstream_401(): manager = MCPServerManager() passthrough_server = MCPServer( server_id="p1", - name="knowledge_qa", + name="sample_docs", url="https://upstream/mcp", transport=MCPTransport.http, auth_type=MCPAuth.none, @@ -409,7 +465,7 @@ async def test_fetch_tools_from_passthrough_raises_on_upstream_401(): assert exc_info.value.www_authenticate == ( 'Bearer resource_metadata="https://upstream"' ) - assert exc_info.value.server_name == "knowledge_qa" + assert exc_info.value.server_name == "sample_docs" mock_client.list_tools.assert_awaited_with(raise_on_error=True) @@ -418,7 +474,7 @@ async def test_fetch_tools_from_passthrough_returns_tools_on_success(): manager = MCPServerManager() passthrough_server = MCPServer( server_id="p1", - name="knowledge_qa", + name="sample_docs", url="https://upstream/mcp", transport=MCPTransport.http, auth_type=MCPAuth.none, @@ -492,12 +548,12 @@ def _make_scope(path: str, headers: list = None) -> dict: "route,expected_metadata_path", [ ( - "/mcp/knowledge_qa", - "/.well-known/oauth-protected-resource/mcp/knowledge_qa", + "/mcp/sample_docs", + "/.well-known/oauth-protected-resource/mcp/sample_docs", ), ( - "/knowledge_qa/mcp", - "/.well-known/oauth-protected-resource/knowledge_qa/mcp", + "/sample_docs/mcp", + "/.well-known/oauth-protected-resource/sample_docs/mcp", ), ], ) @@ -517,9 +573,9 @@ async def test_passthrough_cold_start_emits_401_with_matching_resource_metadata( global_mcp_server_manager.registry.clear() passthrough_server = MCPServer( server_id="pt-cold-start", - name="knowledge_qa", - server_name="knowledge_qa", - alias="knowledge_qa", + name="sample_docs", + server_name="sample_docs", + alias="sample_docs", url="https://upstream.example.com/mcp", transport=MCPTransport.http, auth_type=MCPAuth.none, @@ -534,18 +590,17 @@ async def test_passthrough_cold_start_emits_401_with_matching_resource_metadata( if route.startswith("/mcp/"): scope = _make_scope(route) else: - scope = _make_scope(f"/mcp/knowledge_qa") + scope = _make_scope("/mcp/sample_docs") scope["_original_path"] = route # Verify cold-start detection fires for this path - effective_path = scope.get("_original_path") or scope.get("path", "") servers = _parse_mcp_server_names_from_path( scope.get("path", "") # always /mcp/{name} by the time admission runs ) - assert _is_mcp_passthrough_cold_start(scope, servers) is True + assert _is_mcp_passthrough_cold_start(scope, servers, client_ip=None) is True # Verify resource_metadata_url form selection - server_name = "knowledge_qa" + server_name = "sample_docs" base_url = "http://localhost:4000" path = scope.get("_original_path") or scope.get("path", "") or "" if path.startswith(f"/{server_name}/mcp"): @@ -594,7 +649,7 @@ def test_is_mcp_passthrough_cold_start_false_for_oauth2_server(): global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server scope = _make_scope("/mcp/keycloak_whoami") - result = _is_mcp_passthrough_cold_start(scope, ["keycloak_whoami"]) + result = _is_mcp_passthrough_cold_start(scope, ["keycloak_whoami"], client_ip=None) assert result is False @@ -605,8 +660,8 @@ def test_is_mcp_passthrough_cold_start_false_for_empty_servers(): ) scope = _make_scope("/mcp") - assert _is_mcp_passthrough_cold_start(scope, None) is False - assert _is_mcp_passthrough_cold_start(scope, []) is False + assert _is_mcp_passthrough_cold_start(scope, None, client_ip=None) is False + assert _is_mcp_passthrough_cold_start(scope, [], client_ip=None) is False # -------------------------------------------------------------------------- @@ -617,10 +672,10 @@ def test_is_mcp_passthrough_cold_start_false_for_empty_servers(): @pytest.mark.parametrize( "path,expected", [ - ("/mcp/knowledge_qa", ["knowledge_qa"]), - ("/mcp/knowledge_qa/tools/list", ["knowledge_qa"]), - ("/knowledge_qa/mcp", ["knowledge_qa"]), - ("/knowledge_qa/mcp/tools/list", ["knowledge_qa"]), + ("/mcp/sample_docs", ["sample_docs"]), + ("/mcp/sample_docs/tools/list", ["sample_docs"]), + ("/sample_docs/mcp", ["sample_docs"]), + ("/sample_docs/mcp/tools/list", ["sample_docs"]), ("/mcp", None), ("/mcp/", None), ("/other/path", None), diff --git a/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py b/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py index 3b13ef3641f..2532fe1d107 100644 --- a/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py +++ b/tests/test_litellm/proxy/auth/test_mcp_ip_filtering.py @@ -5,8 +5,9 @@ Tests that internal callers see all MCP servers while external callers only see servers with available_on_public_internet=True. """ -import ipaddress -from unittest.mock import patch +from unittest.mock import MagicMock, patch + +from fastapi import Request from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -58,6 +59,53 @@ class TestIsInternalIp: assert IPAddressUtils.is_internal_ip("not-an-ip") is False +class TestMCPClientIPExtraction: + def test_ignores_xff_without_trusted_proxy_ranges(self): + request = MagicMock(spec=Request) + request.client = MagicMock() + request.client.host = "203.0.113.5" + request.headers = {"x-forwarded-for": "10.0.0.1"} + + result = IPAddressUtils.get_mcp_client_ip( + request, + general_settings={"use_x_forwarded_for": True}, + ) + + assert result == "203.0.113.5" + + def test_honours_xff_from_trusted_proxy(self): + request = MagicMock(spec=Request) + request.client = MagicMock() + request.client.host = "10.0.0.5" + request.headers = {"x-forwarded-for": "192.168.1.10"} + + result = IPAddressUtils.get_mcp_client_ip( + request, + general_settings={ + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + }, + ) + + assert result == "192.168.1.10" + + def test_ignores_xff_from_untrusted_direct_caller(self): + request = MagicMock(spec=Request) + request.client = MagicMock() + request.client.host = "203.0.113.5" + request.headers = {"x-forwarded-for": "10.0.0.1"} + + result = IPAddressUtils.get_mcp_client_ip( + request, + general_settings={ + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + }, + ) + + assert result == "203.0.113.5" + + class TestMCPServerIPFiltering: """Tests that external callers only see public MCP servers.""" 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 01283461de0..265b917ba92 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 @@ -1698,6 +1698,37 @@ class TestTemporaryMCPSessionEndpoints: _, call_kwargs = auth_builder_mock.call_args assert call_kwargs["api_key"] == "Bearer sk-header-key" + def test_mcp_oauth_authorize_token_routes_use_browser_auth_dependency(self): + from fastapi.routing import APIRoute + + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _mcp_oauth_user_api_key_auth, + router, + ) + + oauth_routes = { + route.path: route + for route in router.routes + if isinstance(route, APIRoute) + and route.path + in { + "/v1/mcp/server/oauth/{server_id}/authorize", + "/v1/mcp/server/oauth/{server_id}/token", + } + } + + assert set(oauth_routes) == { + "/v1/mcp/server/oauth/{server_id}/authorize", + "/v1/mcp/server/oauth/{server_id}/token", + } + for route in oauth_routes.values(): + dependency_names = { + dependant.name + for dependant in route.dependant.dependencies + if dependant.call is _mcp_oauth_user_api_key_auth + } + assert dependency_names == {None, "user_api_key_dict"} + @pytest.mark.asyncio async def test_mcp_authorize_proxies_to_discoverable_endpoint(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import (