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)

tests/test_auth_extension.py

@@ -1,7 +1,6 @@
 """Tests for AuthenticationExtensionMiddleware."""
 
 import pytest
-from starlette.datastructures import Headers
 from starlette.requests import Request
 
 from stac_auth_proxy.config import EndpointMethods
@@ -34,10 +33,20 @@ def request_scope():
     }
 
 
-@pytest.fixture(params=["application/json", "application/geo+json"])
-def json_headers(request):
+@pytest.fixture(params=[b"application/json", b"application/geo+json"])
+def initial_message(request):
     """Create headers with JSON content type."""
-    return Headers({"content-type": request.param})
+    return {
+        "type": "http.response.start",
+        "status": 200,
+        "headers": [
+            (b"date", b"Mon, 07 Apr 2025 06:55:37 GMT"),
+            (b"server", b"uvicorn"),
+            (b"content-length", b"27642"),
+            (b"content-type", request.param),
+            (b"x-upstream-time", b"0.063"),
+        ],
+    }
 
 
 @pytest.fixture
@@ -50,7 +59,9 @@ def oidc_metadata():
     }
 
 
-def test_should_transform_response_valid_paths(middleware, request_scope, json_headers):
+def test_should_transform_response_valid_paths(
+    middleware, request_scope, initial_message
+):
     """Test that valid STAC paths are transformed."""
     valid_paths = [
         "/",
@@ -64,11 +75,11 @@ def test_should_transform_response_valid_paths(middleware, request_scope, json_h
     for path in valid_paths:
         request_scope["path"] = path
         request = Request(request_scope)
-        assert middleware.should_transform_response(request, json_headers)
+        assert middleware.should_transform_response(request, initial_message)
 
 
 def test_should_transform_response_invalid_paths(
-    middleware, request_scope, json_headers
+    middleware, request_scope, initial_message
 ):
     """Test that invalid paths are not transformed."""
     invalid_paths = [
@@ -80,14 +91,26 @@ def test_should_transform_response_invalid_paths(
     for path in invalid_paths:
         request_scope["path"] = path
         request = Request(request_scope)
-        assert not middleware.should_transform_response(request, json_headers)
+        assert not middleware.should_transform_response(request, initial_message)
 
 
 def test_should_transform_response_invalid_content_type(middleware, request_scope):
     """Test that non-JSON content types are not transformed."""
     request = Request(request_scope)
-    headers = Headers({"content-type": "text/html"})
-    assert not middleware.should_transform_response(request, headers)
+    assert not middleware.should_transform_response(
+        request,
+        {
+            "type": "http.response.start",
+            "status": 200,
+            "headers": [
+                (b"date", b"Mon, 07 Apr 2025 06:55:37 GMT"),
+                (b"server", b"uvicorn"),
+                (b"content-length", b"27642"),
+                (b"content-type", b"text/html"),
+                (b"x-upstream-time", b"0.063"),
+            ],
+        },
+    )
 
 
 def test_transform_json_catalog(middleware, request_scope, oidc_metadata):

tests/test_middleware.py

@@ -6,7 +6,7 @@
 from starlette.datastructures import Headers
 from starlette.requests import Request
 from starlette.testclient import TestClient
-from starlette.types import ASGIApp
+from starlette.types import ASGIApp, Scope
 
 from stac_auth_proxy.utils.middleware import JsonResponseMiddleware
 
@@ -18,11 +18,9 @@ def __init__(self, app: ASGIApp):
         """Initialize the middleware."""
         self.app = app
 
-    def should_transform_response(
-        self, request: Request, response_headers: Headers
-    ) -> bool:
+    def should_transform_response(self, request: Request, scope: Scope) -> bool:
         """Transform JSON responses based on content type."""
-        return response_headers.get("content-type", "") == "application/json"
+        return Headers(scope=scope).get("content-type", "") == "application/json"
 
     def transform_json(self, data: Any, request: Request) -> Any:
         """Add a test field to the response."""