From 0814951d81cdf6eae01ada14ecf366b5d0109604 Mon Sep 17 00:00:00 2001 From: Thomas Waldmann Date: Mon, 10 Aug 2026 08:48:37 +0200 Subject: [PATCH 1/3] 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 0763ba31737a4c..d2e3f0777ac11b 100644 --- a/Lib/test/test_zstd.py +++ b/Lib/test/test_zstd.py @@ -953,6 +953,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 00000000000000..ec07c3a9d77d86 --- /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 454736d687a430..ce08fd7c031e7d 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); From 431f0f1255f6afb352e047f1da344d5eee7753b7 Mon Sep 17 00:00:00 2001 From: Thomas Waldmann Date: Fri, 6 Nov 2026 19:06:55 +0100 Subject: [PATCH 2/3] gh-155460: Pre-size each whole zstd frame, respect window_log_max On top of the pre-allocation in the previous commit: - Pre-size only when the first input of a frame holds the whole zstd1 frame, and check the recorded size against that frame's compressed size instead of all of the input, so that neither a frame fed in pieces nor the data following a frame leads to a pre-allocation. The 1 GiB limit stays. - Check the frame's window size against window_log_max before pre-sizing: libzstd's single-pass mode, which a pre-sized output enables, doesn't check it. - zstd.decompress() feeds each frame whole, sized with get_frame_size(), so every frame of multi-frame input is pre-sized, not only the first. Co-authored-by: Maik Schulze Co-Authored-By: Claude Opus 5.5 --- Lib/compression/_common/_streams.py | 13 +- Lib/compression/zstd/__init__.py | 11 ++ Lib/test/test_zstd.py | 230 ++++++++++++++++++++++++++++ Modules/_zstd/decompressor.c | 112 +++++++++++--- 4 files changed, 346 insertions(+), 20 deletions(-) diff --git a/Lib/compression/_common/_streams.py b/Lib/compression/_common/_streams.py index 0e3308c0c6e5fd..6bee60b0cf0107 100644 --- a/Lib/compression/_common/_streams.py +++ b/Lib/compression/_common/_streams.py @@ -172,7 +172,8 @@ def tell(self): def _decompress_all(data, decomp_factory, trailing_error=(), *, - stream_magic=(), truncated_error, **decomp_args): + stream_magic=(), truncated_error, stream_size=None, + **decomp_args): """Decompress all of the concatenated streams in *data*. Each stream is decompressed by a new decompressor, made by calling @@ -180,6 +181,9 @@ def _decompress_all(data, decomp_factory, trailing_error=(), *, of class *trailing_error* in a stream after the first one ends the data, unless that stream starts with *stream_magic*. A stream without an end-of-stream marker raises *truncated_error*. + + *stream_size*, if given, is called with the data that starts a stream + and returns the size of that stream, or None if it cannot tell. """ results = [] # The views are released on exit, so that the caller can resize a @@ -190,6 +194,13 @@ def _decompress_all(data, decomp_factory, trailing_error=(), *, # data, to decompress it in one go. chunk_size = len(data) while True: + if stream_size is not None: + # A stream whose size is known is fed whole, in one chunk + # that ends where the stream ends. + with data[pos:] as rest: + size = stream_size(rest) + if size: + chunk_size = size decomp = decomp_factory(**decomp_args) start = pos res = [] diff --git a/Lib/compression/zstd/__init__.py b/Lib/compression/zstd/__init__.py index c87cac26691cf1..6ca5e2c242f9f5 100644 --- a/Lib/compression/zstd/__init__.py +++ b/Lib/compression/zstd/__init__.py @@ -163,9 +163,20 @@ def decompress(data, zstd_dict=None, options=None): """ return _streams._decompress_all(data, ZstdDecompressor, truncated_error=ZstdError, + stream_size=_frame_size, options=options, zstd_dict=zstd_dict) +def _frame_size(data): + # The size of the complete frame at the start of *data*, or None. Fed + # whole, a frame that records its content size is decompressed into an + # output buffer allocated once, at that size. + try: + return get_frame_size(data) + except ZstdError: + return None + + class CompressionParameter(enum.IntEnum): """Compression parameters.""" diff --git a/Lib/test/test_zstd.py b/Lib/test/test_zstd.py index d2e3f0777ac11b..d8a964b2967b61 100644 --- a/Lib/test/test_zstd.py +++ b/Lib/test/test_zstd.py @@ -1053,6 +1053,236 @@ def test_decompress_content_size_known_split_input(self): self.assertEqual(dat, DAT_130K_D) self.assertTrue(d.eof) +# Hand-built frames, independent of what the compressor emits. See RFC 8878. +_BLOCK_SIZE_MAX = 128 * _1K + +def _frame(content_size, blocks, window_log=17): + """A frame whose header records *content_size*, whatever *blocks* hold.""" + return (b'\x28\xb5\x2f\xfd' # magic number + + bytes([0b1100_0000]) # 8-byte content size field + + bytes([(window_log - 10) << 3]) # window size + + content_size.to_bytes(8, 'little') + + b''.join(blocks)) + +def _raw_block(data, last=False): + return (last | 0 << 1 | len(data) << 3).to_bytes(3, 'little') + data + +def _rle_block(byte, size, last=False): + return (last | 1 << 1 | size << 3).to_bytes(3, 'little') + bytes([byte]) + +def _peak_allocation(func, *args, **kwargs): + """Call func() and return the peak of memory allocated by Python + during the call, above what was allocated before it.""" + import tracemalloc + tracing = tracemalloc.is_tracing() + if not tracing: + tracemalloc.start() + try: + before = tracemalloc.get_traced_memory()[0] + tracemalloc.reset_peak() + try: + func(*args, **kwargs) + except ZstdError: + pass + return tracemalloc.get_traced_memory()[1] - before + finally: + if not tracing: + tracemalloc.stop() + + +@support.cpython_only +class DecompressorOutputSizeTestCase(unittest.TestCase): + # A frame that records its content size has its output allocated once, + # at that size. The recorded size is input, though: it must never make + # the decompressor allocate more than the frame can decode to. + + # Small enough not to stress the test machine, large enough that + # allocating it is not mistaken for the noise of a decompress() call. + CLAIM = 64 * _1M + NOISE = 1 * _1M + + def assertAllocatedAbout(self, peak, size): + self.assertGreaterEqual(peak, size) + self.assertLess(peak, size + self.NOISE) + + def test_hand_built_frames(self): + self.assertEqual( + decompress(_frame(5, [_raw_block(b'hello', last=True)])), b'hello') + self.assertEqual( + decompress(_frame(3, [_rle_block(ord('x'), 3, last=True)])), b'xxx') + + def test_output_allocated_once(self): + data = bytes(range(256)) * (8 * _1K) # 2 MiB + frame = compress(data) + self.assertEqual(get_frame_info(frame).decompressed_size, len(data)) + self.assertEqual(decompress(frame), data) + self.assertAllocatedAbout(_peak_allocation(decompress, frame), len(data)) + self.assertAllocatedAbout( + _peak_allocation(ZstdDecompressor().decompress, frame), len(data)) + + def test_output_allocated_once_densest_frame(self): + # RLE blocks are the densest valid encoding: 4 bytes per 128 KiB. + # The plausibility check on the recorded size must accept them. + size = 64 * _BLOCK_SIZE_MAX # 8 MiB + frame = _frame(size, [_rle_block(ord('x'), _BLOCK_SIZE_MAX)] * 63 + + [_rle_block(ord('x'), _BLOCK_SIZE_MAX, last=True)]) + self.assertEqual(len(frame), 14 + 64 * 4) + self.assertEqual(decompress(frame), b'x' * size) + self.assertAllocatedAbout(_peak_allocation(decompress, frame), size) + + def test_output_allocated_once_with_dictionary(self): + data = b''.join(SAMPLES) * 50 + frame = compress(data, zstd_dict=TRAINED_DICT) + self.assertEqual(get_frame_info(frame).decompressed_size, len(data)) + self.assertEqual(decompress(frame, zstd_dict=TRAINED_DICT), data) + self.assertAllocatedAbout( + _peak_allocation(decompress, frame, zstd_dict=TRAINED_DICT), + len(data)) + + def test_output_allocated_per_frame(self): + data1 = bytes(range(256)) * (8 * _1K) # 2 MiB + data2 = b'second frame' * 1000 + frame1, frame2 = compress(data1), compress(data2) + d = ZstdDecompressor() + self.assertEqual(d.decompress(frame1 + frame2), data1) + self.assertEqual(d.unused_data, frame2) + self.assertEqual(decompress(frame1 + frame2), data1 + data2) + self.assertAllocatedAbout( + _peak_allocation(ZstdDecompressor().decompress, frame1 + frame2), + len(data1)) + + def test_output_allocated_once_after_empty_input(self): + data = bytes(range(256)) * (8 * _1K) # 2 MiB + frame = compress(data) + def decompress_after_empty_input(): + d = ZstdDecompressor() + self.assertEqual(d.decompress(b''), b'') + return d.decompress(frame) + self.assertEqual(decompress_after_empty_input(), data) + self.assertAllocatedAbout( + _peak_allocation(decompress_after_empty_input), len(data)) + + def test_output_allocated_once_each_frame(self): + # decompress() feeds every frame to its decompressor whole. + data = bytes(range(256)) * (8 * _1K) # 2 MiB + frame = compress(data) + small = compress(b'small frame') + skippable = (0x184D2A50).to_bytes(4, 'little') + (4).to_bytes(4, 'little') + skippable += b'skip' + for frames, expected in (([small, frame], b'small frame' + data), + ([skippable, frame], data), + ([frame, frame], data + data)): + with self.subTest(frames=[len(f) for f in frames]): + stream = b''.join(frames) + self.assertEqual(decompress(stream), expected) + # The output of each frame, and the joined result. + self.assertAllocatedAbout( + _peak_allocation(decompress, stream), 2 * len(expected)) + + def test_max_length_limits_allocation(self): + size = 64 * _BLOCK_SIZE_MAX # 8 MiB + frame = compress(b'x' * size) + d = ZstdDecompressor() + self.assertEqual(d.decompress(frame, 1000), b'x' * 1000) + self.assertFalse(d.eof) + self.assertLess( + _peak_allocation(ZstdDecompressor().decompress, frame, 1000), + self.NOISE) + + def test_untrusted_content_size(self): + # 27 bytes cannot decode to more than 6 blocks. + frame = _frame(self.CLAIM, [_raw_block(b'0123456789', last=True)]) + self.assertRaises(ZstdError, decompress, frame) + self.assertRaises(ZstdError, ZstdDecompressor().decompress, frame) + self.assertLess(_peak_allocation(decompress, frame), self.NOISE) + self.assertLess( + _peak_allocation(ZstdDecompressor().decompress, frame), self.NOISE) + + def test_untrusted_content_size_huge(self): + for claim in (2**40, 2**62, 2**63 - 1): + with self.subTest(claim=claim): + frame = _frame(claim, [_raw_block(b'0123456789', last=True)]) + # A ZstdError, not a MemoryError or OverflowError + self.assertRaises(ZstdError, decompress, frame) + self.assertLess(_peak_allocation(decompress, frame), self.NOISE) + + def test_untrusted_content_size_trailing_data(self): + # Only the frame itself vouches for its size, not the data after it. + frame = _frame(self.CLAIM, [_raw_block(b'0123456789', last=True)]) + blocks = self.CLAIM // _BLOCK_SIZE_MAX + valid = _frame(self.CLAIM, [_rle_block(ord('x'), _BLOCK_SIZE_MAX)] + * (blocks - 1) + + [_rle_block(ord('x'), _BLOCK_SIZE_MAX, + last=True)]) + for trailing in (bytes(64 * _1K), valid): + with self.subTest(trailing=len(trailing)): + data = frame + trailing + self.assertRaises(ZstdError, decompress, data) + self.assertRaises(ZstdError, ZstdDecompressor().decompress, data) + self.assertLess(_peak_allocation(decompress, data), self.NOISE) + self.assertLess( + _peak_allocation(ZstdDecompressor().decompress, data), + self.NOISE) + + def test_untrusted_content_size_partial_frame(self): + # Until the whole frame is in, nothing vouches for the recorded size. + blocks = [_raw_block(bytes([i]) * _1K) for i in range(64)] + partial = _frame(self.CLAIM, blocks) # no last block + d = ZstdDecompressor() + self.assertEqual(d.decompress(partial), + b''.join(bytes([i]) * _1K for i in range(64))) + self.assertFalse(d.eof) + self.assertTrue(d.needs_input) + self.assertLess( + _peak_allocation(ZstdDecompressor().decompress, partial), self.NOISE) + + def test_window_log_max(self): + # A frame that needs a larger window than allowed is refused, also + # when its output is allocated at once. (zstd doesn't check the + # window of a frame whose output fits the first, 32 KiB, block of the + # output buffer.) + window_log_max = DecompressionParameter.window_log_max + size = 64 * _1K + blocks = [_rle_block(ord('x'), size, last=True)] + frame = _frame(size, blocks, window_log=28) + for options in (None, {window_log_max: 0}): # default: 27 + with self.subTest(options=options): + self.assertRaises(ZstdError, decompress, frame, + options=options) + self.assertRaises(ZstdError, + ZstdDecompressor(options=options).decompress, + frame) + self.assertEqual(decompress(frame, options={window_log_max: 28}), + b'x' * size) + frame = _frame(size, blocks, window_log=20) + options = {window_log_max: 19} + self.assertRaises(ZstdError, decompress, frame, options=options) + self.assertRaises(ZstdError, + ZstdDecompressor(options=options).decompress, frame) + self.assertEqual(decompress(frame, options={window_log_max: 20}), + b'x' * size) + # A single-segment frame needs a window of its content size. + frame = compress(b'x' * _1M) + self.assertRaises(ZstdError, decompress, frame, + options={window_log_max: 19}) + self.assertEqual(decompress(frame, options={window_log_max: 20}), + b'x' * _1M) + + def test_content_size_too_small(self): + # The output must not be cut off at the recorded size. + frame = _frame(_BLOCK_SIZE_MAX, + [_rle_block(ord('x'), _BLOCK_SIZE_MAX), + _rle_block(ord('y'), _BLOCK_SIZE_MAX, last=True)]) + self.assertRaises(ZstdError, decompress, frame) + self.assertRaises(ZstdError, ZstdDecompressor().decompress, frame) + + def test_content_size_too_large(self): + frame = _frame(2 * _BLOCK_SIZE_MAX, + [_rle_block(ord('x'), _BLOCK_SIZE_MAX, last=True)]) + self.assertRaises(ZstdError, decompress, frame) + self.assertRaises(ZstdError, ZstdDecompressor().decompress, frame) + + class DecompressorFlagsTestCase(unittest.TestCase): @classmethod diff --git a/Modules/_zstd/decompressor.c b/Modules/_zstd/decompressor.c index ce08fd7c031e7d..9cd12c4c482f29 100644 --- a/Modules/_zstd/decompressor.c +++ b/Modules/_zstd/decompressor.c @@ -46,6 +46,12 @@ typedef struct { 1 means the end of the first frame has been reached. */ bool eof; + /* 1 until a decompress() call passes the first input of a frame. */ + bool at_frame_start; + + /* The largest window the frame may require, set by window_log_max. */ + unsigned long long max_window_size; + /* Lock to protect the decompression context */ PyMutex lock; } ZstdDecompressor; @@ -80,6 +86,12 @@ _get_DDict(ZstdDict *self) return self->d_dict; } +/* zstd's default window log limit and the largest window a new + decompression context accepts (ZSTD_WINDOWLOG_LIMIT_DEFAULT and + ZSTD_MAXWINDOWSIZE_DEFAULT, in zstd's static-linking-only API). */ +#define WINDOWLOG_LIMIT_DEFAULT 27 +#define DEFAULT_MAX_WINDOW_SIZE ((1ULL << WINDOWLOG_LIMIT_DEFAULT) + 1) + static int _zstd_set_d_parameters(ZstdDecompressor *self, PyObject *options) { @@ -129,6 +141,11 @@ _zstd_set_d_parameters(ZstdDecompressor *self, PyObject *options) set_parameter_error(0, key_v, value_v); return -1; } + if (key_v == ZSTD_d_windowLogMax) { + /* As zstd sets it: 0 means the default window log. */ + self->max_window_size = + 1ULL << (value_v ? value_v : WINDOWLOG_LIMIT_DEFAULT); + } } return 0; } @@ -196,9 +213,63 @@ _zstd_load_d_dict(ZstdDecompressor *self, PyObject *dict) /* 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. */ + fulfilled by the frame and is treated as untrustworthy. */ #define OUTPUT_MAX_EXPANSION 32768 +/* The size to allocate for the whole output of a frame, or -1 to grow the + output buffer block by block. + + If the input holds a whole frame that records a plausible content size, + allocate all of it at once: the data is written once, without the copies + of a growing buffer, and zstd decodes straight into it (its single-pass + mode) instead of through its internal window buffer. */ +static Py_ssize_t +frame_output_size(ZstdDecompressor *self, ZSTD_inBuffer *in) +{ + const unsigned char *src = (const unsigned char *)in->src + in->pos; + size_t avail = in->size - in->pos; + /* Only zstd1 frames: libzstd also accepts legacy (v0.5-v0.7) frames + here, whose headers have another layout and which it never decodes in + a single pass. */ + if (avail < 4 + || ((uint32_t)src[0] | (uint32_t)src[1] << 8 | (uint32_t)src[2] << 16 + | (uint32_t)src[3] << 24) != ZSTD_MAGICNUMBER) + { + return -1; + } + unsigned long long size = ZSTD_getFrameContentSize(src, avail); + if (size == ZSTD_CONTENTSIZE_UNKNOWN || size == ZSTD_CONTENTSIZE_ERROR + || size == 0 || size > (unsigned long long)OUTPUT_PREALLOC_MAX) + { + return -1; + } + /* zstd's single-pass mode doesn't check window_log_max: leave a frame + that needs a larger window to the streaming mode, which refuses it. + The header is complete and not a skippable frame, as it records a + content size other than 0. The window size is the content size for a + single-segment frame, else given by the window descriptor that + follows the frame header descriptor (RFC 8878, 3.1.1.1.2). */ + unsigned long long window_size = size; + if (!(src[4] & 0x20)) { + unsigned char exponent = src[5] >> 3, mantissa = src[5] & 7; + unsigned long long window_base = 1ULL << (10 + exponent); + window_size = window_base + window_base / 8 * mantissa; + } + if (window_size > self->max_window_size) { + return -1; + } + size_t frame_size = ZSTD_findFrameCompressedSize(src, avail); + if (ZSTD_isError(frame_size) || frame_size > avail) { + return -1; + } + /* Only the frame itself vouches for its recorded size, not the data + that follows it. */ + if (size / OUTPUT_MAX_EXPANSION > frame_size) { + return -1; + } + return (Py_ssize_t)size; +} + /* Decompress implementation in pseudo code: @@ -235,25 +306,25 @@ decompress_lock_held(ZstdDecompressor *self, ZSTD_inBuffer *in, /* 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) - { + decompressed size in the frame header. When the first input of a + frame holds the whole frame and it records 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: zstd + fails a frame that decodes to another size than the recorded one, + whether its output was pre-allocated or not. */ + Py_ssize_t init_size = -1; + if (self->at_frame_start && in->pos < in->size) { + /* zstd consumes input from here on, so later calls are not at the + start of the frame, even if this one doesn't pre-size it. */ + self->at_frame_start = 0; + init_size = frame_output_size(self, in); + } + if (init_size > 0) { if (_OutputBuffer_InitWithSize(&buffer, &out, max_length, - (Py_ssize_t)content_size) < 0) { + init_size) < 0) { goto error; } } @@ -330,6 +401,7 @@ decompressor_reset_session_lock_held(ZstdDecompressor *self) /* Reset variables in one operation */ self->needs_input = 1; self->eof = 0; + self->at_frame_start = 1; /* Resetting session is guaranteed to never fail */ ZSTD_DCtx_reset(self->dctx, ZSTD_reset_session_only); @@ -534,6 +606,8 @@ _zstd_ZstdDecompressor_new_impl(PyTypeObject *type, PyObject *zstd_dict, self->in_end = -1; self->unused_data = NULL; self->eof = 0; + self->at_frame_start = 1; + self->max_window_size = DEFAULT_MAX_WINDOW_SIZE; self->dict = NULL; self->lock = (PyMutex){0}; From d81eed6fe3efa8fdf6973895de9631c892a59422 Mon Sep 17 00:00:00 2001 From: Thomas Waldmann Date: Sat, 7 Nov 2026 00:17:37 +0100 Subject: [PATCH 3/3] gh-155460: Make the multi-frame zstd test depend on whole-frame feeding test_output_allocated_once_each_frame passed even without the whole-frame feeding in zstd.decompress(): its 2 MiB frame compresses to 459 bytes, so it fits whole into the 1 KiB that _decompress_all() feeds a later frame first. Use 2 MiB of zeros followed by 4 KiB of random bytes instead: the frame compresses to more than 1 KiB, and most of its output is decoded from its first KiB, which grows the output buffer unless the frame is fed whole. Co-Authored-By: Claude Opus 5.5 --- Lib/test/test_zstd.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/Lib/test/test_zstd.py b/Lib/test/test_zstd.py index d8a964b2967b61..0ff7f5b523cd4d 100644 --- a/Lib/test/test_zstd.py +++ b/Lib/test/test_zstd.py @@ -1163,8 +1163,10 @@ def decompress_after_empty_input(): _peak_allocation(decompress_after_empty_input), len(data)) def test_output_allocated_once_each_frame(self): - # decompress() feeds every frame to its decompressor whole. - data = bytes(range(256)) * (8 * _1K) # 2 MiB + # decompress() feeds every frame to its decompressor whole. This + # frame is larger than the 1 KiB that a later frame would get first + # otherwise, and most of its output is decoded from that first KiB. + data = bytes(2 * _1M) + random.randbytes(4 * _1K) frame = compress(data) small = compress(b'small frame') skippable = (0x184D2A50).to_bytes(4, 'little') + (4).to_bytes(4, 'little')