feat(mcp): aggregate DCR register, authorize, complete, and token flow for the gateway front door - #33189
Conversation
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR implements the full aggregate DCR OAuth front door for the MCP gateway (
Confidence Score: 3/5The stateless DCR design and cookie-bound complete step are sound, but a transient DB outage during the token exchange permanently invalidates the authorization code, making the gateway's own 503 non-retryable. The mark_used call preceding reload_user in _authorization_code_grant causes valid authorization codes to be permanently invalidated on transient DB outages — a present behavioral defect on the new token exchange path. gateway_dcr_flow.py: the ordering of mark_used and reload_user in _authorization_code_grant (lines 495-500) and the empty-prefix _seal call for the flow cookie (line 321).
|
| Filename | Overview |
|---|---|
| litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py | New module implementing the complete aggregate DCR OAuth flow. Core logic is clean; TOCTOU race between already_used/mark_used allows concurrent code replay, and marking the code as used before reload_user permanently invalidates valid codes on transient DB failures. |
| litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py | Wires new DCR flow into existing root endpoints with routing guards ensuring zero behavioral change for non-DCR clients. |
| litellm/proxy/management_endpoints/ui_sso.py | Adds _is_same_origin_return_path helper; check correctly rejects protocol-relative and backslash variants and is well-tested. |
| tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py | Comprehensive tests covering full register→authorize→complete→token walk plus security edges. No real network calls. |
Reviews (1): Last reviewed commit: "feat(mcp): aggregate DCR register, autho..." | Re-trigger Greptile
| if await guard.already_used(parsed.jti): | ||
| return _oauth_error(400, "invalid_grant", "the authorization code was already used") | ||
| await guard.mark_used(parsed.jti) | ||
| failure = await reload_user(parsed.user_id) | ||
| if failure is not None: | ||
| return _reload_failure_response(failure) |
There was a problem hiding this comment.
Single-use guard has a TOCTOU race and marks the code used before the user check
Two concurrent token requests carrying the same code can both get False from already_used before either writes the mark, so both proceed to mint — a narrow but real asyncio-level race since the event loop can switch coroutines between the two awaits. More concretely, mark_used is called on line 497 before reload_user on line 498. When reload_user returns "unavailable" (transient DB outage), the code is already permanently marked; the client gets a 503 and reasonably retries, but then receives invalid_grant — forcing a complete re-authorization cycle through the browser even though the failure was entirely the gateway's own DB blip.
A safer ordering would finish validation (reload_user) before consuming the single-use slot, and ideally use an atomic set-if-not-exists primitive to close the TOCTOU window.
| def _seal(prefix: str, payload: BaseModel) -> str: | ||
| return prefix + encrypt_value_helper(payload.model_dump_json()) |
There was a problem hiding this comment.
Empty-string prefix makes
_open_sealed prefix check vacuous for flow cookies
_seal("", flow) seals with no prefix, so the prefix check in _open_sealed (value.startswith("")) is always true — the actual protection relies entirely on decrypt_value_helper's authenticated encryption plus Pydantic validation. This is functionally safe but means a sealed value of any other type (e.g. a GatewayDcrClient token) would reach Pydantic deserialization as a _ConnectFlow candidate before being rejected by field mismatch.
Using a distinct non-empty prefix for the flow cookie (e.g. "mcp_cf_") would make the routing intention explicit and add an early-exit path consistent with how GATEWAY_DCR_CLIENT_ID_PREFIX and GATEWAY_AUTH_CODE_PREFIX are used.
| failure = await reload_user(opened.principal.user_id) | ||
| if failure is not None: | ||
| return _reload_failure_response(failure) | ||
| return _session_token_pair(opened.principal, keys, now) |
There was a problem hiding this comment.
Medium: Refresh token replay
This mints a new token pair without invalidating the presented refresh token. An attacker who obtains that token can repeatedly mint fresh access and refresh tokens for up to its 14-day lifetime, including after the legitimate client has refreshed. Consume the refresh token's signed jti with an atomic shared-cache claim before minting, and reject subsequent uses.
PR overviewThis PR adds an aggregated MCP gateway front-door flow for dynamic client registration, authorization, completion, and token issuance. The touched code includes the gateway DCR flow handling token refresh behavior. There is one open security issue in the token refresh path: refresh tokens can be reused because the presented token is not consumed or invalidated before minting a new token pair. If a refresh token is obtained by an attacker, it could be replayed during its validity window to maintain access under that client’s scope. No issues have been addressed yet, so the PR still needs a fix for refresh-token rotation/replay protection before the flow is in a safer state. Open issues (1)
Fixed/addressed: 0 · PR risk: 6/10 |
9a81539 to
436673f
Compare
…w for the gateway front door
- atomic single-use guard (async_increment_cache) + reload-before-claim so a transient DB blip does not burn a valid code - PKCE verify over bytes so a non-ASCII code_challenge fails invalid_grant instead of raising a 500; validate code_verifier length (RFC 7636) - flag-off byte-identical for a server literally named mcp (AS well-known delegates to the named-server document) - connect flow is single-use (atomic jti claim) so a double-submit cannot mint two codes - extra=forbid on the sealed models; bound state length; drop unused request param and coarse dict on register - _reload_failure_response exhaustive match+assert_never; dedupe ReloadUserFailure with _KeyResolutionFailure - reject control/whitespace chars in the same-origin return_to
f9b90ee to
90eb828
Compare
186c862
into
litellm_lit3637_session_token
Relevant issues
Stacked on #33188 (UI session cookie exp) -> #33182 (session tokens) -> #33174 (aggregate discovery). The base of this PR is
litellm_lit3637_ui_session_exp, not staging; review and merge after thoseLinear ticket
Part of LIT-3637 (PR 3 of the stack: the aggregate DCR register/authorize/complete/token flow that turns the discovery front door into a working sign-in)
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Proven live end to end at commit head on a real proxy (:4138, Postgres) with
mcp_gateway_dcr: true, driving the exact request sequence a DCR client (Claude Desktop, MCP Inspector) makes. The upstream-server authorization step is the existing per-user vault flow and is not re-exercised here; this PR is the identity front doorType
🆕 New Feature
Changes
PR 1 gave the aggregate
/mcpendpoint OAuth discovery; this PR is the flow that discovery points at, so an OAuth-only DCR client can actually sign in. Everything is gated behindmcp_gateway_dcrand behind thellm_dcrc_client-id prefix, so with the flag off, or for any client that did not register through this flow, the existing per-server authorize/token behavior is byte-identical. The whole flow lives in one new module,gateway_dcr_flow.py, wired into the existing root/register,/authorize, and/tokenhandlers plus one newPOST /authorize/completeRegistration is stateless. The
client_idis the registration: the client's redirect URIs are sealed into it with the repo's authenticated symmetric helper, so nothing is persisted, open registration cannot fill storage, and a forged or tampered client_id simply fails to open. Clients are always registered as public (token_endpoint_auth_method "none") because the gateway issues no client secrets; mandatory S256 PKCE is what protects the code. Redirect URIs must be https, or http on a loopback host for local dev clients (RFC 8252), fragment-free, and bounded in count and length so the sealed id stays smallAuthorize validates the client and redirect URI, requires S256 PKCE, and interposes LiteLLM sign-in. A validation failure answers 400 directly and never redirects, per RFC 6749 4.1.2.1: an unvalidated redirect URI must not receive an error redirect. Without a session cookie the browser is sent through
/sso/key/generatewith a strictly-relativereturn_to(a new same-origin arm in the SSO callback, so login can only ever bounce the browser back to this same authorize request, never off-origin). With a session, the flow parameters and the SSO user are sealed into a per-flow HttpOnly cookie (the same handle-plus-cookie pattern the upstream OAuth relay already uses) and the browser lands on the connect grid, where the user authorizes individual servers before finishingComplete is a deliberate, separate step reached by POST, bound to the HttpOnly flow cookie and to an exact match between the signed-in user and the user sealed into the flow. A GET could be triggered cross-site; a POST with a SameSite=Lax cookie and the identity check cannot mint a code for a victim's session. It seals a short-lived (120s), single-use, PKCE- and client-bound authorization code and redirects back to the registered redirect URI
Token exchanges the code for the identity-only session tokens from the module below it in the stack. The authorization_code grant verifies the code has not expired, is bound to this client and redirect URI, passes PKCE, and has not already been used (a best-effort single-use guard over the shared cache, in-memory or Redis), then re-validates that the litellm user is still active before minting. The refresh_token grant re-validates the same way and rotates the pair, with the refresh token bound to its issuing client (RFC 6749 section 6). Every error is an RFC 6749 5.2 body carrying no token, code, or URL material
The one change outside this module and its wiring is the SSO callback's same-origin
return_toarm, needed so the login round-trip returns the browser to the authorize request. Sealed values open totally (bad input maps to an OAuth error, never a raise); no upstream server credential appears anywhere in this flow, they are vaulted per user by the existing/v1/mcpauthorize endpoints and resolved at egress by user id (unchanged here)QA runbook
Run a proxy with
mcp_gateway_dcr: true, a master key, and a database, then drive the numbered sequence above with curl (the full script is in the proof). Confirm: the login cookie carries an exp; register returns anllm_dcrc_public client; a signed-in authorize 303s to/ui/chat/integrations?connect_flow=...and sets the per-flow cookie; complete 303s back with anllm_gcode_code and the original state; token returnsllm_session_/llm_srefresh_withexpires_in: 3600; the used code replays asinvalid_grant; refresh rotates; a wrong PKCE verifier isinvalid_grant; and an anonymous authorize 303s to/sso/key/generatewith a relativereturn_to. With the flag off, confirm the aggregate well-known routes 404 and/registerreturns its former dummy bodyFinal Attestation