fix: handle compressed OpenAPI responses & ensure paths correctly marked private (#29)

alukach · web-flow · commit 2fe0852c18cb · 2025-03-12T17:17:32.000-07:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -46,8 +46,9 @@ dev = [
   "pytest-asyncio>=0.25.1",
   "pytest-cov>=5.0.0",
   "pytest>=8.3.3",
+  "starlette-cramjam>=0.4.0",
 ]
 
 [tool.pytest.ini_options]
 asyncio_default_fixture_loop_scope = "function"
-asyncio_mode = "auto"
+asyncio_mode = "auto"
diff --git a/src/stac_auth_proxy/app.py b/src/stac_auth_proxy/app.py
@@ -38,6 +38,7 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
             OpenApiMiddleware,
             openapi_spec_path=settings.openapi_spec_endpoint,
             oidc_config_url=str(settings.oidc_discovery_url),
+            public_endpoints=settings.public_endpoints,
             private_endpoints=settings.private_endpoints,
             default_public=settings.default_public,
         )
diff --git a/src/stac_auth_proxy/middleware/UpdateOpenApiMiddleware.py b/src/stac_auth_proxy/middleware/UpdateOpenApiMiddleware.py
@@ -1,15 +1,26 @@
 """Middleware to add auth information to the OpenAPI spec served by upstream API."""
 
+import gzip
 import json
+import re
+import zlib
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Optional
 
+import brotli
+from starlette.datastructures import MutableHeaders
 from starlette.requests import Request
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
 from ..config import EndpointMethods
 from ..utils.requests import dict_to_bytes
 
+ENCODING_HANDLERS = {
+    "gzip": gzip,
+    "deflate": zlib,
+    "br": brotli,
+}
+
 
 @dataclass(frozen=True)
 class OpenApiMiddleware:
@@ -19,6 +30,7 @@ class OpenApiMiddleware:
     openapi_spec_path: str
     oidc_config_url: str
     private_endpoints: EndpointMethods
+    public_endpoints: EndpointMethods
     default_public: bool
     oidc_auth_scheme_name: str = "oidcAuth"
 
@@ -27,26 +39,63 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if scope["type"] != "http" or Request(scope).url.path != self.openapi_spec_path:
             return await self.app(scope, receive, send)
 
-        total_body = b""
+        start_message: Optional[Message] = None
+        body = b""
 
         async def augment_oidc_spec(message: Message):
-            if message["type"] != "http.response.body":
+            nonlocal start_message
+            nonlocal body
+            if message["type"] == "http.response.start":
+                # NOTE: Because we are modifying the response body, we will need to update
+                # the content-length header. However, headers are sent before we see the
+                # body. To handle this, we delay sending the http.response.start message
+                # until after we alter the body.
+                start_message = message
+                return
+            elif message["type"] != "http.response.body":
                 return await send(message)
 
-            # TODO: Make more robust to handle non-JSON responses
-
-            nonlocal total_body
-
-            total_body += message["body"]
+            body += message["body"]
 
-            # Pass empty body chunks until all chunks have been received
+            # Skip body chunks until all chunks have been received
             if message["more_body"]:
-                return await send({**message, "body": b""})
-
+                return
+
+            # Maybe decompress the body
+            headers = MutableHeaders(scope=start_message)
+            content_encoding = headers.get("content-encoding", "").lower()
+            handler = None
+            if content_encoding:
+                handler = ENCODING_HANDLERS.get(content_encoding)
+                assert handler, f"Unsupported content encoding: {content_encoding}"
+                body = (
+                    handler.decompress(body)
+                    if content_encoding != "deflate"
+                    else handler.decompress(body, -zlib.MAX_WBITS)
+                )
+
+            # Augment the spec
+            body = dict_to_bytes(self.augment_spec(json.loads(body)))
+
+            # Maybe re-compress the body
+            if handler:
+                body = handler.compress(body)
+
+            # Update the content-length header
+            headers["content-length"] = str(len(body))
+            assert start_message, "Expected start_message to be set"
+            start_message["headers"] = [
+                (key.encode(), value.encode()) for key, value in headers.items()
+            ]
+
+            # Send http.response.start
+            await send(start_message)
+
+            # Send http.response.body
             await send(
                 {
                     "type": "http.response.body",
-                    "body": dict_to_bytes(self.augment_spec(json.loads(total_body))),
+                    "body": body,
                     "more_body": False,
                 }
             )
