From e051e3612f0a3520b2dae715fd91456ecc73e4d1 Mon Sep 17 00:00:00 2001 From: Thomas Waldmann Date: Mon, 10 Aug 2026 08:48:37 +0200 Subject: [PATCH] gh-155460: Pre-allocate the output buffer in compression.zstd decompression Frames produced by the one-shot compression APIs record the decompressed size in the frame header. When the input starts at such a frame header, read that size with ZSTD_getFrameContentSize() and allocate the output buffer at its exact size, instead of growing it progressively and shrinking it on finish. An exactly filled buffer is then returned without a copy. This is only a sizing hint: decompression does not rely on it, so a hand-crafted header recording a wrong size produces exactly the same results and exceptions as before. Two guards bound the allocation: the recorded size is trusted only up to 1 GiB, and recorded sizes claiming more than a 32768x expansion of the available input -- more than the format can produce -- are ignored. Co-Authored-By: Claude Opus 5 --- Lib/test/test_zstd.py | 100 ++++++++++++++++++ ...-08-09-21-45-00.gh-issue-155460.zstdsz.rst | 4 + Modules/_zstd/decompressor.c | 40 ++++++- 3 files changed, 142 insertions(+), 2 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2026-08-09-21-45-00.gh-issue-155460.zstdsz.rst diff --git a/Lib/test/test_zstd.py b/Lib/test/test_zstd.py index 9225c42e0e711aa..1428a434993ceff 100644 --- a/Lib/test/test_zstd.py +++ b/Lib/test/test_zstd.py @@ -830,6 +830,106 @@ def test_decompress_empty_content_frame(self): self.assertEqual(d.unused_data, b'') self.assertEqual(d.unused_data, b'') # twice + @staticmethod + def _patch_content_size(frame, new_size): + # Rewrite the Frame_Content_Size field of a frame header, see + # RFC 8878 section 3.1.1.1. + frame_header_descriptor = frame[4] + fcs_flag = frame_header_descriptor >> 6 + single_segment = (frame_header_descriptor >> 5) & 1 + did_field_size = (0, 1, 2, 4)[frame_header_descriptor & 3] + offset = 5 + (0 if single_segment else 1) + did_field_size + fcs_field_size = (1 if single_segment else 0, 2, 4, 8)[fcs_flag] + if fcs_field_size == 2: + new_size -= 256 + patched = bytearray(frame) + patched[offset:offset+fcs_field_size] = \ + new_size.to_bytes(fcs_field_size, 'little') + return bytes(patched) + + def test_decompress_wrong_content_size(self): + # The decompressed size recorded in the frame header is used to + # pre-allocate the output buffer, so decompressing frames whose + # recorded size does not match the real one (only possible with + # hand-crafted frames) deserves extra attention. + frame = compress(DAT_130K_D) + self.assertEqual(get_frame_info(frame).decompressed_size, _130_1K) + + # patching the real size back is harmless (checks the patch helper) + patched = self._patch_content_size(frame, _130_1K) + self.assertEqual(patched, frame) + + for lie in (_130_1K + 1000, 1000, 0): + with self.subTest(lie=lie): + patched = self._patch_content_size(frame, lie) + self.assertEqual(get_frame_info(patched).decompressed_size, + lie) + with self.assertRaises(ZstdError): + decompress(patched) + + def test_decompress_absurd_content_size(self): + # A recorded size that is absurdly large for the frame's size, or + # even impossible to produce from it, must not lead to a huge + # pre-allocation. + for data in (DAT_130K_D, # bigger than the compressed frame + b'a' * 66000 # tiny compressed frame + ): + frame = compress(data) + patched = self._patch_content_size(frame, 0xFFFF_FFFF) + with self.subTest(frame_size=len(frame)): + self.assertEqual(get_frame_info(patched).decompressed_size, + 0xFFFF_FFFF) + with self.assertRaises(ZstdError): + decompress(patched) + + def test_decompress_content_size_known(self): + # frames whose header records the decompressed size, with sizes + # around the output buffer block boundaries + big_data = DAT_130K_D * 9 + for size in (1, 100, + 32*_1K - 1, 32*_1K, 32*_1K + 1, + _1M + 17): + with self.subTest(size=size): + data = big_data[:size] + frame = compress(data) + self.assertEqual(get_frame_info(frame).decompressed_size, + size) + self.assertEqual(decompress(frame), data) + + d = ZstdDecompressor() + self.assertEqual(d.decompress(frame), data) + self.assertTrue(d.eof) + + def test_decompress_content_size_unknown(self): + # streaming compression does not record the decompressed size in + # the frame header + c = ZstdCompressor() + frame = c.compress(DAT_130K_D) + c.flush() + self.assertIsNone(get_frame_info(frame).decompressed_size) + self.assertEqual(decompress(frame), DAT_130K_D) + + def test_decompress_content_size_known_max_length(self): + frame = compress(DAT_130K_D) + d = ZstdDecompressor() + dat = d.decompress(frame, max_length=1000) + self.assertEqual(len(dat), 1000) + self.assertFalse(d.needs_input) + while not d.eof: + dat += d.decompress(b'', max_length=32*_1K) + self.assertEqual(dat, DAT_130K_D) + + def test_decompress_content_size_known_split_input(self): + frame = compress(DAT_130K_D) + # a split point of 3 cuts the frame header's magic number, + # 18 cuts right after the (complete) frame header + for split in (3, 18, len(frame) // 2): + with self.subTest(split=split): + d = ZstdDecompressor() + dat = d.decompress(frame[:split]) + dat += d.decompress(frame[split:]) + self.assertEqual(dat, DAT_130K_D) + self.assertTrue(d.eof) + class DecompressorFlagsTestCase(unittest.TestCase): @classmethod diff --git a/Misc/NEWS.d/next/Library/2026-08-09-21-45-00.gh-issue-155460.zstdsz.rst b/Misc/NEWS.d/next/Library/2026-08-09-21-45-00.gh-issue-155460.zstdsz.rst new file mode 100644 index 000000000000000..ec07c3a9d77d865 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-08-09-21-45-00.gh-issue-155460.zstdsz.rst @@ -0,0 +1,4 @@ +Speed up :mod:`compression.zstd` decompression of frames whose header +records the decompressed size (as written by the one-shot compression APIs) +by allocating the output buffer at its exact size up front, instead of +growing it progressively. diff --git a/Modules/_zstd/decompressor.c b/Modules/_zstd/decompressor.c index ca9ea7e3b8f99bb..ab974f90d957cd5 100644 --- a/Modules/_zstd/decompressor.c +++ b/Modules/_zstd/decompressor.c @@ -187,6 +187,18 @@ _zstd_load_d_dict(ZstdDecompressor *self, PyObject *dict) return ret; } +/* Only pre-allocate an output buffer of up to this size based on the + decompressed size recorded in a frame header, so that a hand-crafted + header cannot request an arbitrarily large allocation. Larger outputs + use the progressively growing buffer. */ +#define OUTPUT_PREALLOC_MAX ((Py_ssize_t)1 << 30) + +/* A zstd block cannot expand to more than 128 KiB from less than 4 bytes + of compressed input, so a valid frame never expands by more than 32768x. + A recorded decompressed size claiming a higher ratio than this cannot be + fulfilled by the input and is treated as untrustworthy. */ +#define OUTPUT_MAX_EXPANSION 32768 + /* Decompress implementation in pseudo code: @@ -220,8 +232,32 @@ decompress_lock_held(ZstdDecompressor *self, ZSTD_inBuffer *in, _BlocksOutputBuffer buffer = {.writer = NULL}; PyObject *ret; - /* Initialize the output buffer */ - if (_OutputBuffer_InitAndGrow(&buffer, &out, max_length) < 0) { + /* Initialize the output buffer. + + Frames produced by the one-shot compression APIs record the + decompressed size in the frame header. When *in* starts at a frame + header recording a plausible size, allocate the whole output buffer + at once instead of growing it in blocks: for an exactly-filled + single block, _OutputBuffer_Finish() returns it without a copy. + + This is only a sizing hint, decompression does not rely on it: if + the recorded size turns out to be wrong, the buffer grows further + as needed, or is shrunk to the actual size on finish. */ + size_t avail_in = in->size - in->pos; + unsigned long long content_size = + ZSTD_getFrameContentSize((const char*)in->src + in->pos, avail_in); + if (content_size != ZSTD_CONTENTSIZE_UNKNOWN + && content_size != ZSTD_CONTENTSIZE_ERROR + && 0 < content_size + && content_size <= (unsigned long long)OUTPUT_PREALLOC_MAX + && content_size / OUTPUT_MAX_EXPANSION <= avail_in) + { + if (_OutputBuffer_InitWithSize(&buffer, &out, max_length, + (Py_ssize_t)content_size) < 0) { + goto error; + } + } + else if (_OutputBuffer_InitAndGrow(&buffer, &out, max_length) < 0) { goto error; } assert(out.pos == 0);