Working validation

Author: alukach

pyproject.toml

@@ -8,7 +8,7 @@ classifiers = [
 dependencies = [
   "authlib>=1.3.2",
   "brotli>=1.1.0",
-  "cql2>=0.3.5",
+  "cql2>=0.3.6",
   "fastapi>=0.115.5",
   "httpx[http2]>=0.28.0",
   "jinja2>=3.1.4",

src/stac_auth_proxy/middleware/ApplyCql2FilterMiddleware.py

@@ -1,9 +1,13 @@
 """Middleware to apply CQL2 filters."""
 
 import json
+import re
 from dataclasses import dataclass
 from logging import getLogger
+from typing import Optional
 
+from cql2 import Expr
+from starlette.datastructures import MutableHeaders
 from starlette.requests import Request
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
@@ -28,12 +32,87 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         request = Request(scope)
 
         if request.method == "GET":
-            cql2_filter = getattr(request.state, self.state_key, None)
+            cql2_filter: Optional[Expr] = getattr(request.state, self.state_key, None)
             if cql2_filter:
                 scope["query_string"] = filters.append_qs_filter(
                     request.url.query, cql2_filter
                 )
-            return await self.app(scope, receive, send)
+
+            initial_message = None
+            body = b""
+
+            async def validate_response(message: Message) -> None:
+                nonlocal initial_message
+                nonlocal body
+                headers = MutableHeaders(scope=initial_message)
+                if message["type"] == "http.response.start":
+                    initial_message = message
+                    return
+
+                if message["type"] == "http.response.body":
+                    assert initial_message, "Initial message not set"
+                    assert cql2_filter, "Cql2Filter not set"
+                    body += message["body"]
+                    if message.get("more_body"):
+                        return
+
+                    try:
+                        body = json.loads(body)
+                    except json.JSONDecodeError:
+                        logger.warning("Failed to parse response body as JSON")
+                        not_found_body = json.dumps({"message": "Not found"}).encode(
+                            "utf-8"
+                        )
+                        headers["content-length"] = str(len(not_found_body))
+                        initial_message["status"] = 502
+                        await send(initial_message)
+                        await send(
+                            {
+                                "type": "http.response.body",
+                                "body": not_found_body,
+                                "more_body": False,
+                            }
+                        )
+                        return
+
+                    logger.debug(
+                        "Applying %s filter to %s", cql2_filter.to_text(), body
+                    )
+                    if cql2_filter.matches(body):
+                        await send(initial_message)
+                        await send(
+                            {
+                                "type": "http.response.body",
+                                "body": json.dumps(body).encode("utf-8"),
+                                "more_body": False,
+                            }
+                        )
+                    else:
+                        not_found_body = json.dumps({"message": "Not found"}).encode(
+                            "utf-8"
+                        )
+                        headers["content-length"] = str(len(not_found_body))
+                        initial_message["status"] = 404
+                        await send(initial_message)
+                        await send(
+                            {
+                                "type": "http.response.body",
+                                "body": not_found_body,
+                                "more_body": False,
+                            }
+                        )
+
+                return message
+
+            should_validate_response = cql2_filter and re.match(
+                r"^/collections/([^/]+)/items/([^/]+)$", request.url.path
+            )
+
+            return await self.app(
+                scope,
+                receive,
+                validate_response if should_validate_response else send,
+            )
 
         elif request.method in ["POST", "PUT", "PATCH"]:
 
@@ -55,6 +134,10 @@ async def receive_and_apply_filter() -> Message:
                     message["body"] = json.dumps(new_body).encode("utf-8")
                 return message
 
-            return await self.app(scope, receive_and_apply_filter, send)
+            return await self.app(
+                scope,
+                receive_and_apply_filter,
+                send,
+            )
 
         return await self.app(scope, receive, send)

src/stac_auth_proxy/middleware/BuildCql2FilterMiddleware.py

@@ -1,14 +1,15 @@
 """Middleware to build the Cql2Filter."""
 
 import json
+import re
 from dataclasses import dataclass
 from typing import Callable, Optional
 
 from cql2 import Expr
 from starlette.requests import Request
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
-from ..utils import filters, requests
+from ..utils import requests
 
 
 @dataclass(frozen=True)
@@ -78,11 +79,10 @@ async def receive_build_filter() -> Message:
     def _get_filter(self, path: str) -> Optional[Callable[..., Expr]]:
         """Get the CQL2 filter builder for the given path."""
         endpoint_filters = [
-            (filters.is_collection_endpoint, self.collections_filter),
-            (filters.is_item_endpoint, self.items_filter),
-            (filters.is_search_endpoint, self.items_filter),
+            (r"^/collections(/[^/]+)?$", self.collections_filter),
+            (r"^(/collections/([^/]+)/items(/[^/]+)?$|/search$)", self.items_filter),
         ]
-        for check, builder in endpoint_filters:
-            if check(path):
+        for expr, builder in endpoint_filters:
+            if re.match(expr, path):
                 return builder
         return None

src/stac_auth_proxy/utils/filters.py

@@ -1,7 +1,6 @@
 """Utility functions."""
 
 import json
-import re
 from typing import Optional
 from urllib.parse import parse_qs
 
@@ -32,23 +31,6 @@ def append_body_filter(
     }
 
 
-def is_collection_endpoint(path: str) -> bool:
-    """Check if the path is a collection endpoint."""
-    # TODO: Expand this to cover all cases where a collection filter should be applied
-    return path == "/collections"
-
-
-def is_item_endpoint(path: str) -> bool:
-    """Check if the path is an item endpoint."""
-    # TODO: Expand this to cover all cases where an item filter should be applied
-    return bool(re.compile(r"^(/collections/([^/]+)/items$|/search)").match(path))
-
-
-def is_search_endpoint(path: str) -> bool:
-    """Check if the path is a search endpoint."""
-    return path == "/search"
-
-
 def dict_to_query_string(params: dict) -> str:
     """
     Convert a dictionary to a query string. Dict values are converted to JSON strings,

tests/test_filters_jinja2.py

@@ -307,7 +307,11 @@ def test_item_get(
         "properties": {"private": True},
     }
     response = client.get("/collections/foo/items/bar")
-    expected_status = 404 if is_authenticated else 200
+    expected_status = 200 if is_authenticated else 404
+    expected_body = (
+        {"id": "bar", "properties": {"private": True}}
+        if is_authenticated
+        else {"message": "Not found"}
+    )
     assert response.status_code == expected_status
-    if is_authenticated:
-        assert response.json() == {"id": "bar", "properties": {"private": True}}
+    assert response.json() == expected_body