In progress

Author: alukach

src/stac_auth_proxy/middleware/UpdateOpenApiMiddleware.py

@@ -1,6 +1,9 @@
 """Middleware to add auth information to the OpenAPI spec served by upstream API."""
 
+import brotli
+import gzip
 import json
+import zlib
 from dataclasses import dataclass
 from typing import Any
 
@@ -43,6 +46,21 @@ async def augment_oidc_spec(message: Message):
             if message["more_body"]:
                 return await send({**message, "body": b""})
 
+            # Handle compressed responses
+            # content_encoding = (
+            #     message.get("headers", {})
+            #     .get(b"content-encoding", b"")
+            #     .decode()
+            #     .lower()
+            # )
+            # if content_encoding:
+            #     if "gzip" in content_encoding:
+            #         total_body = gzip.decompress(total_body)
+            #     elif "deflate" in content_encoding:
+            #         total_body = zlib.decompress(total_body)
+            #     elif "br" in content_encoding:
+            #         total_body = brotli.decompress(total_body)
+            # print(f"{message=}")
             await send(
                 {
                     "type": "http.response.body",

tests/conftest.py

@@ -8,7 +8,7 @@
 
 import pytest
 import uvicorn
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from jwcrypto import jwk, jwt
 from utils import single_chunk_async_stream_response
 
@@ -64,61 +64,52 @@ def build_token(payload: dict[str, Any], key=None) -> str:
     return build_token
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture
 def source_api():
     """Create upstream API for testing purposes."""
     app = FastAPI(docs_url="/api.html", openapi_url="/api")
 
+    def default_response_factory(request: Request):
+        """Default response factory."""
+        return {"id": f"Response from {request.method}@{request.url.path}"}
+
+    # Store response factory in app state
+    app.state.response_factory = default_response_factory
+
     for path, methods in {
-        "/": [
-            "GET",
-        ],
-        "/conformance": [
-            "GET",
-        ],
-        "/queryables": [
-            "GET",
-        ],
-        "/search": [
-            "GET",
-            "POST",
-        ],
-        "/collections": [
-            "GET",
-            "POST",
-        ],
-        "/collections/{collection_id}": [
-            "GET",
-            "PUT",
-            "PATCH",
-            "DELETE",
-        ],
-        "/collections/{collection_id}/items": [
-            "GET",
-            "POST",
-        ],
+        "/": ["GET"],
+        "/conformance": ["GET"],
+        "/queryables": ["GET"],
+        "/search": ["GET", "POST"],
+        "/collections": ["GET", "POST"],
+        "/collections/{collection_id}": ["GET", "PUT", "PATCH", "DELETE"],
+        "/collections/{collection_id}/items": ["GET", "POST"],
         "/collections/{collection_id}/items/{item_id}": [
             "GET",
             "PUT",
             "PATCH",
             "DELETE",
         ],
-        "/collections/{collection_id}/bulk_items": [
-            "POST",
-        ],
+        "/collections/{collection_id}/bulk_items": ["POST"],
     }.items():
         for method in methods:
-            # NOTE: declare routes per method separately to avoid warning of "Duplicate Operation ID ... for function <lambda>"
+
+            def endpoint(request: Request):
+                print(f"endpoint: {request}")
+                return app.state.response_factory(request)
+
+            # NOTE: declare routes per method separately to avoid warning of "Duplicate Operation ID"
             app.add_api_route(
                 path,
-                lambda: {"id": f"Response from {method}@{path}"},
+                endpoint,
+                # lambda: {"id": f"Response from {method}@{path}"},
                 methods=[method],
             )
 
     return app
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture
 def source_api_server(source_api):
     """Run the source API in a background thread."""
     host, port = "127.0.0.1", 9119

tests/test_openapi.py

@@ -3,6 +3,8 @@
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 from utils import AppFactory
+import gzip
+from fastapi import Request, Response
 
 app_factory = AppFactory(
     oidc_discovery_url="https://example-stac-api.com/.well-known/openid-configuration"
@@ -59,6 +61,45 @@ def test_oidc_in_openapi_spec(source_api: FastAPI, source_api_server: str):
     assert "oidcAuth" in openapi.get("components", {}).get("securitySchemes", {})
 
 
+def test_oidc_in_openapi_spec_compressed(source_api: FastAPI, source_api_server: str):
+    """When OpenAPI spec endpoint is set, the proxied OpenAPI spec is augmented with oidc details."""
+
+    # Create a compressed response factory
+    def compressed_response_factory(request: Request):
+        assert False
+        # Get the original OpenAPI spec
+        openapi = source_api.openapi()
+        compressed_data = gzip.compress(str(openapi).encode())
+        return Response(
+            content=compressed_data,
+            headers={
+                "Content-Encoding": "gzip",
+                "Content-Type": "application/json",
+            },
+        )
+
+    # Set the custom response factory
+    source_api.state.response_factory = compressed_response_factory
+
+    app = app_factory(
+        upstream_url=source_api_server,
+        openapi_spec_endpoint=source_api.openapi_url,
+    )
+    client = TestClient(app)
+    response = client.get(
+        source_api.openapi_url,
+        headers={"Accept-Encoding": "gzip"},
+    )
+    assert response.status_code == 200
+    assert response.headers.get("content-encoding") == "gzip"
+
+    openapi = response.json()  # TestClient automatically decompresses
+    assert "info" in openapi
+    assert "openapi" in openapi
+    assert "paths" in openapi
+    assert "oidcAuth" in openapi.get("components", {}).get("securitySchemes", {})
+
+
 def test_oidc_in_openapi_spec_private_endpoints(
     source_api: FastAPI, source_api_server: str
 ):