address feedback

Author: patrickkang

examples/example-fastmcp-mcp/src/auth0/__init__.py

@@ -5,13 +5,22 @@
 including token verification, middleware, and scoped tool decorators.
 """
 
+import logging
+import os
+from typing import Callable, Union
+
 from mcp.server.auth.routes import create_protected_resource_routes
 from mcp.server.fastmcp import FastMCP
 from starlette.middleware import Middleware
+from starlette.requests import Request
+from starlette.responses import JSONResponse
 from starlette.routing import Route, Router
 
+from .errors import AuthenticationRequired, InsufficientScope, MalformedAuthorizationRequest
 from .middleware import Auth0Middleware
 
+logger = logging.getLogger(__name__)
+
 
 class Auth0Mcp:
     def __init__(self, name: str, audience: str, domain: str):
@@ -56,3 +65,55 @@ def register_scopes(self, scopes: list[str]) -> None:
         """
         if scopes:
             self._scopes_supported.update(scopes)
+
+    def exception_handlers(self) -> dict[Union[int, type[Exception]], Callable]:
+        return {
+            AuthenticationRequired: self._auth_error_handler,
+            InsufficientScope: self._auth_error_handler,
+            MalformedAuthorizationRequest: self._auth_error_handler,
+            # Generic fallback for any other exceptions
+            Exception: self._generic_exception_handler,
+        }
+
+    def _auth_error_handler(self, request: Request, exc: Exception):
+        """
+        Handle auth errors: malformed authorization requests, missing auth, invalid tokens, and insufficient scopes.
+        """
+        # Include resource metadata parameter for 401 responses per RFC 9728 Section 5.1
+        include_resource_metadata = exc.status_code == 401
+
+        return JSONResponse(
+            {
+                "error": exc.error_code,
+                "error_description": exc.description
+            },
+            status_code=exc.status_code,
+            headers={"WWW-Authenticate": self._build_www_authenticate_header(exc.error_code, exc.description, include_resource_metadata)},
+        )
+
+    def _generic_exception_handler(self, request:Request, exc: Exception):
+        """
+        Fallback handler for all other exceptions.
+        """
+        logger.error(f"Unexpected error in: {exc}", exc_info=exc)
+
+        # Return standard HTTP 500 error
+        return JSONResponse(
+            {
+                "error": "internal_server_error",
+                "error_description": "An unexpected error occurred"
+            },
+            status_code=500,
+        )
+
+    def _build_www_authenticate_header(self, error_code: str, description: str, include_resource_metadata: bool = False) -> str:
+        """
+        Build WWW-Authenticate header according to RFC 9728 Section 5.1.
+        """
+        www_auth_params = [f'error="{error_code}"', f'error_description="{description}"']
+
+        if include_resource_metadata:
+            metadata_url = f"{os.getenv('MCP_SERVER_URL')}/.well-known/oauth-protected-resource"
+            www_auth_params.append(f'resource_metadata="{metadata_url}"')
+
+        return f"Bearer {', '.join(www_auth_params)}"

examples/example-fastmcp-mcp/src/auth0/authz.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Iterable
+from functools import wraps
+
+from mcp.server.fastmcp import Context
+
+from .errors import AuthenticationRequired, InsufficientScope
+
+
+def require_scopes(required_scopes: Iterable[str]):
+    """
+    Decorator that requires scopes on MCP tools.
+
+    Example:
+      @mcp.tool(...)
+      @require_scopes(["tool:greet", "tool:whoami"])
+      def my_tool(name: str, ctx: Context) -> str:
+        return f"Hello {name}!"
+    """
+    required_scopes_list = list(required_scopes)
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            # ctx is passed in either kw or positional
+            ctx: Context | None = (kwargs.get("ctx") if isinstance(kwargs.get("ctx"), Context) else None) or next((arg for arg in args if isinstance(arg, Context)), None)
+            if ctx is None:
+                raise TypeError("ctx: Context is required")
+
+            auth = getattr(ctx.request_context.request.state, "auth", {})
+            if not auth:
+                raise AuthenticationRequired("Authentication required")
+
+            user_scopes = set(auth.get("scopes", []))
+            missing_scopes = [s for s in required_scopes_list if s not in user_scopes]
+            if missing_scopes:
+                raise InsufficientScope(f"Missing required scopes: {missing_scopes}")
+
+            # Call the original function
+            if asyncio.iscoroutinefunction(func):
+                return await func(*args, **kwargs)
+            else:
+                return func(*args, **kwargs)
+        return wrapper
+    return decorator

examples/example-fastmcp-mcp/src/auth0/errors.py

@@ -0,0 +1,47 @@
+class AuthenticationRequired(Exception):
+    """
+    Raised when authentication is required but missing.
+
+    This maps to HTTP 401 Unauthorized status.
+    Indicates the request lacks valid authentication credentials.
+    """
+    status_code = 401
+    error_code = "invalid_token"
+    default_description = "Authentication required"
+
+    def __init__(self, message: str | None = None):
+        self.description = message or self.default_description
+        super().__init__(self.description)
+
+
+class InsufficientScope(Exception):
+    """
+    Raised when user lacks required OAuth scopes.
+
+    This maps to HTTP 403 Forbidden status.
+    Indicates the user is authenticated but doesn't have permission
+    to access the requested resource due to insufficient scopes.
+    """
+    status_code = 403
+    error_code = "insufficient_scope"
+    default_description = "Insufficient scope"
+
+    def __init__(self, message: str | None = None):
+        self.description = message or self.default_description
+        super().__init__(self.description)
+
+
+class MalformedAuthorizationRequest(Exception):
+    """
+    Raised when authorization request is malformed.
+
+    This maps to HTTP 400 Bad Request status.
+    Indicates the authorization header or token format is invalid.
+    """
+    status_code = 400
+    error_code = "invalid_request"
+    default_description = "Malformed authorization request"
+
+    def __init__(self, message: str | None = None):
+        self.description = message or self.default_description
+        super().__init__(self.description)

