diff --git a/src/mcp/client/auth/oauth2.py b/src/mcp/client/auth/oauth2.py index 7dc62b52..d00539b8 100644 --- a/src/mcp/client/auth/oauth2.py +++ b/src/mcp/client/auth/oauth2.py @@ -577,6 +577,92 @@ class OAuthClientProvider(httpx2.Auth): if not check_resource_allowed(requested_resource=default_resource, configured_resource=prm_resource): raise OAuthFlowError(f"Protected resource {prm_resource} does not match expected {default_resource}") + async def _refresh_with_discovery(self) -> AsyncGenerator[httpx2.Request, httpx2.Response]: + """Refresh the token, discovering authorization-server metadata first when needed. + + The token endpoint comes from the AS metadata. On a cold start (a stored refresh + token reused before any 401) that metadata has not been discovered yet, so + ``_refresh_token`` would fall back to ``{origin}/token`` — dropping any issuer + path and 404ing on servers whose token endpoint lives elsewhere. Discovery runs + first, applying the same SEP-2352 issuer-binding checks as the 401 path so stored + credentials are never sent to an authorization server they are not bound to: on a + binding mismatch the credentials and tokens are dropped and the refresh is + skipped, letting the subsequent 401 flow re-register against the new server. + Yields the discovery and refresh requests so they run through the outer httpx + auth flow rather than a side-channel client. + """ + if self.context.oauth_metadata is None: + # Step 1: protected resource metadata -> authorization server URL (SEP-985). + # Best-effort: a legacy server without PRM falls through to the origin + # well-known fallback in the ASM step below. There is no 401 response at + # this point, so no WWW-Authenticate resource_metadata hint is available. + for url in build_protected_resource_metadata_discovery_urls(None, self.context.server_url): + prm = await handle_protected_resource_response((yield create_oauth_metadata_request(url))) + if prm: + # Validate PRM resource matches server URL (RFC 8707) + await self._validate_resource_match(prm) + self.context.protected_resource_metadata = prm + self.context.auth_server_url = str(prm.authorization_servers[0]) + break + else: + logger.debug(f"Protected resource metadata discovery failed: {url}") + + # SEP-2352: stored credentials are bound to the issuer that registered them. + # If the authorization server changed, drop them (and the old tokens) and skip + # the refresh so the 401 flow re-registers instead of presenting another + # server's credentials to the newly discovered one. + if ( + self.context.client_info is not None + and self.context.auth_server_url is not None + and not credentials_match_issuer( + self.context.client_info, self.context.auth_server_url, self.context.client_metadata_url + ) + ): + logger.debug("Authorization server changed; discarding bound credentials and skipping refresh") + self.context.client_info = None + self.context.clear_tokens() + return + + # Step 2: authorization server metadata -> the token endpoint (with fallback + # for legacy servers). + for url in build_oauth_authorization_server_metadata_discovery_urls( + self.context.auth_server_url, self.context.server_url + ): + ok, asm = await handle_auth_metadata_response((yield create_oauth_metadata_request(url))) + if not ok: + break + if asm: + # SEP-2468: metadata issuer must match the discovery issuer + if self.context.auth_server_url is not None: + validate_metadata_issuer(asm, self.context.auth_server_url) + self.context.oauth_metadata = asm + break + else: + logger.debug(f"OAuth metadata discovery failed: {url}") + + # SEP-2352: on the legacy no-PRM path the issuer is only known after ASM + # discovery, so re-evaluate the binding here using the discovered metadata + # issuer (mirroring the 401 path's post-ASM check). + if ( + self.context.client_info is not None + and self.context.auth_server_url is None + and self.context.oauth_metadata is not None + and not credentials_match_issuer( + self.context.client_info, + str(self.context.oauth_metadata.issuer), + self.context.client_metadata_url, + ) + ): + logger.debug("Authorization server changed; discarding bound credentials and skipping refresh") + self.context.client_info = None + self.context.clear_tokens() + return + + refresh_response = yield await self._refresh_token() + if not await self._handle_refresh_response(refresh_response): + # Refresh failed, need full re-authentication + self._initialized = False + async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]: """httpx2 auth flow integration.""" async with self.context.lock: @@ -587,13 +673,17 @@ class OAuthClientProvider(httpx2.Auth): self.context.protocol_version = request.headers.get(MCP_PROTOCOL_VERSION_HEADER) if not self.context.is_token_valid() and self.context.can_refresh_token(): - # Try to refresh token - refresh_request = await self._refresh_token() - refresh_response = yield refresh_request - - if not await self._handle_refresh_response(refresh_response): - # Refresh failed, need full re-authentication - self._initialized = False + # Refresh the token, discovering authorization-server metadata first on a + # cold start (see _refresh_with_discovery). Driven inline so its requests + # run through this httpx auth flow, not a side-channel client. + refresh_flow = self._refresh_with_discovery() + refresh_request = await refresh_flow.__anext__() + while True: + refresh_response = yield refresh_request + try: + refresh_request = await refresh_flow.asend(refresh_response) + except StopAsyncIteration: + break if self.context.is_token_valid(): self._add_auth_header(request) diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index be96cc8e..8da3cbae 100644 --- a/tests/client/test_auth.py +++ b/tests/client/test_auth.py @@ -3253,3 +3253,278 @@ async def test_issuer_is_stamped_when_same_origin_fallback_register_is_on_the_di await auth_flow.asend(httpx2.Response(200, request=final_req)) except StopAsyncIteration: pass + + +@pytest.mark.anyio +async def test_eager_refresh_discovers_token_endpoint_before_refreshing( + oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, valid_tokens: OAuthToken +): + """Regression for #3240/#3250: a cold-start eager refresh discovers the token endpoint. + + On a restart with a stored (expired) token the pre-401 refresh used to POST to the + ``{origin}/token`` fallback because authorization-server metadata had not been + discovered yet, 404ing on servers whose token endpoint lives under a path and + silently clearing the stored tokens. The refresh must run PRM + ASM discovery first + and target the discovered token endpoint. + """ + oauth_provider.context.current_tokens = valid_tokens + oauth_provider.context.token_expiry_time = time.time() - 100 # expired + oauth_provider.context.client_info = OAuthClientInformationFull( + client_id="test_client", + redirect_uris=[AnyUrl("http://localhost:3030/callback")], + token_endpoint_auth_method="none", + ) + oauth_provider._initialized = True + assert oauth_provider.context.oauth_metadata is None + + test_request = httpx2.Request("GET", "https://api.example.com/v1/mcp") + auth_flow = oauth_provider.async_auth_flow(test_request) + + # 1) protected-resource metadata discovery (no WWW-Authenticate hint pre-401) + prm_request = await auth_flow.__anext__() + assert str(prm_request.url) == "https://api.example.com/.well-known/oauth-protected-resource/v1/mcp" + prm_response = httpx2.Response( + 200, + content=( + b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}' + ), + request=prm_request, + ) + + # 2) authorization-server metadata whose token endpoint is NOT {origin}/token + asm_request = await auth_flow.asend(prm_response) + assert str(asm_request.url) == "https://auth.example.com/.well-known/oauth-authorization-server" + asm_response = httpx2.Response( + 200, + content=( + b'{"issuer": "https://auth.example.com", ' + b'"authorization_endpoint": "https://auth.example.com/oauth2/authorize", ' + b'"token_endpoint": "https://auth.example.com/oauth2/api/v1/token"}' + ), + request=asm_request, + ) + + # 3) the refresh targets the discovered token endpoint, not the fallback + refresh_request = await auth_flow.asend(asm_response) + assert refresh_request.method == "POST" + assert str(refresh_request.url) == "https://auth.example.com/oauth2/api/v1/token" + assert "grant_type=refresh_token" in refresh_request.content.decode() + refresh_response = httpx2.Response( + 200, + json={"access_token": "refreshed_token", "token_type": "Bearer", "expires_in": 3600}, + request=refresh_request, + ) + + # 4) the original request goes out with the refreshed token + api_request = await auth_flow.asend(refresh_response) + assert str(api_request.url) == "https://api.example.com/v1/mcp" + assert api_request.headers["Authorization"] == "Bearer refreshed_token" + stored = await mock_storage.get_tokens() + assert stored is not None + assert stored.access_token == "refreshed_token" + + with pytest.raises(StopAsyncIteration): + await auth_flow.asend(httpx2.Response(200, request=api_request)) + + +@pytest.mark.anyio +async def test_eager_refresh_falls_back_to_origin_token_when_no_metadata_published( + oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken +): + """A legacy server publishing no metadata keeps the pre-existing ``{origin}/token`` fallback. + + PRM discovery 404s at both well-known URLs and the legacy origin ASM fallback 404s too, + so the refresh still POSTs to ``{origin}/token`` exactly as before discovery-before-refresh + existed. A failed refresh then clears tokens and lets the request go out unauthenticated. + """ + oauth_provider.context.current_tokens = valid_tokens + oauth_provider.context.token_expiry_time = time.time() - 100 # expired + oauth_provider.context.client_info = OAuthClientInformationFull( + client_id="test_client", + redirect_uris=[AnyUrl("http://localhost:3030/callback")], + token_endpoint_auth_method="none", + ) + oauth_provider._initialized = True + + auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp")) + + # PRM discovery: path-based then root-based, both 404. + prm_request = await auth_flow.__anext__() + assert str(prm_request.url) == "https://api.example.com/.well-known/oauth-protected-resource/v1/mcp" + prm_request = await auth_flow.asend(httpx2.Response(404, request=prm_request)) + assert str(prm_request.url) == "https://api.example.com/.well-known/oauth-protected-resource" + + # ASM discovery: legacy origin fallback, 404 as well. + asm_request = await auth_flow.asend(httpx2.Response(404, request=prm_request)) + assert str(asm_request.url) == "https://api.example.com/.well-known/oauth-authorization-server" + + # Refresh falls back to {origin}/token (pre-existing legacy behavior). + refresh_request = await auth_flow.asend(httpx2.Response(404, request=asm_request)) + assert refresh_request.method == "POST" + assert str(refresh_request.url) == "https://api.example.com/token" + + # The refresh fails; tokens are cleared and the original request goes out unauthenticated. + api_request = await auth_flow.asend(httpx2.Response(401, request=refresh_request)) + assert str(api_request.url) == "https://api.example.com/v1/mcp" + assert "Authorization" not in api_request.headers + assert oauth_provider.context.current_tokens is None + await auth_flow.aclose() + + +@pytest.mark.anyio +async def test_eager_refresh_stops_asm_discovery_on_server_error( + oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken +): + """A non-4XX ASM discovery error stops the fallback chain, mirroring the 401 path. + + The refresh then proceeds against the ``{origin}/token`` fallback rather than + hammering further well-known URLs. + """ + oauth_provider.context.current_tokens = valid_tokens + oauth_provider.context.token_expiry_time = time.time() - 100 # expired + oauth_provider.context.client_info = OAuthClientInformationFull( + client_id="test_client", + redirect_uris=[AnyUrl("http://localhost:3030/callback")], + token_endpoint_auth_method="none", + ) + oauth_provider._initialized = True + + auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp")) + + # PRM discovery succeeds and points at the authorization server. + prm_request = await auth_flow.__anext__() + prm_response = httpx2.Response( + 200, + content=( + b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}' + ), + request=prm_request, + ) + + # ASM discovery hits a 500: stop trying further URLs. + asm_request = await auth_flow.asend(prm_response) + assert str(asm_request.url) == "https://auth.example.com/.well-known/oauth-authorization-server" + refresh_request = await auth_flow.asend(httpx2.Response(500, request=asm_request)) + + assert refresh_request.method == "POST" + assert str(refresh_request.url) == "https://api.example.com/token" + await auth_flow.aclose() + + +@pytest.mark.anyio +async def test_eager_refresh_skips_refresh_when_credentials_bound_to_different_issuer( + oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken +): + """SEP-2352: a cold-start refresh never sends credentials bound to another issuer. + + When PRM discovery reveals an authorization server different from the one the stored + client credentials are bound to, the credentials and tokens are dropped and the + refresh is skipped, so the subsequent 401 flow re-registers against the new server + — mirroring the issuer-binding check on the 401 discovery path. + """ + oauth_provider.context.current_tokens = valid_tokens + oauth_provider.context.token_expiry_time = time.time() - 100 # expired + oauth_provider.context.client_info = OAuthClientInformationFull( + client_id="stale-client", + redirect_uris=[AnyUrl("http://localhost:3030/callback")], + issuer="https://old-as.example.com", + ) + oauth_provider._initialized = True + + auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp")) + + # PRM discovery points at auth.example.com, not the bound old-as.example.com. + prm_request = await auth_flow.__anext__() + prm_response = httpx2.Response( + 200, + content=( + b'{"resource": "https://api.example.com/v1/mcp", "authorization_servers": ["https://auth.example.com"]}' + ), + request=prm_request, + ) + + # No refresh request: the next yield is the original request, unauthenticated. + api_request = await auth_flow.asend(prm_response) + assert str(api_request.url) == "https://api.example.com/v1/mcp" + assert "Authorization" not in api_request.headers + assert oauth_provider.context.client_info is None + assert oauth_provider.context.current_tokens is None + await auth_flow.aclose() + + +@pytest.mark.anyio +async def test_eager_refresh_legacy_path_rechecks_issuer_binding_after_asm( + oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken +): + """SEP-2352 on the legacy no-PRM path: the binding is checked against the ASM issuer. + + PRM discovery fails so the issuer is only known once origin-fallback ASM discovery + succeeds; credentials bound to a different issuer are then dropped and the refresh is + skipped, exactly as on the 401 path's post-ASM re-evaluation. + """ + oauth_provider.context.current_tokens = valid_tokens + oauth_provider.context.token_expiry_time = time.time() - 100 # expired + oauth_provider.context.client_info = OAuthClientInformationFull( + client_id="stale-client", + redirect_uris=[AnyUrl("http://localhost:3030/callback")], + issuer="https://old-as.example.com", + ) + oauth_provider._initialized = True + + auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp")) + + # PRM discovery: both well-known URLs 404. + prm_request = await auth_flow.__anext__() + prm_request = await auth_flow.asend(httpx2.Response(404, request=prm_request)) + + # Origin-fallback ASM discovery succeeds with the resource origin as issuer. + asm_request = await auth_flow.asend(httpx2.Response(404, request=prm_request)) + assert str(asm_request.url) == "https://api.example.com/.well-known/oauth-authorization-server" + asm_response = httpx2.Response( + 200, + content=( + b'{"issuer": "https://api.example.com", ' + b'"authorization_endpoint": "https://api.example.com/authorize", ' + b'"token_endpoint": "https://api.example.com/token"}' + ), + request=asm_request, + ) + + # No refresh request: the next yield is the original request, unauthenticated. + api_request = await auth_flow.asend(asm_response) + assert str(api_request.url) == "https://api.example.com/v1/mcp" + assert "Authorization" not in api_request.headers + assert oauth_provider.context.client_info is None + assert oauth_provider.context.current_tokens is None + # The just-discovered metadata is for the current server and is kept for the 401 flow. + assert oauth_provider.context.oauth_metadata is not None + await auth_flow.aclose() + + +@pytest.mark.anyio +async def test_eager_refresh_skips_discovery_when_metadata_already_known( + oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken +): + """With authorization-server metadata already discovered, the refresh is immediate.""" + oauth_provider.context.current_tokens = valid_tokens + oauth_provider.context.token_expiry_time = time.time() - 100 # expired + oauth_provider.context.client_info = OAuthClientInformationFull( + client_id="test_client", + redirect_uris=[AnyUrl("http://localhost:3030/callback")], + token_endpoint_auth_method="none", + ) + oauth_provider.context.oauth_metadata = OAuthMetadata.model_validate( + { + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/oauth2/authorize", + "token_endpoint": "https://auth.example.com/oauth2/api/v1/token", + } + ) + oauth_provider._initialized = True + + auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp")) + + refresh_request = await auth_flow.__anext__() + assert refresh_request.method == "POST" + assert str(refresh_request.url) == "https://auth.example.com/oauth2/api/v1/token" + await auth_flow.aclose()