Implement RFC 7523 JWT flows

examples/servers/simple-auth/mcp_simple_auth/simple_auth_provider.py

@@ -73,6 +73,8 @@ async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
 
     async def register_client(self, client_info: OAuthClientInformationFull):
         """Register a new OAuth client."""
+        if not client_info.client_id:
+            raise ValueError("No client_id provided")
         self.clients[client_info.client_id] = client_info
 
     async def authorize(self, client: OAuthClientInformationFull, params: AuthorizationParams) -> str:
@@ -209,6 +211,8 @@ async def exchange_authorization_code(
         """Exchange authorization code for tokens."""
         if authorization_code.code not in self.auth_codes:
             raise ValueError("Invalid authorization code")
+        if not client.client_id:
+            raise ValueError("No client_id provided")
 
         # Generate MCP access token
         mcp_token = f"mcp_{secrets.token_hex(32)}"

pyproject.toml

@@ -33,6 +33,7 @@ dependencies = [
     "uvicorn>=0.31.1; sys_platform != 'emscripten'",
     "jsonschema>=4.20.0",
     "pywin32>=310; sys_platform == 'win32'",
+    "pyjwt[crypto]>=2.10.1",
 ]
 
 [project.optional-dependencies]

src/mcp/client/auth/__init__.py

@@ -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

src/mcp/client/auth/extensions/client_credentials.py

@@ -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"}
+        )

src/mcp/client/auth.py → src/mcp/client/auth/oauth2.py

@@ -13,7 +13,7 @@
 import time
 from collections.abc import AsyncGenerator, Awaitable, Callable
 from dataclasses import dataclass, field
-from typing import Protocol
+from typing import Any, Protocol
 from urllib.parse import urlencode, urljoin, urlparse
 
 import anyio
@@ -88,8 +88,8 @@ class OAuthContext:
     server_url: str
     client_metadata: OAuthClientMetadata
     storage: TokenStorage
-    redirect_handler: Callable[[str], Awaitable[None]]
-    callback_handler: Callable[[], Awaitable[tuple[str, str | None]]]
+    redirect_handler: Callable[[str], Awaitable[None]] | None
+    callback_handler: Callable[[], Awaitable[tuple[str, str | None]]] | None
     timeout: float = 300.0
 
     # Discovered metadata
@@ -189,8 +189,8 @@ def __init__(
         server_url: str,
         client_metadata: OAuthClientMetadata,
         storage: TokenStorage,
-        redirect_handler: Callable[[str], Awaitable[None]],
-        callback_handler: Callable[[], Awaitable[tuple[str, str | None]]],
+        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."""
@@ -351,8 +351,21 @@ async def _handle_registration_response(self, response: httpx.Response) -> None:
         except ValidationError as e:
             raise OAuthRegistrationError(f"Invalid registration response: {e}")
 
-    async def _perform_authorization(self) -> tuple[str, str]:
+    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:
@@ -397,24 +410,34 @@ async def _perform_authorization(self) -> tuple[str, str]:
         # Return auth code and code verifier for token exchange
         return auth_code, pkce_params.code_verifier
 
-    async def _exchange_token(self, auth_code: str, code_verifier: str) -> httpx.Request:
-        """Build token exchange request."""
-        if not self.context.client_info:
-            raise OAuthFlowError("Missing client info")
-
+    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_data = {
-            "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,
-        }
+        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):
@@ -430,7 +453,9 @@ async def _exchange_token(self, auth_code: str, code_verifier: str) -> httpx.Req
     async def _handle_token_response(self, response: httpx.Response) -> None:
         """Handle token exchange response."""
         if response.status_code != 200:
-            raise OAuthTokenError(f"Token exchange failed: {response.status_code}")
+            body = await response.aread()
+            body = body.decode("utf-8")
+            raise OAuthTokenError(f"Token exchange failed ({response.status_code}): {body}")
 
         try:
             content = await response.aread()
@@ -577,12 +602,8 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                         registration_response = yield registration_request
                         await self._handle_registration_response(registration_response)
 
-                    # Step 5: Perform authorization
-                    auth_code, code_verifier = await self._perform_authorization()
-
-                    # Step 6: Exchange authorization code for tokens
-                    token_request = await self._exchange_token(auth_code, code_verifier)
-                    token_response = yield token_request
+                    # 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")
@@ -601,17 +622,13 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
                         # Step 2a: Update the required scopes
                         self._select_scopes(response)
 
-                        # Step 2b: Perform (re-)authorization
-                        auth_code, code_verifier = await self._perform_authorization()
-
-                        # Step 2c: Exchange authorization code for tokens
-                        token_request = await self._exchange_token(auth_code, code_verifier)
-                        token_response = yield token_request
+                        # 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
+                # Retry with new tokens
+                self._add_auth_header(request)
+                yield request

src/mcp/shared/auth.py

@@ -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 @@ def validate_scope(self, requested_scope: str | None) -> list[str] | None:
     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