Take care to separate GET and POST middleware behavior

Author: alukach

src/stac_auth_proxy/middleware/Cql2FilterMiddleware.py

@@ -29,6 +29,35 @@ 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 not filter_builder:
+            return await self.app(scope, receive, send)
+
+        async def set_filter(body: Optional[dict] = None) -> None:
+            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": body,
+                    },
+                    **request.state._state,
+                }
+            )
+            cql2_filter.validate()
+            scope["state"][FILTER_STATE_KEY] = cql2_filter
+
+        # For GET requests, we can build the filter immediately
+        # NOTE: It appears that FastAPI will not call receive function for GET requests
+        if request.method == "GET":
+            await set_filter()
+            return await self.app(scope, receive, send)
+
         total_body = b""
 
         async def receive_build_filter() -> Message:
@@ -38,26 +67,7 @@ async def receive_build_filter() -> Message:
             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
+                await set_filter(json.loads(total_body) if total_body else None)
             return message
 
         return await self.app(scope, receive_build_filter, send)
@@ -87,31 +97,35 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if scope["type"] != "http":
             return await self.app(scope, receive, send)
 
-        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
+        request = Request(scope)
 
-            if request.method == "GET":
-                query = filters.insert_qs_filter(qs=query, filter=cql2_filter)
-                # 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")
+        if request.method == "GET":
+            cql2_filter = scope["state"].get(FILTER_STATE_KEY)
+            if cql2_filter:
+                scope["query_string"] = filters.append_qs_filter(
+                    request.url.query, cql2_filter
+                )
+            return await self.app(scope, receive, send)
+        elif request.method in ["POST", "PUT", "PATCH"]:
+            # For methods with bodies, we need to wrap receive to modify the body
+            async def receive_with_filter() -> Message:
+                message = await receive()
+                if message["type"] != "http.request":
+                    return message
+
+                cql2_filter = scope["state"].get(FILTER_STATE_KEY)
+                if cql2_filter:
+                    try:
+                        body = message.get("body", b"{}")
+                    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(json.loads(body), cql2_filter)
+                    message["body"] = json.dumps(new_body).encode("utf-8")
+                return message
 
-            return message
+            return await self.app(scope, receive_with_filter, send)
 
-        return await self.app(scope, apply_filter, send)
+        return await self.app(scope, receive, send)

src/stac_auth_proxy/middleware/EnforceAuthMiddleware.py

@@ -4,7 +4,7 @@
 import logging
 import urllib.request
 
-from fastapi import HTTPException, Security, security, status, Request
+from fastapi import HTTPException, Security, status, Request
 from pydantic import HttpUrl
 from starlette.middleware.base import ASGIApp
 from starlette.responses import JSONResponse

src/stac_auth_proxy/utils/filters.py

@@ -1,32 +1,42 @@
 """Utility functions."""
 
+import json
 import re
+from typing import Optional
 from urllib.parse import parse_qs, urlencode
 
 from cql2 import Expr
 
+from .requests import dict_to_bytes
 
-def append_qs_filter(qs: str, filter: Expr) -> str:
-    """Insert a filter expression into a query string. If a filter already exists, combine them."""
-    qs_dict = parse_qs(qs)
-
-    for qs_filter in qs_dict.get("filter", []):
-        filter += Expr(qs_filter)
-
-    qs_dict["filter"] = filter.to_text()
-    qs_dict["filter-lang"] = "cql2-text"
 
-    return urlencode(qs_dict, doseq=True)
-
-
-def append_body_filter(body: dict, filter: Expr) -> dict:
+def append_qs_filter(qs: str, filter: Expr, filter_lang: Optional[str] = None) -> bytes:
+    """Insert a filter expression into a query string. If a filter already exists, combine them."""
+    qs_dict = {k: v[0] for k, v in parse_qs(qs).items()}
+    new_qs_dict = append_body_filter(
+        qs_dict, filter, filter_lang or qs_dict.get("filter-lang", "cql2-text")
+    )
+    return dict_to_bytes(
+        urlencode(
+            {
+                k: json.dumps(v) if isinstance(v, (list, dict)) else v
+                for k, v in new_qs_dict.items()
+            }
+        )
+    )
+
+
+def append_body_filter(body: dict, filter: Expr, filter_lang: Optional[str]) -> dict:
     """Insert a filter expression into a request body. If a filter already exists, combine them."""
     cur_filter = body.get("filter")
+    filter_lang = filter_lang or body.get("filter-lang", "cql2-json")
     if cur_filter:
         filter = filter + Expr(cur_filter)
-    body["filter"] = filter.to_json()
-    body["filter-lang"] = "cql2-json"
-    return body
+    return {
+        **body,
+        "filter": filter.to_text() if filter_lang == "cql2-text" else filter.to_json(),
+        "filter-lang": filter_lang,
+    }
 
 
 def is_collection_endpoint(path: str) -> bool:

tests/test_filters_jinja2.py

@@ -241,8 +241,11 @@ async def test_search_get(
 
     expected_output = {
         **input_query,
-        "filter": proxy_filter.to_text(),
-        "filter-lang": "cql2-text",
+        "filter": (
+            proxy_filter.to_text()
+            if input_query.get("filter-lang") == "cql2-text"
+            else proxy_filter.to_json()
+        ),
     }
     assert (
         upstream_query == expected_output