diff --git a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py index 48884d82274..2f5973ca371 100644 --- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py @@ -19,10 +19,10 @@ import html as _html_module import time import uuid from typing import Dict, Optional, cast -from urllib.parse import urlencode, urlparse +from urllib.parse import urlencode import jwt -from fastapi import APIRouter, Form, HTTPException, Request +from fastapi import APIRouter, Depends, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse from litellm._logging import verbose_proxy_logger @@ -30,6 +30,11 @@ from litellm.proxy._experimental.mcp_server.db import store_user_credential from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( get_request_base_url, ) +from litellm.proxy._experimental.mcp_server.oauth_utils import ( + TOKEN_NO_CACHE_HEADERS, + validate_loopback_redirect_uri, +) +from litellm.proxy._types import UserAPIKeyAuth # --------------------------------------------------------------------------- # In-memory store for pending authorization codes. @@ -69,6 +74,65 @@ def _purge_expired_codes() -> None: del _byok_auth_codes[k] +def _oauth_token_error(code: str, status: int = 400) -> JSONResponse: + """RFC 6749 §5.2 token-endpoint error body: ``{"error": ""}``. + FastAPI's default ``HTTPException`` renders ``{"detail": ...}`` which + spec-compliant OAuth clients parsing the ``error`` field won't recognize. + """ + return JSONResponse( + status_code=status, content={"error": code}, headers=TOKEN_NO_CACHE_HEADERS + ) + + +def _user_id_from_session_cookie(request: Request) -> Optional[str]: + """Return user_id from the UI ``token`` cookie (HS256-signed with + ``master_key``), or None if missing/invalid. + + The /token endpoint in this file ALSO issues master-key-signed JWTs + (type="byok_session") for MCP-client-side use. They must not be + accepted here as UI sessions — otherwise a leaked byok_session token + could be replayed as a cookie to re-authorize BYOK writes. Distinguish + by requiring a ``login_method`` claim (UI tokens set ``"sso"`` or + ``"username_password"``; byok_session tokens never set it) and + rejecting any token whose ``type`` identifies it as non-UI. + """ + # Inline import avoids a circular dep (proxy_server -> mcp_server router). + from litellm.proxy.proxy_server import master_key + + if not master_key: + return None + token = request.cookies.get("token") + if not token: + return None + try: + 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": + return None + if payload.get("login_method") not in ("sso", "username_password"): + return None + user_id = payload.get("user_id") + return user_id if isinstance(user_id, str) and user_id else None + + +async def _byok_session_auth(request: Request) -> UserAPIKeyAuth: + """Require the UI session cookie. Programmatic BYOK management uses + ``POST /v1/mcp/server/{id}/user-credential`` instead.""" + user_id = _user_id_from_session_cookie(request) + if not user_id: + raise HTTPException(status_code=401, detail="login_required") + return UserAPIKeyAuth(api_key="byok_session_cookie", user_id=user_id) + + def _build_authorize_html( server_name: str, server_initial: str, @@ -582,11 +646,18 @@ async def byok_authorize_get( The MCP client navigates the user here; the user types their API key and clicks "Connect & Authorize", which POSTs back to this same path. + + This GET is intentionally unauthenticated: it only renders HTML with no + state change. The POST handler enforces ``user_api_key_auth`` and pins + the stored credential to the authenticated session. """ if response_type != "code": 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_loopback_redirect_uri(redirect_uri) if not code_challenge: raise HTTPException(status_code=400, detail="code_challenge is required") @@ -636,6 +707,7 @@ async def byok_authorize_post( state: str = Form(default=""), server_id: str = Form(default=""), api_key: str = Form(...), + user_api_key_dict: UserAPIKeyAuth = Depends(_byok_session_auth), ) -> RedirectResponse: """ Process the BYOK API-key form submission. @@ -645,10 +717,7 @@ async def byok_authorize_post( """ _purge_expired_codes() - # Validate redirect_uri scheme to prevent open redirect - parsed_uri = urlparse(redirect_uri) - if parsed_uri.scheme not in ("http", "https"): - raise HTTPException(status_code=400, detail="Invalid redirect_uri scheme") + validate_loopback_redirect_uri(redirect_uri) # Reject new codes if the store is at capacity (prevents memory exhaustion # from a burst of abandoned OAuth flows). @@ -662,13 +731,25 @@ async def byok_authorize_post( status_code=400, detail="Only S256 code_challenge_method is supported" ) + # Identity comes from the authenticated session, not the OAuth client_id + # form field (RFC 6749 §2.2: client_id identifies the client application, + # not the user). We do bind the code to the submitted client_id so the + # /token call must present the same value (RFC 6749 §4.1.3). + user_id = user_api_key_dict.user_id + if not user_id: + raise HTTPException(status_code=401, detail="login_required") + auth_code = str(uuid.uuid4()) _byok_auth_codes[auth_code] = { "api_key": api_key, "server_id": server_id, "code_challenge": code_challenge, "redirect_uri": redirect_uri, - "user_id": client_id, # external client passes LiteLLM user-id as client_id + # RFC 6749 §4.1.3 defense-in-depth: if the authorization request + # declared a client_id, the token request must submit the same + # value. Stored even though we don't pre-register clients. + "client_id": client_id, + "user_id": user_id, "expires_at": time.time() + _AUTH_CODE_TTL_SECONDS, } @@ -704,34 +785,60 @@ async def byok_token( _purge_expired_codes() if grant_type != "authorization_code": - raise HTTPException(status_code=400, detail="unsupported_grant_type") + return _oauth_token_error("unsupported_grant_type") record = _byok_auth_codes.get(code) if record is None: - raise HTTPException(status_code=400, detail="invalid_grant") + return _oauth_token_error("invalid_grant") if time.time() > record["expires_at"]: del _byok_auth_codes[code] - raise HTTPException(status_code=400, detail="invalid_grant") + return _oauth_token_error("invalid_grant") # PKCE verification if not _verify_pkce(code_verifier, record["code_challenge"]): - raise HTTPException(status_code=400, detail="invalid_grant") + return _oauth_token_error("invalid_grant") - # Consume the code (one-time use) - del _byok_auth_codes[code] + # 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 + # /token request MUST authenticate as the same client. We don't + # pre-register clients, so an empty stored client_id skips the check. + if record.get("client_id") and client_id != record["client_id"]: + return _oauth_token_error("invalid_grant") server_id: str = record["server_id"] api_key_value: str = record["api_key"] - # Prefer the user_id that was stored when the code was issued; fall back to - # whatever client_id the token request supplies (they should match). - user_id: str = record.get("user_id") or client_id - + # user_id is stamped by the authenticated /authorize POST. No client_id + # fallback — that fallback was the credential-hijack primitive. The + # token-endpoint client_id is informational per RFC 6749 and is not + # cross-checked against user_id (which identifies the resource owner, + # not the client application). + user_id: str = record.get("user_id") or "" if not user_id: - raise HTTPException( - status_code=400, - detail="Cannot determine user_id; pass LiteLLM user id as client_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: @@ -756,17 +863,12 @@ async def byok_token( server_id, exc, ) - raise HTTPException(status_code=500, detail="Failed to store credential") + return _oauth_token_error("server_error", status=500) else: verbose_proxy_logger.warning( "byok_token: prisma_client is None — credential not persisted" ) - if master_key is None: - raise HTTPException( - status_code=500, detail="Master key not configured; cannot issue token" - ) - now = int(time.time()) payload = { "user_id": user_id, @@ -785,5 +887,6 @@ async def byok_token( "access_token": access_token, "token_type": "bearer", "expires_in": 3600, - } + }, + headers=TOKEN_NO_CACHE_HEADERS, ) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 792a9dace1e..13b19aa2d65 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -10,6 +10,10 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.proxy._experimental.mcp_server.oauth_utils import ( + TOKEN_NO_CACHE_HEADERS, + validate_loopback_redirect_uri, +) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, @@ -322,15 +326,13 @@ async def authorize_with_server( status_code=400, detail="MCP server authorization url is not set" ) + # Loopback-only redirect_uri. The URI is encrypted into the OAuth + # state and decoded on /callback to redirect the user back; a non- + # loopback URI would be an open-redirect + code-theft primitive + # (VERIA-57 root cause B). MCP clients are native apps — loopback is + # the spec-compliant callback pattern. + validate_loopback_redirect_uri(redirect_uri) parsed = urlparse(redirect_uri) - if parsed.scheme not in ("http", "https"): - raise HTTPException( - status_code=400, - detail={ - "error": "invalid_redirect_uri", - "message": "redirect_uri must use http or https scheme", - }, - ) base_url = urlunparse(parsed._replace(query="")) request_base_url = get_request_base_url(request) encoded_state = encode_state_with_base_url( @@ -480,7 +482,8 @@ async def exchange_token_with_server( if "scope" in token_response and token_response["scope"]: result["scope"] = token_response["scope"] - return JSONResponse(result) + # RFC 6749 §5.1: token responses must not be cached. + return JSONResponse(result, headers=TOKEN_NO_CACHE_HEADERS) async def register_client_with_server( @@ -647,20 +650,26 @@ async def token_endpoint( @router.get("/callback") async def callback(code: str, state: str): try: - # Decode the state hash to get base_url, original state, and PKCE params state_data = decode_state_hash(state) base_url = state_data["base_url"] original_state = state_data["original_state"] - # Forward code and original state back to client - params = {"code": code, "state": original_state} + # Re-validate loopback at the sink. /authorize rejects non-loopback + # redirect_uri before encoding into state, but encrypted states + # minted before that check was added have no expiry and remain + # valid indefinitely. Validating here blocks the open-redirect + + # code-theft primitive even for pre-fix states. + validate_loopback_redirect_uri(base_url) - # Forward to client's callback endpoint + params = {"code": code, "state": original_state} complete_returned_url = f"{base_url}?{urlencode(params)}" return RedirectResponse(url=complete_returned_url, status_code=302) + except HTTPException: + # Re-raise so a non-loopback base_url surfaces as 400 instead of + # a generic "authentication incomplete" redirect. + raise except Exception: - # fallback if state hash not found return HTMLResponse( "Authentication incomplete. You can close this window." ) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py new file mode 100644 index 00000000000..b13cf83058c --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -0,0 +1,48 @@ +"""Shared helpers for the MCP OAuth authorization endpoints +(BYOK + discoverable / pass-through OAuth proxy).""" + +from ipaddress import ip_address +from urllib.parse import urlparse + +from fastapi import HTTPException + +# RFC 6749 §5.1 / OAuth 2.1 draft-15 §4.1.3: token-endpoint responses +# must not be cached — both success and error bodies may reveal secrets. +TOKEN_NO_CACHE_HEADERS = {"Cache-Control": "no-store", "Pragma": "no-cache"} + + +def validate_loopback_redirect_uri(redirect_uri: str) -> None: + """Require a loopback ``redirect_uri`` (OAuth 2.1 §4.1.2.1 + RFC 8252 + §7.3 native-app pattern). MCP clients are native apps that listen on + a localhost port; rejecting non-loopback URIs prevents a malicious + client from pointing the callback at its own server to capture the + authorization code — the credential-theft primitive behind VERIA-57 + and pNr1PHa9. + + Accepts the literal ``localhost`` plus any IP in the loopback ranges + (IPv4 ``127.0.0.0/8`` and IPv6 ``::1``). A string match on + ``"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``. + """ + 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 + try: + 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/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5f320655931..c2e998f01e5 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1952,7 +1952,18 @@ if MCP_AVAILABLE: from litellm.proxy.proxy_server import prisma_client if prisma_client is None: - return + # Fail closed on DB unavailability: returning here previously + # bypassed the ownership check and let any proxy-authenticated + # caller invoke BYOK tools during outage windows. + raise HTTPException( + status_code=503, + detail={ + "error": "byok_auth_unavailable", + "server_id": mcp_server.server_id, + "server_name": mcp_server.server_name or mcp_server.name, + "message": "BYOK credential check requires a database connection.", + }, + ) credential = await get_user_credential( prisma_client=prisma_client, 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 4fa676222ba..89992c510f5 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 @@ -11,6 +11,7 @@ Covers: import base64 import hashlib +import json import time import uuid from typing import Any, Optional @@ -61,13 +62,34 @@ def test_verify_pkce_tampered_challenge(): from fastapi import FastAPI +from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( + _byok_session_auth, +) +from litellm.proxy._types import UserAPIKeyAuth + _test_app = FastAPI() _test_app.include_router(router) @pytest.fixture def client(): - return TestClient(_test_app, raise_server_exceptions=False) + """Test client with a fixed authenticated user (bypasses the session + cookie check by overriding the dep).""" + _test_app.dependency_overrides[_byok_session_auth] = lambda: UserAPIKeyAuth( + api_key="hashed", user_id="user-123" + ) + try: + yield TestClient(_test_app, raise_server_exceptions=False) + finally: + _test_app.dependency_overrides.pop(_byok_session_auth, None) + + +@pytest.fixture +def unauthenticated_client(): + """Test client with no dependency override — the real ``_byok_session_auth`` + runs, which checks the ``token`` cookie and falls back to + ``user_api_key_auth``. With neither set, both paths fail → 401.""" + yield TestClient(_test_app, raise_server_exceptions=False) # --------------------------------------------------------------------------- @@ -104,7 +126,7 @@ def test_authorize_get_returns_html(client): "/v1/mcp/oauth/authorize", params={ "client_id": "test-client", - "redirect_uri": "https://client.example.com/callback", + "redirect_uri": "http://127.0.0.1:3000/callback", "response_type": "code", "code_challenge": "abc123", "code_challenge_method": "S256", @@ -138,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", }, @@ -147,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 # --------------------------------------------------------------------------- @@ -160,7 +199,7 @@ def test_authorize_post_creates_code_and_redirects(client): "/v1/mcp/oauth/authorize", data={ "client_id": "user-123", - "redirect_uri": "https://client.example.com/callback", + "redirect_uri": "http://127.0.0.1:3000/callback", "code_challenge": challenge, "code_challenge_method": "S256", "state": "st_abc", @@ -286,10 +325,11 @@ async def test_token_endpoint_success(): ) assert result.status_code == 200 - body = result.body - import json - - data = json.loads(body) + # RFC 6749 §5.1 / OAuth 2.1 draft-15 §4.1.3: token responses MUST NOT + # be cached, as the body contains an access token. + assert result.headers["cache-control"] == "no-store" + assert result.headers["pragma"] == "no-cache" + data = json.loads(result.body) assert "access_token" in data assert data["token_type"] == "bearer" assert data["expires_in"] == 3600 @@ -319,21 +359,20 @@ async def test_token_endpoint_invalid_code(): from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import byok_token mock_request = MagicMock() - with pytest.raises(HTTPException) as exc_info: - with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch("litellm.proxy.proxy_server.master_key", "key"), - ): - await byok_token( - request=mock_request, - grant_type="authorization_code", - code="nonexistent-code", - redirect_uri="", - code_verifier="anything", - client_id="u", - ) - assert exc_info.value.status_code == 400 - assert "invalid_grant" in str(exc_info.value.detail) + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.master_key", "key"), + ): + result = await byok_token( + request=mock_request, + grant_type="authorization_code", + code="nonexistent-code", + redirect_uri="", + code_verifier="anything", + client_id="u", + ) + assert result.status_code == 400 + assert json.loads(result.body) == {"error": "invalid_grant"} @pytest.mark.asyncio @@ -345,27 +384,27 @@ async def test_token_endpoint_expired_code(): server_id="s", user_id="u", challenge=challenge, - redirect_uri="https://cb", + redirect_uri="http://127.0.0.1:3000/cb", ttl=-10, # already expired ) from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import byok_token mock_request = MagicMock() - with pytest.raises(HTTPException) as exc_info: - with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch("litellm.proxy.proxy_server.master_key", "key"), - ): - await byok_token( - request=mock_request, - grant_type="authorization_code", - code=code, - redirect_uri="", - code_verifier=verifier, - client_id="u", - ) - assert exc_info.value.status_code == 400 + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.master_key", "key"), + ): + result = await byok_token( + request=mock_request, + grant_type="authorization_code", + code=code, + redirect_uri="http://127.0.0.1:3000/cb", + code_verifier=verifier, + client_id="u", + ) + assert result.status_code == 400 + assert json.loads(result.body) == {"error": "invalid_grant"} @pytest.mark.asyncio @@ -377,27 +416,26 @@ async def test_token_endpoint_wrong_verifier(): server_id="s", user_id="u", challenge=challenge, - redirect_uri="https://cb", + redirect_uri="http://127.0.0.1:3000/cb", ) from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import byok_token mock_request = MagicMock() - with pytest.raises(HTTPException) as exc_info: - with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch("litellm.proxy.proxy_server.master_key", "key"), - ): - await byok_token( - request=mock_request, - grant_type="authorization_code", - code=code, - redirect_uri="", - code_verifier="wrong_verifier_value_that_wont_match", - client_id="u", - ) - assert exc_info.value.status_code == 400 - assert "invalid_grant" in str(exc_info.value.detail) + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.master_key", "key"), + ): + result = await byok_token( + request=mock_request, + grant_type="authorization_code", + code=code, + redirect_uri="http://127.0.0.1:3000/cb", + code_verifier="wrong_verifier_value_that_wont_match", + client_id="u", + ) + assert result.status_code == 400 + assert json.loads(result.body) == {"error": "invalid_grant"} @pytest.mark.asyncio @@ -405,21 +443,20 @@ async def test_token_endpoint_unsupported_grant_type(): from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import byok_token mock_request = MagicMock() - with pytest.raises(HTTPException) as exc_info: - with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), - patch("litellm.proxy.proxy_server.master_key", "key"), - ): - await byok_token( - request=mock_request, - grant_type="client_credentials", - code="any", - redirect_uri="", - code_verifier="v", - client_id="u", - ) - assert exc_info.value.status_code == 400 - assert "unsupported_grant_type" in str(exc_info.value.detail) + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.master_key", "key"), + ): + result = await byok_token( + request=mock_request, + grant_type="client_credentials", + code="any", + redirect_uri="", + code_verifier="v", + client_id="u", + ) + assert result.status_code == 400 + assert json.loads(result.body) == {"error": "unsupported_grant_type"} # --------------------------------------------------------------------------- @@ -525,3 +562,576 @@ async def test_check_byok_credential_has_credential(): ): # Should not raise await _check_byok_credential(server, user_auth) + + +@pytest.mark.asyncio +async def test_check_byok_credential_db_unavailable_fails_closed(): + """BYOK server with no prisma_client → 503, not silent pass. + + Regression for GHSA-6762: previously returned silently, bypassing the + ownership check during DB outage windows. + """ + from litellm.proxy._experimental.mcp_server.server import _check_byok_credential + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="byok-4", + name="byok-server", + transport=MCPTransport.http, + is_byok=True, + ) + user_auth = UserAPIKeyAuth(user_id="user-55", api_key="sk-test") + + with patch("litellm.proxy.proxy_server.prisma_client", None): + with pytest.raises(HTTPException) as exc_info: + await _check_byok_credential(server, user_auth) + + assert exc_info.value.status_code == 503 + detail: Any = exc_info.value.detail + assert detail["error"] == "byok_auth_unavailable" + assert detail["server_id"] == "byok-4" + + +# --------------------------------------------------------------------------- +# Security regression tests for AO2kf_-9 / GHSA-jg3h: +# Unauthenticated /v1/mcp/oauth/authorize previously allowed an attacker to +# stamp `user_id = client_id` into the auth-code record, overwriting any +# victim's stored BYOK credential at /token. +# --------------------------------------------------------------------------- + + +def test_authorize_post_rejects_unauthenticated(unauthenticated_client): + resp = unauthenticated_client.post( + "/v1/mcp/oauth/authorize", + data={ + "client_id": "victim-user-id", + "redirect_uri": "http://127.0.0.1:3000/cb", + "code_challenge": "abc", + "code_challenge_method": "S256", + "state": "s", + "server_id": "sid", + "api_key": "attacker-key", + }, + follow_redirects=False, + ) + assert resp.status_code == 401 + + +def test_authorize_post_ignores_spec_compliant_client_id(client): + """An MCP client following OAuth 2.1 semantics sends client_id as its + *application* identifier (e.g. "claude-desktop"). The stored user_id + must come from the authenticated session regardless — the form's + client_id is informational only.""" + verifier = "verifier_value_long_enough_to_be_valid_43chars" + challenge = _make_challenge(verifier) + + resp = client.post( + "/v1/mcp/oauth/authorize", + data={ + # Fixture authenticates as "user-123"; the client identifies + # itself as "claude-desktop" per OAuth 2.1 — unrelated to user. + "client_id": "claude-desktop", + "redirect_uri": "http://127.0.0.1:3000/cb", + "code_challenge": challenge, + "code_challenge_method": "S256", + "state": "s", + "server_id": "sid", + "api_key": "upstream-key", + }, + follow_redirects=False, + ) + assert resp.status_code == 302 + + from urllib.parse import parse_qs, urlparse + + qs = parse_qs(urlparse(resp.headers["location"]).query) + code = qs["code"][0] + # Identity = authenticated session, not the form's client_id. + assert _byok_auth_codes[code]["user_id"] == "user-123" + + +def test_authorize_post_binds_code_to_authenticated_user_id(client): + """Ensure the stored auth-code record uses the authenticated user_id, + NOT the form's client_id, as the identity the token endpoint will trust.""" + verifier = "verifier_value_long_enough_to_be_valid_43chars" + challenge = _make_challenge(verifier) + + resp = client.post( + "/v1/mcp/oauth/authorize", + data={ + # client_id omitted — must still bind to the authenticated user. + "redirect_uri": "http://127.0.0.1:3000/cb", + "code_challenge": challenge, + "code_challenge_method": "S256", + "state": "s", + "server_id": "sid", + "api_key": "legit-key", + }, + follow_redirects=False, + ) + assert resp.status_code == 302 + + from urllib.parse import parse_qs, urlparse + + qs = parse_qs(urlparse(resp.headers["location"]).query) + code = qs["code"][0] + assert _byok_auth_codes[code]["user_id"] == "user-123" + + +@pytest.mark.asyncio +async def test_token_endpoint_rejects_missing_user_id_in_code_record(): + """Defense in depth: if a code record somehow lacks user_id (older + format / manual DB write), /token must reject rather than fall back to + the form's client_id.""" + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( + byok_token, + ) + + verifier = "verifier_for_test_missing_user_id_path" + 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", + "user_id": "", # missing / empty + "expires_at": time.time() + 60, + } + + 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="attacker-chosen", + ) + assert result.status_code == 400 + assert json.loads(result.body) == {"error": "invalid_grant"} + + +def _authorize_post_with_cookie(client, cookie_jwt: str, api_key: str = "upstream-key"): + verifier = "verifier_cookie_auth_long_enough_to_be_valid" + return client.post( + "/v1/mcp/oauth/authorize", + data={ + "redirect_uri": "http://127.0.0.1:3000/cb", + "code_challenge": _make_challenge(verifier), + "code_challenge_method": "S256", + "state": "s", + "server_id": "sid", + "api_key": api_key, + }, + cookies={"token": cookie_jwt}, + follow_redirects=False, + ) + + +def test_authorize_post_accepts_ui_session_cookie(unauthenticated_client): + """Browser flow: the native HTML form doesn't add Authorization. Instead + the user's UI session cookie ``token`` carries a master-key-signed JWT + whose ``user_id`` + ``login_method`` claims authenticate the POST.""" + import jwt as _jwt + + with patch("litellm.proxy.proxy_server.master_key", "test-master-key"): + cookie_jwt = _jwt.encode( + { + "user_id": "browser-user-42", + "login_method": "sso", + "exp": int(time.time()) + 3600, + }, + "test-master-key", + algorithm="HS256", + ) + resp = _authorize_post_with_cookie(unauthenticated_client, cookie_jwt) + assert resp.status_code == 302 + + from urllib.parse import parse_qs, urlparse + + qs = parse_qs(urlparse(resp.headers["location"]).query) + code = qs["code"][0] + assert _byok_auth_codes[code]["user_id"] == "browser-user-42" + + +def test_authorize_post_rejects_cookie_signed_with_wrong_key(unauthenticated_client): + """A cookie JWT signed with a different key than the proxy's master_key + must not grant access — otherwise an attacker who can forge a JWT + against any key could impersonate any user.""" + import jwt as _jwt + + with patch("litellm.proxy.proxy_server.master_key", "real-master-key"): + forged = _jwt.encode( + {"user_id": "victim-user", "login_method": "sso"}, + "attacker-key", + algorithm="HS256", + ) + resp = _authorize_post_with_cookie(unauthenticated_client, forged, api_key="k") + assert resp.status_code == 401 + + +def test_authorize_post_rejects_replayed_byok_session_token(unauthenticated_client): + """Regression: the /token endpoint itself issues master-key-signed JWTs + with ``type="byok_session"`` + ``user_id`` (for MCP-client use). Those + tokens must not be accepted here — otherwise an attacker with any + byok_session token could replay it as a Cookie and re-authorize BYOK + writes without a valid UI session.""" + import jwt as _jwt + + with patch("litellm.proxy.proxy_server.master_key", "real-master-key"): + byok_session = _jwt.encode( + { + "user_id": "any-user", + "server_id": "sid", + "type": "byok_session", + }, + "real-master-key", + algorithm="HS256", + ) + resp = _authorize_post_with_cookie( + unauthenticated_client, byok_session, api_key="k" + ) + assert resp.status_code == 401 + + +def test_authorize_post_rejects_non_loopback_redirect_uri(client): + """OAuth 2.1 §4.1.2.1 + RFC 8252: native-app redirects must be loopback. + A public HTTPS callback from an MCP client would let that client capture + the issued code after a legitimate user enters their API key, so we + reject anything that isn't 127.0.0.1/localhost/::1.""" + verifier = "verifier_non_loopback_redirect_uri_test_long" + resp = client.post( + "/v1/mcp/oauth/authorize", + data={ + "redirect_uri": "https://attacker.example.com/cb", + "code_challenge": _make_challenge(verifier), + "code_challenge_method": "S256", + "state": "s", + "server_id": "sid", + "api_key": "k", + }, + follow_redirects=False, + ) + assert resp.status_code == 400 + + +@pytest.mark.asyncio +async def test_token_endpoint_rejects_redirect_uri_mismatch(): + """RFC 6749 §4.1.3 / OAuth 2.1 §4.1.3: if redirect_uri was sent at + /authorize, the /token redirect_uri MUST match exactly.""" + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( + byok_token, + ) + + verifier = "verifier_for_redirect_mismatch_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", + "user_id": "u", + "expires_at": time.time() + 60, + } + + result = await byok_token( + request=MagicMock(), + grant_type="authorization_code", + code=code, + redirect_uri="http://127.0.0.1:9999/cb", # different port + code_verifier=verifier, + client_id="", + ) + assert result.status_code == 400 + assert json.loads(result.body) == {"error": "invalid_grant"} + + +def test_authorize_post_rejects_cookie_missing_login_method(unauthenticated_client): + """Defense in depth: a master-key-signed JWT with only ``user_id`` is not + a valid UI session (UI tokens always carry ``login_method``). Accepting + it would expand the cookie surface to include any master-key-signed + JWT in the system, which is exactly what the byok_session-replay + regression above protects against.""" + import jwt as _jwt + + with patch("litellm.proxy.proxy_server.master_key", "real-master-key"): + malformed = _jwt.encode( + {"user_id": "some-user"}, + "real-master-key", + algorithm="HS256", + ) + resp = _authorize_post_with_cookie( + unauthenticated_client, malformed, api_key="k" + ) + assert resp.status_code == 401 + + +@pytest.mark.asyncio +async def test_token_endpoint_accepts_spec_compliant_client_id(): + """Per OAuth 2.1, the /token client_id is the client application + identifier, not the user. It must not be cross-checked against the + record's user_id — that cross-check would break spec-compliant MCP + clients that pass e.g. client_id="claude-desktop".""" + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( + byok_token, + ) + + verifier = "verifier_for_spec_compliant_long_enough_43chars" + 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", + "user_id": "real-authenticated-user", + "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="http://127.0.0.1:3000/cb", + code_verifier=verifier, + client_id="claude-desktop", # OAuth app id, unrelated to user + ) + # Token should be issued successfully. + assert result.status_code == 200 + + +@pytest.mark.asyncio +async def test_token_endpoint_rejects_client_id_mismatch(): + """RFC 6749 §4.1.3: if the authorization request was bound to a + client_id, the token request must submit the same value. An attacker + who steals a code from another client (different native app) can't + redeem it.""" + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( + byok_token, + ) + + verifier = "verifier_for_client_id_mismatch_long_enough!" + 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": "legitimate-client", + "user_id": "u", + "expires_at": time.time() + 60, + } + + 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="attacker-client", + ) + assert result.status_code == 400 + assert json.loads(result.body) == {"error": "invalid_grant"} + + +def test_authorize_post_accepts_ipv4_loopback_range(client): + """RFC 8252 §7.3 / RFC 5735: ``127.0.0.0/8`` is loopback — a string + match on ``127.0.0.1`` would miss ``127.0.0.2`` and break clients that + pick a loopback alias.""" + verifier = "verifier_for_127002_loopback_test_long_enough" + resp = client.post( + "/v1/mcp/oauth/authorize", + data={ + "redirect_uri": "http://127.0.0.2:3000/cb", + "code_challenge": _make_challenge(verifier), + "code_challenge_method": "S256", + "state": "s", + "server_id": "sid", + "api_key": "k", + }, + follow_redirects=False, + ) + assert resp.status_code == 302 + + +def test_authorize_post_accepts_ipv6_loopback_full_form(client): + """RFC 4291: full IPv6 loopback ``0:0:0:0:0:0:0:1`` must be accepted + equivalently to ``::1``.""" + verifier = "verifier_for_ipv6_full_loopback_test_long_enough" + resp = client.post( + "/v1/mcp/oauth/authorize", + data={ + "redirect_uri": "http://[0:0:0:0:0:0:0:1]:3000/cb", + "code_challenge": _make_challenge(verifier), + "code_challenge_method": "S256", + "state": "s", + "server_id": "sid", + "api_key": "k", + }, + 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() + + +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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index aecb8207735..ae3c16ce8b0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -76,7 +76,7 @@ async def test_authorize_endpoint_includes_response_type(): request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", - redirect_uri="https://client.example.com/callback", + redirect_uri="http://127.0.0.1:60108/callback", state="test_state", ) @@ -139,7 +139,7 @@ async def test_authorize_endpoint_preserves_existing_query_params(): request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", - redirect_uri="https://client.example.com/callback", + redirect_uri="http://127.0.0.1:60108/callback", state="test_state", ) @@ -558,7 +558,7 @@ async def test_authorize_endpoint_respects_x_forwarded_proto(): request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", - redirect_uri="https://client.example.com/callback", + redirect_uri="http://127.0.0.1:60108/callback", state="test_state", ) @@ -855,7 +855,7 @@ async def test_authorize_endpoint_respects_x_forwarded_host(): request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", - redirect_uri="https://client.example.com/callback", + redirect_uri="http://127.0.0.1:60108/callback", state="test_state", ) @@ -1817,3 +1817,191 @@ async def test_token_endpoint_authorization_code_missing_code(): ) assert exc_info.value.status_code == 400 assert "code is required" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_authorize_endpoint_rejects_non_loopback_redirect_uri(): + """VERIA-57 root cause B regression. The client-supplied redirect_uri + is encrypted into the OAuth state and decoded on /callback to 302 the + user back. A non-loopback value is an open-redirect + code-theft + primitive — reject with 400 before encoding anything into state.""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + global_mcp_server_manager.registry.clear() + oauth2_server = MCPServer( + server_id="test_oauth_server", + name="test_oauth", + server_name="test_oauth", + alias="test_oauth", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + client_secret="cs", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + with pytest.raises(HTTPException) as exc_info: + await authorize( + request=mock_request, + client_id="cid", + mcp_server_name="test_oauth", + redirect_uri="https://attacker.example.com/cb", + state="s", + ) + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_authorize_endpoint_accepts_ipv4_loopback_range_and_ipv6_full_form(): + """RFC 8252 §7.3 + RFC 4291: full 127.0.0.0/8 and full-form IPv6 + loopback must be accepted — string match on ``127.0.0.1`` alone + would miss ``127.0.0.2`` and ``0:0:0:0:0:0:0:1``.""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + global_mcp_server_manager.registry.clear() + oauth2_server = MCPServer( + server_id="test_oauth_server", + name="test_oauth", + server_name="test_oauth", + alias="test_oauth", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + client_secret="cs", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + for uri in ( + "http://127.0.0.2:3000/cb", + "http://[0:0:0:0:0:0:0:1]:3000/cb", + "http://localhost:3000/cb", + ): + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" + ) as mock_encrypt: + mock_encrypt.return_value = "mocked_encrypted_state" + response = await authorize( + request=mock_request, + client_id="cid", + mcp_server_name="test_oauth", + redirect_uri=uri, + state="s", + ) + assert response.status_code == 307, f"{uri} should be accepted" + + +@pytest.mark.asyncio +async def test_callback_revalidates_loopback_on_decoded_base_url(): + """VERIA-57 root cause B defense-in-depth: an encrypted state minted + before the /authorize validation was added has no expiry and stays + valid. /callback must re-validate the decoded base_url so those + stale states can't be used as an open-redirect + code-theft + primitive.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + callback, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" + ) as mock_decode: + mock_decode.return_value = { + "base_url": "https://attacker.example.com/cb", + "original_state": "s", + "code_challenge": None, + "code_challenge_method": None, + "client_redirect_uri": "https://attacker.example.com/cb", + } + with pytest.raises(HTTPException) as exc_info: + await callback(code="stolen_code", state="encrypted_stale_state") + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_token_endpoint_sets_no_store_cache_control(): + """RFC 6749 §5.1 / OAuth 2.1 draft-15 §4.1.3: the token response + contains an access token (and possibly a refresh token) — it MUST + NOT be cached by intermediaries or the client.""" + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="t", + name="t", + server_name="t", + alias="t", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + client_secret="cs", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + ) + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + fake_http_response = MagicMock() + fake_http_response.json.return_value = { + "access_token": "tok", + "token_type": "Bearer", + "expires_in": 3600, + } + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ): + response = await exchange_token_with_server( + request=mock_request, + mcp_server=server, + grant_type="authorization_code", + code="c", + redirect_uri="http://127.0.0.1:3000/cb", + client_id="cid", + client_secret=None, + code_verifier=None, + ) + + assert response.headers["cache-control"] == "no-store" + assert response.headers["pragma"] == "no-cache"