@@ -63,9 +112,24 @@ def augment_spec(self, openapi_spec) -> dict[str, Any]:
         }
         for path, method_config in openapi_spec["paths"].items():
             for method, config in method_config.items():
-                for private_method in self.private_endpoints.get(path, []):
-                    if method.casefold() == private_method.casefold():
-                        config.setdefault("security", []).append(
-                            {self.oidc_auth_scheme_name: []}
-                        )
+                requires_auth = (
+                    self.path_matches(path, method, self.private_endpoints)
+                    if self.default_public
+                    else not self.path_matches(path, method, self.public_endpoints)
+                )
+                if requires_auth:
+                    config.setdefault("security", []).append(
+                        {self.oidc_auth_scheme_name: []}
+                    )
         return openapi_spec
+
+    @staticmethod
+    def path_matches(path: str, method: str, endpoints: EndpointMethods) -> bool:
+        """Check if the given path and method match any of the regex patterns and methods in the endpoints."""
+        for pattern, endpoint_methods in endpoints.items():
+            if not re.match(pattern, path):
+                continue
+            for endpoint_method in endpoint_methods:
+                if method.casefold() == endpoint_method.casefold():
+                    return True
+        return False
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -10,6 +10,7 @@
 import uvicorn
 from fastapi import FastAPI
 from jwcrypto import jwk, jwt
+from starlette_cramjam.middleware import CompressionMiddleware
 from utils import single_chunk_async_stream_response
 
 
@@ -69,6 +70,8 @@ def source_api():
     """Create upstream API for testing purposes."""
     app = FastAPI(docs_url="/api.html", openapi_url="/api")
 
