scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

fix(mcp): forward short OAuth state upstream, keep session in a cookie (#32146)

Browse source 
* fix(mcp): forward short OAuth state upstream, keep session in a cookie

Some upstream authorization servers reject the OAuth authorize request with
"state parameter too long" because LiteLLM replaced the client's short state
with its own long encrypted session blob (base_url, original state, PKCE, client
redirect_uri) and sent that upstream as state.

Forward a short random handle as the upstream state instead, and carry the
encrypted session in a per-flow HttpOnly, SameSite=lax cookie bound to that
handle. The browser replays the cookie on /callback, so the session is recovered
without any server-side store and the client still gets its own original state
back. /callback falls back to decoding state directly when no cookie is present,
so flows in flight across a deploy keep working.

Resolves LIT-4197

* test(mcp): cover /callback error path cookie read and clear

The happy-path regression test already asserts the short-handle -> cookie round
trip. Add a focused test for the IdP-error branch of /callback: it must recover
the client's original state from the per-flow cookie (not the short handle),
propagate the error to the client's redirect_uri, and expire the one-time
cookie. Fails if the error path stops reading or clearing the cookie.
This commit is contained in:
tin-berri 2026-07-06 15:47:37 -07:00 • committed by GitHub
parent 101f246fc5
commit fc3c21e837
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 250 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

107
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -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("<html><body>Authentication incomplete. You can close this window.</body></html>")
response = HTMLResponse("<html><body>Authentication incomplete. You can close this window.</body></html>")
_clear_oauth_state_cookie(response, request, state)
return response
# ------------------------------

157
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -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."""