Use classes for handlers

Author: praboud-ant

src/mcp/server/auth/handlers/authorize.py

@@ -5,7 +5,8 @@
 """
 
 import logging
-from typing import Callable, Literal
+from dataclasses import dataclass
+from typing import Literal
 from urllib.parse import urlencode, urlparse, urlunparse
 
 from pydantic import AnyHttpUrl, AnyUrl, BaseModel, Field, RootModel, ValidationError
@@ -117,8 +118,11 @@ class AnyHttpUrlModel(RootModel):
     root: AnyHttpUrl
 
 
-def create_authorization_handler(provider: OAuthServerProvider) -> Callable:
-    async def authorization_handler(request: Request) -> Response:
+@dataclass
+class AuthorizationHandler:
+    provider: OAuthServerProvider
+
+    async def handle(self, request: Request) -> Response:
         # implements authorization requests for grant_type=code;
         # see https://datatracker.ietf.org/doc/html/rfc6749#section-4.1.1
 
@@ -134,7 +138,7 @@ async def error_response(
             if client is None and attempt_load_client:
                 # make last-ditch attempt to load the client
                 client_id = best_effort_extract_string("client_id", params)
-                client = client_id and await provider.clients_store.get_client(
+                client = client_id and await self.provider.clients_store.get_client(
                     client_id
                 )
             if redirect_uri is None and client:
@@ -200,7 +204,9 @@ async def error_response(
                 )
 
             # Get client information
-            client = await provider.clients_store.get_client(auth_request.client_id)
+            client = await self.provider.clients_store.get_client(
+                auth_request.client_id,
+            )
             if not client:
                 # For client_id validation errors, return direct error (no redirect)
                 return await error_response(
@@ -241,7 +247,10 @@ async def error_response(
             response = RedirectResponse(
                 url="", status_code=302, headers={"Cache-Control": "no-store"}
             )
-            response.headers["location"] = await provider.authorize(client, auth_params)
+            response.headers["location"] = await self.provider.authorize(
+                client,
+                auth_params,
+            )
             return response
 
         except Exception as validation_error:
@@ -253,7 +262,6 @@ async def error_response(
                 error="server_error", error_description="An unexpected error occurred"
             )
 
-    return authorization_handler
 
 
 def create_error_redirect(

src/mcp/server/auth/handlers/metadata.py

@@ -4,41 +4,22 @@
 Corresponds to TypeScript file: src/server/auth/handlers/metadata.ts
 """
 
-from typing import Any, Callable
+from dataclasses import dataclass
+from typing import Any
 
 from starlette.requests import Request
 from starlette.responses import JSONResponse, Response
 
 
-def create_metadata_handler(metadata: dict[str, Any]) -> Callable:
-    """
-    Create a handler for OAuth 2.0 Authorization Server Metadata.
+@dataclass
+class MetadataHandler:
+    metadata: dict[str, Any]
 
-    Corresponds to metadataHandler in src/server/auth/handlers/metadata.ts
-
-    Args:
-        metadata: The metadata to return in the response
-
-    Returns:
-        A Starlette endpoint handler function
-    """
-
-    async def metadata_handler(request: Request) -> Response:
-        """
-        Handler for the OAuth 2.0 Authorization Server Metadata endpoint.
-
-        Args:
-            request: The Starlette request
-
-        Returns:
-            JSON response with the authorization server metadata
-        """
+    async def handle(self, request: Request) -> Response:
         # Remove any None values from metadata
-        clean_metadata = {k: v for k, v in metadata.items() if v is not None}
+        clean_metadata = {k: v for k, v in self.metadata.items() if v is not None}
 
         return JSONResponse(
             content=clean_metadata,
             headers={"Cache-Control": "public, max-age=3600"},  # Cache for 1 hour
         )
-
-    return metadata_handler

src/mcp/server/auth/handlers/register.py

@@ -6,7 +6,8 @@
 
 import secrets
 import time
-from typing import Callable, Literal
+from dataclasses import dataclass
+from typing import Literal
 from uuid import uuid4
 
 from pydantic import BaseModel, ValidationError
@@ -29,10 +30,11 @@ class ErrorResponse(BaseModel):
     error_description: str
 
 
