Skip to content

Commit 79259a6

Browse files
committed
fix(client): preserve existing query params on OAuth authorization_endpoint
Closes #2776 The authorization code grant built the redirect URL with `f"{auth_endpoint}?{urlencode(auth_params)}"`, which produces an invalid URL when the server-advertised authorization_endpoint already carries a query string. For example Salesforce advertises `.../services/oauth2/authorize?prompt=select_account`, yielding `...authorize?prompt=select_account?response_type=code&...` (two `?` separators), so the client navigates to a malformed URL and the server rejects the request. Fix: parse the endpoint, merge its existing query params with the flow-generated auth_params (flow params win on conflict), and re-encode into a single well-formed query string. None-valued params are dropped rather than serialized as the literal "None". Tests: add TestAuthorizationEndpointWithQuery covering the helper (no/with/conflicting existing query) plus an end-to-end _perform_authorization_code_grant assertion that the captured redirect URL preserves the server param and stays well-formed. 101 passed.
1 parent 3b78f86 commit 79259a6

2 files changed

Lines changed: 103 additions & 3 deletions

File tree

src/mcp/client/auth/oauth2.py

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,10 @@
99
import secrets
1010
import string
1111
import time
12-
from collections.abc import AsyncGenerator, Awaitable, Callable
12+
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping
1313
from dataclasses import dataclass, field
1414
from typing import Any, Protocol
15-
from urllib.parse import quote, urlencode, urljoin, urlparse
15+
from urllib.parse import parse_qsl, quote, urlencode, urljoin, urlparse, urlunparse
1616

1717
import anyio
1818
import httpx
@@ -59,6 +59,22 @@
5959
logger = logging.getLogger(__name__)
6060

6161

62+
def _build_authorization_url(auth_endpoint: str, auth_params: Mapping[str, str | None]) -> str:
63+
"""Build an authorization URL, preserving any query params already on the endpoint.
64+
65+
Servers may advertise an ``authorization_endpoint`` that already carries query
66+
parameters (e.g. ``https://example.com/authorize?prompt=select_account``).
67+
Naively appending ``?<params>`` would produce an invalid URL with two ``?``
68+
separators, so the existing query is parsed and merged with ``auth_params``.
69+
Flow-generated params take precedence on key conflicts; ``None`` values are
70+
dropped rather than serialized as the literal string ``"None"``.
71+
"""
72+
parsed = urlparse(auth_endpoint)
73+
merged_params = dict(parse_qsl(parsed.query, keep_blank_values=True))
74+
merged_params.update({key: value for key, value in auth_params.items() if value is not None})
75+
return urlunparse(parsed._replace(query=urlencode(merged_params)))
76+
77+
6278
class PKCEParameters(BaseModel):
6379
"""PKCE (Proof Key for Code Exchange) parameters."""
6480

@@ -357,7 +373,7 @@ async def _perform_authorization_code_grant(self) -> tuple[str, str]:
357373
if "offline_access" in self.context.client_metadata.scope.split():
358374
auth_params["prompt"] = "consent"
359375

360-
authorization_url = f"{auth_endpoint}?{urlencode(auth_params)}"
376+
authorization_url = _build_authorization_url(auth_endpoint, auth_params)
361377
await self.context.redirect_handler(authorization_url)
362378

363379
# Wait for callback

tests/client/test_auth.py

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313