+    app.add_middleware(CompressionMiddleware, minimum_size=0, compression_level=1)
+
     for path, methods in {
         "/": [
             "GET",
@@ -118,7 +121,7 @@ def source_api():
     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
diff --git a/tests/test_openapi.py b/tests/test_openapi.py
@@ -1,5 +1,6 @@
 """Tests for OpenAPI spec handling."""
 
+import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 from utils import AppFactory
@@ -40,7 +41,6 @@ def test_no_private_endpoints(source_api_server: str):
     assert "info" in openapi
     assert "openapi" in openapi
     assert "paths" in openapi
-    # assert "oidcAuth" not in openapi.get("components", {}).get("securitySchemes", {})
 
 
 def test_oidc_in_openapi_spec(source_api: FastAPI, source_api_server: str):
@@ -59,37 +59,95 @@ def test_oidc_in_openapi_spec(source_api: FastAPI, source_api_server: str):
     assert "oidcAuth" in openapi.get("components", {}).get("securitySchemes", {})
 
 
+@pytest.mark.parametrize("compression_type", ["gzip", "br", "deflate"])
+def test_oidc_in_openapi_spec_compressed(
+    source_api: FastAPI, source_api_server: str, compression_type: str
+):
+    """When OpenAPI spec endpoint is set, the proxied OpenAPI spec is augmented with oidc details."""
+    app = app_factory(
+        upstream_url=source_api_server,
+        openapi_spec_endpoint=source_api.openapi_url,
+    )
+    client = TestClient(app)
+
+    # Test with gzip acceptance
+    response = client.get(
+        source_api.openapi_url, headers={"Accept-Encoding": compression_type}
+    )
+    assert response.status_code == 200
+    assert response.headers.get("content-encoding") == compression_type
+    assert response.headers.get("content-type") == "application/json"
+    assert response.json()
+
+
 def test_oidc_in_openapi_spec_private_endpoints(
     source_api: FastAPI, source_api_server: str
 ):
     """When OpenAPI spec endpoint is set & endpoints are marked private, those endpoints are marked private in the spec."""
     private_endpoints = {
         # https://github.com/stac-api-extensions/collection-transaction/blob/v1.0.0-beta.1/README.md#methods
+        r"^/collections$": ["POST"],
+        r"^/collections/([^/]+)$": ["PUT", "PATCH", "DELETE"],
+        # https://github.com/stac-api-extensions/transaction/blob/v1.0.0-rc.3/README.md#methods
+        r"^/collections/([^/]+)/items$": ["POST"],
+        r"^/collections/([^/]+)/items/([^/]+)$": ["PUT", "PATCH", "DELETE"],
+        # https://stac-utils.github.io/stac-fastapi/api/stac_fastapi/extensions/third_party/bulk_transactions/#bulktransactionextension
+        r"^/collections/([^/]+)/bulk_items$": ["POST"],
+    }
+    app = app_factory(
+        upstream_url=source_api_server,
+        openapi_spec_endpoint=source_api.openapi_url,
+        default_public=True,
+        private_endpoints=private_endpoints,
+    )
+    client = TestClient(app)
+
+    openapi = client.get(source_api.openapi_url).raise_for_status().json()
+
+    expected_auth = {
         "/collections": ["POST"],
         "/collections/{collection_id}": ["PUT", "PATCH", "DELETE"],
-        # https://github.com/stac-api-extensions/transaction/blob/v1.0.0-rc.3/README.md#methods
         "/collections/{collection_id}/items": ["POST"],
         "/collections/{collection_id}/items/{item_id}": ["PUT", "PATCH", "DELETE"],
-        # https://stac-utils.github.io/stac-fastapi/api/stac_fastapi/extensions/third_party/bulk_transactions/#bulktransactionextension
         "/collections/{collection_id}/bulk_items": ["POST"],
     }
+    for path, method_config in openapi["paths"].items():
+        for method, config in method_config.items():
+            security = config.get("security")
+            path_in_expected_auth = path in expected_auth
+            method_in_expected_auth = any(
+                method.casefold() == m.casefold() for m in expected_auth.get(path, [])
+            )
+            if security:
+                assert path_in_expected_auth
+                assert method_in_expected_auth
+            else:
+                assert not all([path_in_expected_auth, method_in_expected_auth])
+
+
+def test_oidc_in_openapi_spec_public_endpoints(
+    source_api: FastAPI, source_api_server: str
+):
+    """When OpenAPI spec endpoint is set & endpoints are marked public, those endpoints are not marked private in the spec."""
+    public = {r"^/queryables$": ["GET"], r"^/api": ["GET"]}
     app = app_factory(
         upstream_url=source_api_server,
         openapi_spec_endpoint=source_api.openapi_url,
-        private_endpoints=private_endpoints,
+        default_public=False,
+        public_endpoints=public,
     )
     client = TestClient(app)
+
     openapi = client.get(source_api.openapi_url).raise_for_status().json()
-    for path, methods in private_endpoints.items():
-        for method in methods:
-            openapi_path = openapi["paths"].get(path)
-            assert openapi_path, f"Path {path} not found in OpenAPI spec"
-            openapi_path_method = openapi_path.get(method.lower())
-            assert (
-                openapi_path_method
-            ), f"Method {method.lower()!r} not found for path {path!r} in OpenAPI spec for path {path}"
-            security = openapi_path_method.get("security")
-            assert security, f"Security not found for {path!r} {method.lower()!r}"
-            assert any(
-                "oidcAuth" in s for s in security
-            ), f'No "oidcAuth" in security for {path!r} {method.lower()!r}'
+
+    expected_auth = {"/queryables": ["GET"]}
+    for path, method_config in openapi["paths"].items():
+        for method, config in method_config.items():
+            security = config.get("security")
+            if security:
+                assert path not in expected_auth
+            else:
+                assert path in expected_auth
+                assert any(
+                    method.casefold() == m.casefold() for m in expected_auth[path]
+                )
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -38,6 +38,7 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:`
`38`	`38`	`OpenApiMiddleware,`
`39`	`39`	`openapi_spec_path=settings.openapi_spec_endpoint,`
`40`	`40`	`oidc_config_url=str(settings.oidc_discovery_url),`
	`41`	`+ public_endpoints=settings.public_endpoints,`
`41`	`42`	`private_endpoints=settings.private_endpoints,`
`42`	`43`	`default_public=settings.default_public,`
`43`	`44`	`)`