Decompress: guess a multiple of the input for unknown-size frames

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

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 {

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