refactor(cql2-middleware): breakup cql2 middleware into smaller components (#66)

Author: alukach

src/stac_auth_proxy/app.py

@@ -16,9 +16,11 @@
 from .handlers import HealthzHandler, ReverseProxyHandler, SwaggerUI
 from .middleware import (
     AddProcessTimeHeaderMiddleware,
-    ApplyCql2FilterMiddleware,
     AuthenticationExtensionMiddleware,
-    BuildCql2FilterMiddleware,
+    Cql2ApplyFilterBodyMiddleware,
+    Cql2ApplyFilterQueryStringMiddleware,
+    Cql2BuildFilterMiddleware,
+    Cql2ValidateResponseBodyMiddleware,
     EnforceAuthMiddleware,
     OpenApiMiddleware,
     ProcessLinksMiddleware,
@@ -132,11 +134,11 @@ async def lifespan(app: FastAPI):
         )
 
     if settings.items_filter or settings.collections_filter:
+        app.add_middleware(Cql2ValidateResponseBodyMiddleware)
+        app.add_middleware(Cql2ApplyFilterBodyMiddleware)
+        app.add_middleware(Cql2ApplyFilterQueryStringMiddleware)
         app.add_middleware(
-            ApplyCql2FilterMiddleware,
-        )
-        app.add_middleware(
-            BuildCql2FilterMiddleware,
+            Cql2BuildFilterMiddleware,
             items_filter=settings.items_filter() if settings.items_filter else None,
             collections_filter=(
                 settings.collections_filter() if settings.collections_filter else None

src/stac_auth_proxy/middleware/ApplyCql2FilterMiddleware.py

@@ -0,0 +1,98 @@
+"""Middleware to augment the request body with a CQL2 filter for POST/PUT/PATCH requests."""
+
+import json
+from dataclasses import dataclass
+from logging import getLogger
+from typing import Optional
+
+from cql2 import Expr
+from starlette.requests import Request
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from ..utils import filters
+from ..utils.middleware import required_conformance
+
+logger = getLogger(__name__)
+
+
+@required_conformance(
+    r"http://www.opengis.net/spec/cql2/1.0/conf/basic-cql2",
+    r"http://www.opengis.net/spec/cql2/1.0/conf/cql2-text",
+    r"http://www.opengis.net/spec/cql2/1.0/conf/cql2-json",
+)
+@dataclass(frozen=True)
+class Cql2ApplyFilterBodyMiddleware:
+    """Middleware to augment the request body with a CQL2 filter for POST/PUT/PATCH requests."""
+
+    app: ASGIApp
+    state_key: str = "cql2_filter"
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """Apply the CQL2 filter to the request body."""
+        if scope["type"] != "http":
+            return await self.app(scope, receive, send)
+
+        request = Request(scope)
+        cql2_filter: Optional[Expr] = getattr(request.state, self.state_key, None)
+        if not cql2_filter:
+            return await self.app(scope, receive, send)
+
+        if request.method not in ["POST", "PUT", "PATCH"]:
+            return await self.app(scope, receive, send)
+
+        body = b""
+        more_body = True
+        while more_body:
+            message = await receive()
+            if message["type"] == "http.request":
+                body += message.get("body", b"")
+                more_body = message.get("more_body", False)
+
+        try:
+            body_json = json.loads(body) if body else {}
+        except json.JSONDecodeError:
+            logger.warning("Failed to parse request body as JSON")
+            from starlette.responses import JSONResponse
+
+            response = JSONResponse(
+                {
+                    "code": "ParseError",
+                    "description": "Request body must be valid JSON.",
+                },
+                status_code=400,
+            )
+            await response(scope, receive, send)
+            return
+
+        if not isinstance(body_json, dict):
+            logger.warning("Request body must be a JSON object")
+            from starlette.responses import JSONResponse
+
+            response = JSONResponse(
+                {
+                    "code": "TypeError",
+                    "description": "Request body must be a JSON object.",
+                },
+                status_code=400,
+            )
+            await response(scope, receive, send)
+            return
+
+        new_body = json.dumps(
+            filters.append_body_filter(body_json, cql2_filter)
+        ).encode("utf-8")
+
+        # Patch content-length in the headers
+        headers = dict(scope["headers"])
+        headers[b"content-length"] = str(len(new_body)).encode("latin1")
+        scope = dict(scope)
+        scope["headers"] = list(headers.items())
+
+        async def new_receive():
+            return {
+                "type": "http.request",
+                "body": new_body,
+                "more_body": False,
+            }
+
+        await self.app(scope, new_receive, send)

src/stac_auth_proxy/middleware/Cql2ApplyFilterBodyMiddleware.py

@@ -0,0 +1,56 @@
+"""Middleware to inject CQL2 filters into the query string for GET/list endpoints."""
+
+import re
+from dataclasses import dataclass
+from logging import getLogger
+from typing import Optional
+
+from cql2 import Expr
+from starlette.requests import Request
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from ..utils import filters
+from ..utils.middleware import required_conformance
+
+logger = getLogger(__name__)
+
+
+@required_conformance(
+    r"http://www.opengis.net/spec/cql2/1.0/conf/basic-cql2",
+    r"http://www.opengis.net/spec/cql2/1.0/conf/cql2-text",
+    r"http://www.opengis.net/spec/cql2/1.0/conf/cql2-json",
+)
+@dataclass(frozen=True)
+class Cql2ApplyFilterQueryStringMiddleware:
+    """Middleware to inject CQL2 filters into the query string for GET/list endpoints."""
+
+    app: ASGIApp
+    state_key: str = "cql2_filter"
+
+    single_record_endpoints = [
+        r"^/collections/([^/]+)/items/([^/]+)$",
+        r"^/collections/([^/]+)$",
+    ]
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """Apply the CQL2 filter to the query string."""
+        if scope["type"] != "http":
+            return await self.app(scope, receive, send)
+
+        request = Request(scope)
+        cql2_filter: Optional[Expr] = getattr(request.state, self.state_key, None)
+        if not cql2_filter:
+            return await self.app(scope, receive, send)
+
+        # Only handle GET requests that are not single-record endpoints
+        if request.method != "GET":
+            return await self.app(scope, receive, send)
+        if any(
+            re.match(expr, request.url.path) for expr in self.single_record_endpoints
+        ):
+            return await self.app(scope, receive, send)
+
+        # Inject filter into query string
+        scope = dict(scope)
+        scope["query_string"] = filters.append_qs_filter(request.url.query, cql2_filter)
+        return await self.app(scope, receive, send)

src/stac_auth_proxy/middleware/Cql2ApplyFilterQueryStringMiddleware.py

@@ -22,7 +22,7 @@
     "http://www.opengis.net/spec/cql2/1.0/conf/cql2-json",
 )
 @dataclass(frozen=True)
-class BuildCql2FilterMiddleware:
+class Cql2BuildFilterMiddleware:
     """Middleware to build the Cql2Filter."""
 
     app: ASGIApp