refactor: simplify middleware util logic

Author: alukach

src/stac_auth_proxy/middleware/UpdateOpenApiMiddleware.py

@@ -1,8 +1,10 @@
 """Middleware to add auth information to the OpenAPI spec served by upstream API."""
 
+import re
 from dataclasses import dataclass
 from typing import Any
 
+from starlette.datastructures import Headers
 from starlette.requests import Request
 from starlette.types import ASGIApp
 
@@ -23,9 +25,21 @@ class OpenApiMiddleware(JsonResponseMiddleware):
     default_public: bool
     oidc_auth_scheme_name: str = "oidcAuth"
 
-    def should_transform_response(self, request: Request) -> bool:
+    json_content_type_expr: str = r"application/(vnd\.oai\.openapi\+json?|json)"
+
+    def should_transform_response(
+        self, request: Request, response_headers: Headers
+    ) -> bool:
         """Only transform responses for the OpenAPI spec path."""
-        return request.url.path == self.openapi_spec_path
+        return all(
+            [
+                re.match(self.openapi_spec_path, request.url.path),
+                re.match(
+                    self.json_content_type_expr,
+                    response_headers.get("content-type", ""),
+                ),
+            ]
+        )
 
     def transform_json(self, openapi_spec: dict[str, Any]) -> dict[str, Any]:
         """Augment the OpenAPI spec with auth information."""

src/stac_auth_proxy/utils/middleware.py

@@ -1,7 +1,6 @@
 """Utilities for middleware response handling."""
 
 import json
-import re
 from abc import ABC, abstractmethod
 from typing import Any, Optional
 
@@ -14,25 +13,20 @@ class JsonResponseMiddleware(ABC):
     """Base class for middleware that transforms JSON response bodies."""
 
     app: ASGIApp
-    json_content_type_expr: str = (
-        r"application/vnd\.oai\.openapi\+json;.*|application/json|application/geo\+json"
-    )
 
     @abstractmethod
-    def should_transform_response(self, request: Request) -> bool:
+    def should_transform_response(
+        self, request: Request, response_headers: Headers
+    ) -> bool:  # mypy: ignore
         """
-        Determine if this request's response should be transformed.
-
-        Args:
-            request: The incoming request
+        Determine if this response should be transformed. At a minimum, this
+        should check the request's path and content type.
 
         Returns
         -------
             bool: True if the response should be transformed
         """
-        return bool(
-            re.match(self.json_content_type_expr, request.headers.get("accept", ""))
-        )
+        ...
 
     @abstractmethod
     def transform_json(self, data: Any) -> Any:
@@ -46,36 +40,31 @@ def transform_json(self, data: Any) -> Any:
         -------
             The transformed JSON data
         """
-        pass
+        ...
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         """Process the request/response."""
         if scope["type"] != "http":
             return await self.app(scope, receive, send)
 
-        request = Request(scope)
-        if not self.should_transform_response(request):
-            return await self.app(scope, receive, send)
-
         start_message: Optional[Message] = None
         body = b""
-        not_json = False
 
-        async def process_message(message: Message) -> None:
+        async def transform_response(message: Message) -> None:
             nonlocal start_message
             nonlocal body
-            nonlocal not_json
+
             if message["type"] == "http.response.start":
                 # Delay sending start message until we've processed the body
-                if not re.match(
-                    self.json_content_type_expr,
-                    Headers(scope=message).get("content-type", ""),
-                ):
-                    not_json = True
-                    return await send(message)
                 start_message = message
                 return
-            elif message["type"] != "http.response.body" or not_json:
+            assert start_message is not None
+            if not self.should_transform_response(
+                request=Request(scope),
+                response_headers=Headers(scope=start_message),
+            ):
+                return await send(message)
+            if message["type"] != "http.response.body":
                 return await send(message)
 
             body += message["body"]
@@ -109,4 +98,4 @@ async def process_message(message: Message) -> None:
                 }
             )
 
-        return await self.app(scope, receive, process_message)
+        return await self.app(scope, receive, transform_response)