feat: delegated middleware

Author: gao-sun

mcpauth/__init__.py

@@ -1,16 +1,57 @@
-from .exceptions import (
-    MCPAuthException as MCPAuthException,
-    MCPAuthConfigException as MCPAuthConfigException,
-    AuthServerExceptionCode as AuthServerExceptionCode,
-    MCPAuthAuthServerException as MCPAuthAuthServerException,
-    BearerAuthExceptionCode as BearerAuthExceptionCode,
-    MCPAuthBearerAuthExceptionDetails as MCPAuthBearerAuthExceptionDetails,
-    MCPAuthBearerAuthException as MCPAuthBearerAuthException,
-    MCPAuthJwtVerificationExceptionCode as MCPAuthJwtVerificationExceptionCode,
-    MCPAuthJwtVerificationException as MCPAuthJwtVerificationException,
-)
+import logging
+
+from .utils.fetch_server_config import ServerMetadataPaths
+from .config import MCPAuthConfig
+from .exceptions import MCPAuthAuthServerException, AuthServerExceptionCode
+from .utils.validate_server_config import validate_server_config
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.requests import Request
+from starlette.responses import Response, JSONResponse
 
 
 class MCPAuth:
-    def __init__(self):
-        self.config = None
+    def __init__(self, config: MCPAuthConfig):
+        result = validate_server_config(config.server)
+
+        if not result.is_valid:
+            raise MCPAuthAuthServerException(
+                AuthServerExceptionCode.INVALID_SERVER_CONFIG, cause=result
+            )
+
+        if len(result.warnings) > 0:
+            logging.warning("The authorization server configuration has warnings:\n")
+            for warning in result.warnings:
+                logging.warning(f"- {warning}")
+
+        self.config = config
+
+    def delegated_middleware(self) -> type[BaseHTTPMiddleware]:
+        """
+        Returns a middleware that handles OAuth 2.0 Authorization Metadata endpoint
+        (`/.well-known/oauth-authorization-server`) with CORS support (delegated mode).
+
+        :return: A middleware class that can be used in a Starlette or FastAPI application.
+        """
+        server_config = self.config.server
+
+        class DelegatedMiddleware(BaseHTTPMiddleware):
+            async def dispatch(
+                self, request: Request, call_next: RequestResponseEndpoint
+            ) -> Response:
+                path = request.url.path
+                if path == ServerMetadataPaths.OAUTH:
+                    response = JSONResponse(
+                        {
+                            k: v
+                            for k, v in server_config.metadata.model_dump().items()
+                            if v is not None
+                        },
+                        status_code=200,
+                    )
+                    response.headers["Access-Control-Allow-Origin"] = "*"
+                    response.headers["Access-Control-Allow-Methods"] = "GET, OPTIONS"
+                    return response
+                else:
+                    return await call_next(request)
+
+        return DelegatedMiddleware

mcpauth/__init__test.py

