Implement RFC 7523 JWT flows (#1247)
Co-authored-by: Yann Jouanin <yann.jouanin@valueandco.com>
This commit is contained in:
Luca Chang
GitHub
parent
db9e451551
commit
f161149680
4 changed files with 463 additions and 52 deletions
7
src/mcp/client/auth/__init__.py
Normal file
7
src/mcp/client/auth/__init__.py
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
"""
|
||||
OAuth2 Authentication implementation for HTTPX.
|
||||
|
||||
Implements authorization code flow with PKCE and automatic token refresh.
|
||||
"""
|
||||
|
||||
from mcp.client.auth.oauth2 import * # noqa: F403
|
||||
148
src/mcp/client/auth/extensions/client_credentials.py
Normal file
148
src/mcp/client/auth/extensions/client_credentials.py
Normal file
|
|
@ -0,0 +1,148 @@
|
|||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
import httpx
|
||||
import jwt
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from mcp.client.auth import OAuthClientProvider, OAuthFlowError, OAuthTokenError, TokenStorage
|
||||
from mcp.shared.auth import OAuthClientMetadata
|
||||
|
||||
|
||||
class JWTParameters(BaseModel):
|
||||
"""JWT parameters."""
|
||||
|
||||
assertion: str | None = Field(
|
||||
default=None,
|
||||
description="JWT assertion for JWT authentication. "
|
||||
"Will be used instead of generating a new assertion if provided.",
|
||||
)
|
||||
|
||||
issuer: str | None = Field(default=None, description="Issuer for JWT assertions.")
|
||||
subject: str | None = Field(default=None, description="Subject identifier for JWT assertions.")
|
||||
audience: str | None = Field(default=None, description="Audience for JWT assertions.")
|
||||
claims: dict[str, Any] | None = Field(default=None, description="Additional claims for JWT assertions.")
|
||||
jwt_signing_algorithm: str | None = Field(default="RS256", description="Algorithm for signing JWT assertions.")
|
||||
jwt_signing_key: str | None = Field(default=None, description="Private key for JWT signing.")
|
||||
jwt_lifetime_seconds: int = Field(default=300, description="Lifetime of generated JWT in seconds.")
|
||||
|
||||
def to_assertion(self, with_audience_fallback: str | None = None) -> str:
|
||||
if self.assertion is not None:
|
||||
# Prebuilt JWT (e.g. acquired out-of-band)
|
||||
assertion = self.assertion
|
||||
else:
|
||||
if not self.jwt_signing_key:
|
||||
raise OAuthFlowError("Missing signing key for JWT bearer grant")
|
||||
if not self.issuer:
|
||||
raise OAuthFlowError("Missing issuer for JWT bearer grant")
|
||||
if not self.subject:
|
||||
raise OAuthFlowError("Missing subject for JWT bearer grant")
|
||||
|
||||
audience = self.audience if self.audience else with_audience_fallback
|
||||
if not audience:
|
||||
raise OAuthFlowError("Missing audience for JWT bearer grant")
|
||||
|
||||
now = int(time.time())
|
||||
claims: dict[str, Any] = {
|
||||
"iss": self.issuer,
|
||||
"sub": self.subject,
|
||||
"aud": audience,
|
||||
"exp": now + self.jwt_lifetime_seconds,
|
||||
"iat": now,
|
||||
"jti": str(uuid4()),
|
||||
}
|
||||
claims.update(self.claims or {})
|
||||
|
||||
assertion = jwt.encode(
|
||||
claims,
|
||||
self.jwt_signing_key,
|
||||
algorithm=self.jwt_signing_algorithm or "RS256",
|
||||
)
|
||||
return assertion
|
||||
|
||||
|
||||
class RFC7523OAuthClientProvider(OAuthClientProvider):
|
||||
"""OAuth client provider for RFC7532 clients."""
|
||||
|
||||
jwt_parameters: JWTParameters | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server_url: str,
|
||||
client_metadata: OAuthClientMetadata,
|
||||
storage: TokenStorage,
|
||||
redirect_handler: Callable[[str], Awaitable[None]] | None = None,
|
||||
callback_handler: Callable[[], Awaitable[tuple[str, str | None]]] | None = None,
|
||||
timeout: float = 300.0,
|
||||
jwt_parameters: JWTParameters | None = None,
|
||||
) -> None:
|
||||
super().__init__(server_url, client_metadata, storage, redirect_handler, callback_handler, timeout)
|
||||
self.jwt_parameters = jwt_parameters
|
||||
|
||||
async def _exchange_token_authorization_code(
|
||||
self, auth_code: str, code_verifier: str, *, token_data: dict[str, Any] | None = None
|
||||
) -> httpx.Request:
|
||||
"""Build token exchange request for authorization_code flow."""
|
||||
token_data = token_data or {}
|
||||
if self.context.client_metadata.token_endpoint_auth_method == "private_key_jwt":
|
||||
self._add_client_authentication_jwt(token_data=token_data)
|
||||
return await super()._exchange_token_authorization_code(auth_code, code_verifier, token_data=token_data)
|
||||
|
||||
async def _perform_authorization(self) -> httpx.Request:
|
||||
"""Perform the authorization flow."""
|
||||
if "urn:ietf:params:oauth:grant-type:jwt-bearer" in self.context.client_metadata.grant_types:
|
||||
token_request = await self._exchange_token_jwt_bearer()
|
||||
return token_request
|
||||
else:
|
||||
return await super()._perform_authorization()
|
||||
|
||||
def _add_client_authentication_jwt(self, *, token_data: dict[str, Any]):
|
||||
"""Add JWT assertion for client authentication to token endpoint parameters."""
|
||||
if not self.jwt_parameters:
|
||||
raise OAuthTokenError("Missing JWT parameters for private_key_jwt flow")
|
||||
if not self.context.oauth_metadata:
|
||||
raise OAuthTokenError("Missing OAuth metadata for private_key_jwt flow")
|
||||
|
||||
# We need to set the audience to the issuer identifier of the authorization server
|
||||
# https://datatracker.ietf.org/doc/html/draft-ietf-oauth-rfc7523bis-01#name-updates-to-rfc-7523
|
||||
issuer = str(self.context.oauth_metadata.issuer)
|
||||
assertion = self.jwt_parameters.to_assertion(with_audience_fallback=issuer)
|
||||
|
||||
# When using private_key_jwt, in a client_credentials flow, we use RFC 7523 Section 2.2
|
||||
token_data["client_assertion"] = assertion
|
||||
token_data["client_assertion_type"] = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"
|
||||
# We need to set the audience to the resource server, the audience is difference from the one in claims
|
||||
# it represents the resource server that will validate the token
|
||||
token_data["audience"] = self.context.get_resource_url()
|
||||
|
||||
async def _exchange_token_jwt_bearer(self) -> httpx.Request:
|
||||
"""Build token exchange request for JWT bearer grant."""
|
||||
if not self.context.client_info:
|
||||
raise OAuthFlowError("Missing client info")
|
||||
if not self.jwt_parameters:
|
||||
raise OAuthFlowError("Missing JWT parameters")
|
||||
if not self.context.oauth_metadata:
|
||||
raise OAuthTokenError("Missing OAuth metadata")
|
||||
|
||||
# We need to set the audience to the issuer identifier of the authorization server
|
||||
# https://datatracker.ietf.org/doc/html/draft-ietf-oauth-rfc7523bis-01#name-updates-to-rfc-7523
|
||||
issuer = str(self.context.oauth_metadata.issuer)
|
||||
assertion = self.jwt_parameters.to_assertion(with_audience_fallback=issuer)
|
||||
|
||||
token_data = {
|
||||
"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
|
||||
"assertion": assertion,
|
||||
}
|
||||
|
||||
if self.context.should_include_resource_param(self.context.protocol_version):
|
||||
token_data["resource"] = self.context.get_resource_url()
|
||||
|
||||
if self.context.client_metadata.scope:
|
||||
token_data["scope"] = self.context.client_metadata.scope
|
||||
|
||||
token_url = self._get_token_endpoint()
|
||||
return httpx.Request(
|
||||
"POST", token_url, data=token_data, headers={"Content-Type": "application/x-www-form-urlencoded"}
|
||||
)
|
||||
634
src/mcp/client/auth/oauth2.py
Normal file
634
src/mcp/client/auth/oauth2.py
Normal file
|
|
@ -0,0 +1,634 @@
|
|||
"""
|
||||
OAuth2 Authentication implementation for HTTPX.
|
||||
|
||||
Implements authorization code flow with PKCE and automatic token refresh.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
import secrets
|
||||
import string
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Protocol
|
||||
from urllib.parse import urlencode, urljoin, urlparse
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
|
||||
from mcp.client.streamable_http import MCP_PROTOCOL_VERSION
|
||||
from mcp.shared.auth import (
|
||||
OAuthClientInformationFull,
|
||||
OAuthClientMetadata,
|
||||
OAuthMetadata,
|
||||
OAuthToken,
|
||||
ProtectedResourceMetadata,
|
||||
)
|
||||
from mcp.shared.auth_utils import check_resource_allowed, resource_url_from_server_url
|
||||
from mcp.types import LATEST_PROTOCOL_VERSION
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OAuthFlowError(Exception):
|
||||
"""Base exception for OAuth flow errors."""
|
||||
|
||||
|
||||
class OAuthTokenError(OAuthFlowError):
|
||||
"""Raised when token operations fail."""
|
||||
|
||||
|
||||
class OAuthRegistrationError(OAuthFlowError):
|
||||
"""Raised when client registration fails."""
|
||||
|
||||
|
||||
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[tuple[str, str | None]]] | None
|
||||
timeout: float = 300.0
|
||||
|
||||
# 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)
|
||||
|
||||
# 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)
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
# Check if protocol version is 2025-06-18 or later
|
||||
# Version format is YYYY-MM-DD, so string comparison works
|
||||
return protocol_version >= "2025-06-18"
|
||||
|
||||
|
||||
class OAuthClientProvider(httpx.Auth):
|
||||
"""
|
||||
OAuth2 authentication for httpx.
|
||||
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[tuple[str, str | None]]] | None = None,
|
||||
timeout: float = 300.0,
|
||||
):
|
||||
"""Initialize OAuth2 authentication."""
|
||||
self.context = OAuthContext(
|
||||
server_url=server_url,
|
||||
client_metadata=client_metadata,
|
||||
storage=storage,
|
||||
redirect_handler=redirect_handler,
|
||||
callback_handler=callback_handler,
|
||||
timeout=timeout,
|
||||
)
|
||||
self._initialized = False
|
||||
|
||||
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 _discover_protected_resource(self, init_response: httpx.Response) -> httpx.Request:
|
||||
# RFC9728: Try to extract resource_metadata URL from WWW-Authenticate header of the initial response
|
||||
url = self._extract_resource_metadata_from_www_auth(init_response)
|
||||
|
||||
if not url:
|
||||
# Fallback to well-known discovery
|
||||
auth_base_url = self.context.get_authorization_base_url(self.context.server_url)
|
||||
url = urljoin(auth_base_url, "/.well-known/oauth-protected-resource")
|
||||
|
||||
return httpx.Request("GET", url, headers={MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION})
|
||||
|
||||
async def _handle_protected_resource_response(self, response: httpx.Response) -> None:
|
||||
"""Handle discovery response."""
|
||||
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:
|
||||
self.context.auth_server_url = str(metadata.authorization_servers[0])
|
||||
|
||||
except ValidationError:
|
||||
pass
|
||||
else:
|
||||
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:
|
||||
return None
|
||||
|
||||
if self.context.oauth_metadata and self.context.oauth_metadata.registration_endpoint:
|
||||
registration_url = str(self.context.oauth_metadata.registration_endpoint)
|
||||
else:
|
||||
auth_base_url = self.context.get_authorization_base_url(self.context.server_url)
|
||||
registration_url = urljoin(auth_base_url, "/register")
|
||||
|
||||
registration_data = self.context.client_metadata.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
|
||||
return httpx.Request(
|
||||
"POST", registration_url, json=registration_data, headers={"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
async def _handle_registration_response(self, response: httpx.Response) -> None:
|
||||
"""Handle registration response."""
|
||||
if response.status_code not in (200, 201):
|
||||
await response.aread()
|
||||
raise OAuthRegistrationError(f"Registration failed: {response.status_code} {response.text}")
|
||||
|
||||
try:
|
||||
content = await response.aread()
|
||||
client_info = OAuthClientInformationFull.model_validate_json(content)
|
||||
self.context.client_info = client_info
|
||||
await self.context.storage.set_client_info(client_info)
|
||||
except ValidationError as e:
|
||||
raise OAuthRegistrationError(f"Invalid registration response: {e}")
|
||||
|
||||
async def _perform_authorization(self) -> httpx.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")
|
||||
if not self.context.redirect_handler:
|
||||
raise OAuthFlowError("No redirect handler provided for authorization code grant")
|
||||
if not self.context.callback_handler:
|
||||
raise OAuthFlowError("No callback handler provided for authorization code grant")
|
||||
|
||||
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")
|
||||
|
||||
# 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:
|
||||
auth_params["scope"] = self.context.client_metadata.scope
|
||||
|
||||
authorization_url = f"{auth_endpoint}?{urlencode(auth_params)}"
|
||||
await self.context.redirect_handler(authorization_url)
|
||||
|
||||
# Wait for callback
|
||||
auth_code, returned_state = await self.context.callback_handler()
|
||||
|
||||
if returned_state is None or not secrets.compare_digest(returned_state, state):
|
||||
raise OAuthFlowError(f"State parameter mismatch: {returned_state} != {state}")
|
||||
|
||||
if not auth_code:
|
||||
raise OAuthFlowError("No authorization code received")
|
||||
|
||||
# Return auth code and code verifier for token exchange
|
||||
return auth_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, *, token_data: dict[str, Any] | None = {}
|
||||
) -> httpx.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")
|
||||
if not self.context.client_info:
|
||||
raise OAuthFlowError("Missing client info")
|
||||
|
||||
token_url = self._get_token_endpoint()
|
||||
token_data = token_data or {}
|
||||
token_data.update(
|
||||
{
|
||||
"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
|
||||
|
||||
if self.context.client_info.client_secret:
|
||||
token_data["client_secret"] = self.context.client_info.client_secret
|
||||
|
||||
return httpx.Request(
|
||||
"POST", token_url, data=token_data, headers={"Content-Type": "application/x-www-form-urlencoded"}
|
||||
)
|
||||
|
||||
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}")
|
||||
|
||||
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}")
|
||||
|
||||
async def _refresh_token(self) -> httpx.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")
|
||||
|
||||
if not self.context.client_info:
|
||||
raise OAuthTokenError("No client info available")
|
||||
|
||||
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 = {
|
||||
"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
|
||||
|
||||
if self.context.client_info.client_secret:
|
||||
refresh_data["client_secret"] = self.context.client_info.client_secret
|
||||
|
||||
return httpx.Request(
|
||||
"POST", token_url, data=refresh_data, headers={"Content-Type": "application/x-www-form-urlencoded"}
|
||||
)
|
||||
|
||||
async def _handle_refresh_response(self, response: httpx.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)
|
||||
|
||||
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:
|
||||
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: httpx.Request) -> None:
|
||||
"""Add authorization header to request if we have valid tokens."""
|
||||
if self.context.current_tokens and self.context.current_tokens.access_token:
|
||||
request.headers["Authorization"] = f"Bearer {self.context.current_tokens.access_token}"
|
||||
|
||||
def _create_oauth_metadata_request(self, url: str) -> httpx.Request:
|
||||
return httpx.Request("GET", url, headers={MCP_PROTOCOL_VERSION: LATEST_PROTOCOL_VERSION})
|
||||
|
||||
async def _handle_oauth_metadata_response(self, response: httpx.Response) -> None:
|
||||
content = await response.aread()
|
||||
metadata = OAuthMetadata.model_validate_json(content)
|
||||
self.context.oauth_metadata = metadata
|
||||
|
||||
async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
|
||||
"""HTTPX 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)
|
||||
|
||||
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
|
||||
# Step 1: Discover protected resource metadata (RFC9728 with WWW-Authenticate support)
|
||||
discovery_request = await self._discover_protected_resource(response)
|
||||
discovery_response = yield discovery_request
|
||||
await self._handle_protected_resource_response(discovery_response)
|
||||
|
||||
# Step 2: Apply scope selection strategy
|
||||
self._select_scopes(response)
|
||||
|
||||
# Step 3: Discover OAuth metadata (with fallback for legacy servers)
|
||||
discovery_urls = self._get_discovery_urls()
|
||||
for url in discovery_urls:
|
||||
oauth_metadata_request = self._create_oauth_metadata_request(url)
|
||||
oauth_metadata_response = yield oauth_metadata_request
|
||||
|
||||
if oauth_metadata_response.status_code == 200:
|
||||
try:
|
||||
await self._handle_oauth_metadata_response(oauth_metadata_response)
|
||||
break
|
||||
except ValidationError:
|
||||
continue
|
||||
elif oauth_metadata_response.status_code < 400 or oauth_metadata_response.status_code >= 500:
|
||||
break # Non-4XX error, stop trying
|
||||
|
||||
# Step 4: Register client if needed
|
||||
registration_request = await self._register_client()
|
||||
if registration_request:
|
||||
registration_response = yield registration_request
|
||||
await self._handle_registration_response(registration_response)
|
||||
|
||||
# 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 = self._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)
|
||||
|
||||
# Step 2b: Perform (re-)authorization and 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
|
||||
|
|
@ -41,13 +41,13 @@ class OAuthClientMetadata(BaseModel):
|
|||
for the full specification.
|
||||
"""
|
||||
|
||||
redirect_uris: list[AnyUrl] = Field(..., min_length=1)
|
||||
# token_endpoint_auth_method: this implementation only supports none &
|
||||
# client_secret_post;
|
||||
# ie: we do not support client_secret_basic
|
||||
token_endpoint_auth_method: Literal["none", "client_secret_post"] = "client_secret_post"
|
||||
# grant_types: this implementation only supports authorization_code & refresh_token
|
||||
grant_types: list[Literal["authorization_code", "refresh_token"] | str] = [
|
||||
redirect_uris: list[AnyUrl] | None = Field(..., min_length=1)
|
||||
# supported auth methods for the token endpoint
|
||||
token_endpoint_auth_method: Literal["none", "client_secret_post", "private_key_jwt"] = "client_secret_post"
|
||||
# supported grant_types of this implementation
|
||||
grant_types: list[
|
||||
Literal["authorization_code", "refresh_token", "urn:ietf:params:oauth:grant-type:jwt-bearer"] | str
|
||||
] = [
|
||||
"authorization_code",
|
||||
"refresh_token",
|
||||
]
|
||||
|
|
@ -82,10 +82,10 @@ class OAuthClientMetadata(BaseModel):
|
|||
def validate_redirect_uri(self, redirect_uri: AnyUrl | None) -> AnyUrl:
|
||||
if redirect_uri is not None:
|
||||
# Validate redirect_uri against client's registered redirect URIs
|
||||
if redirect_uri not in self.redirect_uris:
|
||||
if self.redirect_uris is None or redirect_uri not in self.redirect_uris:
|
||||
raise InvalidRedirectUriError(f"Redirect URI '{redirect_uri}' not registered for client")
|
||||
return redirect_uri
|
||||
elif len(self.redirect_uris) == 1:
|
||||
elif self.redirect_uris is not None and len(self.redirect_uris) == 1:
|
||||
return self.redirect_uris[0]
|
||||
else:
|
||||
raise InvalidRedirectUriError("redirect_uri must be specified when client has multiple registered URIs")
|
||||
|
|
@ -97,7 +97,7 @@ class OAuthClientInformationFull(OAuthClientMetadata):
|
|||
(client information plus metadata).
|
||||
"""
|
||||
|
||||
client_id: str
|
||||
client_id: str | None = None
|
||||
client_secret: str | None = None
|
||||
client_id_issued_at: int | None = None
|
||||
client_secret_expires_at: int | None = None
|
||||
|
|
|
|||
Loading…
Reference in a new issue