refactor: pull more oauth helpers out

Author: maxisbey

src/mcp/client/auth/oauth2.py

@@ -4,11 +4,8 @@
 Implements authorization code flow with PKCE and automatic token refresh.
 """
 
-import base64
-import hashlib
 import logging
 import secrets
-import string
 import time
 from collections.abc import AsyncGenerator, Awaitable, Callable
 from dataclasses import dataclass, field
@@ -32,6 +29,7 @@
     handle_auth_metadata_response,
     handle_protected_resource_response,
     handle_registration_response,
+    handle_token_response_scopes,
 )
 from mcp.client.streamable_http import MCP_PROTOCOL_VERSION
 from mcp.shared.auth import (
@@ -41,7 +39,12 @@
     OAuthToken,
     ProtectedResourceMetadata,
 )
-from mcp.shared.auth_utils import check_resource_allowed, resource_url_from_server_url
+from mcp.shared.auth_utils import (
+    calculate_token_expiry,
+    check_resource_allowed,
+    generate_pkce_parameters,
+    resource_url_from_server_url,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -54,10 +57,8 @@ class PKCEParameters(BaseModel):
 
     @classmethod
     def generate(cls) -> "PKCEParameters":
-        """Generate new PKCE parameters."""
-        code_verifier = "".join(secrets.choice(string.ascii_letters + string.digits + "-._~") for _ in range(128))
-        digest = hashlib.sha256(code_verifier.encode()).digest()
-        code_challenge = base64.urlsafe_b64encode(digest).decode().rstrip("=")
+        """Generate new PKCE parameters using shared util function."""
+        code_verifier, code_challenge = generate_pkce_parameters(verifier_length=128)
         return cls(code_verifier=code_verifier, code_challenge=code_challenge)
 
 
@@ -114,11 +115,8 @@ def get_authorization_base_url(self, server_url: str) -> str:
         return f"{parsed.scheme}://{parsed.netloc}"
 
     def update_token_expiry(self, token: OAuthToken) -> None:
-        """Update token expiry time."""
-        if token.expires_in:
-            self.token_expiry_time = time.time() + token.expires_in
-        else:
-            self.token_expiry_time = None
+        """Update token expiry time using shared util function."""
+        self.token_expiry_time = calculate_token_expiry(token.expires_in)
 
     def is_token_valid(self) -> bool:
         """Check if current token is valid."""
@@ -364,26 +362,20 @@ async def _handle_token_response(self, response: httpx.Response) -> None:
         """Handle token exchange response."""
         if response.status_code != 200:
             body = await response.aread()
-            body = body.decode("utf-8")
-            raise OAuthTokenError(f"Token exchange failed ({response.status_code}): {body}")
-
-        try:
-            content = await response.aread()
-            token_response = OAuthToken.model_validate_json(content)
-
-            # Validate scopes
-            if token_response.scope and self.context.client_metadata.scope:
-                requested_scopes = set(self.context.client_metadata.scope.split())
-                returned_scopes = set(token_response.scope.split())
-                unauthorized_scopes = returned_scopes - requested_scopes
-                if unauthorized_scopes:
-                    raise OAuthTokenError(f"Server granted unauthorized scopes: {unauthorized_scopes}")
+            body_text = body.decode("utf-8")
+            raise OAuthTokenError(f"Token exchange failed ({response.status_code}): {body_text}")
+
+        # Parse and validate response with scope validation
+        token_response = await handle_token_response_scopes(
+            response,
+            self.context.client_metadata,
+            validate_scope=True,
+        )
 
-            self.context.current_tokens = token_response
-            self.context.update_token_expiry(token_response)
-            await self.context.storage.set_tokens(token_response)
-        except ValidationError as e:
-            raise OAuthTokenError(f"Invalid token response: {e}")
+        # Store tokens in context
+        self.context.current_tokens = token_response
+        self.context.update_token_expiry(token_response)
+        await self.context.storage.set_tokens(token_response)
 
     async def _refresh_token(self) -> httpx.Request:
         """Build token refresh request."""

src/mcp/client/auth/utils.py

@@ -5,9 +5,15 @@
 from httpx import Request, Response
 from pydantic import ValidationError
 
-from mcp.client.auth import OAuthRegistrationError
+from mcp.client.auth import OAuthRegistrationError, OAuthTokenError
 from mcp.client.streamable_http import MCP_PROTOCOL_VERSION
-from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthMetadata, ProtectedResourceMetadata
+from mcp.shared.auth import (
+    OAuthClientInformationFull,
+    OAuthClientMetadata,
+    OAuthMetadata,
+    OAuthToken,
+    ProtectedResourceMetadata,
+)
 from mcp.types import LATEST_PROTOCOL_VERSION
 
 logger = logging.getLogger(__name__)
@@ -220,6 +226,45 @@ async def handle_registration_response(response: Response) -> OAuthClientInforma
         raise OAuthRegistrationError(f"Invalid registration response: {e}")
 
 
+async def handle_token_response_scopes(
+    response: Response,
+    client_metadata: OAuthClientMetadata,
+    validate_scope: bool = True,
+) -> OAuthToken:
+    """Parse and validate token response with optional scope validation.
+
+    Parses token response JSON and validates scopes to prevent scope escalation
+    if requested. Callers should check response.status_code before calling.
+
+    Args:
+        response: HTTP response from token endpoint (status already checked by caller)
+        client_metadata: Client metadata containing requested scopes (if any)
+        validate_scope: Whether to validate scopes (default True). Set False for refresh.
+
+    Returns:
+        Validated OAuthToken model
+
+    Raises:
+        OAuthTokenError: If response JSON is invalid or contains unauthorized scopes
+    """
+    try:
+        content = await response.aread()
+        token_response = OAuthToken.model_validate_json(content)
+
+        # Validate scopes to prevent scope escalation
+        # Only validate during initial token exchange, not during refresh
+        if validate_scope and token_response.scope and client_metadata.scope:
+            requested_scopes = set(client_metadata.scope.split())
+            returned_scopes = set(token_response.scope.split())
+            unauthorized_scopes = returned_scopes - requested_scopes
+            if unauthorized_scopes:
+                raise OAuthTokenError(f"Server granted unauthorized scopes: {unauthorized_scopes}")
+
+        return token_response
+    except ValidationError as e:
+        raise OAuthTokenError(f"Invalid token response: {e}")
+
+
 # async def prm_discovery(
 #     server_url: str,
 #     www_auth_resource_metadata_url: str | None,

src/mcp/shared/auth_utils.py

@@ -1,5 +1,10 @@
-"""Utilities for OAuth 2.0 Resource Indicators (RFC 8707)."""
+"""Utilities for OAuth 2.0 Resource Indicators (RFC 8707) and PKCE (RFC 7636)."""
 
+import base64
+import hashlib
+import secrets
+import string
+import time
 from urllib.parse import urlparse, urlsplit, urlunsplit
 
 from pydantic import AnyUrl, HttpUrl
@@ -67,3 +72,50 @@ def check_resource_allowed(requested_resource: str, configured_resource: str) ->
         configured_path += "/"
 
     return requested_path.startswith(configured_path)
+
+
+def generate_pkce_parameters(verifier_length: int = 128) -> tuple[str, str]:
+    """Generate PKCE verifier and challenge per RFC 7636.
+
+    Generates cryptographically secure code_verifier and code_challenge
+    for OAuth 2.0 PKCE (Proof Key for Code Exchange).
+
+    Args:
+        verifier_length: Length of code_verifier (43-128 chars per RFC 7636, default 128)
+
+    Returns:
+        Tuple of (code_verifier, code_challenge)
+
+    Raises:
+        ValueError: If verifier_length is not between 43 and 128
+    """
+    if not 43 <= verifier_length <= 128:
+        raise ValueError("verifier_length must be between 43 and 128 per RFC 7636")
+
+    # Generate code_verifier using unreserved characters per RFC 7636 Section 4.1
+    # unreserved = ALPHA / DIGIT / "-" / "." / "_" / "~"
+    code_verifier = "".join(
+        secrets.choice(string.ascii_letters + string.digits + "-._~") for _ in range(verifier_length)
+    )
+
+    # Generate code_challenge using S256 method per RFC 7636 Section 4.2
+    # code_challenge = BASE64URL(SHA256(ASCII(code_verifier)))
+    digest = hashlib.sha256(code_verifier.encode("ascii")).digest()
+    code_challenge = base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")
+
+    return code_verifier, code_challenge
+
+
+def calculate_token_expiry(expires_in: int | str | None) -> float | None:
+    """Calculate token expiry timestamp from expires_in seconds.
+
+    Args:
+        expires_in: Seconds until token expiration (may be string from some servers)
+
+    Returns:
+        Unix timestamp when token expires, or None if no expiry specified
+    """
+    if expires_in is None:
+        return None
+    # Defensive: handle servers that return expires_in as string
+    return time.time() + int(expires_in)