Decompress: grow from a small buffer for unknown-size frames with no dst

zstd.go

@@ -145,14 +145,18 @@ func Decompress(dst, src []byte) ([]byte, error) {
 	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.
+	// unknown (the hint is then only an upper bound). Otherwise size from the
+	// hint, or grow from a small buffer when the size is unknown and none was
+	// supplied (see growDecompress).
 	switch {
 	case cap(dst) >= hint:
 		dst = dst[:cap(dst)]
 	case !foundHint && cap(dst) > 0:
 		dst = dst[:cap(dst)]
-	default:
+	case foundHint:
 		dst = make([]byte, hint)
+	default:
+		return growDecompress(src, DecompressInto)
 	}
 
 	written, err := DecompressInto(dst, src)
@@ -169,6 +173,31 @@ func Decompress(dst, src []byte) ([]byte, error) {
 	return ioutil.ReadAll(r)
 }
 
+// growDecompress decompresses src when the frame does not advertise its size and
+// no caller buffer was supplied. It starts small and grows on demand (3x, 6x,
+// 12x the compressed size) before falling back to the streaming reader, so it
+// never allocates the pessimistic decompressSizeBufferLimit bound up front and
+// only ever sizes from the input length, not an attacker-controlled field.
+func growDecompress(src []byte, into func(dst, src []byte) (int, error)) ([]byte, error) {
+	dst := make([]byte, len(src)*3)
+	for attempt := 0; attempt < 3; attempt++ {
+		written, err := into(dst, src)
+		if err == nil {
+			return dst[:written], nil
+		}
+		if !IsDstSizeTooSmallError(err) {
+			return nil, err
+		}
+		if attempt < 2 {
+			dst = make([]byte, len(dst)*2)
+		}
+	}
+
+	r := NewReader(bytes.NewReader(src))
+	defer r.Close()
+	return ioutil.ReadAll(r)
+}
+
 // DecompressInto decompresses src into dst. Unlike Decompress, DecompressInto
 // requires that dst be sufficiently large to hold the decompressed payload.
 // DecompressInto may be used when the caller knows the size of the decompressed

zstd_ctx.go

@@ -118,8 +118,10 @@ func (c *ctx) Decompress(dst, src []byte) ([]byte, error) {
 		dst = dst[:cap(dst)]
 	case !foundHint && cap(dst) > 0:
 		dst = dst[:cap(dst)]
-	default:
+	case foundHint:
 		dst = make([]byte, hint)
+	default:
+		return growDecompress(src, c.DecompressInto)
 	}
 
 	written, err := c.DecompressInto(dst, src)

zstd_nil_buffer_test.go

@@ -0,0 +1,54 @@
+package zstd
+
+import (
+	"bytes"
+	"testing"
+)
+
+// TestDecompressNilBufferUnknownSizeAvoidsBound verifies that decompressing an
+// unknown-size frame with no caller buffer grows from a small size instead of
+// allocating decompressSizeBufferLimit (the pessimistic upper bound).
+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))
+			}
+		})
+	}
+}
+
+// TestDecompressNilBufferKnownSizeExact verifies the known-size path is
+// unaffected: a nil buffer is still sized to the exact content size.
+func TestDecompressNilBufferKnownSizeExact(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) {
+			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) != len(payload) {
+				t.Fatalf("known-size nil decode should allocate exactly: cap %d want %d", cap(out), len(payload))
+			}
+		})
+	}
+}