From ddc2cccc7420486aaffeaa425d45560fe05bb9d2 Mon Sep 17 00:00:00 2001 From: Kris Kennaway Date: Fri, 24 Jul 2026 17:46:36 +0100 Subject: [PATCH] Decompress: reuse the caller-supplied buffer for unknown-size frames Decompress sized its destination from decompressSizeHint(src), which reads the frame's content-size field. Frames that do not carry that field make the hint fall back to a pessimistic upper bound (>= decompressSizeBufferLimit, i.e. 1 MB): this is the case for legacy zstd v0.5 frames and for streaming frames compressed without a pledged source size. When the caller passed a smaller-but-adequate buffer, Decompress discarded it and allocated that bound, so every such decode allocated at least 1 MB regardless of the real payload size. decompressSizeHint now also reports whether the size was read from the frame (foundHint). When it was not, Decompress and ctx.Decompress try the caller-supplied buffer first via DecompressInto -- which reports a too-small buffer without writing past it -- before falling back to the hint-sized allocation and then the stream API. When the size is known (the common case where the same zstd version compressed and decompressed) the original path is unchanged, so callers passing a too-small buffer do not pay for a failed attempt. Co-Authored-By: Claude Opus 4.8 (1M context) --- zstd.go | 46 ++++++++----- zstd_bulk.go | 6 +- zstd_ctx.go | 33 ++++++---- zstd_reuse_buffer_test.go | 132 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 187 insertions(+), 30 deletions(-) create mode 100644 zstd_reuse_buffer_test.go diff --git a/zstd.go b/zstd.go index 8499bf1..f3d7ffb 100644 --- a/zstd.go +++ b/zstd.go @@ -57,31 +57,33 @@ func cCompressBound(srcSize int) int { return int(C.ZSTD_compressBound(C.size_t(srcSize))) } -// decompressSizeHint tries to give a hint on how much of the output buffer size we should have -// based on zstd frame descriptors. To prevent DOS from maliciously-created payloads, limit the size -func decompressSizeHint(src []byte) int { +// 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 + hint = upperBound if len(src) >= zstdFrameHeaderSizeMin { - hint = int(C.ZSTD_getFrameContentSize(unsafe.Pointer(&src[0]), C.size_t(len(src)))) - if hint < 0 { // On error, just use upperBound - hint = upperBound - } - if hint == 0 { // When compressing the empty slice, we need an output of at least 1 to pass down to the C lib - hint = 1 + 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 + hint = 1 + } } } // Take the minimum of both if hint > upperBound { - return upperBound + return upperBound, foundHint } - return hint + return hint, foundHint } // Compress src into dst. If you have a buffer to use, you can pass it to @@ -131,16 +133,26 @@ func CompressLevel(dst, src []byte, level int) ([]byte, error) { // Decompress src into dst. If you have a buffer to use, you can pass it to // prevent allocation. If it is too small, or if nil is passed, a new buffer // will be allocated and returned. +// +// Note: for frames that do not advertise their size (legacy v0.5 or unpledged +// streaming frames) dst may be partially overwritten even if a new slice is +// returned; do not rely on dst's contents after such a call. func Decompress(dst, src []byte) ([]byte, error) { if len(src) == 0 { return []byte{}, ErrEmptySlice } - bound := decompressSizeHint(src) - if cap(dst) >= bound { - dst = dst[0:cap(dst)] - } else { - dst = make([]byte, bound) + hint, foundHint := decompressSizeHint(src) + + // Reuse the caller buffer when it is large enough, or when the size is + // unknown (the hint is then only an upper bound); otherwise allocate the hint. + switch { + case cap(dst) >= hint: + dst = dst[:cap(dst)] + case !foundHint && cap(dst) > 0: + dst = dst[:cap(dst)] + default: + dst = make([]byte, hint) } written, err := DecompressInto(dst, src) diff --git a/zstd_bulk.go b/zstd_bulk.go index 6294a65..31a4009 100644 --- a/zstd_bulk.go +++ b/zstd_bulk.go @@ -110,7 +110,11 @@ func (p *BulkProcessor) Decompress(dst, src []byte) ([]byte, error) { return nil, ErrEmptySlice } - contentSize := decompressSizeHint(src) + // 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 { dst = dst[0:cap(dst)] } else { diff --git a/zstd_ctx.go b/zstd_ctx.go index c4a0889..429c26c 100644 --- a/zstd_ctx.go +++ b/zstd_ctx.go @@ -23,6 +23,10 @@ type Ctx interface { // Decompress src into dst. If you have a buffer to use, you can pass it to // prevent allocation. If it is too small, or if nil is passed, a new buffer // will be allocated and returned. + // + // Note: for frames that do not advertise their size (legacy v0.5 or unpledged + // streaming frames) dst may be partially overwritten even if a new slice is + // returned; do not rely on dst's contents after such a call. Decompress(dst, src []byte) ([]byte, error) // DecompressInto decompresses src into dst. Unlike Decompress, DecompressInto @@ -41,14 +45,14 @@ type ctx struct { } // Create a new ZStd Context. -// When compressing/decompressing many times, it is recommended to allocate a -// context just once, and re-use it for each successive compression operation. -// This will make workload friendlier for system's memory. -// Note : re-using context is just a speed / resource optimization. -// It doesn't change the compression ratio, which remains identical. -// Note 2 : In multi-threaded environments, -// use one different context per thread for parallel execution. // +// When compressing/decompressing many times, it is recommended to allocate a +// context just once, and re-use it for each successive compression operation. +// This will make workload friendlier for system's memory. +// Note : re-using context is just a speed / resource optimization. +// It doesn't change the compression ratio, which remains identical. +// Note 2 : In multi-threaded environments, +// use one different context per thread for parallel execution. func NewCtx() Ctx { c := &ctx{ cctx: C.ZSTD_createCCtx(), @@ -106,11 +110,16 @@ func (c *ctx) Decompress(dst, src []byte) ([]byte, error) { return []byte{}, ErrEmptySlice } - bound := decompressSizeHint(src) - if cap(dst) >= bound { - dst = dst[0:cap(dst)] - } else { - dst = make([]byte, bound) + hint, foundHint := decompressSizeHint(src) + + // See Decompress. + switch { + case cap(dst) >= hint: + dst = dst[:cap(dst)] + case !foundHint && cap(dst) > 0: + dst = dst[:cap(dst)] + default: + dst = make([]byte, hint) } written, err := c.DecompressInto(dst, src) diff --git a/zstd_reuse_buffer_test.go b/zstd_reuse_buffer_test.go new file mode 100644 index 0000000..93c6358 --- /dev/null +++ b/zstd_reuse_buffer_test.go @@ -0,0 +1,132 @@ +package zstd + +import ( + "bytes" + "testing" +) + +// unknownSizeFrame returns a zstd frame that does not advertise its decompressed +// size in the header (as produced by the streaming writer, and by legacy zstd +// v0.5 frames), verifying that is the case so callers exercise the intended path. +func unknownSizeFrame(t *testing.T, payload []byte) []byte { + t.Helper() + var b bytes.Buffer + w := NewWriter(&b) + if _, err := w.Write(payload); err != nil { + t.Fatalf("write: %v", err) + } + if err := w.Close(); err != nil { + t.Fatalf("close: %v", err) + } + frame := b.Bytes() + if _, found := decompressSizeHint(frame); found { + t.Fatal("streaming frame unexpectedly advertises its size") + } + return frame +} + +type namedDecompressor struct { + name string + fn func(dst, src []byte) ([]byte, error) +} + +// decompressors are the entry points that share the caller-buffer-reuse logic. +// ctx.Decompress uses a fresh context per call so tests don't share state. +func decompressors() []namedDecompressor { + return []namedDecompressor{ + {"Decompress", Decompress}, + {"ctx.Decompress", func(dst, src []byte) ([]byte, error) { return NewCtx().Decompress(dst, src) }}, + } +} + +// sameBuffer reports whether out was decoded into buf's backing array. +func sameBuffer(out, buf []byte) bool { + return len(out) > 0 && len(buf) > 0 && &out[0] == &buf[0] +} + +// TestDecompressReusesCallerBufferUnknownSize verifies that for an unknown-size +// frame, an adequate caller buffer is decoded into rather than discarded in +// favour of allocating decompressSizeBufferLimit. +func TestDecompressReusesCallerBufferUnknownSize(t *testing.T) { + payload := bytes.Repeat([]byte("datadog-"), 525) // 4200 bytes + frame := unknownSizeFrame(t, payload) + + for _, d := range decompressors() { + t.Run(d.name, func(t *testing.T) { + buf := make([]byte, 8192) // adequate for the payload, far below the bound + out, err := d.fn(buf, frame) + if err != nil { + t.Fatalf("decompress: %v", err) + } + if !bytes.Equal(out, payload) { + t.Fatalf("round-trip mismatch") + } + if !sameBuffer(out, buf) { + t.Fatalf("caller buffer should have been reused") + } + }) + } +} + +// TestDecompressUnknownSizeTooSmallBuffer verifies the fallback still works when +// the caller buffer is too small for an unknown-size frame. +func TestDecompressUnknownSizeTooSmallBuffer(t *testing.T) { + payload := bytes.Repeat([]byte("datadog-"), 525) + frame := unknownSizeFrame(t, payload) + + for _, d := range decompressors() { + t.Run(d.name, func(t *testing.T) { + out, err := d.fn(make([]byte, 8), frame) // too small; must fall back + if err != nil { + t.Fatalf("decompress: %v", err) + } + if !bytes.Equal(out, payload) { + t.Fatalf("round-trip mismatch on fallback") + } + }) + } +} + +// TestDecompressUnknownSizeNilBuffer verifies nil dst still decompresses. +func TestDecompressUnknownSizeNilBuffer(t *testing.T) { + payload := bytes.Repeat([]byte("datadog-"), 525) + 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 on nil dst") + } + }) + } +} + +// TestDecompressKnownSizeReusesCallerBuffer verifies that for frames that do +// advertise their size, an adequate caller buffer is still reused. +func TestDecompressKnownSizeReusesCallerBuffer(t *testing.T) { + payload := bytes.Repeat([]byte("datadog-"), 525) + frame, err := Compress(nil, payload) + if err != nil { + t.Fatalf("Compress: %v", err) + } + + for _, d := range decompressors() { + t.Run(d.name, func(t *testing.T) { + buf := make([]byte, 8192) + out, err := d.fn(buf, frame) + if err != nil { + t.Fatalf("decompress: %v", err) + } + if !bytes.Equal(out, payload) { + t.Fatalf("round-trip mismatch") + } + if !sameBuffer(out, buf) { + t.Fatalf("caller buffer should have been reused") + } + }) + } +}