From 89125685369e75c847beaf27c9aed0c55500df05 Mon Sep 17 00:00:00 2001 From: Kris Kennaway Date: Mon, 27 Jul 2026 18:26:40 +0100 Subject: [PATCH] Decompress: guess a multiple of the input for unknown-size frames decompressSizeHint returned a flat lower bound of decompressSizeBufferLimit (1 MB) whenever the frame did not carry its decompressed size -- legacy zstd v0.5 frames and streaming frames compressed without a pledged size. Every such Decompress with no adequately sized caller buffer allocated at least 1 MB regardless of the real payload size. Now that decompressSizeHint reports whether the size was found, return 3x the input length for the unknown case instead of the flat bound, letting Decompress fall back to streaming if that is too small. This keeps a single decode path (no extra retry loop) and only ever sizes from the input length, so it does not weaken the zip-bomb guard; 3x is always within the existing 50x cap. BulkProcessor.Decompress has no streaming fallback (the streaming reader ignores the dictionary), so it keeps allocating the upper bound for unknown-size frames rather than the smaller guess, which it could not recover from. Co-Authored-By: Claude Opus 4.8 (1M context) --- zstd.go | 35 ++++++++++++++++++++++------------- zstd_bulk.go | 17 ++++++++++------- zstd_hint_test.go | 45 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 77 insertions(+), 20 deletions(-) create mode 100644 zstd_hint_test.go diff --git a/zstd.go b/zstd.go index f3d7ffb..001cb34 100644 --- a/zstd.go +++ b/zstd.go @@ -57,33 +57,42 @@ func cCompressBound(srcSize int) int { return int(C.ZSTD_compressBound(C.size_t(srcSize))) } +// decompressUpperBound is the largest output size decompressSizeHint will +// suggest: at least decompressSizeBufferLimit, at most 50x the input. It guards +// against zip bombs and bounds the allocation for callers with no streaming +// fallback (see BulkProcessor.Decompress). +func decompressUpperBound(src []byte) int { + b := 50 * len(src) + if b < decompressSizeBufferLimit { + b = decompressSizeBufferLimit + } + return b +} + // decompressSizeHint returns a suggested output size from the frame header, capped to guard against // zip bombs. foundHint is false when the frame does not advertise its size (legacy v0.5 or unpledged // streaming frames); the returned hint is then only a pessimistic upper bound. func decompressSizeHint(src []byte) (hint int, foundHint bool) { - // 1 MB or 50x input size - upperBound := 50 * len(src) - if upperBound < decompressSizeBufferLimit { - upperBound = decompressSizeBufferLimit - } - - hint = upperBound if len(src) >= zstdFrameHeaderSizeMin { contentSize := int(C.ZSTD_getFrameContentSize(unsafe.Pointer(&src[0]), C.size_t(len(src)))) if contentSize >= 0 { // a negative value means the size is unknown or the header is in error foundHint = true hint = contentSize - if hint == 0 { // When compressing the empty slice, we need an output of at least 1 to pass down to the C lib + if hint == 0 { // an empty payload still needs an output of at least 1 for the C lib hint = 1 } + if upper := decompressUpperBound(src); hint > upper { + hint = upper + } + return hint, true } } - // Take the minimum of both - if hint > upperBound { - return upperBound, foundHint - } - return hint, foundHint + // The frame does not advertise its size (legacy v0.5 or unpledged streaming + // frames). Guess a small multiple of the input rather than the flat upper + // bound; Decompress falls back to streaming if it is too small. 3x is always + // within decompressUpperBound (which is at least 50x the input). + return 3 * len(src), false } // Compress src into dst. If you have a buffer to use, you can pass it to diff --git a/zstd_bulk.go b/zstd_bulk.go index 31a4009..e46f2fa 100644 --- a/zstd_bulk.go +++ b/zstd_bulk.go @@ -110,15 +110,18 @@ func (p *BulkProcessor) Decompress(dst, src []byte) ([]byte, error) { return nil, ErrEmptySlice } - // Unlike Decompress, this always sizes from the hint and does not reuse a - // too-small caller buffer for unknown-size frames: there is no streaming - // fallback here (the streaming reader ignores the dictionary), so a - // too-small buffer could not be recovered and would fail the decode. - contentSize, _ := decompressSizeHint(src) - if cap(dst) >= contentSize { + // Unlike Decompress, there is no streaming fallback here (the streaming + // reader ignores the dictionary), so a too-small buffer can't be recovered. + // Size to the exact length when the frame advertises it, otherwise to the + // pessimistic upper bound rather than the small unknown-size guess. + size, foundHint := decompressSizeHint(src) + if !foundHint { + size = decompressUpperBound(src) + } + if cap(dst) >= size { dst = dst[0:cap(dst)] } else { - dst = make([]byte, contentSize) + dst = make([]byte, size) } if len(dst) == 0 { diff --git a/zstd_hint_test.go b/zstd_hint_test.go new file mode 100644 index 0000000..4dfaca0 --- /dev/null +++ b/zstd_hint_test.go @@ -0,0 +1,45 @@ +package zstd + +import ( + "bytes" + "testing" +) + +// TestDecompressSizeHintUnknownIsMultipleOfInput verifies that when the frame +// does not advertise its size, the hint is a small multiple of the input rather +// than the flat upper bound. +func TestDecompressSizeHintUnknownIsMultipleOfInput(t *testing.T) { + payload := bytes.Repeat([]byte("datadog-"), 525) + frame := unknownSizeFrame(t, payload) + + hint, found := decompressSizeHint(frame) + if found { + t.Fatal("streaming frame should not advertise its size") + } + if want := 3 * len(frame); hint != want { + t.Fatalf("unknown-size hint = %d, want %d (3x input)", hint, want) + } +} + +// TestDecompressNilBufferUnknownSizeAvoidsBound verifies that decompressing an +// unknown-size frame with no caller buffer does not allocate the flat upper +// bound (decompressSizeBufferLimit). +func TestDecompressNilBufferUnknownSizeAvoidsBound(t *testing.T) { + payload := bytes.Repeat([]byte("datadog-"), 525) // 4200 bytes, well under 1 MB + frame := unknownSizeFrame(t, payload) + + for _, d := range decompressors() { + t.Run(d.name, func(t *testing.T) { + out, err := d.fn(nil, frame) + if err != nil { + t.Fatalf("decompress: %v", err) + } + if !bytes.Equal(out, payload) { + t.Fatalf("round-trip mismatch") + } + if cap(out) >= decompressSizeBufferLimit { + t.Fatalf("nil-buffer unknown-size decode allocated the bound (cap %d)", cap(out)) + } + }) + } +}