From f5b45644663d0dacd9559a5398073610c972783d Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 23 Apr 2026 03:39:41 +0000 Subject: [PATCH] fix(mcp): OAuth 2.1 spec compliance + cookie exp + loopback edge cases MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address codex review P1 + P2 findings: - BYOK /token now accepts OAuth 2.1 clients that omit redirect_uri (draft-15 §4.1.3 dropped the requirement). When the client does submit a value, equality is still enforced vs the /authorize record. PKCE + client_id binding cover the security role redirect_uri played under RFC 6749. - _user_id_from_session_cookie requires the ``exp`` claim on the UI session JWT (PyJWT options={"require": ["exp"]}) so leaked cookies have a bounded lifetime. - validate_loopback_redirect_uri rejects URIs with a fragment (RFC 6749 §3.1.2) and catches malformed-URI ValueError so unparseable input surfaces as 400 invalid_request instead of 500. --- .../mcp_server/byok_oauth_endpoints.py | 26 ++++- .../_experimental/mcp_server/oauth_utils.py | 12 +- .../mcp_server/test_byok_oauth_endpoints.py | 110 +++++++++++++++++- 3 files changed, 141 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 ab753bae4fb..2f5973ca371 100644 --- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py @@ -105,7 +105,15 @@ def _user_id_from_session_cookie(request: Request) -> Optional[str]: if not token: return None try: - payload = jwt.decode(token, master_key, algorithms=["HS256"]) + payload = jwt.decode( + token, + master_key, + algorithms=["HS256"], + # Require an expiry claim so a leaked UI session cookie has a + # bounded lifetime. PyJWT verifies exp by default when present; + # require=["exp"] additionally rejects tokens that omit it. + options={"require": ["exp"]}, + ) except jwt.InvalidTokenError: return None if payload.get("type") == "byok_session": @@ -791,10 +799,18 @@ async def byok_token( if not _verify_pkce(code_verifier, record["code_challenge"]): return _oauth_token_error("invalid_grant") - # RFC 6749 §4.1.3 / OAuth 2.1 §4.1.3: if redirect_uri was sent with the - # authorization request, the token request MUST include the identical - # value. Enforce exact match. - if record.get("redirect_uri") and redirect_uri != record["redirect_uri"]: + # RFC 6749 §4.1.3: if redirect_uri was sent with the authorization + # request, the token request MUST include the identical value. + # OAuth 2.1 draft-15 §4.1.3 drops this requirement — strict OAuth 2.1 + # clients will omit it. Enforce equality ONLY when the client + # actually submitted a value, so we stay RFC 6749-backward-compatible + # without breaking OAuth 2.1 clients. PKCE + client_id binding + # (checked below) cover the security role redirect_uri played. + if ( + record.get("redirect_uri") + and redirect_uri + and redirect_uri != record["redirect_uri"] + ): return _oauth_token_error("invalid_grant") # RFC 6749 §4.1.3: if the client was identified at /authorize, the diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index ad1b97b1ac0..b13cf83058c 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -24,9 +24,17 @@ def validate_loopback_redirect_uri(redirect_uri: str) -> None: ``"127.0.0.1"`` alone would miss ``127.0.0.2`` and the full-form IPv6 loopback ``0:0:0:0:0:0:0:1``. """ - parsed = urlparse(redirect_uri) + try: + parsed = urlparse(redirect_uri) + except ValueError: + raise HTTPException(status_code=400, detail="invalid_request") if parsed.scheme not in ("http", "https"): raise HTTPException(status_code=400, detail="invalid_request") + # Fragments are not allowed in OAuth redirect URIs (RFC 6749 §3.1.2) + # — rejecting them prevents a ``http://127.0.0.1/cb#frag?code=...`` + # from silently eating the authorization code. + if parsed.fragment: + raise HTTPException(status_code=400, detail="invalid_request") host = (parsed.hostname or "").lower() if host == "localhost": return @@ -34,5 +42,7 @@ def validate_loopback_redirect_uri(redirect_uri: str) -> None: if ip_address(host).is_loopback: return except ValueError: + # Unparseable host (malformed IPv6, etc.) — treat as invalid, + # don't let it bubble up as a 500. pass raise HTTPException(status_code=400, detail="invalid_request") 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 c298731648a..48dfc1715f7 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 @@ -708,7 +708,11 @@ def test_authorize_post_accepts_ui_session_cookie(unauthenticated_client): with patch("litellm.proxy.proxy_server.master_key", "test-master-key"): cookie_jwt = _jwt.encode( - {"user_id": "browser-user-42", "login_method": "sso"}, + { + "user_id": "browser-user-42", + "login_method": "sso", + "exp": int(time.time()) + 3600, + }, "test-master-key", algorithm="HS256", ) @@ -998,3 +1002,107 @@ async def test_token_endpoint_missing_master_key_preserves_code_and_db(): assert code in _byok_auth_codes # Credential never written — no inconsistent DB state. mock_store.assert_not_awaited() + + +def test_authorize_post_rejects_cookie_without_exp(unauthenticated_client): + """Defense-in-depth: UI session cookies must carry an ``exp`` claim + so a leaked cookie has a bounded lifetime. A master-key-signed JWT + without ``exp`` is rejected at decode time (PyJWT + ``options={"require": ["exp"]}``).""" + import jwt as _jwt + + with patch("litellm.proxy.proxy_server.master_key", "real-master-key"): + no_exp = _jwt.encode( + {"user_id": "u", "login_method": "sso"}, + "real-master-key", + algorithm="HS256", + ) + resp = _authorize_post_with_cookie(unauthenticated_client, no_exp, api_key="k") + assert resp.status_code == 401 + + +def test_authorize_post_rejects_expired_cookie(unauthenticated_client): + """An expired UI session cookie is rejected, not accepted as valid.""" + import jwt as _jwt + + with patch("litellm.proxy.proxy_server.master_key", "real-master-key"): + expired = _jwt.encode( + { + "user_id": "u", + "login_method": "sso", + "exp": int(time.time()) - 60, # 1 minute ago + }, + "real-master-key", + algorithm="HS256", + ) + resp = _authorize_post_with_cookie(unauthenticated_client, expired, api_key="k") + assert resp.status_code == 401 + + +@pytest.mark.asyncio +async def test_token_endpoint_accepts_oauth21_client_omitting_redirect_uri(): + """OAuth 2.1 draft-15 §4.1.3 dropped the redirect_uri requirement at + the token endpoint. A strict OAuth 2.1 client will omit the value — + LiteLLM must accept that and rely on PKCE + client_id binding for + the security role redirect_uri played under RFC 6749. + + Enforcement still fires when the client DOES submit a value that + disagrees with the record (see test_token_endpoint_rejects_redirect_uri_mismatch). + """ + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( + byok_token, + ) + + verifier = "verifier_for_oauth21_no_redirect_uri_omit_ok!" + 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": "claude-desktop", + "user_id": "u", + "expires_at": time.time() + 60, + } + + with ( + patch( + "litellm.proxy._experimental.mcp_server.byok_oauth_endpoints.store_user_credential", + new=AsyncMock(), + ), + patch("litellm.proxy.proxy_server.master_key", "test-master"), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + ): + result = await byok_token( + request=MagicMock(), + grant_type="authorization_code", + code=code, + redirect_uri="", # OAuth 2.1 client omits it + code_verifier=verifier, + client_id="claude-desktop", + ) + assert result.status_code == 200 + + +def test_validate_loopback_redirect_uri_rejects_fragment(): + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_loopback_redirect_uri, + ) + + with pytest.raises(HTTPException) as exc: + validate_loopback_redirect_uri("http://127.0.0.1:3000/cb#code=1") + assert exc.value.status_code == 400 + + +def test_validate_loopback_redirect_uri_rejects_malformed_cleanly(): + """Malformed / unparseable URIs should surface as 400 invalid_request, + not a 500 from an unhandled exception inside ip_address().""" + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_loopback_redirect_uri, + ) + + # Netloc that parses but whose host is neither "localhost" nor a valid IP. + with pytest.raises(HTTPException) as exc: + validate_loopback_redirect_uri("http://[not-an-ip]/cb") + assert exc.value.status_code == 400