refactor: decompress response within reverse-proxy tooling

Author: alukach

pyproject.toml

@@ -14,6 +14,7 @@ dependencies = [
   "jinja2>=3.1.4",
   "pydantic-settings>=2.6.1",
   "pyjwt>=2.10.1",
+  "starlette-cramjam>=0.4.0",
   "uvicorn>=0.32.1",
 ]
 description = "STAC authentication proxy with FastAPI"

src/stac_auth_proxy/app.py

@@ -9,6 +9,7 @@
 from typing import Optional
 
 from fastapi import FastAPI
+from starlette_cramjam.middleware import CompressionMiddleware
 
 from .config import Settings
 from .handlers import HealthzHandler, ReverseProxyHandler
@@ -56,7 +57,7 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
 
     app.add_api_route(
         "/{path:path}",
-        ReverseProxyHandler(upstream=str(settings.upstream_url)).stream,
+        ReverseProxyHandler(upstream=str(settings.upstream_url)).proxy_request,
         methods=["GET", "POST", "PUT", "PATCH", "DELETE"],
     )
 
@@ -90,6 +91,10 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
         oidc_config_url=settings.oidc_discovery_internal_url,
     )
 
+    app.add_middleware(
+        CompressionMiddleware,
+    )
+
     app.add_middleware(
         AddProcessTimeHeaderMiddleware,
     )

src/stac_auth_proxy/handlers/reverse_proxy.py

@@ -6,9 +6,8 @@
 
 import httpx
 from fastapi import Request
-from starlette.background import BackgroundTask
 from starlette.datastructures import MutableHeaders
-from starlette.responses import StreamingResponse
+from starlette.responses import Response
 
 logger = logging.getLogger(__name__)
 
@@ -28,7 +27,7 @@ def __post_init__(self):
             timeout=self.timeout,
         )
 
-    async def proxy_request(self, request: Request) -> httpx.Response:
+    async def proxy_request(self, request: Request) -> Response:
         """Proxy a request to the upstream STAC API."""
         headers = MutableHeaders(request.headers)
         headers.setdefault("X-Forwarded-For", request.client.host)
@@ -60,14 +59,15 @@ async def proxy_request(self, request: Request) -> httpx.Response:
             f"Received response status {rp_resp.status_code!r} from {rp_req.url} in {proxy_time:.3f}s"
         )
         rp_resp.headers["X-Upstream-Time"] = f"{proxy_time:.3f}"
-        return rp_resp
 
-    async def stream(self, request: Request) -> StreamingResponse:
-        """Transparently proxy a request to the upstream STAC API."""
-        rp_resp = await self.proxy_request(request)
-        return StreamingResponse(
-            rp_resp.aiter_raw(),
+        # We read the content here to make use of HTTPX's decompression, ensuring we have
+        # non-compressed content for the middleware to work with.
+        content = await rp_resp.aread()
+        if rp_resp.headers.get("Content-Encoding"):
+            del rp_resp.headers["Content-Encoding"]
+
+        return Response(
+            content=content,
             status_code=rp_resp.status_code,
-            headers=rp_resp.headers,
-            background=BackgroundTask(rp_resp.aclose),
+            headers=dict(rp_resp.headers),
         )

src/stac_auth_proxy/middleware/UpdateOpenApiMiddleware.py

@@ -1,25 +1,16 @@
 """Middleware to add auth information to the OpenAPI spec served by upstream API."""
 
-import gzip
 import json
-import zlib
 from dataclasses import dataclass
 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, find_match
 
-ENCODING_HANDLERS = {
-    "gzip": gzip,
-    "deflate": zlib,
-    "br": brotli,
-}
-
 
 @dataclass(frozen=True)
 class OpenApiMiddleware:
@@ -57,29 +48,15 @@ async def augment_oidc_spec(message: Message):
             body += message["body"]
 
             # Skip body chunks until all chunks have been received
-            if message["more_body"]:
+            if message.get("more_body"):
                 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"