Refactor API spec augmentation to middleware

Author: alukach

src/stac_auth_proxy/app.py

@@ -8,12 +8,16 @@
 import logging
 from typing import Optional
 
-from fastapi import FastAPI, Security
+from fastapi import FastAPI
 
 from .auth import OpenIdConnectAuth
 from .config import Settings
-from .handlers import ReverseProxyHandler, build_openapi_spec_handler
-from .middleware import AddProcessTimeHeaderMiddleware
+from .handlers import ReverseProxyHandler
+from .middleware import (
+    AddProcessTimeHeaderMiddleware,
+    EnforceAuthMiddleware,
+    OpenApiMiddleware,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -25,7 +29,17 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
     app = FastAPI(
         openapi_url=None,  # Disable OpenAPI schema endpoint, we want to serve upstream's schema
     )
+
     app.add_middleware(AddProcessTimeHeaderMiddleware)
+    if settings.openapi_spec_endpoint:
+        app.add_middleware(
+            OpenApiMiddleware,
+            openapi_spec_path=settings.openapi_spec_endpoint,
+            oidc_config_url=str(settings.oidc_discovery_url),
+            private_endpoints=settings.private_endpoints,
+            default_public=settings.default_public,
+        )
+    app.add_middleware(EnforceAuthMiddleware)
 
     if settings.debug:
         app.add_api_route(
@@ -44,37 +58,33 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
         collections_filter=settings.collections_filter,
         items_filter=settings.items_filter,
     )
-    openapi_handler = build_openapi_spec_handler(
-        proxy=proxy_handler,
-        oidc_config_url=str(settings.oidc_discovery_url),
-    )
 
-    # Configure security dependency for explicitely specified endpoints
-    for path_methods, dependencies in [
-        (settings.private_endpoints, [Security(auth_scheme.validated_user)]),
-        (settings.public_endpoints, []),
-    ]:
-        for path, methods in path_methods.items():
-            endpoint = (
-                openapi_handler
-                if path == settings.openapi_spec_endpoint
-                else proxy_handler.stream
-            )
-            app.add_api_route(
-                path,
-                endpoint=endpoint,
-                methods=methods,
-                dependencies=dependencies,
-            )
+    # # Configure security dependency for explicitely specified endpoints
+    # for path_methods, dependencies in [
+    #     (settings.private_endpoints, [Security(auth_scheme.validated_user)]),
+    #     (settings.public_endpoints, []),
+    # ]:
+    #     for path, methods in path_methods.items():
+    #         endpoint = (
+    #             openapi_handler
+    #             if path == settings.openapi_spec_endpoint
+    #             else proxy_handler.stream
+    #         )
+    #         app.add_api_route(
+    #             path,
+    #             endpoint=endpoint,
+    #             methods=methods,
+    #             dependencies=dependencies,
+    #         )
 
     # Catchall for remainder of the endpoints
     app.add_api_route(
         "/{path:path}",
         proxy_handler.stream,
         methods=["GET", "POST", "PUT", "PATCH", "DELETE"],
-        dependencies=(
-            [] if settings.default_public else [Security(auth_scheme.validated_user)]
-        ),
+        # dependencies=(
+        #     [] if settings.default_public else [Security(auth_scheme.validated_user)]
+        # ),
     )
 
     return app

src/stac_auth_proxy/handlers/__init__.py

@@ -1,6 +1,5 @@
 """Handlers to process requests."""
 
-from .open_api_spec import build_openapi_spec_handler
 from .reverse_proxy import ReverseProxyHandler
 
-__all__ = ["build_openapi_spec_handler", "ReverseProxyHandler"]
+__all__ = ["ReverseProxyHandler"]

src/stac_auth_proxy/handlers/open_api_spec.py

@@ -1,9 +1,15 @@
 """Custom middleware."""
 
+from dataclasses import dataclass
+from typing import Any
+import json
 import time
 
 from fastapi import Request, Response
 from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import ASGIApp, Message, Scope, Receive, Send
+
+from .config import EndpointMethods
 
 
 class AddProcessTimeHeaderMiddleware(BaseHTTPMiddleware):
@@ -16,3 +22,66 @@ async def dispatch(self, request: Request, call_next) -> Response:
         process_time = time.perf_counter() - start_time
         response.headers["X-Process-Time"] = f"{process_time:.3f}"
         return response
+
+
+@dataclass(frozen=True)
+class OpenApiMiddleware:
+    """Middleware to add the OpenAPI spec to the response."""
+
+    app: ASGIApp
+    openapi_spec_path: str
+    oidc_config_url: str
+    private_endpoints: EndpointMethods
+    default_public: bool
+    oidc_auth_scheme_name: str = "oidcAuth"
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """Add the OpenAPI spec to the response."""
+        if scope["type"] != "http":
+            return await self.app(scope, receive, send)
+
+        async def augment_oidc_spec(message: Message):
+            if message["type"] != "http.response.body":
+                return await send(message)
+
+            # TODO: Make more robust to handle non-JSON responses
+            body = json.loads(message["body"])
+
+            await send(
+                {
+                    "type": "http.response.body",
+                    "body": dict_to_bytes(self.augment_spec(body)),
+                }
+            )
+
+        return await self.app(scope, receive, augment_oidc_spec)
+
+    def augment_spec(self, openapi_spec) -> dict[str, Any]:
+        components = openapi_spec.setdefault("components", {})
+        securitySchemes = components.setdefault("securitySchemes", {})
+        securitySchemes[self.oidc_auth_scheme_name] = {
+            "type": "openIdConnect",
+            "openIdConnectUrl": self.oidc_config_url,
+        }
+        for path, method_config in openapi_spec["paths"].items():
+            for method, config in method_config.items():
+                for private_method in self.private_endpoints.get(path, []):
+                    if method.lower() == private_method.lower():
+                        config.setdefault("security", []).append(
+                            {self.oidc_auth_scheme_name: []}
+                        )
+        return openapi_spec
+
+
+# TODO
+class EnforceAuthMiddleware(BaseHTTPMiddleware):
+    """Middleware to enforce authentication."""
+
+    async def dispatch(self, request: Request, call_next) -> Response:
+        """Enforce authentication."""
+        return await call_next(request)
+
+
+def dict_to_bytes(d: dict) -> bytes:
+    """Convert a dictionary to a body."""
+    return json.dumps(d, separators=(",", ":")).encode("utf-8")