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)) + } + }) + } +}