@@ -0,0 +1,137 @@
+import pytest
+from unittest.mock import AsyncMock, patch, MagicMock
+from starlette.requests import Request
+from starlette.responses import Response
+from mcpauth import MCPAuth, MCPAuthAuthServerException, AuthServerExceptionCode
+from mcpauth.config import MCPAuthConfig
+from mcpauth.models.auth_server import AuthServerConfig, AuthServerType
+from mcpauth.models.oauth import AuthorizationServerMetadata
+
+
+class TestMCPAuth:
+    def test_init_with_valid_config(self):
+        # Setup
+        server_config = AuthServerConfig(
+            type=AuthServerType.OAUTH,
+            metadata=AuthorizationServerMetadata(
+                issuer="https://example.com",
+                authorization_endpoint="https://example.com/oauth/authorize",
+                token_endpoint="https://example.com/oauth/token",
+                response_types_supported=["code"],
+                grant_types_supported=["authorization_code"],
+                code_challenge_methods_supported=["S256"],
+            ),
+        )
+        config = MCPAuthConfig(server=server_config)
+
+        # Exercise
+        auth = MCPAuth(config)
+
+        # Verify
+        assert auth.config == config
+
+    def test_init_with_invalid_config(self):
+        # Setup
+        server_config = AuthServerConfig(
+            type=AuthServerType.OAUTH,
+            metadata=AuthorizationServerMetadata(
+                issuer="https://example.com",
+                authorization_endpoint="https://example.com/oauth/authorize",
+                token_endpoint="https://example.com/oauth/token",
+                response_types_supported=["token"],  # Invalid response type
+            ),
+        )
+        config = MCPAuthConfig(server=server_config)
+
+        # Exercise & Verify
+        with pytest.raises(MCPAuthAuthServerException) as exc_info:
+            MCPAuth(config)
+
+        assert exc_info.value.code == AuthServerExceptionCode.INVALID_SERVER_CONFIG
+
+    @patch("mcpauth.logging.warning")
+    def test_init_with_warnings(self, mock_warning: MagicMock):
+        # Setup
+        server_config = AuthServerConfig(
+            type=AuthServerType.OAUTH,
+            metadata=AuthorizationServerMetadata(
+                issuer="https://example.com",
+                authorization_endpoint="https://example.com/oauth/authorize",
+                token_endpoint="https://example.com/oauth/token",
+                response_types_supported=["code"],
+                grant_types_supported=["authorization_code"],
+                code_challenge_methods_supported=["S256"],
+                # Missing registration_endpoint will cause a warning
+            ),
+        )
+        config = MCPAuthConfig(server=server_config)
+
+        # Exercise
+        MCPAuth(config)
+
+        # Verify
+        assert mock_warning.called
+
+    @pytest.mark.asyncio
+    async def test_delegated_middleware_oauth_endpoint(self):
+        # Setup
+        server_config = AuthServerConfig(
+            type=AuthServerType.OAUTH,
+            metadata=AuthorizationServerMetadata(
+                issuer="https://example.com",
+                authorization_endpoint="https://example.com/oauth/authorize",
+                token_endpoint="https://example.com/oauth/token",
+                response_types_supported=["code"],
+                grant_types_supported=["authorization_code"],
+                code_challenge_methods_supported=["S256"],
+            ),
+        )
+        config = MCPAuthConfig(server=server_config)
+        auth = MCPAuth(config)
+
+        middleware_class = auth.delegated_middleware()
+        middleware = middleware_class(app=MagicMock())
+
+        mock_request = MagicMock(spec=Request)
+        mock_request.url.path = "/.well-known/oauth-authorization-server"
+
+        # Exercise
+        response = await middleware.dispatch(mock_request, call_next=AsyncMock())
+
+        # Verify
+        assert response.status_code == 200
+        assert response.headers["Access-Control-Allow-Origin"] == "*"
+        assert response.headers["Access-Control-Allow-Methods"] == "GET, OPTIONS"
+
+    @pytest.mark.asyncio
+    async def test_delegated_middleware_other_endpoint(self):
+        # Setup
+        server_config = AuthServerConfig(
+            type=AuthServerType.OAUTH,
+            metadata=AuthorizationServerMetadata(
+                issuer="https://example.com",
+                authorization_endpoint="https://example.com/oauth/authorize",
+                token_endpoint="https://example.com/oauth/token",
+                response_types_supported=["code"],
+                grant_types_supported=["authorization_code"],
+                code_challenge_methods_supported=["S256"],
+            ),
+        )
+        config = MCPAuthConfig(server=server_config)
+        auth = MCPAuth(config)
+
+        middleware_class = auth.delegated_middleware()
+        middleware = middleware_class(app=MagicMock())
+
+        mock_request = MagicMock(spec=Request)
+        mock_request.url.path = "/some-other-path"
+
+        mock_response = Response(content="Test response")
+        mock_call_next = AsyncMock(return_value=mock_response)
+
+        # Exercise
+        response = await middleware.dispatch(mock_request, call_next=mock_call_next)
+
+        # Verify
+        assert mock_call_next.called
+        assert response == mock_response