diff --git a/src/debug_toolbar/litestar/middleware.py b/src/debug_toolbar/litestar/middleware.py index 4bf6d60..8b04efa 100644 --- a/src/debug_toolbar/litestar/middleware.py +++ b/src/debug_toolbar/litestar/middleware.py @@ -2,6 +2,7 @@ from __future__ import annotations +import contextlib import gzip import logging import re @@ -492,26 +493,21 @@ def _inject_toolbar(self, body: bytes, context: RequestContext, content_encoding If gzip was decompressed, returns uncompressed body with empty encoding. """ # Handle gzip-compressed responses - # Track whether we successfully decompressed the body - decompressed = False + # Store original body in case we need to return it + original_body = body encodings = [e.strip() for e in content_encoding.lower().split(",")] if content_encoding else [] if "gzip" in encodings: - try: + # Try to decompress, fall back to treating as uncompressed if invalid + with contextlib.suppress(gzip.BadGzipFile): body = gzip.decompress(body) - decompressed = True - except gzip.BadGzipFile: - # Not valid gzip, try to decode as-is - pass try: html = body.decode("utf-8") except UnicodeDecodeError: - # Can't decode. If we successfully decompressed gzip, return the - # decompressed body with no content-encoding. Otherwise, return - # the body as-is with the original encoding. - if decompressed: - return body, "" - return body, content_encoding + # Can't decode as UTF-8. Return the original body with original + # content-encoding to preserve HTTP semantics (client expects the + # format they requested via Accept-Encoding). + return original_body, content_encoding toolbar_data = self.toolbar.get_toolbar_data(context) toolbar_html = self._render_toolbar(toolbar_data) diff --git a/tests/integration/test_litestar_middleware.py b/tests/integration/test_litestar_middleware.py index 8a92396..4540bbd 100644 --- a/tests/integration/test_litestar_middleware.py +++ b/tests/integration/test_litestar_middleware.py @@ -5,11 +5,11 @@ import gzip import pytest +from litestar.status_codes import HTTP_200_OK from litestar.testing import TestClient from debug_toolbar.litestar import DebugToolbarPlugin, LitestarDebugToolbarConfig from litestar import Litestar, MediaType, Response, get -from litestar.status_codes import HTTP_200_OK @get("/", media_type=MediaType.HTML) @@ -277,7 +277,7 @@ def test_works_with_before_after_request(self) -> None: Note: We only verify before_request hook is called. The after_request hook timing varies in CI environments due to async execution order. """ - from litestar import Request, Response + from litestar import Request hook_state: dict[str, bool] = {"before": False, "after": False} @@ -378,7 +378,8 @@ def test_gzip_decompressed_data_fails_utf8_decoding(self) -> None: """Test handling of valid gzip data that fails UTF-8 decoding after decompression. When gzipped data decompresses successfully but contains non-UTF-8 bytes, - the middleware should return the original compressed data. + the middleware should return the original compressed data with the original + content-encoding header to preserve HTTP semantics. """ @get("/binary-gzip", media_type=MediaType.HTML) @@ -403,8 +404,9 @@ async def binary_gzip_handler() -> Response: with TestClient(app) as client: response = client.get("/binary-gzip") assert response.status_code == 200 - # Should return decompressed binary data since UTF-8 decode failed - # The middleware has removed the gzip encoding, so we check for the raw binary content + # Should return original compressed data since UTF-8 decode failed + # TestClient automatically decompresses gzip responses, so we verify + # the decompressed content matches the original binary data assert response.content == b"\x80\x81\x82\x83\x84\x85\x86\x87\x88\x89" def test_gzip_header_case_insensitive(self) -> None: