refactor: enhance CQL2 filter middleware

- Removed outdated conformance URLs from ApplyCql2FilterMiddleware.
- Refactored error response handling to include error codes in BuildCql2FilterMiddleware.
- Improved logging and validation processes for response bodies in Cql2ResponseBodyValidator.
- Added required conformance checks in BuildCql2FilterMiddleware based on filter functions.

Author: alukach

src/stac_auth_proxy/middleware/ApplyCql2FilterMiddleware.py

@@ -21,8 +21,6 @@
     r"http://www.opengis.net/spec/cql2/1.0/conf/basic-cql2",
     r"http://www.opengis.net/spec/cql2/1.0/conf/cql2-text",
     r"http://www.opengis.net/spec/cql2/1.0/conf/cql2-json",
-    r"http://www.opengis.net/spec/ogcapi-features-3/1.0/conf/features-filter",
-    r"https://api.stacspec.org/v1\.\d+\.\d+(?:-[\w\.]+)?/item-search#filter",
 )
 @dataclass(frozen=True)
 class ApplyCql2FilterMiddleware:
@@ -31,6 +29,11 @@ class ApplyCql2FilterMiddleware:
     app: ASGIApp
     state_key: str = "cql2_filter"
 
+    single_record_endpoints = [
+        r"^/collections/([^/]+)/items/([^/]+)$",
+        r"^/collections/([^/]+)$",
+    ]
+
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         """Add the Cql2Filter to the request."""
         if scope["type"] != "http":
@@ -52,11 +55,9 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
             return await req_body_handler(scope, receive, send)
 
         # Handle single record requests (ie non-filterable endpoints)
-        single_record_endpoints = [
-            r"^/collections/([^/]+)/items/([^/]+)$",
-            r"^/collections/([^/]+)$",
-        ]
-        if any(re.match(expr, request.url.path) for expr in single_record_endpoints):
+        if any(
+            re.match(expr, request.url.path) for expr in self.single_record_endpoints
+        ):
             res_body_validator = Cql2ResponseBodyValidator(
                 app=self.app,
                 cql2_filter=cql2_filter,
@@ -130,18 +131,22 @@ async def __call__(self, scope: Scope, send: Send, receive: Receive) -> None:
         body = b""
         initial_message: Optional[Message] = None
 
-        async def _send_error_response(status: int, message: str) -> None:
+        async def _send_error_response(status: int, code: str, 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")
+            response_dict = {
+                "code": code,
+                "description": message,
+            }
+            response_bytes = json.dumps(response_dict).encode("utf-8")
             headers = MutableHeaders(scope=initial_message)
-            headers["content-length"] = str(len(error_body))
+            headers["content-length"] = str(len(response_bytes))
             initial_message["status"] = status
             await send(initial_message)
             await send(
                 {
                     "type": "http.response.body",
-                    "body": error_body,
+                    "body": response_bytes,
                     "more_body": False,
                 }
             )
@@ -150,40 +155,48 @@ async def buffered_send(message: Message) -> None:
             """Process a response message and apply filtering if needed."""
             nonlocal body
             nonlocal initial_message
+            initial_message = initial_message or message
+            # NOTE: to avoid data-leak, we process 404s so their responses are the same as rejected 200s
+            should_process = initial_message["status"] in [200, 404]
+
+            if not should_process:
+                return await send(message)
 
             if message["type"] == "http.response.start":
-                initial_message = message
+                # Hold off on sending response headers until we've validated the response body
                 return
 
-            assert initial_message, "Initial message not set"
-
             body += message["body"]
             if message.get("more_body"):
                 return
 
             try:
                 body_json = json.loads(body)
             except json.JSONDecodeError:
-                logger.warning("Failed to parse response body as JSON")
-                await _send_error_response(502, "Not found")
+                msg = "Failed to parse response body as JSON"
+                logger.warning(msg)
+                await _send_error_response(status=502, code="ParseError", message=msg)
                 return
 
-            logger.debug(
-                "Applying %s filter to %s", self.cql2_filter.to_text(), body_json
-            )
             try:
-                if self.cql2_filter.matches(body_json):
-                    await send(initial_message)
-                    return await send(
-                        {
-                            "type": "http.response.body",
-                            "body": json.dumps(body_json).encode("utf-8"),
-                            "more_body": False,
-                        }
-                    )
+                cql2_matches = self.cql2_filter.matches(body_json)
             except Exception as e:
+                cql2_matches = False
                 logger.warning("Failed to apply filter: %s", e)
 
-            return await _send_error_response(404, "Not found")
+            if cql2_matches:
+                logger.debug("Response matches filter, returning record")
+                await send(initial_message)
+                return await send(
+                    {
+                        "type": "http.response.body",
+                        "body": json.dumps(body_json).encode("utf-8"),
+                        "more_body": False,
+                    }
+                )
+            logger.debug("Response did not match filter, returning 404")
+            return await _send_error_response(
+                status=404, code="NotFoundError", message="Record not found."
+            )
 
         return await self.app(scope, receive, buffered_send)

src/stac_auth_proxy/middleware/BuildCql2FilterMiddleware.py

@@ -11,10 +11,16 @@
 from starlette.types import ASGIApp, Receive, Scope, Send
 
 from ..utils import requests
+from ..utils.middleware import required_conformance
 
 logger = logging.getLogger(__name__)
 
 
+@required_conformance(
+    "http://www.opengis.net/spec/cql2/1.0/conf/basic-cql2",
+    "http://www.opengis.net/spec/cql2/1.0/conf/cql2-text",
+    "http://www.opengis.net/spec/cql2/1.0/conf/cql2-json",
+)
 @dataclass(frozen=True)
 class BuildCql2FilterMiddleware:
     """Middleware to build the Cql2Filter."""
@@ -29,6 +35,34 @@ class BuildCql2FilterMiddleware:
     items_filter: Optional[Callable] = None
     items_filter_path: str = r"^(/collections/([^/]+)/items(/[^/]+)?$|/search$)"
 
+    def __post_init__(self):
+        """Set required conformances based on the filter functions."""
+        required_conformances = set()
+        if self.collections_filter:
+            logger.debug("Appending required conformance for collections filter")
+            required_conformances.update(
+                [
+                    "http://www.opengis.net/spec/ogcapi-features-3/1.0/conf/filter",
+                    "http://www.opengis.net/spec/cql2/1.0/conf/basic-cql2",
+                    r"https://api.stacspec.org/v1\.0\.0(?:-[\w\.]+)?/item-search#filter",
+                    "http://www.opengis.net/spec/ogcapi-features-3/1.0/conf/features-filter",
+                ]
+            )
+        if self.items_filter:
+            logger.debug("Appending required conformance for items filter")
+            required_conformances.update(
+                [
+                    "https://api.stacspec.org/v1.0.0/core",
+                    r"https://api.stacspec.org/v1\.0\.0(?:-[\w\.]+)?/collection-search#filter",
+                    "http://www.opengis.net/spec/ogcapi-common-2/1.0/conf/simple-query",
+                ]
+            )
+
+        # Must set required conformances on class
+        self.__class__.__required_conformances__ = required_conformances.union(
+            getattr(self.__class__, "__required_conformances__", [])
+        )
+
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         """Build the CQL2 filter, place on the request state."""
         if scope["type"] != "http":