Support compressed bodies

Author: alukach

src/stac_auth_proxy/middleware/UpdateOpenApiMiddleware.py

@@ -7,13 +7,21 @@
 from dataclasses import dataclass
 from typing import Any
 
+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:
     """Middleware to add the OpenAPI spec to the response."""
@@ -30,41 +38,56 @@ 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 = 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
+            body += message["body"]
 
-            nonlocal total_body
+            # Skip body chunks until all chunks have been received
+            if message["more_body"]:
+                return
 
-            total_body += message["body"]
+            # 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)
 
-            # Pass empty body chunks until all chunks have been received
-            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=}")
+            # 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))
+            start_message["headers"] = 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,
                 }
             )