examples/example-fastmcp-mcp/src/auth0/middleware.py

@@ -1,13 +1,13 @@
 import logging
-import os
 
 from auth0_api_python import ApiClient, ApiClientOptions
 from auth0_api_python.errors import VerifyAccessTokenError
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.requests import Request
-from starlette.responses import JSONResponse
 from starlette.types import ASGIApp
 
+from .errors import AuthenticationRequired, MalformedAuthorizationRequest
+
 logger = logging.getLogger(__name__)
 
 class Auth0Middleware(BaseHTTPMiddleware):
@@ -29,13 +29,9 @@ async def dispatch(self, request: Request, call_next):
         # Extract Authorization header
         auth_header = request.headers.get("authorization")
         if not auth_header:
-            return self._return_auth_error_response(status_code=401, error="Authentication required", description="Missing Authorization header")
+            raise AuthenticationRequired("Missing Authorization header")
         if not auth_header.lower().startswith("bearer "):
-            return self._return_auth_error_response(
-                status_code=401,
-                error="Authentication required",
-                description="Invalid Authorization header format"
-            )
+            raise MalformedAuthorizationRequest("Invalid Authorization header format")
 
         # Extract and verify token
         token = auth_header[7:] # Remove "Bearer " prefix
@@ -72,25 +68,8 @@ async def dispatch(self, request: Request, call_next):
             return await call_next(request)
         except VerifyAccessTokenError as e:
             logger.error(f"Token verification failed: {str(e)}")
-            return self._return_auth_error_response(
-                status_code=401,
-                error="Authentication failed",
-                description="Invalid token"
-            )
+            raise AuthenticationRequired("Invalid token")
         except Exception as e:
             logger.error(f"Unexpected error in middleware: {str(e)}")
-            return self._return_auth_error_response(
-                status_code=500,
-                error="Internal Server Error",
-                description="Internal Server Error"
-            )
-
-    def _return_auth_error_response(self, status_code: int, error: str, description: str) -> JSONResponse:
-        www_auth_parts = [f'error="{error}"', f'error_description="{description}"', f'resource_metadata="{os.getenv("MCP_SERVER_URL")}"']
-        www_authenticate = f"Bearer {', '.join(www_auth_parts)}"
-
-        return JSONResponse(
-            status_code=status_code,
-            content={"error": error, "error_description": description},
-            headers={"WWW-Authenticate": www_authenticate}
-        )
+            # Re-raise unexpected errors to be handled by generic exception handler
+            raise

examples/example-fastmcp-mcp/src/auth0/tools.py

@@ -44,6 +44,7 @@ async def lifespan(app: Starlette) -> AsyncIterator[None]:
         ),
     ],
     lifespan=lifespan,
+    exception_handlers=auth0_mcp.exception_handlers(),
 )
 
 # Wrap ASGI application with CORS middleware to expose Mcp-Session-Id header

examples/example-fastmcp-mcp/src/server.py

@@ -2,16 +2,16 @@
 
 from mcp.server.fastmcp import Context
 
-from .auth0.tools import create_scoped_tool_decorator, get_auth_info
+from .auth0.authz import require_scopes
 
 
 def register_tools(auth0Mcp):
     """
     Register all tools with the MCP server.
     """
     mcp = auth0Mcp.mcp
-    # Create a scoped_tool decorator bound to the mcp instance
-    scoped_tool = create_scoped_tool_decorator(auth0Mcp)
+    # Register scopes used by tools for Protected Resource Metadata
+    auth0Mcp.register_scopes(["tool:greet", "tool:whoami"])
 
     # Tool without required scopes
     @mcp.tool()
@@ -20,31 +20,29 @@ def echo(text: str) -> str:
         return text
 
     # A MCP tool with required scopes
-    @scoped_tool(
-        required_scopes=["tool:greet"],
+    @mcp.tool(
         name="greet",
         title="Greet Tool",
         description="Greets a user",
         annotations={"readOnlyHint": True}
     )
+    @require_scopes(["tool:greet"])
     def greet(name: str, ctx: Context) -> str:
         name = (name or "").strip() or "world"
-        request = ctx.request_context.request
-        auth_info = get_auth_info(request)
+        auth_info = ctx.request_context.request.state.auth
         user_id = auth_info.get("extra", {}).get("sub")
         return f"Hello, {name}! You are authenticated as {user_id}"
 
     # A MCP tool with required scopes
-    @scoped_tool(
-        required_scopes=["tool:whoami"],
+    @mcp.tool(
         name="whoami",
         title="Who Am I Tool",
         description="Returns information about the authenticated user",
         annotations={"readOnlyHint": True}
     )
+    @require_scopes(["tool:whoami"])
     def whoami(ctx: Context) -> str:
-        request = ctx.request_context.request
-        auth_info = get_auth_info(request)
+        auth_info = ctx.request_context.request.state.auth
 
         response_data = {
             "user": auth_info.get("extra", {}),