diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 70e4bac01b5..02184effc95 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -145,10 +145,16 @@ every session-token claim set. Four 256-character ASCII URIs seal to roughly 1.5 the encoded client_id is checked against its own cap before registration succeeds. VS Code registers four callbacks for its web and desktop environments.""" -MAX_STATE_LENGTH: Final = 1024 -"""Bound on the client ``state`` sealed into the flow cookie and echoed on the auth-code -redirect. An unbounded ``state`` can push the sealed cookie past the browser's ~4KB cap -(silently dropped, breaking the flow); spec clients send a short opaque value.""" +MAX_STATE_LENGTH: Final = 2048 +"""Coarse bound on the client ``state`` sealed into the flow cookie and echoed on the +auth-code redirect. RFC 6749 sets no limit and real clients send long opaque values +(Gemini web sends ~1.1KB), so this only rejects clearly abusive input; the sealed cookie +itself is checked against ``MAX_FLOW_COOKIE_BYTES``.""" + +MAX_FLOW_COOKIE_BYTES: Final = 4000 +"""Bound on the flow cookie's name plus sealed value. Browsers drop cookies over ~4096 +bytes silently, which would break the flow at the connect page with an opaque "unknown or +expired connect flow", so an oversized authorize request is rejected up front instead.""" MIN_CODE_VERIFIER_LENGTH: Final = 43 MAX_CODE_VERIFIER_LENGTH: Final = 128 @@ -758,8 +764,11 @@ def aggregate_authorize( upstream_scopes=upstream_scopes, ) connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),)) + sealed_flow: Final = _seal(_UNPREFIXED, flow) + if _flow_cookie_too_large(handle, sealed_flow): + return _flow_cookie_too_large_error() response: Final = RedirectResponse(connect_url, status_code=303) - _set_flow_cookie(response, request, handle, flow) + _set_flow_cookie(response, request, handle, sealed_flow) return response @@ -804,6 +813,9 @@ async def native_client_authorize( resource_server_ids=None, audience=PROXY_API_AUDIENCE, ) + sealed_flow: Final = _seal(_UNPREFIXED, flow) + if _flow_cookie_too_large(handle, sealed_flow): + return _flow_cookie_too_large_error() page: Final = render_native_client_consent_page( client_origin=_origin_only(redirect_uri), user_id=session_user_id, @@ -812,7 +824,7 @@ async def native_client_authorize( complete_url=f"{base_url}/authorize/complete", ) response: Final = HTMLResponse(page, headers=_CONSENT_PAGE_HEADERS) - _set_flow_cookie(response, request, handle, flow) + _set_flow_cookie(response, request, handle, sealed_flow) return response @@ -936,11 +948,19 @@ def _new_connect_flow( ) -def _set_flow_cookie(response: Response, request: Request, handle: str, flow: _ConnectFlow) -> None: +def _flow_cookie_too_large(handle: str, sealed_flow: str) -> bool: + return len(_flow_cookie_name(handle)) + len(sealed_flow) > MAX_FLOW_COOKIE_BYTES + + +def _flow_cookie_too_large_error() -> Response: + return _oauth_error(400, "invalid_request", "authorization request is too large (shorten state or client_id)") + + +def _set_flow_cookie(response: Response, request: Request, handle: str, sealed_flow: str) -> None: path, secure = _cookie_path_and_secure(request) response.set_cookie( key=_flow_cookie_name(handle), - value=_seal(_UNPREFIXED, flow), + value=sealed_flow, max_age=CONNECT_FLOW_TTL_SECONDS, path=path, secure=secure, diff --git a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 530667fdfdd..7536816871b 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -23,6 +23,8 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( JWT_BEARER_GRANT_TYPE, MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS, MAX_CLIENT_ID_LENGTH, + MAX_FLOW_COOKIE_BYTES, + MAX_STATE_LENGTH, SUBJECT_TOKEN_TYPES, TOKEN_EXCHANGE_GRANT_TYPE, AssertionIdentity, @@ -631,7 +633,7 @@ async def test_authorize_rejects_over_long_state(): request=_request(query=f"client_id={client_id}"), client_id=client_id, redirect_uri=REDIRECT_URI, - state="s" * 2000, + state="s" * (MAX_STATE_LENGTH + 1), code_challenge=CODE_CHALLENGE, code_challenge_method="S256", response_type="code", @@ -641,6 +643,41 @@ async def test_authorize_rejects_over_long_state(): assert json.loads(response.body)["error"] == "invalid_request" +def _authorize_with_state(client_id, redirect_uri, state): + return aggregate_authorize( + request=_request(query=f"client_id={client_id}"), + client_id=client_id, + redirect_uri=redirect_uri, + state=state, + code_challenge=CODE_CHALLENGE, + code_challenge_method="S256", + response_type="code", + session_user_id="u1", + ) + + +@pytest.mark.asyncio +async def test_authorize_accepts_long_opaque_state_from_gemini(): + """Gemini web sends a ~1.1KB opaque state (#46414); RFC 6749 sets no limit on it.""" + client_id = (await _register([REDIRECT_URI]))["client_id"] + response = _authorize_with_state(client_id, REDIRECT_URI, "s" * 1112) + assert response.status_code == 303 + flow_cookie = next(header.decode() for name, header in response.raw_headers if name == b"set-cookie") + name_and_value = flow_cookie.split(";", 1)[0] + assert len(name_and_value) - 1 <= MAX_FLOW_COOKIE_BYTES + + +@pytest.mark.asyncio +async def test_authorize_rejects_request_whose_flow_cookie_would_exceed_browser_limit(): + """A within-cap state plus a maximal registration must not set a cookie the browser + would silently drop; the request is rejected up front instead.""" + client_id = (await _register(list(MAX_LENGTH_REDIRECT_URIS)))["client_id"] + response = _authorize_with_state(client_id, MAX_LENGTH_REDIRECT_URIS[0], "s" * MAX_STATE_LENGTH) + assert response.status_code == 400 + assert json.loads(response.body)["error"] == "invalid_request" + assert not any(name == b"set-cookie" for name, _ in response.raw_headers) + + @pytest.mark.asyncio async def test_non_ascii_code_challenge_fails_grant_not_500(): """A non-ASCII code_challenge (unvalidated from the client) must yield a clean