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 <noreply@anthropic.com>

Author: ThomasWaldmann

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

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.

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);