Refactor codebase to use middleware pattern (#20)

* Refactor API spec augmentation to middleware

* Breakout middleware into separate files

* Rename middleware

* Reorg

* chore: reorg imports

* Mv auth to middleware, rework url patterns to regex

* fix: pass empty chunks while waiting for entirety of body

* Rm scope scheck in auth enforcement

* typing cleanup

* fix: extract user from request state

* Rm auth class in favor of middleware

* Rework template tooling as dataclass

* cleanup: rm unused code

* Cleanup

* CQL2 Middleware: in progress

* In progress

* Update tests

* Take care to separate GET and POST middleware behavior

* Conditionally enable CQL2 filters

+ cleanup

* Add formatting helpers

* Test tweaks

* Fix querystring parse test util

* bugfix: place auth middleware after cql2 middleware to ensure correct order

* Allow customizing state key

* Minor update to auth token

* pre-commit cleanup

* Rm unused code

Author: alukach

.vscode/launch.json

@@ -16,6 +16,7 @@
       ],
       "jinja": true,
       "cwd": "${workspaceFolder}/src",
+      "justMyCode": false
     }
   ]
 }

pyproject.toml

@@ -50,4 +50,4 @@ dev = [
 
 [tool.pytest.ini_options]
 asyncio_default_fixture_loop_scope = "function"
-asyncio_mode = "strict"
+asyncio_mode = "auto"

src/stac_auth_proxy/app.py

@@ -8,12 +8,17 @@
 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,
+    ApplyCql2FilterMiddleware,
+    BuildCql2FilterMiddleware,
+    EnforceAuthMiddleware,
+    OpenApiMiddleware,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -25,56 +30,47 @@ 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,
+        )
+
+    if settings.items_filter:
+        app.add_middleware(ApplyCql2FilterMiddleware)
+        app.add_middleware(
+            BuildCql2FilterMiddleware,
+            # collections_filter=settings.collections_filter,
+            items_filter=settings.items_filter(),
+        )
+
+    app.add_middleware(
+        EnforceAuthMiddleware,
+        public_endpoints=settings.public_endpoints,
+        private_endpoints=settings.private_endpoints,
+        default_public=settings.default_public,
+        oidc_config_url=settings.oidc_discovery_url,
+    )
+
     if settings.debug:
         app.add_api_route(
             "/_debug",
             lambda: {"settings": settings},
             methods=["GET"],
         )
 
-    # Tooling
-    auth_scheme = OpenIdConnectAuth(
-        openid_configuration_url=settings.oidc_discovery_url
-    )
-    proxy_handler = ReverseProxyHandler(
-        upstream=str(settings.upstream_url),
-        auth_dependency=auth_scheme.maybe_validated_user,
-        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,
-            )
-
-    # Catchall for remainder of the endpoints
+    # Catchall for any endpoint
+    proxy_handler = ReverseProxyHandler(upstream=str(settings.upstream_url))
     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)]
-        ),
     )
 
     return app

src/stac_auth_proxy/config.py

@@ -37,26 +37,26 @@ class Settings(BaseSettings):
     default_public: bool = False
     private_endpoints: EndpointMethods = {
         # https://github.com/stac-api-extensions/collection-transaction/blob/v1.0.0-beta.1/README.md#methods
-        "/collections": ["POST"],
-        "/collections/{collection_id}": ["PUT", "PATCH", "DELETE"],
+        r"^/collections$": ["POST"],
+        r"^/collections/([^/]+)$": ["PUT", "PATCH", "DELETE"],
         # https://github.com/stac-api-extensions/transaction/blob/v1.0.0-rc.3/README.md#methods
-        "/collections/{collection_id}/items": ["POST"],
-        "/collections/{collection_id}/items/{item_id}": ["PUT", "PATCH", "DELETE"],
+        r"^/collections/([^/]+)/items$": ["POST"],
+        r"^/collections/([^/]+)/items/([^/]+)$": ["PUT", "PATCH", "DELETE"],
         # https://stac-utils.github.io/stac-fastapi/api/stac_fastapi/extensions/third_party/bulk_transactions/#bulktransactionextension
-        "/collections/{collection_id}/bulk_items": ["POST"],
+        r"^/collections/([^/]+)/bulk_items$": ["POST"],
     }
-    public_endpoints: EndpointMethods = {"/api.html": ["GET"], "/api": ["GET"]}
+    public_endpoints: EndpointMethods = {r"^/api.html$": ["GET"], r"^/api$": ["GET"]}
     openapi_spec_endpoint: Optional[str] = None
 
-    collections_filter: Optional[ClassInput] = None
-    collections_filter_endpoints: Optional[EndpointMethods] = {
-        "/collections": ["GET"],
-        "/collections/{collection_id}": ["GET"],
-    }
+    # collections_filter: Optional[ClassInput] = None
+    # collections_filter_endpoints: Optional[EndpointMethods] = {
+    #     r"^/collections$": ["GET"],
+    #     r"^/collections$/([^/]+)": ["GET"],
+    # }
     items_filter: Optional[ClassInput] = None
     items_filter_endpoints: Optional[EndpointMethods] = {
-        "/search": ["POST"],
-        "/collections/{collection_id}/items": ["GET", "POST"],
+        r"^/search$": ["POST"],
+        r"^/collections/([^/]+)/items$": ["GET", "POST"],
     }
 
     model_config = SettingsConfigDict(env_prefix="STAC_AUTH_PROXY_")

src/stac_auth_proxy/filters/template.py

@@ -1,42 +1,27 @@
 """Generate CQL2 filter expressions via Jinja2 templating."""
 
-from typing import Annotated, Any
+from dataclasses import dataclass, field
+from typing import Any
 
 from cql2 import Expr
-from fastapi import Request
 from jinja2 import BaseLoader, Environment
 
-from ..utils.requests import extract_variables
 
-
-def Template(template_str: str):
+@dataclass
+class Template:
     """Generate CQL2 filter expressions via Jinja2 templating."""
-    env = Environment(loader=BaseLoader).from_string(template_str)
 
-    async def dependency(
-        request: Request,
-        auth_token: Annotated[dict[str, Any], ...],
-    ) -> Expr:
+    template_str: str
+    env: Environment = field(init=False)
+
+    def __post_init__(self):
+        """Initialize the Jinja2 environment."""
+        self.env = Environment(loader=BaseLoader).from_string(self.template_str)
+
+    async def __call__(self, context: dict[str, Any]) -> Expr:
         """Render a CQL2 filter expression with the request and auth token."""
         # TODO: How to handle the case where auth_token is null?
-        context = {
-            "req": {
-                "path": request.url.path,
-                "method": request.method,
-                "query_params": dict(request.query_params),
-                "path_params": extract_variables(request.url.path),
-                "headers": dict(request.headers),
-                "body": (
-                    await request.json()
-                    if request.headers.get("content-type") == "application/json"
-                    else (await request.body()).decode()
-                ),
-            },
-            "token": auth_token,
-        }
-        cql2_str = env.render(**context).strip()
+        cql2_str = self.env.render(**context).strip()
         cql2_expr = Expr(cql2_str)
         cql2_expr.validate()
         return cql2_expr
-
-    return dependency

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"]