tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-14 02:22:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mcp): harden BYOK OAuth authorize/token endpoints

Browse Source 
- POST /v1/mcp/oauth/authorize now requires an authenticated UI session
  cookie. The authenticated user_id — not the OAuth client_id form
  field — is stamped onto the authorization code record (RFC 6749 §2.2:
  client_id identifies the client, not the user).
- redirect_uri is restricted to loopback per RFC 8252 §7.3 (localhost
  plus any ipaddress.is_loopback IP, covering 127.0.0.0/8 and IPv6
  loopback forms).
- POST /v1/mcp/oauth/token enforces exact-match of the redirect_uri and
  client_id submitted at /authorize (RFC 6749 §4.1.3).
- /token error responses use the RFC 6749 §5.2 format ({"error":
  "<code>"}), and all /token responses set Cache-Control: no-store +
  Pragma: no-cache (RFC 6749 §5.1).
This commit is contained in:
user
2026-04-22 20:44:01 +00:00
parent c95cac5d46
commit bd1c3ea94d
2 changed files with 606 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py
+131 -23
View File
@@ -18,11 +18,12 @@ import hashlib
import html as _html_module
import time
import uuid
from ipaddress import ip_address
from typing import Dict, Optional, cast
from urllib.parse import urlencode, urlparse
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 +31,7 @@ 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._types import UserAPIKeyAuth
# ---------------------------------------------------------------------------
# In-memory store for pending authorization codes.
@@ -69,6 +71,88 @@ def _purge_expired_codes() -> None:
del _byok_auth_codes[k]
# 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 _oauth_token_error(code: str, status: int = 400) -> JSONResponse:
"""RFC 6749 §5.2 token-endpoint error body: ``{"error": "<code>"}``.
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 _validate_redirect_uri(redirect_uri: str) -> None:
"""Require a loopback redirect_uri (OAuth 2.1 §4.1.2.1 + RFC 8252
native-app pattern). MCP clients are native apps that listen on a
localhost port; rejecting non-loopback URIs prevents a malicious MCP
client from pointing the callback at its own server to capture the
code after a legitimate user enters their upstream API key.
Accepts the literal ``localhost`` plus any IP in the loopback ranges
(IPv4 ``127.0.0.0/8`` and IPv6 ``::1``) per RFC 8252 §7.3 — a string
match on ``"127.0.0.1"`` would miss ``127.0.0.2`` and full-form IPv6
(``0:0:0:0:0:0:0:1``).
"""
parsed = urlparse(redirect_uri)
if parsed.scheme not in ("http", "https"):
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:
pass
raise HTTPException(status_code=400, detail="invalid_request")
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"])
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,6 +666,10 @@ 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'")
@@ -636,6 +724,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 +734,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_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 +748,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 +802,45 @@ 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")
# 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"]:
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")
# Consume the code (one-time use)
del _byok_auth_codes[code]
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")
# Persist the BYOK credential
if prisma_client is not None:
@@ -756,16 +865,14 @@ 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"
)
return _oauth_token_error("server_error", status=500)
now = int(time.time())
payload = {
@@ -785,5 +892,6 @@ async def byok_token(
"access_token": access_token,
"token_type": "bearer",
"expires_in": 3600,
}
},
headers=_TOKEN_NO_CACHE_HEADERS,
)
tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py
+475 -68
View File
@@ -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",
@@ -160,7 +182,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 +308,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 +342,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 +367,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 +399,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 +426,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 +545,390 @@ async def test_check_byok_credential_has_credential():
):
# Should not raise
await _check_byok_credential(server, user_auth)
# ---------------------------------------------------------------------------
# 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"},
"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