1414
from mcp.client.auth import OAuthClientProvider, PKCEParameters
1515
from mcp.client.auth.exceptions import OAuthFlowError, OAuthTokenError
16+
from mcp.client.auth.oauth2 import _build_authorization_url
1617
from mcp.client.auth.utils import (
1718
build_oauth_authorization_server_metadata_discovery_urls,
1819
build_protected_resource_metadata_discovery_urls,
@@ -3169,3 +3170,86 @@ async def echo_callback() -> AuthorizationCodeResult:
31693170
await auth_flow.asend(httpx.Response(200, request=final_req))
31703171
except StopAsyncIteration:
31713172
pass
3173+
3174+
3175+
class TestAuthorizationEndpointWithQuery:
3176+
"""Regression tests for #2776 - authorization_endpoint carrying query params."""
3177+
3178+
def test_build_authorization_url_no_existing_query(self):
3179+
url = _build_authorization_url(
3180+
"https://auth.example.com/authorize",
3181+
{"response_type": "code", "client_id": "abc"},
3182+
)
3183+
parsed = urlparse(url)
3184+
params = parse_qs(parsed.query)
3185+
assert parsed.path == "/authorize"
3186+
assert params["response_type"] == ["code"]
3187+
assert params["client_id"] == ["abc"]
3188+
# No malformed double "?" separator.
3189+
assert url.count("?") == 1
3190+
3191+
def test_build_authorization_url_preserves_existing_query(self):
3192+
# e.g. Salesforce advertises .../authorize?prompt=select_account
3193+
url = _build_authorization_url(
3194+
"https://test.salesforce.com/services/oauth2/authorize?prompt=select_account",
3195+
{"response_type": "code", "client_id": "abc"},
3196+
)
3197+
parsed = urlparse(url)
3198+
params = parse_qs(parsed.query)
3199+
assert parsed.path == "/services/oauth2/authorize"
3200+
# The server-provided param survives...
3201+
assert params["prompt"] == ["select_account"]
3202+
# ...alongside the flow-generated params.
3203+
assert params["response_type"] == ["code"]
3204+
assert params["client_id"] == ["abc"]
3205+
# Exactly one "?" - the old f-string produced "...?prompt=...?response_type=...".
3206+
assert url.count("?") == 1
3207+
3208+
def test_build_authorization_url_flow_params_win_on_conflict(self):
3209+
url = _build_authorization_url(
3210+
"https://auth.example.com/authorize?response_type=token",
3211+
{"response_type": "code"},
3212+
)
3213+
params = parse_qs(urlparse(url).query)
3214+
assert params["response_type"] == ["code"]
3215+
3216+
@pytest.mark.anyio
3217+
async def test_perform_authorization_preserves_endpoint_query(self, oauth_provider: OAuthClientProvider):
3218+
"""End-to-end: redirect URL stays valid when the endpoint has a query string."""
3219+
oauth_provider.context.oauth_metadata = OAuthMetadata(
3220+
issuer=AnyHttpUrl("https://test.salesforce.com"),
3221+
authorization_endpoint=AnyHttpUrl(
3222+
"https://test.salesforce.com/services/oauth2/authorize?prompt=select_account"
3223+
),
3224+
token_endpoint=AnyHttpUrl("https://test.salesforce.com/services/oauth2/token"),
3225+
)
3226+
oauth_provider.context.client_info = OAuthClientInformationFull(
3227+
client_id="test_client_id",
3228+
client_secret="test_client_secret",
3229+
redirect_uris=[AnyUrl("http://localhost:3030/callback")],
3230+
)
3231+
3232+
captured_url: str | None = None
3233+
captured_state: str | None = None
3234+
3235+
async def capture_redirect(url: str) -> None:
3236+
nonlocal captured_url, captured_state
3237+
captured_url = url
3238+
captured_state = parse_qs(urlparse(url).query).get("state", [None])[0]
3239+
3240+
async def mock_callback() -> AuthorizationCodeResult:
3241+
return AuthorizationCodeResult(code="test_auth_code", state=captured_state)
3242+
3243+
oauth_provider.context.redirect_handler = capture_redirect
3244+
oauth_provider.context.callback_handler = mock_callback
3245+
3246+
await oauth_provider._perform_authorization_code_grant()
3247+
3248+
assert captured_url is not None
3249+
parsed = urlparse(captured_url)
3250+
params = parse_qs(parsed.query)
3251+
assert parsed.path == "/services/oauth2/authorize"
3252+
assert params["prompt"] == ["select_account"]
3253+
assert params["response_type"] == ["code"]
3254+
assert params["client_id"] == ["test_client_id"]
3255+
assert captured_url.count("?") == 1

0 commit comments

Comments
 (0)