pass in request rather than scope, use state_key

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, Scope
+from starlette.types import ASGIApp
 
 from ..config import EndpointMethods
 from ..utils.middleware import JsonResponseMiddleware
@@ -36,6 +36,8 @@ class AuthenticationExtensionMiddleware(JsonResponseMiddleware):
 
     json_content_type_expr: str = r"(application/json|geo\+json)"
 
+    state_key: str = "oidc_metadata"
+
     def should_transform_response(
         self, request: Request, response_headers: Headers
     ) -> bool:
@@ -55,9 +57,9 @@ def should_transform_response(
             ]
         )
 
-    def transform_json(self, doc: dict[str, Any], scope: Scope) -> dict[str, Any]:
+    def transform_json(self, data: dict[str, Any], request: Request) -> dict[str, Any]:
         """Augment the STAC Item with auth information."""
-        extensions = doc.setdefault("stac_extensions", [])
+        extensions = data.setdefault("stac_extensions", [])
         if self.extension_url not in extensions:
             extensions.append(self.extension_url)
 
@@ -70,30 +72,30 @@ def transform_json(self, doc: dict[str, Any], scope: Scope) -> dict[str, Any]:
         # - Collections
         # - Item Properties
 
-        if "oidc_metadata" not in scope:
+        if self.state_key not in request.state:
             logger.error(
                 "OIDC metadata not found in scope. "
                 "Skipping authentication extension."
             )
-            return doc
+            return data
 
-        scheme_loc = doc["properties"] if "properties" in doc else doc
+        scheme_loc = data["properties"] if "properties" in data else data
         schemes = scheme_loc.setdefault("auth:schemes", {})
         schemes[self.auth_scheme_name] = self.parse_oidc_config(
-            scope.get("oidc_metadata", {})
+            request.state.get(self.state_key, {})
         )
 
         # auth:refs
         # ---
         # Annotate links with "auth:refs": [auth_scheme]
         links = chain(
             # Item/Collection
-            doc.get("links", []),
+            data.get("links", []),
             # Collections/Items/Search
             (
                 link
                 for prop in ["features", "collections"]
-                for object_with_links in doc.get(prop, [])
+                for object_with_links in data.get(prop, [])
                 for link in object_with_links.get("links", [])
             ),
         )
@@ -111,7 +113,7 @@ def transform_json(self, doc: dict[str, Any], scope: Scope) -> dict[str, Any]:
             if match.is_private:
                 link.setdefault("auth:refs", []).append(self.auth_scheme_name)
 
-        return doc
+        return data
 
     def parse_oidc_config(self, oidc_config: dict[str, Any]) -> dict[str, Any]:
         """Parse the OIDC configuration."""

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, Scope
+from starlette.types import ASGIApp
 
 from ..config import EndpointMethods
 from ..utils.middleware import JsonResponseMiddleware
@@ -41,17 +41,15 @@ def should_transform_response(
             ]
         )
 
-    def transform_json(
-        self, openapi_spec: dict[str, Any], scope: Scope
-    ) -> dict[str, Any]:
+    def transform_json(self, data: dict[str, Any], request: Request) -> dict[str, Any]:
         """Augment the OpenAPI spec with auth information."""
-        components = openapi_spec.setdefault("components", {})
+        components = data.setdefault("components", {})
         securitySchemes = components.setdefault("securitySchemes", {})
         securitySchemes[self.oidc_auth_scheme_name] = {
             "type": "openIdConnect",
             "openIdConnectUrl": self.oidc_config_url,
         }
-        for path, method_config in openapi_spec["paths"].items():
+        for path, method_config in data["paths"].items():
             for method, config in method_config.items():
                 match = find_match(
                     path,
@@ -64,4 +62,4 @@ def transform_json(
                     config.setdefault("security", []).append(
                         {self.oidc_auth_scheme_name: match.required_scopes}
                     )
-        return openapi_spec
+        return data

src/stac_auth_proxy/utils/middleware.py

@@ -29,7 +29,7 @@ def should_transform_response(
         ...
 
     @abstractmethod
-    def transform_json(self, data: Any, scope: Scope) -> Any:
+    def transform_json(self, data: Any, request: Request) -> Any:
         """
         Transform the JSON data.
 
@@ -56,9 +56,10 @@ async def transform_response(message: Message) -> None:
 
             start_message = start_message or message
             headers = MutableHeaders(scope=start_message)
+            request = Request(scope)
 
             if not self.should_transform_response(
-                request=Request(scope),
+                request=request,
                 response_headers=headers,
             ):
                 # For non-JSON responses, send the start message immediately
@@ -78,7 +79,7 @@ async def transform_response(message: Message) -> None:
             # Transform the JSON body
             if body:
                 data = json.loads(body)
-                transformed = self.transform_json(data, scope=scope)
+                transformed = self.transform_json(data, request=request)
                 body = json.dumps(transformed).encode()
 
             # Update content-length header

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, Scope
+from starlette.types import ASGIApp
 
 from stac_auth_proxy.utils.middleware import JsonResponseMiddleware
 
@@ -24,7 +24,7 @@ def should_transform_response(
         """Transform JSON responses based on content type."""
         return response_headers.get("content-type", "") == "application/json"
 
-    def transform_json(self, data: Any, scope: Scope) -> Any:
+    def transform_json(self, data: Any, request: Request) -> Any:
         """Add a test field to the response."""
         if isinstance(data, dict):
             data["transformed"] = True