refactor: pull out oauth helper functions

Author: maxisbey

src/mcp/client/auth/oauth2.py

@@ -7,7 +7,6 @@
 import base64
 import hashlib
 import logging
-import re
 import secrets
 import string
 import time
@@ -20,6 +19,14 @@
 import httpx
 from pydantic import BaseModel, Field, ValidationError
 
+from mcp.client.auth.utils import (
+    build_protected_resource_discovery_urls,
+    extract_field_from_www_auth,
+    extract_resource_metadata_from_www_auth,
+    extract_scope_from_www_auth,
+    get_client_metadata_scopes,
+    get_discovery_urls,
+)
 from mcp.client.streamable_http import MCP_PROTOCOL_VERSION
 from mcp.shared.auth import (
     OAuthClientInformationFull,
@@ -200,85 +207,6 @@ def __init__(
         )
         self._initialized = False
 
-    def _build_protected_resource_discovery_urls(self, init_response: httpx.Response) -> list[str]:
-        """
-        Build ordered list of URLs to try for protected resource metadata discovery.
-
-        Per SEP-985, the client MUST:
-        1. Try resource_metadata from WWW-Authenticate header (if present)
-        2. Fall back to path-based well-known URI: /.well-known/oauth-protected-resource/{path}
-        3. Fall back to root-based well-known URI: /.well-known/oauth-protected-resource
-
-        Args:
-            init_response: The initial 401 response from the server
-
-        Returns:
-            Ordered list of URLs to try for discovery
-        """
-        urls: list[str] = []
-
-        # Priority 1: WWW-Authenticate header with resource_metadata parameter
-        www_auth_url = self._extract_resource_metadata_from_www_auth(init_response)
-        if www_auth_url:
-            urls.append(www_auth_url)
-
-        # Priority 2-3: Well-known URIs (RFC 9728)
-        parsed = urlparse(self.context.server_url)
-        base_url = f"{parsed.scheme}://{parsed.netloc}"
-
-        # Priority 2: Path-based well-known URI (if server has a path component)
-        if parsed.path and parsed.path != "/":
-            path_based_url = urljoin(base_url, f"/.well-known/oauth-protected-resource{parsed.path}")
-            urls.append(path_based_url)
-
-        # Priority 3: Root-based well-known URI
-        root_based_url = urljoin(base_url, "/.well-known/oauth-protected-resource")
-        urls.append(root_based_url)
-
-        return urls
-
-    def _extract_field_from_www_auth(self, init_response: httpx.Response, field_name: str) -> str | None:
-        """
-        Extract field from WWW-Authenticate header.
-
-        Returns:
-            Field value if found in WWW-Authenticate header, None otherwise
-        """
-        www_auth_header = init_response.headers.get("WWW-Authenticate")
-        if not www_auth_header:
-            return None
-
-        # Pattern matches: field_name="value" or field_name=value (unquoted)
-        pattern = rf'{field_name}=(?:"([^"]+)"|([^\s,]+))'
-        match = re.search(pattern, www_auth_header)
-
-        if match:
-            # Return quoted value if present, otherwise unquoted value
-            return match.group(1) or match.group(2)
-
-        return None
-
-    def _extract_resource_metadata_from_www_auth(self, init_response: httpx.Response) -> str | None:
-        """
-        Extract protected resource metadata URL from WWW-Authenticate header as per RFC9728.
-
-        Returns:
-            Resource metadata URL if found in WWW-Authenticate header, None otherwise
-        """
-        if not init_response or init_response.status_code != 401:
-            return None
-
-        return self._extract_field_from_www_auth(init_response, "resource_metadata")
-
-    def _extract_scope_from_www_auth(self, init_response: httpx.Response) -> str | None:
-        """
-        Extract scope parameter from WWW-Authenticate header as per RFC6750.
-
-        Returns:
-            Scope string if found in WWW-Authenticate header, None otherwise
-        """
-        return self._extract_field_from_www_auth(init_response, "scope")
-
     async def _handle_protected_resource_response(self, response: httpx.Response) -> bool:
         """
         Handle protected resource metadata discovery response.
@@ -309,54 +237,6 @@ async def _handle_protected_resource_response(self, response: httpx.Response) ->
             # Other error - fail immediately
             raise OAuthFlowError(f"Protected Resource Metadata request failed: {response.status_code}")
 
-    def _select_scopes(self, init_response: httpx.Response) -> None:
-        """Select scopes as outlined in the 'Scope Selection Strategy in the MCP spec."""
-        # Per MCP spec, scope selection priority order:
-        # 1. Use scope from WWW-Authenticate header (if provided)
-        # 2. Use all scopes from PRM scopes_supported (if available)
-        # 3. Omit scope parameter if neither is available
-        #
-        www_authenticate_scope = self._extract_scope_from_www_auth(init_response)
-        if www_authenticate_scope is not None:
-            # Priority 1: WWW-Authenticate header scope
-            self.context.client_metadata.scope = www_authenticate_scope
-        elif (
-            self.context.protected_resource_metadata is not None
-            and self.context.protected_resource_metadata.scopes_supported is not None
-        ):
-            # Priority 2: PRM scopes_supported
-            self.context.client_metadata.scope = " ".join(self.context.protected_resource_metadata.scopes_supported)
-        else:
-            # Priority 3: Omit scope parameter
-            self.context.client_metadata.scope = None
-
-    def _get_discovery_urls(self) -> list[str]:
-        """Generate ordered list of (url, type) tuples for discovery attempts."""
-        urls: list[str] = []
-        auth_server_url = self.context.auth_server_url or self.context.server_url
-        parsed = urlparse(auth_server_url)
-        base_url = f"{parsed.scheme}://{parsed.netloc}"
-
-        # RFC 8414: Path-aware OAuth discovery
-        if parsed.path and parsed.path != "/":
-            oauth_path = f"/.well-known/oauth-authorization-server{parsed.path.rstrip('/')}"
-            urls.append(urljoin(base_url, oauth_path))
-
-        # OAuth root fallback
-        urls.append(urljoin(base_url, "/.well-known/oauth-authorization-server"))
-
-        # RFC 8414 section 5: Path-aware OIDC discovery
-        # See https://www.rfc-editor.org/rfc/rfc8414.html#section-5
-        if parsed.path and parsed.path != "/":
-            oidc_path = f"/.well-known/openid-configuration{parsed.path.rstrip('/')}"
-            urls.append(urljoin(base_url, oidc_path))
-
-        # OIDC 1.0 fallback (appends to full URL per OIDC spec)
-        oidc_fallback = f"{auth_server_url.rstrip('/')}/.well-known/openid-configuration"
-        urls.append(oidc_fallback)
-
-        return urls
-
     async def _register_client(self) -> httpx.Request | None:
         """Build registration request or skip if already registered."""
         if self.context.client_info:
@@ -610,8 +490,12 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                 # Perform full OAuth flow
                 try:
                     # OAuth flow must be inline due to generator constraints
+                    www_auth_resource_metadata_url = extract_resource_metadata_from_www_auth(response)
+
                     # Step 1: Discover protected resource metadata (SEP-985 with fallback support)
-                    discovery_urls = self._build_protected_resource_discovery_urls(response)
+                    discovery_urls = build_protected_resource_discovery_urls(
+                        www_auth_resource_metadata_url, self.context.server_url
+                    )
                     discovery_success = False
                     for url in discovery_urls:
                         discovery_request = httpx.Request(
@@ -626,10 +510,12 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                         raise OAuthFlowError("Protected resource metadata discovery failed: no valid metadata found")
 
                     # Step 2: Apply scope selection strategy
-                    self._select_scopes(response)
+                    self.context.client_metadata.scope = get_client_metadata_scopes(
+                        www_auth_resource_metadata_url, self.context.protected_resource_metadata
+                    )
 
                     # Step 3: Discover OAuth metadata (with fallback for legacy servers)
-                    discovery_urls = self._get_discovery_urls()
+                    discovery_urls = get_discovery_urls(self.context.auth_server_url or self.context.server_url)
                     for url in discovery_urls:
                         oauth_metadata_request = self._create_oauth_metadata_request(url)
                         oauth_metadata_response = yield oauth_metadata_request
@@ -661,13 +547,15 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                 yield request
             elif response.status_code == 403:
                 # Step 1: Extract error field from WWW-Authenticate header
-                error = self._extract_field_from_www_auth(response, "error")
+                error = extract_field_from_www_auth(response, "error")
 
                 # Step 2: Check if we need to step-up authorization
                 if error == "insufficient_scope":
                     try:
                         # Step 2a: Update the required scopes
-                        self._select_scopes(response)
+                        self.context.client_metadata.scope = get_client_metadata_scopes(
+                            extract_scope_from_www_auth(response), self.context.protected_resource_metadata
+                        )
 
                         # Step 2b: Perform (re-)authorization and token exchange
                         token_response = yield await self._perform_authorization()

src/mcp/client/auth/utils.py

@@ -0,0 +1,136 @@
+import re
+from urllib.parse import urljoin, urlparse
+
+from httpx import Response
+
+from mcp.shared.auth import ProtectedResourceMetadata
+
+
+def extract_field_from_www_auth(response: Response, field_name: str) -> str | None:
+    """
+    Extract field from WWW-Authenticate header.
+
+    Returns:
+        Field value if found in WWW-Authenticate header, None otherwise
+    """
+    www_auth_header = response.headers.get("WWW-Authenticate")
+    if not www_auth_header:
+        return None
+
+    # Pattern matches: field_name="value" or field_name=value (unquoted)
+    pattern = rf'{field_name}=(?:"([^"]+)"|([^\s,]+))'
+    match = re.search(pattern, www_auth_header)
+
+    if match:
+        # Return quoted value if present, otherwise unquoted value
+        return match.group(1) or match.group(2)
+
+    return None
+
+
+def extract_scope_from_www_auth(response: Response) -> str | None:
+    """
+    Extract scope parameter from WWW-Authenticate header as per RFC6750.
+
+    Returns:
+        Scope string if found in WWW-Authenticate header, None otherwise
+    """
+    return extract_field_from_www_auth(response, "scope")
+
+
+def extract_resource_metadata_from_www_auth(response: Response) -> str | None:
+    """
+    Extract protected resource metadata URL from WWW-Authenticate header as per RFC9728.
+
+    Returns:
+        Resource metadata URL if found in WWW-Authenticate header, None otherwise
+    """
+    if not response or response.status_code != 401:
+        return None
+
+    return extract_field_from_www_auth(response, "resource_metadata")
+
+
+def build_protected_resource_discovery_urls(www_auth_url: str | None, server_url: str) -> list[str]:
+    """
+    Build ordered list of URLs to try for protected resource metadata discovery.
+
+    Per SEP-985, the client MUST:
+    1. Try resource_metadata from WWW-Authenticate header (if present)
+    2. Fall back to path-based well-known URI: /.well-known/oauth-protected-resource/{path}
+    3. Fall back to root-based well-known URI: /.well-known/oauth-protected-resource
+
+    Args:
+        www_auth_url: optional resource_metadata url extracted from the WWW-Authenticate header
+        server_url: server url
+
+    Returns:
+        Ordered list of URLs to try for discovery
+    """
+    urls: list[str] = []
+
+    # Priority 1: WWW-Authenticate header with resource_metadata parameter
+    if www_auth_url:
+        urls.append(www_auth_url)
+
+    # Priority 2-3: Well-known URIs (RFC 9728)
+    parsed = urlparse(server_url)
+    base_url = f"{parsed.scheme}://{parsed.netloc}"
+
+    # Priority 2: Path-based well-known URI (if server has a path component)
+    if parsed.path and parsed.path != "/":
+        path_based_url = urljoin(base_url, f"/.well-known/oauth-protected-resource{parsed.path}")
+        urls.append(path_based_url)
+
+    # Priority 3: Root-based well-known URI
+    root_based_url = urljoin(base_url, "/.well-known/oauth-protected-resource")
+    urls.append(root_based_url)
+
+    return urls
+
+
+def get_client_metadata_scopes(
+    www_authenticate_scope: str | None, protected_resource_metadata: ProtectedResourceMetadata | None
+) -> str | None:
+    """Select scopes as outlined in the 'Scope Selection Strategy' in the MCP spec."""
+    # Per MCP spec, scope selection priority order:
+    # 1. Use scope from WWW-Authenticate header (if provided)
+    # 2. Use all scopes from PRM scopes_supported (if available)
+    # 3. Omit scope parameter if neither is available
+
+    if www_authenticate_scope is not None:
+        # Priority 1: WWW-Authenticate header scope
+        return www_authenticate_scope
+    elif protected_resource_metadata is not None and protected_resource_metadata.scopes_supported is not None:
+        # Priority 2: PRM scopes_supported
+        return " ".join(protected_resource_metadata.scopes_supported)
+    else:
+        # Priority 3: Omit scope parameter
+        return None
+
+
+def get_discovery_urls(auth_server_url: str) -> list[str]:
+    """Generate ordered list of (url, type) tuples for discovery attempts."""
+    urls: list[str] = []
+    parsed = urlparse(auth_server_url)
+    base_url = f"{parsed.scheme}://{parsed.netloc}"
+
+    # RFC 8414: Path-aware OAuth discovery
+    if parsed.path and parsed.path != "/":
+        oauth_path = f"/.well-known/oauth-authorization-server{parsed.path.rstrip('/')}"
+        urls.append(urljoin(base_url, oauth_path))
+
+    # OAuth root fallback
+    urls.append(urljoin(base_url, "/.well-known/oauth-authorization-server"))
+
+    # RFC 8414 section 5: Path-aware OIDC discovery
+    # See https://www.rfc-editor.org/rfc/rfc8414.html#section-5
+    if parsed.path and parsed.path != "/":
+        oidc_path = f"/.well-known/openid-configuration{parsed.path.rstrip('/')}"
+        urls.append(urljoin(base_url, oidc_path))
+
+    # OIDC 1.0 fallback (appends to full URL per OIDC spec)
+    oidc_fallback = f"{auth_server_url.rstrip('/')}/.well-known/openid-configuration"
+    urls.append(oidc_fallback)
+
+    return urls