In progress

Author: alukach

src/stac_auth_proxy/app.py

@@ -14,6 +14,8 @@
     OpenApiMiddleware,
     AddProcessTimeHeaderMiddleware,
     EnforceAuthMiddleware,
+    BuildCql2FilterMiddleware,
+    ApplyCql2FilterMiddleware,
 )
 
 from .config import Settings
@@ -48,6 +50,13 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
             default_public=settings.default_public,
         )
 
+    app.add_middleware(ApplyCql2FilterMiddleware)
+    app.add_middleware(
+        BuildCql2FilterMiddleware,
+        # collections_filter=settings.collections_filter,
+        items_filter=settings.items_filter(),
+    )
+
     if settings.debug:
         app.add_api_route(
             "/_debug",

src/stac_auth_proxy/config.py

@@ -55,8 +55,8 @@ class Settings(BaseSettings):
     }
     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,13 +1,12 @@
 """Generate CQL2 filter expressions via Jinja2 templating."""
 
+from typing import Any
 from dataclasses import dataclass, field
 
 from cql2 import Expr
 from fastapi import Request
 from jinja2 import BaseLoader, Environment
 
-from ..utils.requests import extract_variables
-
 
 @dataclass
 class Template:
@@ -20,24 +19,9 @@ def __post_init__(self):
         """Initialize the Jinja2 environment."""
         self.env = Environment(loader=BaseLoader).from_string(self.template_str)
 
-    async def __call__(self, request: Request) -> Expr:
+    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": request.state.user,
-        }
         cql2_str = self.env.render(**context).strip()
         cql2_expr = Expr(cql2_str)
         cql2_expr.validate()

src/stac_auth_proxy/middleware/Cql2FilterMiddleware.py

@@ -1,3 +1,4 @@
+from logging import getLogger
 import json
 from dataclasses import dataclass
 from typing import Annotated, Callable, Optional
@@ -9,6 +10,7 @@
 from ..config import EndpointMethods
 from ..utils import di, filters, requests
 
+logger = getLogger(__name__)
 
 FILTER_STATE_KEY = "cql2_filter"
 
@@ -27,17 +29,38 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if scope["type"] != "http":
             return await self.app(scope, receive, send)
 
-        request = Request(scope)
-        filter_builder = self._get_filter(request.url.path)
-        if filter_builder:
-            cql2_filter = await di.call_with_injected_dependencies(
-                func=filter_builder,
-                request=request,
-            )
-            cql2_filter.validate()
-            scope["state"][FILTER_STATE_KEY] = cql2_filter
+        total_body = b""
+
+        async def receive_build_filter() -> Message:
+            nonlocal total_body
 
-        return await self.app(scope, receive, send)
+            message = await receive()
+            total_body += message.get("body", b"")
+
+            if not message.get("more_body"):
+                request = Request(scope)
+                filter_builder = self._get_filter(request.url.path)
+                if filter_builder:
+                    cql2_filter = await filter_builder(
+                        {
+                            "req": {
+                                "path": request.url.path,
+                                "method": request.method,
+                                "query_params": dict(request.query_params),
+                                "path_params": requests.extract_variables(
+                                    request.url.path
+                                ),
+                                "headers": dict(request.headers),
+                                "body": json.loads(total_body),
+                            },
+                            **request.state._state,
+                        }
+                    )
+                    cql2_filter.validate()
+                    scope["state"][FILTER_STATE_KEY] = cql2_filter
+            return message
+
+        return await self.app(scope, receive_build_filter, send)
 
     def _get_filter(self, path: str) -> Optional[Callable[..., Expr]]:
         """Get the CQL2 filter builder for the given path."""
@@ -55,51 +78,40 @@ def _get_filter(self, path: str) -> Optional[Callable[..., Expr]]:
 
 @dataclass(frozen=True)
 class ApplyCql2FilterMiddleware:
-    """Middleware to add the OpenAPI spec to the response."""
+    """Middleware to apply the Cql2Filter to the request."""
 
     app: ASGIApp
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         """Add the Cql2Filter to the request."""
-        request = Request(scope)
-        cql2_filter = request.state.get(FILTER_STATE_KEY)
-
-        if scope["type"] != "http" or not cql2_filter:
+        if scope["type"] != "http":
             return await self.app(scope, receive, send)
 
-        # Apply filter if applicable
+        async def apply_filter() -> Message:
+            message = await receive()
+            request = Request(scope)
+            cql2_filter = getattr(request.state, FILTER_STATE_KEY, None)
+            if not cql2_filter:
+                logger.debug("No cql2 filter found on message.")
+                return message
 
-        total_body = b""
-
-        async def receive_with_filter(message: Message):
-            query = request.url.query
-
-            # TODO: How do we handle querystrings in middleware?
             if request.method == "GET":
                 query = filters.insert_qs_filter(qs=query, filter=cql2_filter)
-
-            if message["type"] == "http.response.body":
-                nonlocal total_body
-                total_body += message["body"]
-                if message["more_body"]:
-                    return await receive({**message, "body": b""})
-
-                # TODO: Only on search, not on create or update...
-                if request.method in ["POST", "PUT"]:
-                    return await receive(
-                        {
-                            "type": "http.response.body",
-                            "body": requests.dict_to_bytes(
-                                filters.append_body_filter(
-                                    json.loads(total_body), cql2_filter
-                                )
-                            ),
-                            "more_body": False,
-                        }
+                # Get the original query string
+                original_qs = scope["query_string"].decode("utf-8")
+                # Apply the filter to query string
+                new_qs = filters.append_qs_filter(original_qs, cql2_filter)
+                # Update the scope with new query string
+                # scope["query_string"] = new_qs.encode("utf-8")
+            elif request.method in ["POST", "PUT", "PATCH"]:
+                # TODO: Apply the filter to the body
+                message["body"] = json.dumps(
+                    filters.append_body_filter(
+                        body=json.loads(message.get("body", "{}")),
+                        filter=cql2_filter,
                     )
+                ).encode("utf-8")
 
-                return await receive(message)
-
-            await receive(message)
+            return message
 
-        return await self.app(scope, receive_with_filter, send)
+        return await self.app(scope, apply_filter, send)

src/stac_auth_proxy/middleware/__init__.py

@@ -3,9 +3,12 @@
 from .UpdateOpenApiMiddleware import OpenApiMiddleware
 from .AddProcessTimeHeaderMiddleware import AddProcessTimeHeaderMiddleware
 from .EnforceAuthMiddleware import EnforceAuthMiddleware
+from .Cql2FilterMiddleware import BuildCql2FilterMiddleware, ApplyCql2FilterMiddleware
 
 __all__ = [
-    UpdateOpenApiMiddleware,
+    OpenApiMiddleware,
     AddProcessTimeHeaderMiddleware,
     EnforceAuthMiddleware,
+    BuildCql2FilterMiddleware,
+    ApplyCql2FilterMiddleware,
 ]