fix: only transform non-errors

Author: alukach

src/stac_auth_proxy/middleware/AuthenticationExtensionMiddleware.py

@@ -9,7 +9,7 @@
 
 from starlette.datastructures import Headers
 from starlette.requests import Request
-from starlette.types import ASGIApp
+from starlette.types import ASGIApp, Scope
 
 from ..config import EndpointMethods
 from ..utils.middleware import JsonResponseMiddleware
@@ -38,24 +38,27 @@ class AuthenticationExtensionMiddleware(JsonResponseMiddleware):
 
     state_key: str = "oidc_metadata"
 
-    def should_transform_response(
-        self, request: Request, response_headers: Headers
-    ) -> bool:
+    def should_transform_response(self, request: Request, scope: Scope) -> bool:
         """Determine if the response should be transformed."""
         # Match STAC catalog, collection, or item URLs with a single regex
-        return all(
-            re.match(expr, val)
-            for expr, val in [
+        return (
+            all(
                 (
-                    # catalog, collections, collection, items, item, search
-                    r"^(/|/collections(/[^/]+(/items(/[^/]+)?)?)?|/search)$",
-                    request.url.path,
+                    re.match(expr, val)
+                    for expr, val in [
+                        (
+                            # catalog, collections, collection, items, item, search
+                            r"^(/|/collections(/[^/]+(/items(/[^/]+)?)?)?|/search)$",
+                            request.url.path,
+                        ),
+                        (
+                            self.json_content_type_expr,
+                            Headers(scope=scope).get("content-type", ""),
+                        ),
+                    ]
                 ),
-                (
-                    self.json_content_type_expr,
-                    response_headers.get("content-type", ""),
-                ),
-            ]
+            )
+            and 200 >= scope["status"] < 300
         )
 
     def transform_json(self, data: dict[str, Any], request: Request) -> dict[str, Any]:

src/stac_auth_proxy/middleware/UpdateOpenApiMiddleware.py

@@ -6,7 +6,7 @@
 
 from starlette.datastructures import Headers
 from starlette.requests import Request
-from starlette.types import ASGIApp
+from starlette.types import ASGIApp, Scope
 
 from ..config import EndpointMethods
 from ..utils.middleware import JsonResponseMiddleware
@@ -27,19 +27,20 @@ class OpenApiMiddleware(JsonResponseMiddleware):
 
     json_content_type_expr: str = r"application/(vnd\.oai\.openapi\+json?|json)"
 
-    def should_transform_response(
-        self, request: Request, response_headers: Headers
-    ) -> bool:
+    def should_transform_response(self, request: Request, scope: Scope) -> bool:
         """Only transform responses for the OpenAPI spec path."""
-        return all(
-            re.match(expr, val)
-            for expr, val in [
-                (self.openapi_spec_path, request.url.path),
-                (
-                    self.json_content_type_expr,
-                    response_headers.get("content-type", ""),
-                ),
-            ]
+        return (
+            all(
+                re.match(expr, val)
+                for expr, val in [
+                    (self.openapi_spec_path, request.url.path),
+                    (
+                        self.json_content_type_expr,
+                        Headers(scope=scope).get("content-type", ""),
+                    ),
+                ]
+            )
+            and 200 >= scope["status"] < 300
         )
 
     def transform_json(self, data: dict[str, Any], request: Request) -> dict[str, Any]:

src/stac_auth_proxy/utils/middleware.py

@@ -4,7 +4,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, Optional
 
-from starlette.datastructures import Headers, MutableHeaders
+from starlette.datastructures import MutableHeaders
 from starlette.requests import Request
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
@@ -16,7 +16,7 @@ class JsonResponseMiddleware(ABC):
 
     @abstractmethod
     def should_transform_response(
-        self, request: Request, response_headers: Headers
+        self, request: Request, scope: Scope
     ) -> bool:  # mypy: ignore
         """
         Determine if this response should be transformed. At a minimum, this
@@ -60,7 +60,7 @@ async def transform_response(message: Message) -> None:
 
             if not self.should_transform_response(
                 request=request,
-                response_headers=headers,
+                scope=start_message,
             ):
                 # For non-JSON responses, send the start message immediately
                 await send(message)