From c279f4e3efd1d44ba26855c2dea81520d89e1b4a Mon Sep 17 00:00:00 2001 From: gym-cmd <186399764+gym-cmd@users.noreply.github.com> Date: Fri, 15 May 2026 13:56:09 +0100 Subject: [PATCH] fix(mcp): validate oauth callback redirect sink --- .../mcp_server/discoverable_endpoints.py | 10 ++-- litellm/proxy/auth/handle_jwt.py | 2 +- .../mcp_management_endpoints.py | 10 +++- .../mcp_server/test_discoverable_endpoints.py | 49 +++++++++++++++++-- 4 files changed, 60 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 19db31d54d2..303c9d2e5fa 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -661,7 +661,11 @@ async def token_endpoint( async def callback(request: Request, code: str, state: str): try: state_data = decode_state_hash(state) - base_url = state_data["base_url"] + 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 at the sink. /authorize rejects untrusted @@ -670,10 +674,10 @@ async def callback(request: Request, code: str, state: str): # valid indefinitely. Validating here (same-origin OR loopback) # blocks the open-redirect + code-theft primitive even for pre-fix # states while allowing the UI's same-origin callback to work. - validate_trusted_redirect_uri(request, base_url) + validate_trusted_redirect_uri(request, redirect_uri) params = {"code": code, "state": original_state} - complete_returned_url = f"{base_url}?{urlencode(params)}" + complete_returned_url = _append_query_params(redirect_uri, params) return RedirectResponse(url=complete_returned_url, status_code=302) except HTTPException: diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 041766894f2..36ccf58a40d 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -953,7 +953,7 @@ class JWTHandler: raise Exception(f"Validation fails: {str(e)}") async def auth_jwt(self, token: str) -> dict: - decode_kwargs = self._build_decode_kwargs() + decode_kwargs = self._build_decode_kwargs() header = jwt.get_unverified_header(token) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index eb42cd4912d..de81449fd41 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -152,8 +152,14 @@ if MCP_AVAILABLE: UserAPIKeyAuth, UserMCPManagementMode, ) - from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - from litellm.proxy.common_utils.http_parsing_utils import _read_request_body + from litellm.proxy.auth.user_api_key_auth import ( + _user_api_key_auth_builder, + user_api_key_auth, + ) + from litellm.proxy.common_utils.http_parsing_utils import ( + _read_request_body, + populate_request_with_path_params, + ) from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.types.mcp import MCPCredentials 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 e4521ffb964..6cf1e8184f5 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 @@ -1893,11 +1893,11 @@ async def test_discovery_root_does_not_expose_private_server_for_external_client "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.get_mcp_client_ip", return_value="198.51.100.10", ): - authorization_response = _build_oauth_authorization_server_response( + authorization_response = await _build_oauth_authorization_server_response( request=mock_request, mcp_server_name=None, ) - resource_response = _build_oauth_protected_resource_response( + resource_response = await _build_oauth_protected_resource_response( request=mock_request, mcp_server_name=None, use_standard_pattern=False, @@ -1922,6 +1922,10 @@ 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", @@ -1938,6 +1942,7 @@ async def test_oauth_callback_redirects_with_state(): # Call callback endpoint with code and state response = await callback( + request=mock_request, code="test_authorization_code_12345", state="encrypted_state_value", ) @@ -1967,6 +1972,10 @@ 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: @@ -1981,6 +1990,7 @@ async def test_oauth_callback_preserves_client_redirect_uri_query(): } response = await callback( + request=mock_request, code="test_authorization_code_12345", state="encrypted_state_value", ) @@ -2003,6 +2013,10 @@ 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" @@ -2011,6 +2025,7 @@ async def test_oauth_callback_handles_invalid_state(): # Call callback endpoint with invalid state response = await callback( + request=mock_request, code="test_code", state="invalid_encrypted_state", ) @@ -2390,6 +2405,10 @@ 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: @@ -2401,7 +2420,11 @@ async def test_callback_revalidates_loopback_on_decoded_base_url(): "client_redirect_uri": "https://attacker.example.com/cb", } with pytest.raises(HTTPException) as exc_info: - await callback(code="stolen_code", state="encrypted_stale_state") + await callback( + request=mock_request, + code="stolen_code", + state="encrypted_stale_state", + ) assert exc_info.value.status_code == 400 @@ -2412,6 +2435,10 @@ 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: @@ -2423,7 +2450,11 @@ async def test_callback_revalidates_loopback_on_decoded_client_redirect_uri(): "client_redirect_uri": "https://attacker.example.com/cb", } with pytest.raises(HTTPException) as exc_info: - await callback(code="stolen_code", state="encrypted_stale_state") + await callback( + request=mock_request, + code="stolen_code", + state="encrypted_stale_state", + ) assert exc_info.value.status_code == 400 @@ -2434,6 +2465,10 @@ 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: @@ -2443,7 +2478,11 @@ async def test_callback_rejects_state_missing_redirect_uri(): "code_challenge_method": None, } with pytest.raises(HTTPException) as exc_info: - await callback(code="code", state="encrypted_malformed_state") + await callback( + request=mock_request, + code="code", + state="encrypted_malformed_state", + ) assert exc_info.value.status_code == 400