diff --git a/scrapy/downloadermiddlewares/httpcompression.py b/scrapy/downloadermiddlewares/httpcompression.py index 91748e57e..8ee1d95a6 100644 --- a/scrapy/downloadermiddlewares/httpcompression.py +++ b/scrapy/downloadermiddlewares/httpcompression.py @@ -1,10 +1,14 @@ -import io import warnings from scrapy.exceptions import IgnoreRequest, NotConfigured from scrapy.http import Response, TextResponse from scrapy.responsetypes import responsetypes -from scrapy.utils._compression import _DecompressionMaxSizeExceeded, _inflate, _unbrotli +from scrapy.utils._compression import ( + _DecompressionMaxSizeExceeded, + _inflate, + _unbrotli, + _unzstd, +) from scrapy.utils.deprecate import ScrapyDeprecationWarning from scrapy.utils.gz import gunzip @@ -18,7 +22,7 @@ else: ACCEPTED_ENCODINGS.append(b"br") try: - import zstandard + import zstandard # noqa: F401 except ImportError: pass else: @@ -102,8 +106,5 @@ class HttpCompressionMiddleware: if encoding == b"br" and b"br" in ACCEPTED_ENCODINGS: return _unbrotli(body, max_size=self._max_size) if encoding == b"zstd" and b"zstd" in ACCEPTED_ENCODINGS: - # Using its streaming API since its simple API could handle only cases - # where there is content size data embedded in the frame - reader = zstandard.ZstdDecompressor().stream_reader(io.BytesIO(body)) - return reader.read() + return _unzstd(body, max_size=self._max_size) return body diff --git a/scrapy/utils/_compression.py b/scrapy/utils/_compression.py index 9a32ce4f0..93aa254b2 100644 --- a/scrapy/utils/_compression.py +++ b/scrapy/utils/_compression.py @@ -7,6 +7,11 @@ try: except ImportError: pass +try: + import zstandard +except ImportError: + pass + class _DecompressionMaxSizeExceeded(ValueError): pass @@ -65,3 +70,23 @@ def _unbrotli(data: bytes, *, max_size: int = 0) -> bytes: ) output_list.append(output_chunk) return b"".join(output_list) + + +def _unzstd(data: bytes, *, max_size: int = 0) -> bytes: + decompressor = zstandard.ZstdDecompressor() + stream_reader = decompressor.stream_reader(BytesIO(data)) + output_list: List[bytes] = [] + output_chunk = b"." + decompressed_size = 0 + CHUNK_SIZE = 8196 + while output_chunk: + output_chunk = stream_reader.read(CHUNK_SIZE) + decompressed_size += len(output_chunk) + if max_size and decompressed_size > max_size: + raise _DecompressionMaxSizeExceeded( + f"The number of bytes decompressed so far " + f"({decompressed_size}B) exceed the specified maximum " + f"({max_size}B)." + ) + output_list.append(output_chunk) + return b"".join(output_list) diff --git a/tests/sample_data/compressed/bomb-zstd.bin b/tests/sample_data/compressed/bomb-zstd.bin new file mode 100644 index 000000000..4b0efa8a4 Binary files /dev/null and b/tests/sample_data/compressed/bomb-zstd.bin differ diff --git a/tests/test_downloadermiddleware_httpcompression.py b/tests/test_downloadermiddleware_httpcompression.py index 8858916bc..7babd1318 100644 --- a/tests/test_downloadermiddleware_httpcompression.py +++ b/tests/test_downloadermiddleware_httpcompression.py @@ -41,7 +41,7 @@ FORMAT = { "br", # 34 → 11 511 612 "deflate", # 27 968 → 11 511 612 "gzip", # 27 988 → 11 511 612 - # "zstd", + "zstd", # 1 096 → 11 511 612 ) }, } @@ -396,6 +396,9 @@ class HttpCompressionTest(TestCase): def test_compression_bomb_gzip(self): self._test_compression_bomb("gzip") + def test_compression_bomb_zstd(self): + self._test_compression_bomb("zstd") + class HttpCompressionSubclassTest(TestCase): def test_init_missing_stats(self):