-def create_registration_handler(
-    clients_store: OAuthRegisteredClientsStore, client_secret_expiry_seconds: int | None
-) -> Callable:
-    async def registration_handler(request: Request) -> Response:
+@dataclass
+class RegistrationHandler:
+    clients_store: OAuthRegisteredClientsStore
+    client_secret_expiry_seconds: int | None
+    async def handle(self, request: Request) -> Response:
         # Implements dynamic client registration as defined in https://datatracker.ietf.org/doc/html/rfc7591#section-3.1
         try:
             # Parse request body as JSON
@@ -55,8 +57,8 @@ async def registration_handler(request: Request) -> Response:
 
         client_id_issued_at = int(time.time())
         client_secret_expires_at = (
-            client_id_issued_at + client_secret_expiry_seconds
-            if client_secret_expiry_seconds is not None
+            client_id_issued_at + self.client_secret_expiry_seconds
+            if self.client_secret_expiry_seconds is not None
             else None
         )
 
@@ -83,9 +85,7 @@ async def registration_handler(request: Request) -> Response:
             software_version=client_metadata.software_version,
         )
         # Register client
-        await clients_store.register_client(client_info)
+        await self.clients_store.register_client(client_info)
 
         # Return client information
         return PydanticJSONResponse(content=client_info, status_code=201)
-
-    return registration_handler

src/mcp/server/auth/handlers/revoke.py

@@ -4,7 +4,7 @@
 Corresponds to TypeScript file: src/server/auth/handlers/revoke.ts
 """
 
-from typing import Callable
+from dataclasses import dataclass
 
 from pydantic import ValidationError
 from starlette.requests import Request
@@ -27,10 +27,12 @@ class RevocationRequest(OAuthTokenRevocationRequest, ClientAuthRequest):
     pass
 
 
-def create_revocation_handler(
-    provider: OAuthServerProvider, client_authenticator: ClientAuthenticator
-) -> Callable:
-    async def revocation_handler(request: Request) -> Response:
+@dataclass
+class RevocationHandler:
+    provider: OAuthServerProvider
+    client_authenticator: ClientAuthenticator
+
+    async def handle(self, request: Request) -> Response:
         """
         Handler for the OAuth 2.0 Token Revocation endpoint.
         """
@@ -48,13 +50,13 @@ async def revocation_handler(request: Request) -> Response:
 
         # Authenticate client
         try:
-            client_auth_result = await client_authenticator(revocation_request)
+            client_auth_result = await self.client_authenticator(revocation_request)
         except InvalidClientError as e:
             return PydanticJSONResponse(status_code=401, content=e.error_response())
 
         # Revoke token
-        if provider.revoke_token:
-            await provider.revoke_token(client_auth_result, revocation_request)
+        if self.provider.revoke_token:
+            await self.provider.revoke_token(client_auth_result, revocation_request)
 
         # Return successful empty response
         return Response(
@@ -64,5 +66,3 @@ async def revocation_handler(request: Request) -> Response:
                 "Pragma": "no-cache",
             },
         )
-
-    return revocation_handler

src/mcp/server/auth/handlers/token.py

@@ -7,7 +7,8 @@
 import base64
 import hashlib
 import time
-from typing import Annotated, Callable, Literal
+from dataclasses import dataclass
+from typing import Annotated, Literal
 
 from pydantic import AnyHttpUrl, Field, RootModel, ValidationError
 from starlette.requests import Request
@@ -52,10 +53,12 @@ class TokenRequest(RootModel):
     ]
 
 
-def create_token_handler(
-    provider: OAuthServerProvider, client_authenticator: ClientAuthenticator
-) -> Callable:
-    def response(obj: TokenSuccessResponse | TokenErrorResponse | ErrorResponse):
+@dataclass
+class TokenHandler:
+    provider: OAuthServerProvider
+    client_authenticator: ClientAuthenticator
+
+    def response(self, obj: TokenSuccessResponse | TokenErrorResponse | ErrorResponse):
         status_code = 200
         if isinstance(obj, TokenErrorResponse):
             status_code = 400
@@ -69,25 +72,25 @@ def response(obj: TokenSuccessResponse | TokenErrorResponse | ErrorResponse):
             },
         )
 
-    async def token_handler(request: Request):
+    async def handle(self, request: Request):
         try:
             form_data = await request.form()
             token_request = TokenRequest.model_validate(dict(form_data)).root
         except ValidationError as validation_error:
-            return response(
+            return self.response(
                 TokenErrorResponse(
                     error="invalid_request",
                     error_description=stringify_pydantic_error(validation_error),
                 )
             )
 
         try:
-            client_info = await client_authenticator(token_request)
+            client_info = await self.client_authenticator(token_request)
         except InvalidClientError as e:
-            return response(e.error_response())
+            return self.response(e.error_response())
 
         if token_request.grant_type not in client_info.grant_types:
-            return response(
+            return self.response(
                 TokenErrorResponse(
                     error="unsupported_grant_type",
                     error_description=(
@@ -101,12 +104,12 @@ async def token_handler(request: Request):
 
         match token_request:
             case AuthorizationCodeRequest():
-                auth_code = await provider.load_authorization_code(
+                auth_code = await self.provider.load_authorization_code(
                     client_info, token_request.code
                 )
                 if auth_code is None or auth_code.client_id != token_request.client_id:
                     # if code belongs to different client, pretend it doesn't exist
-                    return response(
+                    return self.response(
                         TokenErrorResponse(
                             error="invalid_grant",
                             error_description="authorization code does not exist",
@@ -116,7 +119,7 @@ async def token_handler(request: Request):
                 # make auth codes expire after a deadline
                 # see https://datatracker.ietf.org/doc/html/rfc6749#section-10.5
                 if auth_code.expires_at < time.time():
-                    return response(
+                    return self.response(
                         TokenErrorResponse(
                             error="invalid_grant",
                             error_description="authorization code has expired",
@@ -126,7 +129,7 @@ async def token_handler(request: Request):
                 # verify redirect_uri doesn't change between /authorize and /tokens
                 # see https://datatracker.ietf.org/doc/html/rfc6749#section-10.6
                 if token_request.redirect_uri != auth_code.redirect_uri:
-                    return response(
+                    return self.response(
                         TokenErrorResponse(
                             error="invalid_request",
                             error_description=(
@@ -144,28 +147,28 @@ async def token_handler(request: Request):
 
                 if hashed_code_verifier != auth_code.code_challenge:
                     # see https://datatracker.ietf.org/doc/html/rfc7636#section-4.6
-                    return response(
+                    return self.response(
                         TokenErrorResponse(
                             error="invalid_grant",
                             error_description="incorrect code_verifier",
                         )
                     )
 
                 # Exchange authorization code for tokens
-                tokens = await provider.exchange_authorization_code(
+                tokens = await self.provider.exchange_authorization_code(
                     client_info, auth_code
                 )
 
             case RefreshTokenRequest():
-                refresh_token = await provider.load_refresh_token(
+                refresh_token = await self.provider.load_refresh_token(
                     client_info, token_request.refresh_token
                 )
                 if (
                     refresh_token is None
                     or refresh_token.client_id != token_request.client_id
                 ):
                     # if token belongs to different client, pretend it doesn't exist
-                    return response(
+                    return self.response(
                         TokenErrorResponse(
                             error="invalid_grant",
                             error_description="refresh token does not exist",
@@ -174,7 +177,7 @@ async def token_handler(request: Request):
 
                 if refresh_token.expires_at and refresh_token.expires_at < time.time():
                     # if the refresh token has expired, pretend it doesn't exist
-                    return response(
+                    return self.response(
                         TokenErrorResponse(
                             error="invalid_grant",
                             error_description="refresh token has expired",
@@ -190,20 +193,19 @@ async def token_handler(request: Request):
 
                 for scope in scopes:
                     if scope not in refresh_token.scopes:
-                        return response(
+                        return self.response(
                             TokenErrorResponse(
                                 error="invalid_scope",
                                 error_description=(
-                        f"cannot request scope `{scope}` not provided by refresh token"
-                    ),
+                                    f"cannot request scope `{scope}` "
+                                    "not provided by refresh token"
+                                ),
                             )
                         )
 
                 # Exchange refresh token for new tokens
-                tokens = await provider.exchange_refresh_token(
+                tokens = await self.provider.exchange_refresh_token(
                     client_info, refresh_token, scopes
                 )
 
-        return response(tokens)
-
-    return token_handler
+        return self.response(tokens)