Decompress: reuse the caller-supplied buffer for unknown-size frames

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)

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 {

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)

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