786 lines
38 KiB
Python
786 lines
38 KiB
Python
"""OAuth2 Authentication implementation for httpx2.
|
|
|
|
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
|
|
from typing import Any, Protocol, get_args
|
|
from urllib.parse import quote, urlencode, urljoin, urlparse
|
|
|
|
import anyio
|
|
import httpx2
|
|
from mcp_types.version import is_version_at_least
|
|
from pydantic import BaseModel, Field, ValidationError
|
|
|
|
from mcp.client.auth.exceptions import OAuthFlowError, OAuthRegistrationError, OAuthTokenError
|
|
from mcp.client.auth.utils import (
|
|
build_oauth_authorization_server_metadata_discovery_urls,
|
|
build_protected_resource_metadata_discovery_urls,
|
|
create_client_info_from_metadata_url,
|
|
create_client_registration_request,
|
|
create_oauth_metadata_request,
|
|
credentials_match_issuer,
|
|
extract_field_from_www_auth,
|
|
extract_resource_metadata_from_www_auth,
|
|
extract_scope_from_www_auth,
|
|
get_client_metadata_scopes,
|
|
handle_auth_metadata_response,
|
|
handle_protected_resource_response,
|
|
handle_registration_response,
|
|
handle_token_response_scopes,
|
|
is_valid_client_metadata_url,
|
|
should_use_client_metadata_url,
|
|
union_scopes,
|
|
validate_authorization_response_iss,
|
|
validate_metadata_issuer,
|
|
)
|
|
from mcp.shared.auth import (
|
|
AuthorizationCodeResult,
|
|
OAuthClientInformationFull,
|
|
OAuthClientMetadata,
|
|
OAuthMetadata,
|
|
OAuthToken,
|
|
ProtectedResourceMetadata,
|
|
TokenEndpointAuthMethod,
|
|
)
|
|
from mcp.shared.auth_utils import (
|
|
calculate_token_expiry,
|
|
check_resource_allowed,
|
|
resource_url_from_server_url,
|
|
)
|
|
from mcp.shared.inbound import MCP_PROTOCOL_VERSION_HEADER
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Methods a registered client's record may carry without a token request being an error,
|
|
# derived from the set the SDK is willing to request so the two cannot drift. `None`/"none"
|
|
# send no client secret. `private_key_jwt` sends none from here either: only
|
|
# `PrivateKeyJWTOAuthProvider` signs the assertion, and only in its client-credentials
|
|
# exchange, so its inherited refresh path must pass through here without raising - a refresh
|
|
# the server then rejects falls back to a fresh client-credentials exchange, which signs.
|
|
# Anything else is a method no client here can apply.
|
|
_KNOWN_TOKEN_ENDPOINT_AUTH_METHODS: tuple[str | None, ...] = (None, *get_args(TokenEndpointAuthMethod))
|
|
|
|
# Methods that authenticate the token request with the minted `client_secret`; a
|
|
# registration assigning one is only usable if the server issued that secret.
|
|
_SECRET_TOKEN_ENDPOINT_AUTH_METHODS = ("client_secret_post", "client_secret_basic")
|
|
|
|
# Methods a registration completed by the authorization-code flow can act on. That flow
|
|
# authenticates the token request with the minted client secret (or nothing); it holds no key
|
|
# to sign a `private_key_jwt` assertion, so a server assigning that method has registered a
|
|
# client this flow cannot use. `PrivateKeyJWTOAuthProvider` never registers dynamically.
|
|
_REGISTRATION_USABLE_TOKEN_ENDPOINT_AUTH_METHODS: tuple[str | None, ...] = tuple(
|
|
method for method in _KNOWN_TOKEN_ENDPOINT_AUTH_METHODS if method != "private_key_jwt"
|
|
)
|
|
|
|
|
|
def check_registration_usable(client_info: OAuthClientInformationFull) -> None:
|
|
"""Confirm a registration this flow completed is one it can act on.
|
|
|
|
RFC 7591 §3.2.1 lets the authorization server replace requested metadata and leaves it to
|
|
the client to "check the values in the response to determine if the registration is
|
|
sufficient for use". Two substitutions make the minted credentials unusable, and both are
|
|
judged here - before the record is persisted or any interactive authorization begins -
|
|
rather than surfacing later as an opaque failure at the token endpoint: a token-endpoint
|
|
auth method the authorization-code flow cannot apply (one it does not implement, or
|
|
`private_key_jwt`, whose assertion this flow has no key to sign), and a secret-based
|
|
method the flow could apply but for which the server issued no `client_secret`.
|
|
|
|
Raises:
|
|
OAuthRegistrationError: The server registered the client with a
|
|
`token_endpoint_auth_method` this flow cannot apply, or with a secret-based
|
|
method but no `client_secret`.
|
|
"""
|
|
method = client_info.token_endpoint_auth_method
|
|
if method not in _REGISTRATION_USABLE_TOKEN_ENDPOINT_AUTH_METHODS:
|
|
raise OAuthRegistrationError(
|
|
f"Authorization server registered the client with unsupported token_endpoint_auth_method {method!r}"
|
|
)
|
|
if method in _SECRET_TOKEN_ENDPOINT_AUTH_METHODS and client_info.client_secret is None:
|
|
raise OAuthRegistrationError(
|
|
f"Authorization server registered the client for {method!r} but issued no client_secret"
|
|
)
|
|
|
|
|
|
class PKCEParameters(BaseModel):
|
|
"""PKCE (Proof Key for Code Exchange) parameters."""
|
|
|
|
code_verifier: str = Field(..., min_length=43, max_length=128)
|
|
code_challenge: str = Field(..., min_length=43, max_length=128)
|
|
|
|
@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("=")
|
|
return cls(code_verifier=code_verifier, code_challenge=code_challenge)
|
|
|
|
|
|
class TokenStorage(Protocol):
|
|
"""Protocol for token storage implementations."""
|
|
|
|
async def get_tokens(self) -> OAuthToken | None:
|
|
"""Get stored tokens."""
|
|
...
|
|
|
|
async def set_tokens(self, tokens: OAuthToken) -> None:
|
|
"""Store tokens."""
|
|
...
|
|
|
|
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."""
|
|
...
|
|
|
|
|
|
@dataclass
|
|
class OAuthContext:
|
|
"""OAuth flow context."""
|
|
|
|
server_url: str
|
|
client_metadata: OAuthClientMetadata
|
|
storage: TokenStorage
|
|
redirect_handler: Callable[[str], Awaitable[None]] | None
|
|
callback_handler: Callable[[], Awaitable[AuthorizationCodeResult]] | None
|
|
client_metadata_url: str | None = None
|
|
|
|
# Discovered metadata
|
|
protected_resource_metadata: ProtectedResourceMetadata | None = None
|
|
oauth_metadata: OAuthMetadata | None = None
|
|
auth_server_url: str | None = None
|
|
protocol_version: str | None = None
|
|
|
|
# Client registration
|
|
client_info: OAuthClientInformationFull | None = None
|
|
|
|
# Token management
|
|
current_tokens: OAuthToken | None = None
|
|
token_expiry_time: float | None = None
|
|
|
|
# State
|
|
lock: anyio.Lock = field(default_factory=anyio.Lock)
|
|
|
|
def get_authorization_base_url(self, server_url: str) -> str:
|
|
"""Extract base URL by removing path component."""
|
|
parsed = urlparse(server_url)
|
|
return f"{parsed.scheme}://{parsed.netloc}"
|
|
|
|
def update_token_expiry(self, token: OAuthToken) -> 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."""
|
|
return bool(
|
|
self.current_tokens
|
|
and self.current_tokens.access_token
|
|
and (not self.token_expiry_time or time.time() <= self.token_expiry_time)
|
|
)
|
|
|
|
def can_refresh_token(self) -> bool:
|
|
"""Check if token can be refreshed."""
|
|
return bool(self.current_tokens and self.current_tokens.refresh_token and self.client_info)
|
|
|
|
def clear_tokens(self) -> None:
|
|
"""Clear current tokens."""
|
|
self.current_tokens = None
|
|
self.token_expiry_time = None
|
|
|
|
def get_resource_url(self) -> str:
|
|
"""Get resource URL for RFC 8707.
|
|
|
|
Uses PRM resource if it's a valid parent, otherwise uses canonical server URL.
|
|
"""
|
|
resource = resource_url_from_server_url(self.server_url)
|
|
|
|
# If PRM provides a resource that's a valid parent, use it
|
|
if self.protected_resource_metadata and self.protected_resource_metadata.resource:
|
|
prm_resource = str(self.protected_resource_metadata.resource)
|
|
if check_resource_allowed(requested_resource=resource, configured_resource=prm_resource):
|
|
resource = prm_resource
|
|
|
|
return resource
|
|
|
|
def should_include_resource_param(self, protocol_version: str | None = None) -> bool:
|
|
"""Determine if the resource parameter should be included in OAuth requests.
|
|
|
|
Returns True if:
|
|
- Protected resource metadata is available, OR
|
|
- MCP-Protocol-Version header is 2025-06-18 or later
|
|
"""
|
|
# If we have protected resource metadata, include the resource param
|
|
if self.protected_resource_metadata is not None:
|
|
return True
|
|
|
|
# If no protocol version provided, don't include resource param
|
|
if not protocol_version:
|
|
return False
|
|
|
|
return is_version_at_least(protocol_version, "2025-06-18")
|
|
|
|
def prepare_token_auth(
|
|
self, data: dict[str, str], headers: dict[str, str] | None = None
|
|
) -> tuple[dict[str, str], dict[str, str]]:
|
|
"""Prepare authentication for token requests.
|
|
|
|
Args:
|
|
data: The form data to send
|
|
headers: Optional headers dict to update
|
|
|
|
Returns:
|
|
Tuple of (updated_data, updated_headers)
|
|
|
|
Raises:
|
|
OAuthTokenError: The client record carries a `token_endpoint_auth_method` this
|
|
client does not know. A dynamic registration assigning an unusable method is
|
|
rejected earlier, by `check_registration_usable`; this fires for a stored or
|
|
pre-registered record that reaches a token request with such a method.
|
|
"""
|
|
if headers is None:
|
|
headers = {} # pragma: no cover
|
|
|
|
if not self.client_info:
|
|
return data, headers
|
|
|
|
auth_method = self.client_info.token_endpoint_auth_method
|
|
|
|
if auth_method == "client_secret_basic" and self.client_info.client_secret:
|
|
# URL-encode client ID and secret per RFC 6749 Section 2.3.1
|
|
encoded_id = quote(self.client_info.client_id, safe="")
|
|
encoded_secret = quote(self.client_info.client_secret, safe="")
|
|
credentials = f"{encoded_id}:{encoded_secret}"
|
|
encoded_credentials = base64.b64encode(credentials.encode()).decode()
|
|
headers["Authorization"] = f"Basic {encoded_credentials}"
|
|
# Don't include client_secret in body for basic auth
|
|
data = {k: v for k, v in data.items() if k != "client_secret"}
|
|
elif auth_method == "client_secret_post" and self.client_info.client_secret:
|
|
# Include client_id and client_secret in request body (RFC 6749 §2.3.1)
|
|
data["client_id"] = self.client_info.client_id
|
|
data["client_secret"] = self.client_info.client_secret
|
|
elif auth_method not in _KNOWN_TOKEN_ENDPOINT_AUTH_METHODS:
|
|
raise OAuthTokenError(f"Registered client uses unsupported token_endpoint_auth_method {auth_method!r}")
|
|
# For "none" (or absent), don't add any client_secret; "private_key_jwt" adds its
|
|
# assertion in the provider that implements it, not here.
|
|
|
|
return data, headers
|
|
|
|
|
|
class OAuthClientProvider(httpx2.Auth):
|
|
"""OAuth2 authentication for httpx2.
|
|
|
|
Handles OAuth flow with automatic client registration and token storage.
|
|
"""
|
|
|
|
requires_response_body = True
|
|
|
|
def __init__(
|
|
self,
|
|
server_url: str,
|
|
client_metadata: OAuthClientMetadata,
|
|
storage: TokenStorage,
|
|
redirect_handler: Callable[[str], Awaitable[None]] | None = None,
|
|
callback_handler: Callable[[], Awaitable[AuthorizationCodeResult]] | None = None,
|
|
client_metadata_url: str | None = None,
|
|
validate_resource_url: Callable[[str, str | None], Awaitable[None]] | None = None,
|
|
):
|
|
"""Initialize OAuth2 authentication.
|
|
|
|
Args:
|
|
server_url: The MCP server URL.
|
|
client_metadata: OAuth client metadata for registration.
|
|
storage: Token storage implementation.
|
|
redirect_handler: Handler for authorization redirects.
|
|
callback_handler: Handler for authorization callbacks.
|
|
client_metadata_url: URL-based client ID. When provided and the server
|
|
advertises client_id_metadata_document_supported=True, this URL will be
|
|
used as the client_id instead of performing dynamic client registration.
|
|
Must be a valid HTTPS URL with a non-root pathname.
|
|
validate_resource_url: Optional callback to override resource URL validation.
|
|
Called with (server_url, prm_resource) where prm_resource is the resource
|
|
from Protected Resource Metadata (or None if not present). If not provided,
|
|
default validation rejects mismatched resources per RFC 8707.
|
|
|
|
Raises:
|
|
ValueError: If client_metadata_url is provided but not a valid HTTPS URL
|
|
with a non-root pathname.
|
|
"""
|
|
# Validate client_metadata_url if provided
|
|
if client_metadata_url is not None and not is_valid_client_metadata_url(client_metadata_url):
|
|
raise ValueError(
|
|
f"client_metadata_url must be a valid HTTPS URL with a non-root pathname, got: {client_metadata_url}"
|
|
)
|
|
|
|
self.context = OAuthContext(
|
|
server_url=server_url,
|
|
client_metadata=client_metadata,
|
|
storage=storage,
|
|
redirect_handler=redirect_handler,
|
|
callback_handler=callback_handler,
|
|
client_metadata_url=client_metadata_url,
|
|
)
|
|
self._validate_resource_url_callback = validate_resource_url
|
|
self._initialized = False
|
|
|
|
async def _handle_protected_resource_response(self, response: httpx2.Response) -> bool:
|
|
"""Handle protected resource metadata discovery response.
|
|
|
|
Per SEP-985, supports fallback when discovery fails at one URL.
|
|
|
|
Returns:
|
|
True if metadata was successfully discovered, False if we should try next URL
|
|
"""
|
|
if response.status_code == 200:
|
|
try:
|
|
content = await response.aread()
|
|
metadata = ProtectedResourceMetadata.model_validate_json(content)
|
|
self.context.protected_resource_metadata = metadata
|
|
if metadata.authorization_servers: # pragma: no branch
|
|
self.context.auth_server_url = str(metadata.authorization_servers[0])
|
|
return True
|
|
|
|
except ValidationError: # pragma: no cover
|
|
# Invalid metadata - try next URL
|
|
logger.warning(f"Invalid protected resource metadata at {response.request.url}")
|
|
return False
|
|
elif response.status_code == 404: # pragma: no cover
|
|
# Not found - try next URL in fallback chain
|
|
logger.debug(f"Protected resource metadata not found at {response.request.url}, trying next URL")
|
|
return False
|
|
else:
|
|
# Other error - fail immediately
|
|
raise OAuthFlowError(
|
|
f"Protected Resource Metadata request failed: {response.status_code}"
|
|
) # pragma: no cover
|
|
|
|
async def _perform_authorization(self) -> httpx2.Request:
|
|
"""Perform the authorization flow."""
|
|
auth_code, code_verifier = await self._perform_authorization_code_grant()
|
|
token_request = await self._exchange_token_authorization_code(auth_code, code_verifier)
|
|
return token_request
|
|
|
|
async def _perform_authorization_code_grant(self) -> tuple[str, str]:
|
|
"""Perform the authorization redirect and get auth code."""
|
|
if self.context.client_metadata.redirect_uris is None:
|
|
raise OAuthFlowError("No redirect URIs provided for authorization code grant") # pragma: no cover
|
|
if not self.context.redirect_handler:
|
|
raise OAuthFlowError("No redirect handler provided for authorization code grant") # pragma: no cover
|
|
if not self.context.callback_handler:
|
|
raise OAuthFlowError("No callback handler provided for authorization code grant") # pragma: no cover
|
|
|
|
if self.context.oauth_metadata and self.context.oauth_metadata.authorization_endpoint:
|
|
auth_endpoint = str(self.context.oauth_metadata.authorization_endpoint)
|
|
else:
|
|
auth_base_url = self.context.get_authorization_base_url(self.context.server_url)
|
|
auth_endpoint = urljoin(auth_base_url, "/authorize")
|
|
|
|
if not self.context.client_info:
|
|
raise OAuthFlowError("No client info available for authorization") # pragma: no cover
|
|
|
|
# Generate PKCE parameters
|
|
pkce_params = PKCEParameters.generate()
|
|
state = secrets.token_urlsafe(32)
|
|
|
|
auth_params = {
|
|
"response_type": "code",
|
|
"client_id": self.context.client_info.client_id,
|
|
"redirect_uri": str(self.context.client_metadata.redirect_uris[0]),
|
|
"state": state,
|
|
"code_challenge": pkce_params.code_challenge,
|
|
"code_challenge_method": "S256",
|
|
}
|
|
|
|
# Only include resource param if conditions are met
|
|
if self.context.should_include_resource_param(self.context.protocol_version):
|
|
auth_params["resource"] = self.context.get_resource_url() # RFC 8707
|
|
|
|
if self.context.client_metadata.scope: # pragma: no branch
|
|
auth_params["scope"] = self.context.client_metadata.scope
|
|
|
|
# OIDC requires prompt=consent when offline_access is requested
|
|
# https://openid.net/specs/openid-connect-core-1_0.html#OfflineAccess
|
|
if "offline_access" in self.context.client_metadata.scope.split():
|
|
auth_params["prompt"] = "consent"
|
|
|
|
authorization_url = f"{auth_endpoint}?{urlencode(auth_params)}"
|
|
await self.context.redirect_handler(authorization_url)
|
|
|
|
# Wait for callback
|
|
result = await self.context.callback_handler()
|
|
|
|
if result.state is None or not secrets.compare_digest(result.state, state):
|
|
raise OAuthFlowError(f"State parameter mismatch: {result.state} != {state}")
|
|
|
|
# RFC 9207: validate the authorization-response issuer
|
|
validate_authorization_response_iss(result.iss, self.context.oauth_metadata)
|
|
|
|
if not result.code:
|
|
raise OAuthFlowError("No authorization code received")
|
|
|
|
# Return auth code and code verifier for token exchange
|
|
return result.code, pkce_params.code_verifier
|
|
|
|
def _get_token_endpoint(self) -> str:
|
|
if self.context.oauth_metadata and self.context.oauth_metadata.token_endpoint:
|
|
token_url = str(self.context.oauth_metadata.token_endpoint)
|
|
else:
|
|
auth_base_url = self.context.get_authorization_base_url(self.context.server_url)
|
|
token_url = urljoin(auth_base_url, "/token")
|
|
return token_url
|
|
|
|
async def _exchange_token_authorization_code(self, auth_code: str, code_verifier: str) -> httpx2.Request:
|
|
"""Build token exchange request for authorization_code flow."""
|
|
if self.context.client_metadata.redirect_uris is None:
|
|
raise OAuthFlowError("No redirect URIs provided for authorization code grant") # pragma: no cover
|
|
if not self.context.client_info:
|
|
raise OAuthFlowError("Missing client info") # pragma: no cover
|
|
|
|
token_url = self._get_token_endpoint()
|
|
token_data: dict[str, Any] = {
|
|
"grant_type": "authorization_code",
|
|
"code": auth_code,
|
|
"redirect_uri": str(self.context.client_metadata.redirect_uris[0]),
|
|
"client_id": self.context.client_info.client_id,
|
|
"code_verifier": code_verifier,
|
|
}
|
|
|
|
# Only include resource param if conditions are met
|
|
if self.context.should_include_resource_param(self.context.protocol_version):
|
|
token_data["resource"] = self.context.get_resource_url() # RFC 8707
|
|
|
|
# Prepare authentication based on preferred method
|
|
headers = {"Content-Type": "application/x-www-form-urlencoded"}
|
|
token_data, headers = self.context.prepare_token_auth(token_data, headers)
|
|
|
|
return httpx2.Request("POST", token_url, data=token_data, headers=headers)
|
|
|
|
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()
|
|
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)
|
|
|
|
# RFC 6749 §5.1: an omitted scope means the granted scope equals the requested
|
|
# scope. Record it explicitly so the persisted token is self-describing — the
|
|
# SEP-2350 step-up union reads it after a restart, when client_metadata.scope
|
|
# has reverted to its constructor value.
|
|
if token_response.scope is None:
|
|
token_response.scope = self.context.client_metadata.scope
|
|
|
|
# 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) -> httpx2.Request:
|
|
"""Build token refresh request."""
|
|
if not self.context.current_tokens or not self.context.current_tokens.refresh_token:
|
|
raise OAuthTokenError("No refresh token available") # pragma: no cover
|
|
|
|
if not self.context.client_info or not self.context.client_info.client_id:
|
|
raise OAuthTokenError("No client info available") # pragma: no cover
|
|
|
|
if self.context.oauth_metadata and self.context.oauth_metadata.token_endpoint:
|
|
token_url = str(self.context.oauth_metadata.token_endpoint)
|
|
else:
|
|
auth_base_url = self.context.get_authorization_base_url(self.context.server_url)
|
|
token_url = urljoin(auth_base_url, "/token")
|
|
|
|
refresh_data: dict[str, str] = {
|
|
"grant_type": "refresh_token",
|
|
"refresh_token": self.context.current_tokens.refresh_token,
|
|
"client_id": self.context.client_info.client_id,
|
|
}
|
|
|
|
# Only include resource param if conditions are met
|
|
if self.context.should_include_resource_param(self.context.protocol_version):
|
|
refresh_data["resource"] = self.context.get_resource_url() # RFC 8707
|
|
|
|
# Prepare authentication based on preferred method
|
|
headers = {"Content-Type": "application/x-www-form-urlencoded"}
|
|
refresh_data, headers = self.context.prepare_token_auth(refresh_data, headers)
|
|
|
|
return httpx2.Request("POST", token_url, data=refresh_data, headers=headers)
|
|
|
|
async def _handle_refresh_response(self, response: httpx2.Response) -> bool:
|
|
"""Handle token refresh response. Returns True if successful."""
|
|
if response.status_code != 200:
|
|
logger.warning(f"Token refresh failed: {response.status_code}")
|
|
self.context.clear_tokens()
|
|
return False
|
|
|
|
try:
|
|
content = await response.aread()
|
|
token_response = OAuthToken.model_validate_json(content)
|
|
|
|
# RFC 6749 §6: a refresh response may omit scope (unchanged) and refresh_token
|
|
# (the AS does not rotate). Carry both forward so the persisted token stays
|
|
# self-describing for the SEP-2350 step-up union and the next expiry can
|
|
# still refresh instead of forcing a full re-authorization.
|
|
prior = self.context.current_tokens
|
|
if token_response.scope is None and prior is not None:
|
|
token_response.scope = prior.scope
|
|
if token_response.refresh_token is None and prior is not None:
|
|
token_response.refresh_token = prior.refresh_token
|
|
|
|
self.context.current_tokens = token_response
|
|
self.context.update_token_expiry(token_response)
|
|
await self.context.storage.set_tokens(token_response)
|
|
|
|
return True
|
|
except ValidationError: # pragma: no cover
|
|
logger.exception("Invalid refresh response")
|
|
self.context.clear_tokens()
|
|
return False
|
|
|
|
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()
|
|
self._initialized = True
|
|
|
|
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
|
|
request.headers["Authorization"] = f"Bearer {self.context.current_tokens.access_token}"
|
|
|
|
async def _handle_oauth_metadata_response(self, response: httpx2.Response) -> None:
|
|
content = await response.aread()
|
|
metadata = OAuthMetadata.model_validate_json(content)
|
|
self.context.oauth_metadata = metadata
|
|
|
|
async def _validate_resource_match(self, prm: ProtectedResourceMetadata) -> None:
|
|
"""Validate that PRM resource matches the server URL per RFC 8707."""
|
|
prm_resource = str(prm.resource) if prm.resource else None
|
|
|
|
if self._validate_resource_url_callback is not None:
|
|
await self._validate_resource_url_callback(self.context.server_url, prm_resource)
|
|
return
|
|
|
|
if not prm_resource:
|
|
return # pragma: no cover
|
|
default_resource = resource_url_from_server_url(self.context.server_url)
|
|
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 async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
|
|
"""httpx2 auth flow integration."""
|
|
async with self.context.lock:
|
|
if not self._initialized:
|
|
await self._initialize()
|
|
|
|
# Capture protocol version from request headers
|
|
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
|
|
|
|
if self.context.is_token_valid():
|
|
self._add_auth_header(request)
|
|
|
|
response = yield request
|
|
|
|
if response.status_code == 401:
|
|
# 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)
|
|
prm_discovery_urls = build_protected_resource_metadata_discovery_urls(
|
|
www_auth_resource_metadata_url, self.context.server_url
|
|
)
|
|
|
|
for url in prm_discovery_urls: # pragma: no branch
|
|
discovery_request = create_oauth_metadata_request(url)
|
|
|
|
discovery_response = yield discovery_request # sending request
|
|
|
|
prm = await handle_protected_resource_response(discovery_response)
|
|
if prm:
|
|
# Validate PRM resource matches server URL (RFC 8707)
|
|
await self._validate_resource_match(prm)
|
|
self.context.protected_resource_metadata = prm
|
|
|
|
# todo: try all authorization_servers to find the OASM
|
|
assert (
|
|
len(prm.authorization_servers) > 0
|
|
) # this is always true as authorization_servers has a min length of 1
|
|
|
|
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) so the
|
|
# flow re-registers instead of presenting another server's credentials.
|
|
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 re-registering")
|
|
self.context.client_info = None
|
|
self.context.clear_tokens()
|
|
# Any cached AS metadata is for the old server; drop it so a failed
|
|
# rediscovery cannot leak the old registration/token endpoints into Step 4.
|
|
self.context.oauth_metadata = None
|
|
|
|
asm_discovery_urls = build_oauth_authorization_server_metadata_discovery_urls(
|
|
self.context.auth_server_url, self.context.server_url
|
|
)
|
|
|
|
# Step 2: Discover OAuth Authorization Server Metadata (OASM) (with fallback for legacy servers)
|
|
for url in asm_discovery_urls: # pragma: no branch
|
|
oauth_metadata_request = create_oauth_metadata_request(url)
|
|
oauth_metadata_response = yield oauth_metadata_request
|
|
|
|
ok, asm = await handle_auth_metadata_response(oauth_metadata_response)
|
|
if not ok:
|
|
break
|
|
if ok and 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 bound_issuer fallback in Step 4).
|
|
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 re-registering")
|
|
self.context.client_info = None
|
|
self.context.clear_tokens()
|
|
|
|
# Step 3: Apply scope selection strategy
|
|
self.context.client_metadata.scope = get_client_metadata_scopes(
|
|
extract_scope_from_www_auth(response),
|
|
self.context.protected_resource_metadata,
|
|
self.context.oauth_metadata,
|
|
self.context.client_metadata.grant_types,
|
|
)
|
|
|
|
# Step 4: Register client or use URL-based client ID (CIMD)
|
|
if not self.context.client_info:
|
|
# SEP-2352: the issuer to bind these credentials to, when known.
|
|
discovered_issuer: str | None = None
|
|
if self.context.oauth_metadata is not None:
|
|
discovered_issuer = self.context.auth_server_url or str(self.context.oauth_metadata.issuer)
|
|
|
|
if should_use_client_metadata_url(
|
|
self.context.oauth_metadata, self.context.client_metadata_url
|
|
):
|
|
# Use URL-based client ID (CIMD). CIMD records are portable across
|
|
# authorization servers, so the issuer stamp is informational.
|
|
logger.debug(f"Using URL-based client ID (CIMD): {self.context.client_metadata_url}")
|
|
client_information = create_client_info_from_metadata_url(
|
|
self.context.client_metadata_url, # type: ignore[arg-type]
|
|
redirect_uris=self.context.client_metadata.redirect_uris,
|
|
)
|
|
client_information.issuer = discovered_issuer
|
|
self.context.client_info = client_information
|
|
await self.context.storage.set_client_info(client_information)
|
|
else:
|
|
# Fallback to Dynamic Client Registration
|
|
fallback_base = self.context.get_authorization_base_url(self.context.server_url)
|
|
registration_request = create_client_registration_request(
|
|
self.context.oauth_metadata, self.context.client_metadata, fallback_base
|
|
)
|
|
registration_response = yield registration_request
|
|
client_information = await handle_registration_response(registration_response)
|
|
check_registration_usable(client_information)
|
|
# Only record the issuer when the registration above actually targeted
|
|
# the discovered AS — either via its published registration_endpoint,
|
|
# or because the resource-origin /register fallback is on the issuer's
|
|
# own host (legacy same-origin embedded AS). Otherwise the fallback hit
|
|
# a different server and recording a binding to the PRM-advertised AS
|
|
# would persist a binding that was never established.
|
|
if (
|
|
self.context.oauth_metadata is not None
|
|
and discovered_issuer is not None
|
|
and (
|
|
self.context.oauth_metadata.registration_endpoint is not None
|
|
or self.context.get_authorization_base_url(discovered_issuer) == fallback_base
|
|
)
|
|
):
|
|
client_information.issuer = discovered_issuer
|
|
self.context.client_info = client_information
|
|
await self.context.storage.set_client_info(client_information)
|
|
|
|
# Step 5: Perform authorization and complete token exchange
|
|
token_response = yield await self._perform_authorization()
|
|
await self._handle_token_response(token_response)
|
|
except Exception:
|
|
logger.exception("OAuth flow error")
|
|
raise
|
|
|
|
# Retry with new tokens
|
|
self._add_auth_header(request)
|
|
yield request
|
|
elif response.status_code == 403:
|
|
# Step 1: Extract error field from WWW-Authenticate header
|
|
error = extract_field_from_www_auth(response, "error")
|
|
|
|
# Step 2: Check if we need to step-up authorization
|
|
if error == "insufficient_scope": # pragma: no branch
|
|
try:
|
|
# Step 2a: Union previously requested scopes with the newly challenged
|
|
# scopes (SEP-2350) so escalating one operation keeps the others' grants.
|
|
# Fold in the stored token's scope too: on a restart the token is reloaded
|
|
# but client_metadata.scope is not, so it would otherwise be the only basis.
|
|
challenged_scope = get_client_metadata_scopes(
|
|
extract_scope_from_www_auth(response),
|
|
self.context.protected_resource_metadata,
|
|
self.context.oauth_metadata,
|
|
self.context.client_metadata.grant_types,
|
|
)
|
|
granted_scope = self.context.current_tokens.scope if self.context.current_tokens else None
|
|
prior_scope = union_scopes(self.context.client_metadata.scope, granted_scope)
|
|
self.context.client_metadata.scope = union_scopes(prior_scope, challenged_scope)
|
|
|
|
# Step 2b: Perform (re-)authorization and token exchange
|
|
token_response = yield await self._perform_authorization()
|
|
await self._handle_token_response(token_response)
|
|
except Exception: # pragma: no cover
|
|
logger.exception("OAuth flow error")
|
|
raise
|
|
|
|
# Retry with new tokens
|
|
self._add_auth_header(request)
|
|
yield request
|