diff --git a/docs_src/identity_assertion/tutorial001.py b/docs_src/identity_assertion/tutorial001.py index afcd537896..1954e34111 100644 --- a/docs_src/identity_assertion/tutorial001.py +++ b/docs_src/identity_assertion/tutorial001.py @@ -20,13 +20,13 @@ def __init__(self) -> None: async def get_tokens(self) -> OAuthToken | None: return self.tokens - async def set_tokens(self, tokens: OAuthToken) -> None: + async def set_tokens(self, tokens: OAuthToken | None) -> None: self.tokens = tokens async def get_client_info(self) -> OAuthClientInformationFull | None: return self.client_info - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: self.client_info = client_info diff --git a/docs_src/oauth_clients/tutorial001.py b/docs_src/oauth_clients/tutorial001.py index d150dc5da6..ef062fd2e1 100644 --- a/docs_src/oauth_clients/tutorial001.py +++ b/docs_src/oauth_clients/tutorial001.py @@ -17,13 +17,13 @@ def __init__(self) -> None: async def get_tokens(self) -> OAuthToken | None: return self.tokens - async def set_tokens(self, tokens: OAuthToken) -> None: + async def set_tokens(self, tokens: OAuthToken | None) -> None: self.tokens = tokens async def get_client_info(self) -> OAuthClientInformationFull | None: return self.client_info - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: self.client_info = client_info diff --git a/docs_src/oauth_clients/tutorial002.py b/docs_src/oauth_clients/tutorial002.py index dd4105f937..6eaf27b9a6 100644 --- a/docs_src/oauth_clients/tutorial002.py +++ b/docs_src/oauth_clients/tutorial002.py @@ -14,13 +14,13 @@ def __init__(self) -> None: async def get_tokens(self) -> OAuthToken | None: return self.tokens - async def set_tokens(self, tokens: OAuthToken) -> None: + async def set_tokens(self, tokens: OAuthToken | None) -> None: self.tokens = tokens async def get_client_info(self) -> OAuthClientInformationFull | None: return self.client_info - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: self.client_info = client_info diff --git a/examples/clients/simple-auth-client/mcp_simple_auth_client/main.py b/examples/clients/simple-auth-client/mcp_simple_auth_client/main.py index a190b89970..558b741edb 100644 --- a/examples/clients/simple-auth-client/mcp_simple_auth_client/main.py +++ b/examples/clients/simple-auth-client/mcp_simple_auth_client/main.py @@ -37,13 +37,13 @@ def __init__(self): async def get_tokens(self) -> OAuthToken | None: return self._tokens - async def set_tokens(self, tokens: OAuthToken) -> None: + async def set_tokens(self, tokens: OAuthToken | None) -> None: self._tokens = tokens async def get_client_info(self) -> OAuthClientInformationFull | None: return self._client_info - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: self._client_info = client_info diff --git a/examples/snippets/clients/identity_assertion_client.py b/examples/snippets/clients/identity_assertion_client.py index 19cde274c5..ec0913e91d 100644 --- a/examples/snippets/clients/identity_assertion_client.py +++ b/examples/snippets/clients/identity_assertion_client.py @@ -34,13 +34,13 @@ def __init__(self) -> None: async def get_tokens(self) -> OAuthToken | None: return self.tokens - async def set_tokens(self, tokens: OAuthToken) -> None: + async def set_tokens(self, tokens: OAuthToken | None) -> None: self.tokens = tokens async def get_client_info(self) -> OAuthClientInformationFull | None: return self.client_info - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: self.client_info = client_info diff --git a/examples/snippets/clients/oauth_client.py b/examples/snippets/clients/oauth_client.py index 58c542ea43..3e094a9801 100644 --- a/examples/snippets/clients/oauth_client.py +++ b/examples/snippets/clients/oauth_client.py @@ -29,7 +29,7 @@ async def get_tokens(self) -> OAuthToken | None: """Get stored tokens.""" return self.tokens - async def set_tokens(self, tokens: OAuthToken) -> None: + async def set_tokens(self, tokens: OAuthToken | None) -> None: """Store tokens.""" self.tokens = tokens @@ -37,7 +37,7 @@ async def get_client_info(self) -> OAuthClientInformationFull | None: """Get stored client information.""" return self.client_info - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: """Store client information.""" self.client_info = client_info diff --git a/examples/stories/_shared/auth.py b/examples/stories/_shared/auth.py index 35e997242f..18a22839fa 100644 --- a/examples/stories/_shared/auth.py +++ b/examples/stories/_shared/auth.py @@ -38,13 +38,13 @@ class InMemoryTokenStorage: async def get_tokens(self) -> OAuthToken | None: return self.tokens - async def set_tokens(self, tokens: OAuthToken) -> None: + async def set_tokens(self, tokens: OAuthToken | None) -> None: self.tokens = tokens async def get_client_info(self) -> OAuthClientInformationFull | None: return self.client_info - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: self.client_info = client_info diff --git a/src/mcp/client/auth/oauth2.py b/src/mcp/client/auth/oauth2.py index 7dc62b52b9..21ae427dd1 100644 --- a/src/mcp/client/auth/oauth2.py +++ b/src/mcp/client/auth/oauth2.py @@ -5,13 +5,14 @@ import base64 import hashlib +import json import logging import secrets import string import time from collections.abc import AsyncGenerator, Awaitable, Callable from dataclasses import dataclass, field -from typing import Any, Protocol, get_args +from typing import Any, Protocol, cast, get_args from urllib.parse import quote, urlencode, urljoin, urlparse import anyio @@ -109,6 +110,21 @@ def check_registration_usable(client_info: OAuthClientInformationFull) -> None: ) +def _is_expired_client_secret(client_info: OAuthClientInformationFull) -> bool: + """Return whether a stored registration reports an expired client secret.""" + expires_at = client_info.client_secret_expires_at + return expires_at is not None and expires_at != 0 and expires_at <= int(time.time()) + + +def _is_invalid_client_response(body: bytes) -> bool: + """Identify RFC 6749 ``invalid_client`` responses independent of HTTP status.""" + try: + payload: Any = json.loads(body) + except (json.JSONDecodeError, UnicodeDecodeError): + return False + return isinstance(payload, dict) and cast(dict[str, Any], payload).get("error") == "invalid_client" + + class PKCEParameters(BaseModel): """PKCE (Proof Key for Code Exchange) parameters.""" @@ -131,16 +147,16 @@ async def get_tokens(self) -> OAuthToken | None: """Get stored tokens.""" ... - async def set_tokens(self, tokens: OAuthToken) -> None: - """Store tokens.""" + async def set_tokens(self, tokens: OAuthToken | None) -> None: + """Store tokens, or clear them when ``tokens`` is ``None``.""" ... async def get_client_info(self) -> OAuthClientInformationFull | None: """Get stored client information.""" ... - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: - """Store client information.""" + async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: + """Store client information, or clear it when ``client_info`` is ``None``.""" ... @@ -468,6 +484,8 @@ async def _handle_token_response(self, response: httpx2.Response) -> None: """Handle token exchange response.""" if response.status_code not in {200, 201}: body = await response.aread() + if _is_invalid_client_response(body): + await self._clear_stored_credentials() body_text = body.decode("utf-8") raise OAuthTokenError(f"Token exchange failed ({response.status_code}): {body_text}") @@ -519,8 +537,12 @@ async def _refresh_token(self) -> httpx2.Request: async def _handle_refresh_response(self, response: httpx2.Response) -> bool: """Handle token refresh response. Returns True if successful.""" if response.status_code != 200: + body = await response.aread() logger.warning(f"Token refresh failed: {response.status_code}") - self.context.clear_tokens() + if _is_invalid_client_response(body): + await self._clear_stored_credentials() + else: + self.context.clear_tokens() return False try: @@ -551,8 +573,18 @@ async def _initialize(self) -> None: """Load stored tokens and client info.""" self.context.current_tokens = await self.context.storage.get_tokens() self.context.client_info = await self.context.storage.get_client_info() + if self.context.client_info and _is_expired_client_secret(self.context.client_info): + logger.info("Stored client registration has expired; clearing credentials and re-registering") + await self._clear_stored_credentials() self._initialized = True + async def _clear_stored_credentials(self) -> None: + """Clear the in-memory and persisted credentials bound to a client registration.""" + self.context.client_info = None + self.context.clear_tokens() + await self.context.storage.set_client_info(None) + await self.context.storage.set_tokens(None) + def _add_auth_header(self, request: httpx2.Request) -> None: """Add authorization header to request if we have valid tokens.""" if self.context.current_tokens and self.context.current_tokens.access_token: # pragma: no branch diff --git a/tests/client/auth/extensions/test_client_credentials.py b/tests/client/auth/extensions/test_client_credentials.py index 16336f8002..fbb9d48591 100644 --- a/tests/client/auth/extensions/test_client_credentials.py +++ b/tests/client/auth/extensions/test_client_credentials.py @@ -27,13 +27,13 @@ def __init__(self): async def get_tokens(self) -> OAuthToken | None: return self._tokens - async def set_tokens(self, tokens: OAuthToken) -> None: # pragma: no cover + async def set_tokens(self, tokens: OAuthToken | None) -> None: # pragma: no cover self._tokens = tokens async def get_client_info(self) -> OAuthClientInformationFull | None: # pragma: no cover return self._client_info - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: # pragma: no cover + async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: # pragma: no cover self._client_info = client_info diff --git a/tests/client/auth/extensions/test_identity_assertion.py b/tests/client/auth/extensions/test_identity_assertion.py index d4c965edc5..aeb786b1ec 100644 --- a/tests/client/auth/extensions/test_identity_assertion.py +++ b/tests/client/auth/extensions/test_identity_assertion.py @@ -29,13 +29,13 @@ def __init__(self, tokens: OAuthToken | None = None) -> None: async def get_tokens(self) -> OAuthToken | None: return self.tokens - async def set_tokens(self, tokens: OAuthToken) -> None: + async def set_tokens(self, tokens: OAuthToken | None) -> None: self.tokens = tokens async def get_client_info(self) -> OAuthClientInformationFull | None: raise NotImplementedError - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: raise NotImplementedError diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index be96cc8eec..a38e44c280 100644 --- a/tests/client/test_auth.py +++ b/tests/client/test_auth.py @@ -13,6 +13,7 @@ from mcp.client.auth import OAuthClientProvider, PKCEParameters from mcp.client.auth.exceptions import OAuthFlowError, OAuthRegistrationError, OAuthTokenError +from mcp.client.auth.oauth2 import _is_invalid_client_response from mcp.client.auth.utils import ( build_oauth_authorization_server_metadata_discovery_urls, build_protected_resource_metadata_discovery_urls, @@ -53,13 +54,13 @@ def __init__(self): async def get_tokens(self) -> OAuthToken | None: return self._tokens - async def set_tokens(self, tokens: OAuthToken) -> None: + async def set_tokens(self, tokens: OAuthToken | None) -> None: self._tokens = tokens async def get_client_info(self) -> OAuthClientInformationFull | None: return self._client_info - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: self._client_info = client_info @@ -78,6 +79,20 @@ def client_metadata(): ) +@pytest.mark.parametrize( + ("body", "expected"), + [ + (b"not json", False), + (b"\xff", False), + (b"[]", False), + (b'{"error": "invalid_grant"}', False), + (b'{"error": "invalid_client"}', True), + ], +) +def test_invalid_client_response_detection(body: bytes, expected: bool) -> None: + assert _is_invalid_client_response(body) is expected + + @pytest.fixture def valid_tokens(): return OAuthToken( @@ -264,6 +279,26 @@ def test_clear_tokens(self, oauth_provider: OAuthClientProvider, valid_tokens: O assert context.current_tokens is None assert context.token_expiry_time is None + @pytest.mark.anyio + async def test_initialize_discards_expired_client_registration( + self, oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, valid_tokens: OAuthToken + ): + expired_client = OAuthClientInformationFull( + client_id="expired-client", + client_secret="expired-secret", + client_secret_expires_at=int(time.time()) - 1, + redirect_uris=[AnyUrl("http://localhost:3030/callback")], + ) + await mock_storage.set_client_info(expired_client) + await mock_storage.set_tokens(valid_tokens) + + await oauth_provider._initialize() + + assert oauth_provider.context.client_info is None + assert oauth_provider.context.current_tokens is None + assert await mock_storage.get_client_info() is None + assert await mock_storage.get_tokens() is None + class TestOAuthFlow: """Test OAuth flow methods.""" @@ -2955,6 +2990,34 @@ async def test_handle_token_response_raises_on_non_2xx_with_body(oauth_provider: await oauth_provider._handle_token_response(response) +@pytest.mark.anyio +async def test_handle_token_response_invalid_client_clears_stored_credentials( + oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, valid_tokens: OAuthToken +): + client_info = OAuthClientInformationFull( + client_id="stale-client", + client_secret="stale-secret", + redirect_uris=[AnyUrl("http://localhost:3030/callback")], + ) + oauth_provider.context.client_info = client_info + oauth_provider.context.current_tokens = valid_tokens + await mock_storage.set_client_info(client_info) + await mock_storage.set_tokens(valid_tokens) + response = httpx2.Response( + 401, + json={"error": "invalid_client"}, + request=httpx2.Request("POST", "https://auth.example.com/token"), + ) + + with pytest.raises(OAuthTokenError, match=r"Token exchange failed \(401\).*invalid_client"): + await oauth_provider._handle_token_response(response) + + assert oauth_provider.context.client_info is None + assert oauth_provider.context.current_tokens is None + assert await mock_storage.get_client_info() is None + assert await mock_storage.get_tokens() is None + + @pytest.mark.anyio async def test_handle_refresh_response_carries_prior_scope_and_refresh_token_when_omitted( oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage @@ -3007,6 +3070,34 @@ async def test_handle_refresh_response_adopts_rotated_refresh_token_when_returne assert stored.refresh_token == "rotated" +@pytest.mark.anyio +async def test_handle_refresh_response_invalid_client_clears_stored_credentials( + oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, valid_tokens: OAuthToken +): + client_info = OAuthClientInformationFull( + client_id="stale-client", + client_secret="stale-secret", + redirect_uris=[AnyUrl("http://localhost:3030/callback")], + ) + oauth_provider.context.client_info = client_info + oauth_provider.context.current_tokens = valid_tokens + await mock_storage.set_client_info(client_info) + await mock_storage.set_tokens(valid_tokens) + response = httpx2.Response( + 400, + json={"error": "invalid_client"}, + request=httpx2.Request("POST", "https://auth.example.com/token"), + ) + + ok = await oauth_provider._handle_refresh_response(response) + + assert ok is False + assert oauth_provider.context.client_info is None + assert oauth_provider.context.current_tokens is None + assert await mock_storage.get_client_info() is None + assert await mock_storage.get_tokens() is None + + @pytest.mark.anyio async def test_issuer_binding_re_evaluated_after_asm_when_prm_discovery_failed( oauth_provider: OAuthClientProvider, diff --git a/tests/client/test_scope_bug_1630.py b/tests/client/test_scope_bug_1630.py index 0782b4a037..f0c75ffdee 100644 --- a/tests/client/test_scope_bug_1630.py +++ b/tests/client/test_scope_bug_1630.py @@ -29,13 +29,13 @@ def __init__(self) -> None: async def get_tokens(self) -> OAuthToken | None: return self._tokens # pragma: no cover - async def set_tokens(self, tokens: OAuthToken) -> None: + async def set_tokens(self, tokens: OAuthToken | None) -> None: self._tokens = tokens async def get_client_info(self) -> OAuthClientInformationFull | None: return self._client_info # pragma: no cover - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: self._client_info = client_info # pragma: no cover diff --git a/tests/interaction/auth/_harness.py b/tests/interaction/auth/_harness.py index 856a1fe9a8..b667452ce2 100644 --- a/tests/interaction/auth/_harness.py +++ b/tests/interaction/auth/_harness.py @@ -116,13 +116,13 @@ def __init__(self, *, client_info: OAuthClientInformationFull | None = None) -> async def get_tokens(self) -> OAuthToken | None: return self.tokens - async def set_tokens(self, tokens: OAuthToken) -> None: + async def set_tokens(self, tokens: OAuthToken | None) -> None: self.tokens = tokens async def get_client_info(self) -> OAuthClientInformationFull | None: return self.client_info - async def set_client_info(self, client_info: OAuthClientInformationFull) -> None: + async def set_client_info(self, client_info: OAuthClientInformationFull | None) -> None: self.client_info = client_info