diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 0ccd7f9b2f9..89d645b6f8a 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,6 +1,7 @@ import asyncio import html as _html import json +import secrets import time from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple @@ -8,7 +9,7 @@ from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import httpx from fastapi import APIRouter, Form, HTTPException, Request -from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse +from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response from pydantic import BaseModel, ValidationError from litellm._logging import verbose_logger @@ -137,6 +138,72 @@ def decode_state_hash(encrypted_state: str) -> dict: return state_data +# LIT-4197: some upstream authorization servers reject an over-long ``state`` +# (the encrypted OAuth session blob routinely exceeds their limit). The upstream +# only needs an opaque value it echoes back on ``/callback``, so we forward a +# short random handle and keep the encrypted session in a per-flow HttpOnly +# cookie bound to that handle. The browser carries the cookie across the +# upstream round trip, so the flow stays correct with no server-side session +# store (works across proxy replicas, unlike an in-process map). +_OAUTH_STATE_COOKIE_PREFIX = "mcp_oauth_state_" +_OAUTH_STATE_COOKIE_TTL_SECONDS = 600 +_OAUTH_STATE_HANDLE_BYTES = 32 + + +def _oauth_state_cookie_name(relay_state: str) -> str: + return f"{_OAUTH_STATE_COOKIE_PREFIX}{relay_state}" + + +def _oauth_state_cookie_path_and_secure(request: Request) -> tuple[str, bool]: + parsed = urlparse(get_request_base_url(request)) + return parsed.path or "/", parsed.scheme == "https" + + +def _set_oauth_state_cookie( + response: Response, + request: Request, + relay_state: str, + encoded_state: str, +) -> None: + path, secure = _oauth_state_cookie_path_and_secure(request) + response.set_cookie( + key=_oauth_state_cookie_name(relay_state), + value=encoded_state, + max_age=_OAUTH_STATE_COOKIE_TTL_SECONDS, + path=path, + secure=secure, + httponly=True, + samesite="lax", + ) + + +def _resolve_encoded_oauth_state(request: Request, state: str) -> str: + """Return the encrypted OAuth session for a ``/callback`` request. + + New flows carry it in a per-flow cookie keyed by the short handle we + forwarded upstream (the IdP echoes that handle back as ``state``). Flows + started before this change - or in flight across a deploy - carry the + encrypted blob directly in ``state``, so fall back to it when the cookie + is absent. + """ + cookie_value = request.cookies.get(_oauth_state_cookie_name(state)) + return cookie_value if cookie_value else state + + +def _clear_oauth_state_cookie(response: Response, request: Request, state: str) -> None: + cookie_name = _oauth_state_cookie_name(state) + if cookie_name not in request.cookies: + return + path, secure = _oauth_state_cookie_path_and_secure(request) + response.delete_cookie( + key=cookie_name, + path=path, + secure=secure, + httponly=True, + samesite="lax", + ) + + 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. @@ -462,11 +529,12 @@ async def authorize_with_server( code_challenge_method=code_challenge_method, client_redirect_uri=redirect_uri, ) + relay_state = secrets.token_urlsafe(_OAUTH_STATE_HANDLE_BYTES) params = { "client_id": mcp_server.client_id if mcp_server.client_id else client_id, "redirect_uri": f"{request_base_url}/callback", - "state": encoded_state, + "state": relay_state, "response_type": response_type or "code", } if scope: @@ -483,7 +551,9 @@ async def authorize_with_server( existing_params = dict(parse_qsl(parsed_auth_url.query)) existing_params.update(params) final_url = urlunparse(parsed_auth_url._replace(query=urlencode(existing_params))) - return RedirectResponse(final_url) + response = RedirectResponse(final_url) + _set_oauth_state_cookie(response, request, relay_state, encoded_state) + return response async def exchange_token_with_server( @@ -1017,17 +1087,19 @@ async def callback( error_description, ) if state: + encoded_state = _resolve_encoded_oauth_state(request, state) try: - state_data = decode_state_hash(state) + state_data = decode_state_hash(encoded_state) original_state = state_data.get("original_state") redirect_uri = _get_validated_client_redirect_uri(request, state_data) - except HTTPException: - # Untrusted/invalid client redirect_uri — surface inline rather - # than blindly forwarding the error to an attacker-controlled URL. - return _render_oauth_error_html(error, error_description) except Exception: - # State could not be decrypted (expired key, tampered, etc.). - return _render_oauth_error_html(error, error_description) + # Untrusted/invalid client redirect_uri (HTTPException), or an + # undecryptable state (expired key, tampered): surface the IdP + # error inline rather than forwarding it to an attacker-controlled + # URL, and drop the one-time cookie we can no longer consume. + response = _render_oauth_error_html(error, error_description) + _clear_oauth_state_cookie(response, request, state) + return response params: Dict[str, str] = {"error": error} if error_description: @@ -1037,7 +1109,9 @@ async def callback( if original_state is not None: params["state"] = original_state complete_returned_url = _append_query_params(redirect_uri, params) - return RedirectResponse(url=complete_returned_url, status_code=302) + response = RedirectResponse(url=complete_returned_url, status_code=302) + _clear_oauth_state_cookie(response, request, state) + return response # No state — nothing to round-trip to. Show the user the error. return _render_oauth_error_html(error, error_description) @@ -1053,7 +1127,8 @@ async def callback( # 3. Successful authorization response. try: - state_data = decode_state_hash(state) + encoded_state = _resolve_encoded_oauth_state(request, state) + state_data = decode_state_hash(encoded_state) original_state = state_data["original_state"] # Re-validate the client redirect URI at the sink. /authorize @@ -1066,14 +1141,18 @@ async def callback( params = {"code": code, "state": original_state} complete_returned_url = _append_query_params(redirect_uri, params) - return RedirectResponse(url=complete_returned_url, status_code=302) + response = RedirectResponse(url=complete_returned_url, status_code=302) + _clear_oauth_state_cookie(response, request, state) + return response except HTTPException: # Re-raise so a non-loopback base_url surfaces as 400 instead of # a generic "authentication incomplete" redirect. raise except Exception: - return HTMLResponse("Authentication incomplete. You can close this window.") + response = HTMLResponse("Authentication incomplete. You can close this window.") + _clear_oauth_state_cookie(response, request, state) + return response # ------------------------------ 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 8ee8f658f6e..fe90fd45856 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 @@ -35,6 +35,7 @@ def _mock_callback_request(base_url: str = "http://localhost:3000/"): req = MagicMock() req.base_url = base_url req.headers = {} + req.cookies = {} return req @@ -2636,6 +2637,162 @@ async def test_oauth_callback_accepts_same_origin_ui_redirect(): assert "state=state-123" in response.headers["location"] +@pytest.mark.asyncio +async def test_authorize_forwards_short_state_and_round_trips_via_cookie(monkeypatch): + """LIT-4197: the ``state`` sent to the upstream authorization server must be + a short opaque handle, not the long encrypted OAuth session (some IdPs + reject an over-long state). The session must instead ride in a per-flow + HttpOnly cookie so ``/callback`` still recovers the client's original state + and redirects back to the client's redirect_uri.""" + from http.cookies import SimpleCookie + from urllib.parse import parse_qs, urlparse + + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _oauth_state_cookie_name, + authorize_with_server, + callback, + decode_state_hash, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + # Real encryption so the cookie value is a genuine encrypted session. + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-for-LIT-4197") + + client_state = "ee230e3dfd4f19c7441941684f39c8a4e0e2c3c61a088e33403df5662b4047b8" + client_redirect_uri = "http://127.0.0.1:6274/oauth/callback/debug" + + server = MCPServer( + server_id="leanix_server", + name="leanix", + server_name="leanix", + alias="leanix", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="upstream-client-id", + authorization_url="https://idp.example.com/oauth/authorize", + token_url="https://idp.example.com/oauth/token", + ) + + authorize_request = MagicMock(spec=Request) + authorize_request.base_url = "https://proxy.example.com/" + authorize_request.headers = {} + + authorize_response = await authorize_with_server( + request=authorize_request, + mcp_server=server, + client_id="upstream-client-id", + redirect_uri=client_redirect_uri, + state=client_state, + code_challenge="challenge", + code_challenge_method="S256", + ) + + location = authorize_response.headers["location"] + upstream_state = parse_qs(urlparse(location).query)["state"][0] + + # The upstream must receive a short handle, not the encrypted session blob. + assert len(upstream_state) <= 64 + assert upstream_state != client_state + + # The encrypted session rides in a per-flow HttpOnly cookie bound to it. + jar = SimpleCookie() + jar.load(authorize_response.headers["set-cookie"]) + cookie_name = _oauth_state_cookie_name(upstream_state) + assert cookie_name in jar + morsel = jar[cookie_name] + assert morsel["httponly"] + assert morsel["samesite"].lower() == "lax" + assert len(morsel.value) > len(upstream_state) + session = decode_state_hash(morsel.value) + assert session["original_state"] == client_state + assert session["client_redirect_uri"] == client_redirect_uri + + # /callback recovers the original state from the cookie (not the handle) and + # redirects back to the client with the client's own state. + callback_request = MagicMock(spec=Request) + callback_request.base_url = "https://proxy.example.com/" + callback_request.headers = {} + callback_request.cookies = {cookie_name: morsel.value} + + callback_response = await callback( + request=callback_request, + code="upstream-auth-code", + state=upstream_state, + ) + + assert callback_response.status_code == 302 + cb_query = parse_qs(urlparse(callback_response.headers["location"]).query) + assert callback_response.headers["location"].startswith(client_redirect_uri) + assert cb_query["code"] == ["upstream-auth-code"] + assert cb_query["state"] == [client_state] + + # The one-time cookie is expired on the callback response so it cannot be replayed. + cleared = SimpleCookie() + cleared.load(callback_response.headers["set-cookie"]) + assert cookie_name in cleared + assert cleared[cookie_name].value == "" + assert cleared[cookie_name]["max-age"] == "0" + + +@pytest.mark.asyncio +async def test_callback_error_path_reads_cookie_and_clears_it(monkeypatch): + """LIT-4197: an IdP error routed through /callback must recover the client's + original state from the cookie (not the short handle), propagate the error to + the client's redirect_uri, and expire the one-time cookie.""" + from http.cookies import SimpleCookie + from urllib.parse import parse_qs, urlparse + + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _oauth_state_cookie_name, + callback, + encode_state_with_base_url, + ) + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-for-LIT-4197") + + client_state = "client-original-state-abc" + client_redirect_uri = "http://127.0.0.1:6274/oauth/callback/debug" + handle = "shortRelayHandle123" + encoded_state = encode_state_with_base_url( + base_url=client_redirect_uri, + original_state=client_state, + client_redirect_uri=client_redirect_uri, + ) + cookie_name = _oauth_state_cookie_name(handle) + + request = MagicMock(spec=Request) + request.base_url = "https://proxy.example.com/" + request.headers = {} + request.cookies = {cookie_name: encoded_state} + + response = await callback( + request=request, + error="access_denied", + error_description="User declined access", + state=handle, + ) + + assert response.status_code == 302 + location = response.headers["location"] + assert location.startswith(client_redirect_uri) + query = parse_qs(urlparse(location).query) + assert query["error"] == ["access_denied"] + # The client's own state is echoed back, recovered from the cookie. + assert query["state"] == [client_state] + + cleared = SimpleCookie() + cleared.load(response.headers["set-cookie"]) + assert cookie_name in cleared + assert cleared[cookie_name].value == "" + assert cleared[cookie_name]["max-age"] == "0" + + @pytest.mark.asyncio async def test_oauth_authorize_includes_scopes_from_server_config(): """Test that authorize endpoint includes scopes from server configuration."""