Litestar Brotli Decompression Middleware v1

decompression.py

from typing import ClassVar, cast

import brotli
from litestar.datastructures import Headers, MutableScopeHeaders
from litestar.enums import CompressionEncoding
from litestar.middleware import AbstractMiddleware
from litestar.types import HTTPRequestEvent, Receive, ReceiveMessage, Scope, Send


class BrotliDecompressionMiddleware(AbstractMiddleware):
    """Middleware to decompress incoming brotli encoded requests.

    The available standard litestar middleware only provides brotli/gzip *response* compression. Whilst any framework
    agnostic ASGI middleware could implement decompression, there appears to be no decompression middleware for any ASGI
    framework available currently.
    """

    # class var compiled to regex to match paths that will NOT use this BrotliDecompressionMiddleware
    # exclude = ...

    # route decorators can use this kwarg to exclude themselves from this middleware
    # e.g. @get("/foo", no_decompression=True) ... a very minor performance benefit
    exclude_opt_key = "no_decompression"

    async def __call__(
        self,
        scope: "Scope",
        receive: "Receive",
        send: "Send",
    ) -> None:
        self.scope = scope
        content_encoding = Headers.from_scope(scope).get("content-encoding", "").lower().strip()
        if content_encoding == "br":
            receive_wrapper = self.create_br_decompression_receive_wrapper(receive=receive)

            # app is an asgi app, next in the chain, not the LiteStar
            await self.app(scope, receive_wrapper, send)
            return

        await self.app(scope, receive, send)
        return None

    def create_br_decompression_receive_wrapper(self, receive: Receive) -> Receive:
        brotli_decompressor = brotli.Decompressor()

        async def receive_wrapper() -> ReceiveMessage:
            nonlocal brotli_decompressor

            message = await receive()

            if message["type"] == "http.request":
                http_req_event = cast("HTTPRequestEvent", message)  # TypedDict
                if encoded_body := http_req_event["body"]:
                    decoded_body = brotli_decompressor.process(encoded_body)
                    http_req_event["body"] = decoded_body

                if not http_req_event["more_body"]:
                    # finished, update content length header
                    headers = MutableScopeHeaders.from_message(message=self.scope)
                    headers["content-length"] = str(len(http_req_event["body"]))
                    self.scope["headers"] = headers.headers  # RawHeadersList
                    assert brotli_decompressor.is_finished()

            return message

        return receive_wrapper