mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
feat(mcp): dcr_bridge authorize and token relay redirect handling with mandatory S256
This commit is contained in:
parent
bca3e88c5f
commit
aa4f585e4c
2 changed files with 328 additions and 1 deletions
|
|
@ -516,6 +516,67 @@ def _raise_unless_oauth2_discovery_server(
|
|||
)
|
||||
|
||||
|
||||
def _dcr_bridge_relays_client_registration(mcp_server: MCPServer) -> bool:
|
||||
"""True when a DCR-bridge server relays client registration to the upstream authorization
|
||||
server instead of short-circuiting to an admin-configured OAuth client. In the relay arm the
|
||||
upstream holds each client's own registration, so the authorize and token relays pass the
|
||||
client's ``client_id`` and ``redirect_uri`` through verbatim and the authorization code
|
||||
returns directly to the client's redirect URI without transiting the gateway. Gateway-side
|
||||
redirect trust and the ``/callback`` state relay therefore only apply to the short-circuit
|
||||
arm, where the upstream only knows the gateway's own callback."""
|
||||
return mcp_server.is_dcr_bridge and bool(mcp_server.registration_url) and not mcp_server.client_id
|
||||
|
||||
|
||||
def _require_s256_pkce(
|
||||
code_challenge: Optional[str],
|
||||
code_challenge_method: Optional[str],
|
||||
) -> Tuple[str, str]:
|
||||
"""DCR-bridge servers serve unauthenticated public OAuth clients, so the PKCE downgrade
|
||||
paths (no challenge, or a non-S256 method; RFC 7636 defaults a missing method to ``plain``)
|
||||
are rejected at the gateway instead of relying on upstream enforcement. Returns the
|
||||
validated pair so callers get non-optional values."""
|
||||
if code_challenge and code_challenge_method == "S256":
|
||||
return code_challenge, code_challenge_method
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": (
|
||||
"This server requires PKCE: send code_challenge with "
|
||||
"code_challenge_method=S256 on the authorization request"
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _redirect_to_upstream_authorize(
|
||||
*,
|
||||
mcp_server: MCPServer,
|
||||
client_id: str,
|
||||
redirect_uri: str,
|
||||
state: str,
|
||||
code_challenge: str,
|
||||
code_challenge_method: str,
|
||||
response_type: Optional[str],
|
||||
scope: Optional[str],
|
||||
) -> RedirectResponse:
|
||||
"""The bridge relay arm's authorize redirect: every client-supplied parameter passes through
|
||||
to the upstream authorize endpoint verbatim, no relay state cookie is set, and the upstream
|
||||
enforces its own registered redirect binding for the client."""
|
||||
scope_value = scope or (" ".join(mcp_server.scopes) if mcp_server.scopes else None)
|
||||
passthrough_params = {
|
||||
"client_id": client_id,
|
||||
"redirect_uri": redirect_uri,
|
||||
"state": state,
|
||||
"response_type": response_type or "code",
|
||||
"code_challenge": code_challenge,
|
||||
"code_challenge_method": code_challenge_method,
|
||||
**({"scope": scope_value} if scope_value else {}),
|
||||
}
|
||||
parsed_auth_url = urlparse(mcp_server.authorization_url or "")
|
||||
merged_params = {**dict(parse_qsl(parsed_auth_url.query)), **passthrough_params}
|
||||
return RedirectResponse(urlunparse(parsed_auth_url._replace(query=urlencode(merged_params))))
|
||||
|
||||
|
||||
async def authorize_with_server(
|
||||
request: Request,
|
||||
mcp_server: MCPServer,
|
||||
|
|
@ -531,6 +592,20 @@ async def authorize_with_server(
|
|||
if mcp_server.authorization_url is None:
|
||||
raise HTTPException(status_code=400, detail="MCP server authorization url is not set")
|
||||
|
||||
if mcp_server.is_dcr_bridge:
|
||||
bridge_challenge, bridge_method = _require_s256_pkce(code_challenge, code_challenge_method)
|
||||
if _dcr_bridge_relays_client_registration(mcp_server):
|
||||
return _redirect_to_upstream_authorize(
|
||||
mcp_server=mcp_server,
|
||||
client_id=client_id,
|
||||
redirect_uri=redirect_uri,
|
||||
state=state,
|
||||
code_challenge=bridge_challenge,
|
||||
code_challenge_method=bridge_method,
|
||||
response_type=response_type,
|
||||
scope=scope,
|
||||
)
|
||||
|
||||
# Trusted redirect_uri: same-origin, loopback, or ops-allowlisted.
|
||||
# The URI is encrypted into the OAuth state and decoded on
|
||||
# /callback to redirect the user back; a non-trusted URI would be
|
||||
|
|
@ -626,11 +701,22 @@ async def exchange_token_with_server(
|
|||
status_code=400,
|
||||
detail="code is required for authorization_code grant",
|
||||
)
|
||||
if _dcr_bridge_relays_client_registration(mcp_server) and not redirect_uri:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(
|
||||
"redirect_uri is required for the authorization_code grant on this server; "
|
||||
"send the same redirect_uri used on the authorization request"
|
||||
),
|
||||
)
|
||||
proxy_base_url = get_request_base_url(request)
|
||||
resolved_redirect_uri = (
|
||||
redirect_uri if _dcr_bridge_relays_client_registration(mcp_server) else f"{proxy_base_url}/callback"
|
||||
)
|
||||
token_data = {
|
||||
"grant_type": "authorization_code",
|
||||
"code": code,
|
||||
"redirect_uri": f"{proxy_base_url}/callback",
|
||||
"redirect_uri": resolved_redirect_uri,
|
||||
**client_auth.body,
|
||||
}
|
||||
if code_verifier:
|
||||
|
|
|
|||
|
|
@ -3577,6 +3577,247 @@ async def test_token_exchange_passes_through_upstream_expires_in():
|
|||
assert body["expires_in"] == 43200
|
||||
|
||||
|
||||
_BRIDGE_CLIENT_REDIRECT = "https://claude.ai/api/mcp/auth_callback"
|
||||
|
||||
|
||||
def _bridge_server(**overrides):
|
||||
from litellm.proxy._types import MCPTransport
|
||||
from litellm.types.mcp import MCPAuth
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
fields = {
|
||||
"server_id": "bridge_srv",
|
||||
"name": "bridge_srv",
|
||||
"server_name": "bridge_srv",
|
||||
"alias": "bridge_srv",
|
||||
"transport": MCPTransport.http,
|
||||
"auth_type": MCPAuth.true_passthrough,
|
||||
"dcr_bridge": True,
|
||||
"authorization_url": "https://provider.com/oauth/authorize",
|
||||
"token_url": "https://provider.com/oauth/token",
|
||||
"registration_url": "https://provider.com/oauth/register",
|
||||
**overrides,
|
||||
}
|
||||
return MCPServer(**fields)
|
||||
|
||||
|
||||
def _bridge_mock_request():
|
||||
from fastapi import Request
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.base_url = "https://litellm.example.com/"
|
||||
mock_request.headers = {}
|
||||
return mock_request
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("auth_type_value", ["true_passthrough", "oauth_delegate"])
|
||||
async def test_authorize_bridge_relay_passes_client_params_verbatim(auth_type_value):
|
||||
"""The bridge relay arm (registration relayed upstream, no admin-configured client) passes the
|
||||
client's client_id, redirect_uri, state, and PKCE through verbatim: the code returns straight
|
||||
to the client's own redirect URI, so the gateway sets no state cookie, injects no /callback,
|
||||
and applies no gateway-side redirect trust (the upstream enforces its registered binding)."""
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
authorize_with_server,
|
||||
)
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
response = await authorize_with_server(
|
||||
request=_bridge_mock_request(),
|
||||
mcp_server=_bridge_server(auth_type=MCPAuth(auth_type_value)),
|
||||
client_id="dcr-client-123",
|
||||
redirect_uri=_BRIDGE_CLIENT_REDIRECT,
|
||||
state="client-state",
|
||||
code_challenge="chal",
|
||||
code_challenge_method="S256",
|
||||
)
|
||||
|
||||
assert response.status_code == 307
|
||||
location = response.headers["location"]
|
||||
assert location.startswith("https://provider.com/oauth/authorize")
|
||||
query = parse_qs(urlparse(location).query)
|
||||
assert query["client_id"] == ["dcr-client-123"]
|
||||
assert query["redirect_uri"] == [_BRIDGE_CLIENT_REDIRECT]
|
||||
assert query["state"] == ["client-state"]
|
||||
assert query["code_challenge"] == ["chal"]
|
||||
assert query["code_challenge_method"] == ["S256"]
|
||||
assert "litellm.example.com" not in location
|
||||
assert "set-cookie" not in {key.lower() for key in response.headers.keys()}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"code_challenge,code_challenge_method",
|
||||
[(None, None), ("chal", None), ("chal", "plain"), (None, "S256")],
|
||||
)
|
||||
async def test_authorize_bridge_requires_s256_pkce(code_challenge, code_challenge_method):
|
||||
"""Bridge servers serve unauthenticated public clients, so the PKCE downgrade paths (missing
|
||||
challenge, or a method that is not S256; RFC 7636 defaults a missing method to plain) are
|
||||
rejected at the gateway on both bridge arms."""
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
authorize_with_server,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await authorize_with_server(
|
||||
request=_bridge_mock_request(),
|
||||
mcp_server=_bridge_server(),
|
||||
client_id="dcr-client-123",
|
||||
redirect_uri=_BRIDGE_CLIENT_REDIRECT,
|
||||
state="s",
|
||||
code_challenge=code_challenge,
|
||||
code_challenge_method=code_challenge_method,
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 400
|
||||
assert "S256" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorize_bridge_short_circuit_keeps_callback_and_redirect_trust():
|
||||
"""The bridge short-circuit arm (admin-configured OAuth client, upstream only knows the
|
||||
gateway callback) keeps the /callback state relay and the gateway redirect trust: a public
|
||||
client redirect target is rejected unless ops allowlist it, and a trusted target still routes
|
||||
through the gateway callback with the state cookie."""
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
authorize_with_server,
|
||||
)
|
||||
|
||||
short_circuit_server = _bridge_server(client_id="admin-client", registration_url=None)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await authorize_with_server(
|
||||
request=_bridge_mock_request(),
|
||||
mcp_server=short_circuit_server,
|
||||
client_id="ignored",
|
||||
redirect_uri=_BRIDGE_CLIENT_REDIRECT,
|
||||
state="s",
|
||||
code_challenge="chal",
|
||||
code_challenge_method="S256",
|
||||
)
|
||||
assert exc.value.status_code in (400, 403)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper",
|
||||
return_value="mocked_encrypted_state",
|
||||
):
|
||||
response = await authorize_with_server(
|
||||
request=_bridge_mock_request(),
|
||||
mcp_server=short_circuit_server,
|
||||
client_id="ignored",
|
||||
redirect_uri="http://127.0.0.1:60108/callback",
|
||||
state="s",
|
||||
code_challenge="chal",
|
||||
code_challenge_method="S256",
|
||||
)
|
||||
|
||||
query = parse_qs(urlparse(response.headers["location"]).query)
|
||||
assert query["redirect_uri"] == ["https://litellm.example.com/callback"]
|
||||
assert query["client_id"] == ["admin-client"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorize_non_bridge_client_forwarded_keeps_pre_bridge_contract():
|
||||
"""A client-forwarded server without dcr_bridge keeps the pre-bridge behavior: no PKCE
|
||||
requirement and the gateway /callback relay (this is the browser-only Authorize path)."""
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
authorize_with_server,
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper",
|
||||
return_value="mocked_encrypted_state",
|
||||
):
|
||||
response = await authorize_with_server(
|
||||
request=_bridge_mock_request(),
|
||||
mcp_server=_bridge_server(dcr_bridge=None),
|
||||
client_id="cid",
|
||||
redirect_uri="http://127.0.0.1:60108/callback",
|
||||
state="s",
|
||||
)
|
||||
|
||||
assert response.status_code == 307
|
||||
query = parse_qs(urlparse(response.headers["location"]).query)
|
||||
assert query["redirect_uri"] == ["https://litellm.example.com/callback"]
|
||||
|
||||
|
||||
async def _bridge_token_post_data(server, redirect_uri):
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
exchange_token_with_server,
|
||||
)
|
||||
|
||||
fake_http_response = MagicMock()
|
||||
fake_http_response.json.return_value = {"access_token": "tok", "token_type": "Bearer"}
|
||||
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,
|
||||
):
|
||||
await exchange_token_with_server(
|
||||
request=_bridge_mock_request(),
|
||||
mcp_server=server,
|
||||
grant_type="authorization_code",
|
||||
code="auth-code",
|
||||
redirect_uri=redirect_uri,
|
||||
client_id="dcr-client-123",
|
||||
client_secret=None,
|
||||
code_verifier="verifier",
|
||||
)
|
||||
return fake_http_client.post.call_args.kwargs["data"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_bridge_relay_posts_client_redirect_uri():
|
||||
"""The bridge relay arm's token exchange sends the client's own redirect_uri upstream (it must
|
||||
match the authorize leg) with the caller's public client_id and PKCE verifier."""
|
||||
data = await _bridge_token_post_data(_bridge_server(), redirect_uri=_BRIDGE_CLIENT_REDIRECT)
|
||||
|
||||
assert data["redirect_uri"] == _BRIDGE_CLIENT_REDIRECT
|
||||
assert data["client_id"] == "dcr-client-123"
|
||||
assert data["code_verifier"] == "verifier"
|
||||
assert "client_secret" not in data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_bridge_relay_requires_redirect_uri():
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
exchange_token_with_server,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await exchange_token_with_server(
|
||||
request=_bridge_mock_request(),
|
||||
mcp_server=_bridge_server(),
|
||||
grant_type="authorization_code",
|
||||
code="auth-code",
|
||||
redirect_uri=None,
|
||||
client_id="dcr-client-123",
|
||||
client_secret=None,
|
||||
code_verifier="verifier",
|
||||
)
|
||||
|
||||
assert exc.value.status_code == 400
|
||||
assert "redirect_uri" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_non_bridge_keeps_gateway_callback():
|
||||
"""Without dcr_bridge the token exchange keeps posting the gateway callback as redirect_uri,
|
||||
pinning the pre-bridge contract for the browser-only Authorize path."""
|
||||
data = await _bridge_token_post_data(_bridge_server(dcr_bridge=None), redirect_uri=_BRIDGE_CLIENT_REDIRECT)
|
||||
|
||||
assert data["redirect_uri"] == "https://litellm.example.com/callback"
|
||||
|
||||
|
||||
async def _exchange_persistence_attempted_for_auth_type(auth_type) -> bool:
|
||||
"""Run exchange_token_with_server for a server of ``auth_type`` and report whether it attempted
|
||||
to persist the exchanged token server-side. The client-forwarded token modes must not persist:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue