feat: augment openapi spec

Author: alukach

src/stac_auth_proxy/app.py

@@ -12,6 +12,7 @@
 from .proxy import ReverseProxy
 from .config import Settings
 from .middleware import AddProcessTimeHeaderMiddleware
+from .handlers import OpenApiSpecHandler
 
 
 def create_app(settings: Optional[Settings] = None) -> FastAPI:
@@ -41,12 +42,21 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
     }.items():
         app.add_api_route(
             path,
-            proxy.passthrough,
+            proxy.stream,
             methods=methods,
-            dependencies=[Depends(open_id_connect_scheme)],
+        )
+
+    # Endpoint with special OpenAPI transformation functionality
+    if settings.openapi_spec_endpoint:
+        app.add_api_route(
+            settings.openapi_spec_endpoint,
+            OpenApiSpecHandler(
+                proxy=proxy, oidc_config_url=str(settings.oidc_discovery_url)
+            ).dispatch,
+            methods=["GET"],
         )
 
     # Catchall proxy
-    app.add_route("/{path:path}", proxy.passthrough)
+    app.add_route("/{path:path}", proxy.stream)
 
     return app

src/stac_auth_proxy/config.py

@@ -1,10 +1,12 @@
+from typing import Optional
 from pydantic.networks import HttpUrl
 from pydantic_settings import BaseSettings
 
 
 class Settings(BaseSettings):
     upstream_url: HttpUrl = "https://earth-search.aws.element84.com/v1"
     oidc_discovery_url: HttpUrl
+    openapi_spec_endpoint: Optional[str] = None
 
     class Config:
         env_prefix = "STAC_AUTH_PROXY_"

src/stac_auth_proxy/handlers.py

@@ -0,0 +1,65 @@
+from dataclasses import dataclass
+import logging
+
+from fastapi import Request, Response
+from fastapi.routing import APIRoute
+
+from .proxy import ReverseProxy
+from .utils import safe_headers
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OpenApiSpecHandler:
+    proxy: ReverseProxy
+    oidc_config_url: str
+    auth_scheme_name: str = "oidcAuth"
+
+    async def dispatch(self, req: Request, res: Response):
+        """
+        Proxy the OpenAPI spec from the upstream STAC API, updating it with OIDC security
+        requirements.
+        """
+        oidc_spec_response = await self.proxy.proxy_request(req)
+        openapi_spec = oidc_spec_response.json()
+
+        # Pass along the response headers
+        res.headers.update(safe_headers(oidc_spec_response.headers))
+
+        # Add the OIDC security scheme to the components
+        openapi_spec.setdefault("components", {}).setdefault("securitySchemes", {})[
+            self.auth_scheme_name
+        ] = {
+            "type": "openIdConnect",
+            "openIdConnectUrl": self.oidc_config_url,
+        }
+
+        proxy_auth_routes = [
+            r
+            for r in req.app.routes
+            # Ignore non-APIRoutes (we can't check their security dependencies)
+            if isinstance(r, APIRoute)
+            # Ignore routes that don't have security requirements
+            and (
+                r.dependant.security_requirements
+                or any(d.security_requirements for d in r.dependant.dependencies)
+            )
+        ]
+
+        # Update the paths with the specified security requirements
+        for path, method_config in openapi_spec["paths"].items():
+            for method, config in method_config.items():
+                for route in proxy_auth_routes:
+                    match, _ = route.matches(
+                        {"type": "http", "method": method.upper(), "path": path}
+                    )
+                    if match.name != "FULL":
+                        continue
+                    # Add the OIDC security requirement
+                    config.setdefault("security", []).append(
+                        [{self.auth_scheme_name: []}]
+                    )
+                    break
+
+        return openapi_spec

src/stac_auth_proxy/proxy.py

@@ -1,15 +1,13 @@
 import logging
 import time
 from dataclasses import dataclass
-from urllib.parse import urlparse
 
+import httpx
 from fastapi import Request
-
 from starlette.datastructures import MutableHeaders
 from starlette.responses import StreamingResponse
 from starlette.background import BackgroundTask
 
-import httpx
 
 logger = logging.getLogger(__name__)
 
@@ -25,9 +23,7 @@ def __post_init__(self):
             timeout=httpx.Timeout(timeout=15.0),
         )
 
-    async def passthrough(self, request: Request):
-        """Transparently proxy a request to the upstream STAC API."""
-
+    async def proxy_request(self, request: Request, *, stream=False) -> httpx.Response:
         headers = MutableHeaders(request.headers)
 
         # https://github.com/fastapi/fastapi/discussions/7382#discussioncomment-5136466
@@ -43,14 +39,18 @@ async def passthrough(self, request: Request):
         logger.debug(f"Proxying request to {rp_req.url}")
 
         start_time = time.perf_counter()
-        rp_resp = await self.client.send(rp_req, stream=True)
+        rp_resp = await self.client.send(rp_req, stream=stream)
         proxy_time = time.perf_counter() - start_time
 
         logger.debug(
             f"Received response status {rp_resp.status_code!r} from {rp_req.url} in {proxy_time:.3f}s"
         )
         rp_resp.headers["X-Upstream-Time"] = f"{proxy_time:.3f}"
+        return rp_resp
 
+    async def stream(self, request: Request) -> StreamingResponse:
+        """Transparently proxy a request to the upstream STAC API."""
+        rp_resp = await self.proxy_request(request, stream=True)
         return StreamingResponse(
             rp_resp.aiter_raw(),
             status_code=rp_resp.status_code,

src/stac_auth_proxy/utils.py

@@ -0,0 +1,16 @@
+from httpx import Headers
+
+
+def safe_headers(headers: Headers) -> dict[str, str]:
+    """
+    Scrub headers that should not be proxied to the client.
+    """
+    excluded_headers = [
+        "content-length",
+        "content-encoding",
+    ]
+    return {
+        key: value
+        for key, value in headers.items()
+        if key.lower() not in excluded_headers
+    }