In progress...

Author: alukach

src/stac_auth_proxy/middleware/ApplyCql2FilterMiddleware.py

@@ -3,9 +3,8 @@
 import json
 import re
 from dataclasses import dataclass, field
-from functools import partial
 from logging import getLogger
-from typing import Callable, Optional
+from typing import Optional
 
 from cql2 import Expr
 from starlette.datastructures import MutableHeaders, State
@@ -39,25 +38,60 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
 
         request = Request(scope)
 
-        get_cql2_filter: Callable[[], Optional[Expr]] = partial(
-            getattr, request.state, self.state_key, None
-        )
+        cql2_filter: Optional[Expr] = getattr(request.state, self.state_key, None)
+
+        if not cql2_filter:
+            return await self.app(scope, receive, send)
 
         # Handle POST, PUT, PATCH
         if request.method in ["POST", "PUT", "PATCH"]:
-            return await self.app(
-                scope,
-                Cql2RequestBodyAugmentor(
-                    receive=receive,
-                    state=request.state,
-                    get_cql2_filter=get_cql2_filter,
-                ),
-                send,
+            body = b""
+            more_body = True
+            receive_ = receive
+
+            async def buffered_receive():
+                nonlocal body, more_body
+                if more_body:
+                    message = await receive_()
+                    if message["type"] == "http.request":
+                        body += message.get("body", b"")
+                        more_body = message.get("more_body", False)
+                        return message
+                return {"type": "http.request", "body": b"", "more_body": False}
+
+            while more_body:
+                await buffered_receive()
+
+            # Modify body
+            # modified_body = body + b"\nAppended content."
+            try:
+                body = json.loads(body)
+            except json.JSONDecodeError as e:
+                logger.warning("Failed to parse request body as JSON")
+                # TODO: Return a 400 error
+                raise e
+
+            new_body = json.dumps(filters.append_body_filter(body, cql2_filter)).encode(
+                "utf-8"
             )
 
-        cql2_filter = get_cql2_filter()
-        if not cql2_filter:
-            return await self.app(scope, receive, send)
+            # Override the receive function with a generator sending our modified body
+            async def new_receive():
+                nonlocal new_body
+                chunk = new_body
+                new_body = b""
+                return {
+                    "type": "http.request",
+                    "body": chunk,
+                    "more_body": False,
+                }
+
+            # Patch content-length in the headers
+            headers = dict(scope["headers"])
+            headers[b"content-length"] = str(len(new_body)).encode("latin1")
+            scope["headers"] = list(headers.items())
+
+            await self.app(scope, new_receive, send)
 
         if re.match(r"^/collections/([^/]+)/items/([^/]+)$", request.url.path):
             return await self.app(
@@ -76,27 +110,27 @@ class Cql2RequestBodyAugmentor:
 
     receive: Receive
     state: State
-    get_cql2_filter: Callable[[], Optional[Expr]]
+    cql2_filter: Expr
+    initial_message: Optional[Message] = field(init=False)
 
     async def __call__(self) -> Message:
         """Process a request body and augment with a CQL2 filter if available."""
         message = await self.receive()
-        if message["type"] != "http.request":
-            return message
+        if message["type"] == "http.response.start":
+            self.initial_message = message
+            return
 
-        # NOTE: Can only get cql2 filter _after_ calling self.receive()
-        cql2_filter = self.get_cql2_filter()
-        if not cql2_filter:
+        if not message.get("body"):
             return message
 
         try:
-            body = json.loads(message.get("body", b"{}"))
+            body = json.loads(message.get("body"))
         except json.JSONDecodeError as e:
             logger.warning("Failed to parse request body as JSON")
             # TODO: Return a 400 error
             raise e
 
-        new_body = filters.append_body_filter(body, cql2_filter)
+        new_body = filters.append_body_filter(body, self.cql2_filter)
         message["body"] = json.dumps(new_body).encode("utf-8")
         return message
 
@@ -116,33 +150,30 @@ async def __call__(self, message: Message) -> None:
             self.initial_message = message
             return
 
-        if message["type"] == "http.response.body":
-            assert self.initial_message, "Initial message not set"
+        assert self.initial_message, "Initial message not set"
 
-            self.body += message["body"]
-            if message.get("more_body"):
-                return
+        self.body += message["body"]
+        if message.get("more_body"):
+            return
 
-            try:
-                body_json = json.loads(self.body)
-            except json.JSONDecodeError:
-                logger.warning("Failed to parse response body as JSON")
-                await self._send_error_response(502, "Not found")
-                return
-
-            logger.debug(
-                "Applying %s filter to %s", self.cql2_filter.to_text(), body_json
+        try:
+            body_json = json.loads(self.body)
+        except json.JSONDecodeError:
+            logger.warning("Failed to parse response body as JSON")
+            await self._send_error_response(502, "Not found")
+            return
+
+        logger.debug("Applying %s filter to %s", self.cql2_filter.to_text(), body_json)
+        if self.cql2_filter.matches(body_json):
+            await self.send(self.initial_message)
+            return await self.send(
+                {
+                    "type": "http.response.body",
+                    "body": json.dumps(body_json).encode("utf-8"),
+                    "more_body": False,
+                }
             )
-            if self.cql2_filter.matches(body_json):
-                await self.send(self.initial_message)
-                return await self.send(
-                    {
-                        "type": "http.response.body",
-                        "body": json.dumps(body_json).encode("utf-8"),
-                        "more_body": False,
-                    }
-                )
-            return await self._send_error_response(404, "Not found")
+        return await self._send_error_response(404, "Not found")
 
     async def _send_error_response(self, status: int, message: str) -> None:
         """Send an error response with the given status and message."""

src/stac_auth_proxy/middleware/BuildCql2FilterMiddleware.py

@@ -1,16 +1,19 @@
 """Middleware to build the Cql2Filter."""
 
-import json
+import logging
 import re
 from dataclasses import dataclass
 from typing import Any, Awaitable, Callable, Optional
 
-from cql2 import Expr
+from cql2 import Expr, ValidationError
 from starlette.requests import Request
-from starlette.types import ASGIApp, Message, Receive, Scope, Send
+from starlette.responses import Response
+from starlette.types import ASGIApp, Receive, Scope, Send
 
 from ..utils import requests
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass(frozen=True)
 class BuildCql2FilterMiddleware:
@@ -35,47 +38,27 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if not filter_builder:
             return await self.app(scope, receive, send)
 
-        async def set_filter(body: Optional[dict] = None) -> None:
-            assert filter_builder is not None
-            filter_expr = 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": body,
-                    },
-                    **scope["state"],
-                }
-            )
-            cql2_filter = Expr(filter_expr)
+        filter_expr = 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),
+                },
+                **scope["state"],
+            }
+        )
+        cql2_filter = Expr(filter_expr)
+        try:
             cql2_filter.validate()
-            setattr(request.state, self.state_key, cql2_filter)
-
-        # For GET requests, we can build the filter immediately
-        if request.method == "GET":
-            await set_filter()
-            return await self.app(scope, receive, send)
-
-        total_body = b""
-
-        async def receive_build_filter() -> Message:
-            """
-            Receive the body of the request and build the filter.
-            NOTE: This is not called for GET requests.
-            """
-            nonlocal total_body
-
-            message = await receive()
-            total_body += message.get("body", b"")
-
-            if not message.get("more_body"):
-                await set_filter(json.loads(total_body) if total_body else None)
-            return message
+        except ValidationError:
+            logger.exception("Invalid CQL2 filter: %s", filter_expr)
+            return await Response(status_code=502, content="Invalid CQL2 filter")
+        setattr(request.state, self.state_key, cql2_filter)
 
-        return await self.app(scope, receive_build_filter, send)
+        return await self.app(scope, receive, send)
 
     def _get_filter(
         self, path: str