From 862bad363ef672e5ca130699767cd6f8eed7a665 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Wed, 22 Apr 2026 21:26:27 +0000 Subject: [PATCH] fix(mcp): validate preconditions before consuming auth code MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address two Greptile P2 findings on the BYOK OAuth endpoints: - GET /v1/mcp/oauth/authorize now runs _validate_redirect_uri up front so a non-loopback redirect_uri is rejected before the HTML form is rendered. Previously the user typed an API key, submitted, and got a 400 with no form state. - POST /v1/mcp/oauth/token moves the master_key guard ahead of the code-consumption and credential-store steps. Without this, a proxy with master_key unset would burn the code and persist the credential but return an error — leaving the user with no way to retrieve a session token without restarting the whole flow. --- .../mcp_server/byok_oauth_endpoints.py | 19 ++++-- .../mcp_server/test_byok_oauth_endpoints.py | 68 ++++++++++++++++++- 2 files changed, 80 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py index 2992a34dded..7b8716f70f2 100644 --- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py @@ -675,6 +675,9 @@ async def byok_authorize_get( raise HTTPException(status_code=400, detail="response_type must be 'code'") if not redirect_uri: raise HTTPException(status_code=400, detail="redirect_uri is required") + # Validate here too so the user sees the rejection before typing their + # API key into the HTML form (the POST handler also validates). + _validate_redirect_uri(redirect_uri) if not code_challenge: raise HTTPException(status_code=400, detail="code_challenge is required") @@ -828,9 +831,6 @@ async def byok_token( if record.get("client_id") and client_id != record["client_id"]: return _oauth_token_error("invalid_grant") - # Consume the code (one-time use) - del _byok_auth_codes[code] - server_id: str = record["server_id"] api_key_value: str = record["api_key"] # user_id is stamped by the authenticated /authorize POST. No client_id @@ -842,6 +842,16 @@ async def byok_token( if not user_id: return _oauth_token_error("invalid_grant") + # Verify preconditions that would fail token issuance BEFORE consuming + # the code or writing to the DB — otherwise a misconfigured proxy + # (missing master_key) silently persists the user's credential without + # ever returning an access token, and the user has no way to recover. + if master_key is None: + return _oauth_token_error("server_error", status=500) + + # Consume the code (one-time use) + del _byok_auth_codes[code] + # Persist the BYOK credential if prisma_client is not None: try: @@ -871,9 +881,6 @@ async def byok_token( "byok_token: prisma_client is None — credential not persisted" ) - if master_key is None: - return _oauth_token_error("server_error", status=500) - now = int(time.time()) payload = { "user_id": user_id, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py index f0ae33f2956..c298731648a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py @@ -160,7 +160,7 @@ def test_authorize_get_wrong_response_type(client): resp = client.get( "/v1/mcp/oauth/authorize", params={ - "redirect_uri": "https://example.com/cb", + "redirect_uri": "http://127.0.0.1:3000/cb", "response_type": "token", "code_challenge": "abc", }, @@ -169,6 +169,23 @@ def test_authorize_get_wrong_response_type(client): assert resp.status_code == 400 +def test_authorize_get_rejects_non_loopback_redirect_uri(client): + """GET /authorize validates redirect_uri up front so the user sees + the rejection before typing an API key into the HTML form — matches + the POST handler's rule and avoids the ``user fills form → POST 400 + with no form state`` UX.""" + resp = client.get( + "/v1/mcp/oauth/authorize", + params={ + "redirect_uri": "https://attacker.example.com/cb", + "response_type": "code", + "code_challenge": "abc", + }, + follow_redirects=False, + ) + assert resp.status_code == 400 + + # --------------------------------------------------------------------------- # Authorization POST endpoint # --------------------------------------------------------------------------- @@ -932,3 +949,52 @@ def test_authorize_post_accepts_ipv6_loopback_full_form(client): follow_redirects=False, ) assert resp.status_code == 302 + + +@pytest.mark.asyncio +async def test_token_endpoint_missing_master_key_preserves_code_and_db(): + """If master_key is unset, /token must reject BEFORE consuming the + code or writing the credential — otherwise a misconfigured deploy + burns the code and persists the key with no way for the user to + retrieve a session token without restarting the flow.""" + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( + byok_token, + ) + + verifier = "verifier_for_missing_master_key_test_long_!" + challenge = _make_challenge(verifier) + code = str(uuid.uuid4()) + _byok_auth_codes[code] = { + "api_key": "k", + "server_id": "sid", + "code_challenge": challenge, + "redirect_uri": "http://127.0.0.1:3000/cb", + "client_id": "", + "user_id": "u", + "expires_at": time.time() + 60, + } + + mock_store = AsyncMock() + with ( + patch( + "litellm.proxy._experimental.mcp_server.byok_oauth_endpoints.store_user_credential", + mock_store, + ), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.master_key", None), + ): + result = await byok_token( + request=MagicMock(), + grant_type="authorization_code", + code=code, + redirect_uri="http://127.0.0.1:3000/cb", + code_verifier=verifier, + client_id="", + ) + + assert result.status_code == 500 + assert json.loads(result.body) == {"error": "server_error"} + # Code still present — user can retry once master_key is configured. + assert code in _byok_auth_codes + # Credential never written — no inconsistent DB state. + mock_store.assert_not_awaited()