Use decompressor max_length parameter (#11898) (#11918)

(cherry picked from commit 92477c5)

---------

Co-authored-by: J. Nick Koston <[email protected]>

Author: Dreamsorcerer

CHANGES/11898.breaking.rst

@@ -0,0 +1,2 @@
+``Brotli`` and ``brotlicffi`` minimum version is now 1.2.
+Decompression now has a default maximum output size of 32MiB per decompress call -- by :user:`Dreamsorcerer`.

aiohttp/compression_utils.py

@@ -1,6 +1,7 @@
 import asyncio
 import sys
 import zlib
+from abc import ABC, abstractmethod
 from concurrent.futures import Executor
 from typing import Any, Final, Optional, Protocol, TypedDict, cast
 
@@ -32,7 +33,12 @@
     HAS_ZSTD = False
 
 
-MAX_SYNC_CHUNK_SIZE = 1024
+MAX_SYNC_CHUNK_SIZE = 4096
+DEFAULT_MAX_DECOMPRESS_SIZE = 2**25  # 32MiB
+
+# Unlimited decompression constants - different libraries use different conventions
+ZLIB_MAX_LENGTH_UNLIMITED = 0  # zlib uses 0 to mean unlimited
+ZSTD_MAX_LENGTH_UNLIMITED = -1  # zstd uses -1 to mean unlimited
 
 
 class ZLibCompressObjProtocol(Protocol):
@@ -144,19 +150,37 @@ def encoding_to_mode(
     return -ZLibBackend.MAX_WBITS if suppress_deflate_header else ZLibBackend.MAX_WBITS
 
 
-class ZlibBaseHandler:
+class DecompressionBaseHandler(ABC):
     def __init__(
         self,
-        mode: int,
         executor: Optional[Executor] = None,
         max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE,
     ):
-        self._mode = mode
+        """Base class for decompression handlers."""
         self._executor = executor
         self._max_sync_chunk_size = max_sync_chunk_size
 
+    @abstractmethod
+    def decompress_sync(
+        self, data: bytes, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED
+    ) -> bytes:
+        """Decompress the given data."""
+
+    async def decompress(
+        self, data: bytes, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED
+    ) -> bytes:
+        """Decompress the given data."""
+        if (
+            self._max_sync_chunk_size is not None
+            and len(data) > self._max_sync_chunk_size
+        ):
+            return await asyncio.get_event_loop().run_in_executor(
+                self._executor, self.decompress_sync, data, max_length
+            )
+        return self.decompress_sync(data, max_length)
+
 
-class ZLibCompressor(ZlibBaseHandler):
+class ZLibCompressor:
     def __init__(
         self,
         encoding: Optional[str] = None,
@@ -167,14 +191,12 @@ def __init__(
         executor: Optional[Executor] = None,
         max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE,
     ):
-        super().__init__(
-            mode=(
-                encoding_to_mode(encoding, suppress_deflate_header)
-                if wbits is None
-                else wbits
-            ),
-            executor=executor,
-            max_sync_chunk_size=max_sync_chunk_size,
+        self._executor = executor
+        self._max_sync_chunk_size = max_sync_chunk_size
+        self._mode = (
+            encoding_to_mode(encoding, suppress_deflate_header)
+            if wbits is None
+            else wbits
         )
         self._zlib_backend: Final = ZLibBackendWrapper(ZLibBackend._zlib_backend)
 
@@ -233,41 +255,24 @@ def flush(self, mode: Optional[int] = None) -> bytes:
         )
 
 
-class ZLibDecompressor(ZlibBaseHandler):
+class ZLibDecompressor(DecompressionBaseHandler):
     def __init__(
         self,
         encoding: Optional[str] = None,
         suppress_deflate_header: bool = False,
         executor: Optional[Executor] = None,
         max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE,
     ):
-        super().__init__(
-            mode=encoding_to_mode(encoding, suppress_deflate_header),
-            executor=executor,
-            max_sync_chunk_size=max_sync_chunk_size,
-        )
+        super().__init__(executor=executor, max_sync_chunk_size=max_sync_chunk_size)
+        self._mode = encoding_to_mode(encoding, suppress_deflate_header)
         self._zlib_backend: Final = ZLibBackendWrapper(ZLibBackend._zlib_backend)
         self._decompressor = self._zlib_backend.decompressobj(wbits=self._mode)
 
-    def decompress_sync(self, data: bytes, max_length: int = 0) -> bytes:
+    def decompress_sync(
+        self, data: Buffer, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED
+    ) -> bytes:
         return self._decompressor.decompress(data, max_length)
 
-    async def decompress(self, data: bytes, max_length: int = 0) -> bytes:
-        """Decompress the data and return the decompressed bytes.
-
-        If the data size is large than the max_sync_chunk_size, the decompression
-        will be done in the executor. Otherwise, the decompression will be done
-        in the event loop.
-        """
-        if (
-            self._max_sync_chunk_size is not None
-            and len(data) > self._max_sync_chunk_size
-        ):
-            return await asyncio.get_running_loop().run_in_executor(
-                self._executor, self._decompressor.decompress, data, max_length
-            )
-        return self.decompress_sync(data, max_length)
-
     def flush(self, length: int = 0) -> bytes:
         return (
             self._decompressor.flush(length)
@@ -280,40 +285,64 @@ def eof(self) -> bool:
         return self._decompressor.eof
 
 
-class BrotliDecompressor:
+class BrotliDecompressor(DecompressionBaseHandler):
     # Supports both 'brotlipy' and 'Brotli' packages
     # since they share an import name. The top branches
     # are for 'brotlipy' and bottom branches for 'Brotli'
-    def __init__(self) -> None:
+    def __init__(
+        self,
+        executor: Optional[Executor] = None,
+        max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE,
+    ) -> None:
+        """Decompress data using the Brotli library."""
         if not HAS_BROTLI:
             raise RuntimeError(
                 "The brotli decompression is not available. "
                 "Please install `Brotli` module"
             )
         self._obj = brotli.Decompressor()
+        super().__init__(executor=executor, max_sync_chunk_size=max_sync_chunk_size)
 
-    def decompress_sync(self, data: bytes) -> bytes:
+    def decompress_sync(
+        self, data: Buffer, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED
+    ) -> bytes:
+        """Decompress the given data."""
         if hasattr(self._obj, "decompress"):
-            return cast(bytes, self._obj.decompress(data))
-        return cast(bytes, self._obj.process(data))
+            return cast(bytes, self._obj.decompress(data, max_length))
+        return cast(bytes, self._obj.process(data, max_length))
 
     def flush(self) -> bytes:
+        """Flush the decompressor."""
         if hasattr(self._obj, "flush"):
             return cast(bytes, self._obj.flush())
         return b""
 
 
-class ZSTDDecompressor:
-    def __init__(self) -> None:
+class ZSTDDecompressor(DecompressionBaseHandler):
+    def __init__(
+        self,
+        executor: Optional[Executor] = None,
+        max_sync_chunk_size: Optional[int] = MAX_SYNC_CHUNK_SIZE,
+    ) -> None:
         if not HAS_ZSTD:
             raise RuntimeError(
                 "The zstd decompression is not available. "
                 "Please install `backports.zstd` module"
             )
         self._obj = ZstdDecompressor()
-
-    def decompress_sync(self, data: bytes) -> bytes:
-        return self._obj.decompress(data)
+        super().__init__(executor=executor, max_sync_chunk_size=max_sync_chunk_size)
+
+    def decompress_sync(
+        self, data: bytes, max_length: int = ZLIB_MAX_LENGTH_UNLIMITED
+    ) -> bytes:
+        # zstd uses -1 for unlimited, while zlib uses 0 for unlimited
+        # Convert the zlib convention (0=unlimited) to zstd convention (-1=unlimited)
+        zstd_max_length = (
+            ZSTD_MAX_LENGTH_UNLIMITED
+            if max_length == ZLIB_MAX_LENGTH_UNLIMITED
+            else max_length
+        )
+        return self._obj.decompress(data, zstd_max_length)
 
     def flush(self) -> bytes:
         return b""

aiohttp/http_exceptions.py

@@ -74,6 +74,10 @@ class ContentLengthError(PayloadEncodingError):
     """Not enough data to satisfy content length header."""
 
 
+class DecompressSizeError(PayloadEncodingError):
+    """Decompressed size exceeds the configured limit."""
+
+
 class LineTooLong(BadHttpMessage):
     def __init__(
         self, line: str, limit: str = "Unknown", actual_size: str = "Unknown"

aiohttp/http_parser.py

@@ -27,6 +27,7 @@
 from . import hdrs
 from .base_protocol import BaseProtocol
 from .compression_utils import (
+    DEFAULT_MAX_DECOMPRESS_SIZE,
     HAS_BROTLI,
     HAS_ZSTD,
     BrotliDecompressor,
@@ -48,6 +49,7 @@
     BadStatusLine,
     ContentEncodingError,
     ContentLengthError,
+    DecompressSizeError,
     InvalidHeader,
     InvalidURLError,
     LineTooLong,
@@ -963,7 +965,12 @@ class DeflateBuffer:
 
     decompressor: Any
 
-    def __init__(self, out: StreamReader, encoding: Optional[str]) -> None:
+    def __init__(
+        self,
+        out: StreamReader,
+        encoding: Optional[str],
+        max_decompress_size: int = DEFAULT_MAX_DECOMPRESS_SIZE,
+    ) -> None:
         self.out = out
         self.size = 0
         out.total_compressed_bytes = self.size
@@ -988,6 +995,8 @@ def __init__(self, out: StreamReader, encoding: Optional[str]) -> None:
         else:
             self.decompressor = ZLibDecompressor(encoding=encoding)
 
+        self._max_decompress_size = max_decompress_size
+
     def set_exception(
         self,
         exc: BaseException,
@@ -1017,14 +1026,24 @@ def feed_data(self, chunk: bytes, size: int) -> None:
             )
 
         try:
-            chunk = self.decompressor.decompress_sync(chunk)
+            # Decompress with limit + 1 so we can detect if output exceeds limit
+            chunk = self.decompressor.decompress_sync(
+                chunk, max_length=self._max_decompress_size + 1
+            )
         except Exception:
             raise ContentEncodingError(
                 "Can not decode content-encoding: %s" % self.encoding
             )
 
         self._started_decoding = True
 
+        # Check if decompression limit was exceeded
+        if len(chunk) > self._max_decompress_size:
+            raise DecompressSizeError(
+                "Decompressed data exceeds the configured limit of %d bytes"
+                % self._max_decompress_size
+            )
+
         if chunk:
             self.out.feed_data(chunk, len(chunk))
 

aiohttp/multipart.py

@@ -25,7 +25,12 @@
 
 from multidict import CIMultiDict, CIMultiDictProxy
 
-from .compression_utils import ZLibCompressor, ZLibDecompressor
+from .abc import AbstractStreamWriter
+from .compression_utils import (
+    DEFAULT_MAX_DECOMPRESS_SIZE,
+    ZLibCompressor,
+    ZLibDecompressor,
+)
 from .hdrs import (
     CONTENT_DISPOSITION,
     CONTENT_ENCODING,
@@ -273,6 +278,7 @@ def __init__(
         *,
         subtype: str = "mixed",
         default_charset: Optional[str] = None,
+        max_decompress_size: int = DEFAULT_MAX_DECOMPRESS_SIZE,
     ) -> None:
         self.headers = headers
         self._boundary = boundary
@@ -289,6 +295,7 @@ def __init__(
         self._prev_chunk: Optional[bytes] = None
         self._content_eof = 0
         self._cache: Dict[str, Any] = {}
+        self._max_decompress_size = max_decompress_size
 
     def __aiter__(self: Self) -> Self:
         return self
@@ -318,7 +325,7 @@ async def read(self, *, decode: bool = False) -> bytes:
         while not self._at_eof:
             data.extend(await self.read_chunk(self.chunk_size))
         if decode:
-            return self.decode(data)
+            return await self.decode(data)
         return data
 
     async def read_chunk(self, size: int = chunk_size) -> bytes:
@@ -496,7 +503,7 @@ def at_eof(self) -> bool:
         """Returns True if the boundary was reached or False otherwise."""
         return self._at_eof
 
-    def decode(self, data: bytes) -> bytes:
+    async def decode(self, data: bytes) -> bytes:
         """Decodes data.
 
         Decoding is done according the specified Content-Encoding
@@ -506,18 +513,18 @@ def decode(self, data: bytes) -> bytes:
             data = self._decode_content_transfer(data)
         # https://datatracker.ietf.org/doc/html/rfc7578#section-4.8
         if not self._is_form_data and CONTENT_ENCODING in self.headers:
-            return self._decode_content(data)
+            return await self._decode_content(data)
         return data
 
-    def _decode_content(self, data: bytes) -> bytes:
+    async def _decode_content(self, data: bytes) -> bytes:
         encoding = self.headers.get(CONTENT_ENCODING, "").lower()
         if encoding == "identity":
             return data
         if encoding in {"deflate", "gzip"}:
-            return ZLibDecompressor(
+            return await ZLibDecompressor(
                 encoding=encoding,
                 suppress_deflate_header=True,
-            ).decompress_sync(data)
+            ).decompress(data, max_length=self._max_decompress_size)
 
         raise RuntimeError(f"unknown content encoding: {encoding}")
 
@@ -588,11 +595,11 @@ async def as_bytes(self, encoding: str = "utf-8", errors: str = "strict") -> byt
         """
         raise TypeError("Unable to read body part as bytes. Use write() to consume.")
 
-    async def write(self, writer: Any) -> None:
+    async def write(self, writer: AbstractStreamWriter) -> None:
         field = self._value
         chunk = await field.read_chunk(size=2**16)
         while chunk:
-            await writer.write(field.decode(chunk))
+            await writer.write(await field.decode(chunk))
             chunk = await field.read_chunk(size=2**16)
 
 
@@ -1032,7 +1039,9 @@ async def as_bytes(self, encoding: str = "utf-8", errors: str = "strict") -> byt
 
         return b"".join(parts)
 
-    async def write(self, writer: Any, close_boundary: bool = True) -> None:
+    async def write(
+        self, writer: AbstractStreamWriter, close_boundary: bool = True
+    ) -> None:
         """Write body."""
         for part, encoding, te_encoding in self._parts:
             if self._is_form_data:
@@ -1086,7 +1095,7 @@ async def close(self) -> None:
 
 
 class MultipartPayloadWriter:
-    def __init__(self, writer: Any) -> None:
+    def __init__(self, writer: AbstractStreamWriter) -> None:
         self._writer = writer
         self._encoding: Optional[str] = None
         self._compress: Optional[ZLibCompressor] = None