feat: create bearer auth middleware (#1)

Author: gao-sun

mcpauth/exceptions.py

@@ -34,7 +34,15 @@ def to_json(self, show_cause: bool = False) -> Record:
         data: Record = {
             "error": self.code.value if isinstance(self.code, Enum) else self.code,
             "error_description": self.message,
-            "cause": self.cause if show_cause and hasattr(self, "cause") else None,
+            "cause": (
+                (
+                    {k: v for k, v in self.cause.model_dump().items() if v is not None}
+                    if isinstance(self.cause, BaseModel)
+                    else str(self.cause)
+                )
+                if show_cause and hasattr(self, "cause")
+                else None
+            ),
         }
         return {k: v for k, v in data.items() if v is not None}
 
@@ -99,8 +107,8 @@ class MCPAuthBearerAuthExceptionDetails(BaseModel):
     cause: Any = None
     uri: Optional[str] = None
     missing_scopes: Optional[List[str]] = None
-    expected: Optional[Union[str, Record]] = None
-    actual: Optional[Union[str, Record]] = None
+    expected: Any = None
+    actual: Any = None
 
 
 class MCPAuthBearerAuthException(MCPAuthException):
@@ -124,15 +132,15 @@ def __init__(
 
     def to_json(self, show_cause: bool = False) -> Dict[str, Optional[str]]:
         # Matches the OAuth 2.0 exception response format at best effort
-        result = super().to_json(show_cause)
+        data = super().to_json(show_cause)
         if self.cause:
-            result.update(
+            data.update(
                 {
                     "error_uri": self.cause.uri,
                     "missing_scopes": self.cause.missing_scopes,
                 }
             )
-        return result
+        return {k: v for k, v in data.items() if v is not None}
 
 
 class MCPAuthJwtVerificationExceptionCode(str, Enum):

mcpauth/exceptioins_test.py renamed to mcpauth/exceptions_test.py

@@ -28,7 +28,7 @@ def test_to_json(self):
         assert mcp_exception.to_json(show_cause=True) == {
             "error": "test_code",
             "error_description": "Test message",
-            "cause": exception,
+            "cause": str(exception),
         }
 
     def test_properties(self):

mcpauth/middleware/create_bearer_auth.py

@@ -0,0 +1,213 @@
+from typing import Any, Dict, List, Optional
+from urllib.parse import urlparse
+import logging
+from pydantic import BaseModel
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.requests import Request
+from starlette.responses import Response, JSONResponse
+from starlette.datastructures import Headers
+
+from ..exceptions import (
+    MCPAuthBearerAuthException,
+    MCPAuthJwtVerificationException,
+    MCPAuthAuthServerException,
+    MCPAuthConfigException,
+    BearerAuthExceptionCode,
+    MCPAuthBearerAuthExceptionDetails,
+)
+from ..types import VerifyAccessTokenFunction, Record
+
+
+class BearerAuthConfig(BaseModel):
+    """
+    Configuration for the Bearer auth handler.
+
+    Attributes:
+      issuer: The expected issuer of the access token.
+      audience: The expected audience of the access token.
+      required_scopes: An array of required scopes that the access token must have.
+      show_error_details: Whether to show detailed error information in the response.
+    """
+
+    issuer: str
+    audience: Optional[str] = None
+    required_scopes: Optional[List[str]] = None
+    show_error_details: bool = False
+
+
+def get_bearer_token_from_headers(headers: Headers) -> str:
+    """
+    Extract the Bearer token from the request headers.
+
+    Args:
+      headers: The HTTP request headers.
+
+    Returns:
+      The Bearer token.
+
+    Raises:
+      MCPAuthBearerAuthException: If the Authorization header is missing or invalid.
+    """
+
+    auth_header = headers.get("authorization") or headers.get("Authorization")
+
+    print(f"Authorization header: {auth_header}")
+
+    if not auth_header:
+        raise MCPAuthBearerAuthException(BearerAuthExceptionCode.MISSING_AUTH_HEADER)
+
+    parts = auth_header.split(" ")
+    if len(parts) != 2 or parts[0].lower() != "bearer":
+        raise MCPAuthBearerAuthException(
+            BearerAuthExceptionCode.INVALID_AUTH_HEADER_FORMAT
+        )
+
+    token = parts[1]
+    if not token:
+        raise MCPAuthBearerAuthException(BearerAuthExceptionCode.MISSING_BEARER_TOKEN)
+
+    return token
+
+
+def _handle_error(
+    error: Exception, show_error_details: bool = False
+) -> tuple[int, Dict[str, Any]]:
+    """
+    Handle errors from the Bearer auth process.
+
+    Args:
+      error: The exception that was caught.
+      show_error_details: Whether to include detailed error information in the response.
+
+    Returns:
+      A tuple of (status_code, response_body).
+    """
+    if isinstance(error, MCPAuthJwtVerificationException):
+        return 401, error.to_json(show_error_details)
+
+    if isinstance(error, MCPAuthBearerAuthException):
+        if error.code == BearerAuthExceptionCode.MISSING_REQUIRED_SCOPES:
+            return 403, error.to_json(show_error_details)
+        return 401, error.to_json(show_error_details)
+
+    if isinstance(error, (MCPAuthAuthServerException, MCPAuthConfigException)):
+        response: Record = {
+            "error": "server_error",
+            "error_description": "An error occurred with the authorization server.",
+        }
+        if show_error_details:
+            response["cause"] = error.to_json()
+        return 500, response
+
+    # Re-raise other errors
+    raise error
+
+
+def create_bearer_auth(
+    verify_access_token: VerifyAccessTokenFunction, config: BearerAuthConfig
+) -> type[BaseHTTPMiddleware]:
+    """
+    Creates a middleware function for handling Bearer auth.
+
+    This middleware extracts the Bearer token from the `Authorization` header, verifies it using the
+    provided `verify_access_token` function, and checks the issuer, audience, and required scopes.
+
+    Args:
+      verify_access_token: A function that takes a Bearer token and returns an `AuthInfo` object.
+      config: Configuration for the Bearer auth handler.
+
+    Returns:
+      A middleware class that handles Bearer auth.
+    """
+
+    if not callable(verify_access_token):
+        raise TypeError(
+            "`verify_access_token` must be a function that takes a token and returns an `AuthInfo` object."
+        )
+
+    try:
+        result = urlparse(config.issuer)
+        if not all([result.scheme, result.netloc]):
+            raise ValueError("Invalid URL")
+    except:
+        raise TypeError("`issuer` must be a valid URL.")
+
+    class BearerAuthMiddleware(BaseHTTPMiddleware):
+        """
+        Middleware class that handles Bearer auth.
+
+        This class is used to wrap the request handling process and apply Bearer auth checks.
+        """
+
+        async def dispatch(
+            self, request: Request, call_next: RequestResponseEndpoint
+        ) -> Response:
+            """
+            Dispatch method that processes the request and applies Bearer auth checks.
+
+            Args:
+              request: The HTTP request.
+              call_next: The next middleware or route handler to call.
+
+            Returns:
+              The HTTP response after processing the request.
+            """
+            try:
+                token = get_bearer_token_from_headers(request.headers)
+                auth_info = verify_access_token(token)
+
+                if auth_info.issuer != config.issuer:
+                    details = MCPAuthBearerAuthExceptionDetails(
+                        expected=config.issuer, actual=auth_info.issuer
+                    )
+                    raise MCPAuthBearerAuthException(
+                        BearerAuthExceptionCode.INVALID_ISSUER, cause=details
+                    )
+
+                if config.audience:
+                    audience_matches = (
+                        config.audience == auth_info.audience
+                        if isinstance(auth_info.audience, str)
+                        else (
+                            isinstance(auth_info.audience, list)
+                            and config.audience in auth_info.audience
+                        )
+                    )
+                    if not audience_matches:
+                        details = MCPAuthBearerAuthExceptionDetails(
+                            expected=config.audience, actual=auth_info.audience
+                        )
+                        raise MCPAuthBearerAuthException(
+                            BearerAuthExceptionCode.INVALID_AUDIENCE, cause=details
+                        )
+
+                if config.required_scopes:
+                    missing_scopes = [
+                        scope
+                        for scope in config.required_scopes
+                        if scope not in auth_info.scopes
+                    ]
+                    if missing_scopes:
+                        details = MCPAuthBearerAuthExceptionDetails(
+                            missing_scopes=missing_scopes
+                        )
+                        raise MCPAuthBearerAuthException(
+                            BearerAuthExceptionCode.MISSING_REQUIRED_SCOPES,
+                            cause=details,
+                        )
+
+                # Attach auth info to the request
+                request.state.auth = auth_info
+
+                # Call the next middleware or route handler
+                response = await call_next(request)
+                return response
+
+            except Exception as error:
+                logging.error(f"Error during Bearer auth: {error}")
+                status_code, response_body = _handle_error(
+                    error, config.show_error_details
+                )
+                return JSONResponse(status_code=status_code, content=response_body)
+
+    return BearerAuthMiddleware