diff --git a/src/mcp/client/auth.py b/src/mcp/client/auth.py index c174385e..359e0585 100644 --- a/src/mcp/client/auth.py +++ b/src/mcp/client/auth.py @@ -106,6 +106,10 @@ class OAuthContext: # State lock: anyio.Lock = field(default_factory=anyio.Lock) + # Discovery state for fallback support + discovery_base_url: str | None = None + discovery_pathname: str | None = None + def get_authorization_base_url(self, server_url: str) -> str: """Extract base URL by removing path component.""" parsed = urlparse(server_url) @@ -197,18 +201,53 @@ class OAuthClientProvider(httpx.Auth): except ValidationError: pass - async def _discover_oauth_metadata(self) -> httpx.Request: - """Build OAuth metadata discovery request.""" - if self.context.auth_server_url: - base_url = self.context.get_authorization_base_url(self.context.auth_server_url) - else: - base_url = self.context.get_authorization_base_url(self.context.server_url) + def _build_well_known_path(self, pathname: str) -> str: + """Construct well-known path for OAuth metadata discovery.""" + well_known_path = f"/.well-known/oauth-authorization-server{pathname}" + if pathname.endswith("/"): + # Strip trailing slash from pathname to avoid double slashes + well_known_path = well_known_path[:-1] + return well_known_path - url = urljoin(base_url, "/.well-known/oauth-authorization-server") + def _should_attempt_fallback(self, response_status: int, pathname: str) -> bool: + """Determine if fallback to root discovery should be attempted.""" + return response_status == 404 and pathname != "/" + + async def _try_metadata_discovery(self, url: str) -> httpx.Request: + """Build metadata discovery request for a specific URL.""" return httpx.Request("GET", url, headers={MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION}) - async def _handle_oauth_metadata_response(self, response: httpx.Response) -> None: - """Handle OAuth metadata response.""" + async def _discover_oauth_metadata(self) -> httpx.Request: + """Build OAuth metadata discovery request with fallback support.""" + if self.context.auth_server_url: + auth_server_url = self.context.auth_server_url + else: + auth_server_url = self.context.server_url + + # Per RFC 8414, try path-aware discovery first + parsed = urlparse(auth_server_url) + well_known_path = self._build_well_known_path(parsed.path) + base_url = f"{parsed.scheme}://{parsed.netloc}" + url = urljoin(base_url, well_known_path) + + # Store fallback info for use in response handler + self.context.discovery_base_url = base_url + self.context.discovery_pathname = parsed.path + + return await self._try_metadata_discovery(url) + + async def _discover_oauth_metadata_fallback(self) -> httpx.Request: + """Build fallback OAuth metadata discovery request for legacy servers.""" + base_url = getattr(self.context, "discovery_base_url", "") + if not base_url: + raise OAuthFlowError("No base URL available for fallback discovery") + + # Fallback to root discovery for legacy servers + url = urljoin(base_url, "/.well-known/oauth-authorization-server") + return await self._try_metadata_discovery(url) + + async def _handle_oauth_metadata_response(self, response: httpx.Response, is_fallback: bool = False) -> bool: + """Handle OAuth metadata response. Returns True if handled successfully.""" if response.status_code == 200: try: content = await response.aread() @@ -217,9 +256,18 @@ class OAuthClientProvider(httpx.Auth): # Apply default scope if none specified if self.context.client_metadata.scope is None and metadata.scopes_supported is not None: self.context.client_metadata.scope = " ".join(metadata.scopes_supported) + return True except ValidationError: pass + # Check if we should attempt fallback (404 on path-aware discovery) + if not is_fallback and self._should_attempt_fallback( + response.status_code, getattr(self.context, "discovery_pathname", "/") + ): + return False # Signal that fallback should be attempted + + return True # Signal no fallback needed (either success or non-404 error) + async def _register_client(self) -> httpx.Request | None: """Build registration request or skip if already registered.""" if self.context.client_info: @@ -418,10 +466,16 @@ class OAuthClientProvider(httpx.Auth): discovery_response = yield discovery_request await self._handle_protected_resource_response(discovery_response) - # Step 2: Discover OAuth metadata + # Step 2: Discover OAuth metadata (with fallback for legacy servers) oauth_request = await self._discover_oauth_metadata() oauth_response = yield oauth_request - await self._handle_oauth_metadata_response(oauth_response) + handled = await self._handle_oauth_metadata_response(oauth_response, is_fallback=False) + + # If path-aware discovery failed with 404, try fallback to root + if not handled: + fallback_request = await self._discover_oauth_metadata_fallback() + fallback_response = yield fallback_request + await self._handle_oauth_metadata_response(fallback_response, is_fallback=True) # Step 3: Register client if needed registration_request = await self._register_client() @@ -464,10 +518,16 @@ class OAuthClientProvider(httpx.Auth): discovery_response = yield discovery_request await self._handle_protected_resource_response(discovery_response) - # Step 2: Discover OAuth metadata + # Step 2: Discover OAuth metadata (with fallback for legacy servers) oauth_request = await self._discover_oauth_metadata() oauth_response = yield oauth_request - await self._handle_oauth_metadata_response(oauth_response) + handled = await self._handle_oauth_metadata_response(oauth_response, is_fallback=False) + + # If path-aware discovery failed with 404, try fallback to root + if not handled: + fallback_request = await self._discover_oauth_metadata_fallback() + fallback_response = yield fallback_request + await self._handle_oauth_metadata_response(fallback_response, is_fallback=True) # Step 3: Register client if needed registration_request = await self._register_client() diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index 8dee687a..d87410d0 100644 --- a/tests/client/test_auth.py +++ b/tests/client/test_auth.py @@ -208,10 +208,159 @@ class TestOAuthFlow: """Test OAuth metadata discovery request building.""" request = await oauth_provider._discover_oauth_metadata() + assert request.method == "GET" + assert str(request.url) == "https://api.example.com/.well-known/oauth-authorization-server/v1/mcp" + assert "mcp-protocol-version" in request.headers + + @pytest.mark.anyio + async def test_discover_oauth_metadata_request_no_path(self, client_metadata, mock_storage): + """Test OAuth metadata discovery request building when server has no path.""" + + async def redirect_handler(url: str) -> None: + pass + + async def callback_handler() -> tuple[str, str | None]: + return "test_auth_code", "test_state" + + provider = OAuthClientProvider( + server_url="https://api.example.com", + client_metadata=client_metadata, + storage=mock_storage, + redirect_handler=redirect_handler, + callback_handler=callback_handler, + ) + + request = await provider._discover_oauth_metadata() + assert request.method == "GET" assert str(request.url) == "https://api.example.com/.well-known/oauth-authorization-server" assert "mcp-protocol-version" in request.headers + @pytest.mark.anyio + async def test_discover_oauth_metadata_request_trailing_slash(self, client_metadata, mock_storage): + """Test OAuth metadata discovery request building when server path has trailing slash.""" + + async def redirect_handler(url: str) -> None: + pass + + async def callback_handler() -> tuple[str, str | None]: + return "test_auth_code", "test_state" + + provider = OAuthClientProvider( + server_url="https://api.example.com/v1/mcp/", + client_metadata=client_metadata, + storage=mock_storage, + redirect_handler=redirect_handler, + callback_handler=callback_handler, + ) + + request = await provider._discover_oauth_metadata() + + assert request.method == "GET" + assert str(request.url) == "https://api.example.com/.well-known/oauth-authorization-server/v1/mcp" + assert "mcp-protocol-version" in request.headers + + +class TestOAuthFallback: + """Test OAuth discovery fallback behavior for legacy (act as AS not RS) servers.""" + + @pytest.mark.anyio + async def test_fallback_discovery_request(self, client_metadata, mock_storage): + """Test fallback discovery request building.""" + + async def redirect_handler(url: str) -> None: + pass + + async def callback_handler() -> tuple[str, str | None]: + return "test_auth_code", "test_state" + + provider = OAuthClientProvider( + server_url="https://api.example.com/v1/mcp", + client_metadata=client_metadata, + storage=mock_storage, + redirect_handler=redirect_handler, + callback_handler=callback_handler, + ) + + # Set up discovery state manually as if path-aware discovery was attempted + provider.context.discovery_base_url = "https://api.example.com" + provider.context.discovery_pathname = "/v1/mcp" + + # Test fallback request building + request = await provider._discover_oauth_metadata_fallback() + + assert request.method == "GET" + assert str(request.url) == "https://api.example.com/.well-known/oauth-authorization-server" + assert "mcp-protocol-version" in request.headers + + @pytest.mark.anyio + async def test_should_attempt_fallback(self, oauth_provider): + """Test fallback decision logic.""" + # Should attempt fallback on 404 with non-root path + assert oauth_provider._should_attempt_fallback(404, "/v1/mcp") + + # Should NOT attempt fallback on 404 with root path + assert not oauth_provider._should_attempt_fallback(404, "/") + + # Should NOT attempt fallback on other status codes + assert not oauth_provider._should_attempt_fallback(200, "/v1/mcp") + assert not oauth_provider._should_attempt_fallback(500, "/v1/mcp") + + @pytest.mark.anyio + async def test_handle_metadata_response_success(self, oauth_provider): + """Test successful metadata response handling.""" + # Create minimal valid OAuth metadata + content = b"""{ + "issuer": "https://auth.example.com", + "authorization_endpoint": "https://auth.example.com/authorize", + "token_endpoint": "https://auth.example.com/token" + }""" + response = httpx.Response(200, content=content) + + # Should return True (success) and set metadata + result = await oauth_provider._handle_oauth_metadata_response(response, is_fallback=False) + assert result is True + assert oauth_provider.context.oauth_metadata is not None + assert str(oauth_provider.context.oauth_metadata.issuer) == "https://auth.example.com/" + + @pytest.mark.anyio + async def test_handle_metadata_response_404_needs_fallback(self, oauth_provider): + """Test 404 response handling that should trigger fallback.""" + # Set up discovery state for non-root path + oauth_provider.context.discovery_base_url = "https://api.example.com" + oauth_provider.context.discovery_pathname = "/v1/mcp" + + # Mock 404 response + response = httpx.Response(404) + + # Should return False (needs fallback) + result = await oauth_provider._handle_oauth_metadata_response(response, is_fallback=False) + assert result is False + + @pytest.mark.anyio + async def test_handle_metadata_response_404_no_fallback_needed(self, oauth_provider): + """Test 404 response handling when no fallback is needed.""" + # Set up discovery state for root path + oauth_provider.context.discovery_base_url = "https://api.example.com" + oauth_provider.context.discovery_pathname = "/" + + # Mock 404 response + response = httpx.Response(404) + + # Should return True (no fallback needed) + result = await oauth_provider._handle_oauth_metadata_response(response, is_fallback=False) + assert result is True + + @pytest.mark.anyio + async def test_handle_metadata_response_404_fallback_attempt(self, oauth_provider): + """Test 404 response handling during fallback attempt.""" + # Mock 404 response during fallback + response = httpx.Response(404) + + # Should return True (fallback attempt complete, no further action needed) + result = await oauth_provider._handle_oauth_metadata_response(response, is_fallback=True) + assert result is True + @pytest.mark.anyio async def test_register_client_request(self, oauth_provider): """Test client registration request building."""