fix: Rework body augmentor to avoid error on empty POSTs

Was seeing `h11._util.LocalProtocolError: Too much data for declared Content-Length.`

Author: alukach

src/stac_auth_proxy/middleware/ApplyCql2FilterMiddleware.py

@@ -2,13 +2,12 @@
 
 import json
 import re
-from dataclasses import dataclass, field
-from functools import partial
+from dataclasses import dataclass
 from logging import getLogger
-from typing import Callable, Optional
+from typing import Optional
 
 from cql2 import Expr
-from starlette.datastructures import MutableHeaders, State
+from starlette.datastructures import MutableHeaders
 from starlette.requests import Request
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
@@ -39,32 +38,25 @@ 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,
+            req_body_handler = Cql2RequestBodyAugmentor(
+                app=self.app,
+                cql2_filter=cql2_filter,
             )
-
-        cql2_filter = get_cql2_filter()
-        if not cql2_filter:
-            return await self.app(scope, receive, send)
+            return await req_body_handler(scope, receive, send)
 
         if re.match(r"^/collections/([^/]+)/items/([^/]+)$", request.url.path):
-            return await self.app(
-                scope,
-                receive,
-                Cql2ResponseBodyValidator(cql2_filter=cql2_filter, send=send),
+            res_body_validator = Cql2ResponseBodyValidator(
+                app=self.app,
+                cql2_filter=cql2_filter,
             )
+            return await res_body_validator(scope, send, receive)
 
         scope["query_string"] = filters.append_qs_filter(request.url.query, cql2_filter)
         return await self.app(scope, receive, send)
@@ -74,88 +66,115 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
 class Cql2RequestBodyAugmentor:
     """Handler to augment the request body with a CQL2 filter."""
 
-    receive: Receive
-    state: State
-    get_cql2_filter: Callable[[], Optional[Expr]]
-
-    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
-
-        # NOTE: Can only get cql2 filter _after_ calling self.receive()
-        cql2_filter = self.get_cql2_filter()
-        if not cql2_filter:
-            return message
+    app: ASGIApp
+    cql2_filter: Expr
 
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """Augment the request body with a CQL2 filter."""
+        body = b""
+        more_body = True
+
+        # Read the body
+        while more_body:
+            message = await receive()
+            if message["type"] == "http.request":
+                body += message.get("body", b"")
+                more_body = message.get("more_body", False)
+
+        # Modify body
         try:
-            body = json.loads(message.get("body", b"{}"))
+            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 = filters.append_body_filter(body, cql2_filter)
-        message["body"] = json.dumps(new_body).encode("utf-8")
-        return message
+        # Augment the body
+        assert isinstance(body, dict), "Request body must be a JSON object"
+        new_body = json.dumps(
+            filters.append_body_filter(body, self.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["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)
 
 
 @dataclass
 class Cql2ResponseBodyValidator:
     """Handler to validate response body with CQL2."""
 
-    send: Send
+    app: ASGIApp
     cql2_filter: Expr
-    initial_message: Optional[Message] = field(init=False)
-    body: bytes = field(init=False, default_factory=bytes)
 
-    async def __call__(self, message: Message) -> None:
+    async def __call__(self, scope: Scope, send: Send, receive: Receive) -> None:
         """Process a response message and apply filtering if needed."""
-        if message["type"] == "http.response.start":
-            self.initial_message = message
-            return
+        if scope["type"] != "http":
+            return await self.app(scope, send, receive)
+
+        body = b""
+        initial_message: Optional[Message] = None
+
+        async def _send_error_response(status: int, message: str) -> None:
+            """Send an error response with the given status and message."""
+            assert initial_message, "Initial message not set"
+            error_body = json.dumps({"message": message}).encode("utf-8")
+            headers = MutableHeaders(scope=initial_message)
+            headers["content-length"] = str(len(error_body))
+            initial_message["status"] = status
+            await send(initial_message)
+            await send(
+                {
+                    "type": "http.response.body",
+                    "body": error_body,
+                    "more_body": False,
+                }
+            )
 
-        if message["type"] == "http.response.body":
-            assert self.initial_message, "Initial message not set"
+        async def buffered_send(message: Message) -> None:
+            """Process a response message and apply filtering if needed."""
+            nonlocal body
+            nonlocal initial_message
 
-            self.body += message["body"]
+            if message["type"] == "http.response.start":
+                initial_message = message
+                return
+
+            assert initial_message, "Initial message not set"
+
+            body += message["body"]
             if message.get("more_body"):
                 return
 
             try:
-                body_json = json.loads(self.body)
+                body_json = json.loads(body)
             except json.JSONDecodeError:
                 logger.warning("Failed to parse response body as JSON")
-                await self._send_error_response(502, "Not found")
+                await _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(
+                await send(initial_message)
+                return await 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")
-
-    async def _send_error_response(self, status: int, message: str) -> None:
-        """Send an error response with the given status and message."""
-        assert self.initial_message, "Initial message not set"
-        error_body = json.dumps({"message": message}).encode("utf-8")
-        headers = MutableHeaders(scope=self.initial_message)
-        headers["content-length"] = str(len(error_body))
-        self.initial_message["status"] = status
-        await self.send(self.initial_message)
-        await self.send(
-            {
-                "type": "http.response.body",
-                "body": error_body,
-                "more_body": False,
-            }
-        )
+            return await _send_error_response(404, "Not found")
+
+        return await self.app(scope, receive, buffered_send)

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

tests/test_filters_jinja2.py

@@ -314,3 +314,26 @@ def test_item_get(
     else:
         assert response.status_code == 404
         assert response.json() == {"message": "Not found"}
+
+
+@pytest.mark.parametrize("is_authenticated", [True, False], ids=["auth", "anon"])
+async def test_search_post_empty_body(
+    source_api_server,
+    is_authenticated,
+    token_builder,
+):
+    """Test that POST /search with empty body."""
+    client = _build_client(
+        src_api_server=source_api_server,
+        template_expr="(properties.private = false)",
+        is_authenticated=is_authenticated,
+        token_builder=token_builder,
+    )
+
+    # Send request with Content-Length header that doesn't match actual body size
+    response = client.post(
+        "/search",
+        json={},
+    )
+
+    assert response.status_code == 200