From 43c36f4ae540226e246083e8293f93d7e465bcd7 Mon Sep 17 00:00:00 2001
From: Vianney Tran <vianney.tran@datadoghq.com>
Date: Thu, 2 Jun 2022 10:42:19 -0400
Subject: [PATCH 1/2] [zstd] Fix unneededly allocate large decompression buffer

Fix #118
With the introduction of #115 to prevent zipbombs, the check was too strict and was checking a too large size boundary. This fixes it and adds a test
---
 zstd.go      |  4 ++--
 zstd_test.go | 19 +++++++++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/zstd.go b/zstd.go
index f792116..7f3d29d 100644
--- a/zstd.go
+++ b/zstd.go
@@ -36,7 +36,7 @@ const (
 	// decompressed, err := zstd.Decompress(dst, src)
 	decompressSizeBufferLimit = 1000 * 1000
 
-	zstdFrameHeaderSizeMax = 18 // From zstd.h. Since it's experimental API, hardcoding it
+	zstdFrameHeaderSizeMin = 2 // From zstd.h. Since it's experimental API, hardcoding it
 )
 
 // CompressBound returns the worst case size needed for a destination buffer,
@@ -67,7 +67,7 @@ func decompressSizeHint(src []byte) int {
 	}
 
 	hint := upperBound
-	if len(src) >= zstdFrameHeaderSizeMax {
+	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
diff --git a/zstd_test.go b/zstd_test.go
index b7299cc..0253537 100644
--- a/zstd_test.go
+++ b/zstd_test.go
@@ -293,6 +293,25 @@ func TestBadPayloadZipBomb(t *testing.T) {
 	}
 }
 
+func TestSmallPayload(t *testing.T) {
+	// Test that we can compress really small payloads and this doesn't generate a huge output buffer
+	compressed, err := Compress(nil, []byte("a"))
+	if err != nil {
+		t.Fatalf("failed to compress: %s", err)
+	}
+
+	preAllocated := make([]byte, 1, 64) // Don't use more than that
+	decompressed, err := Decompress(preAllocated, compressed)
+	if err != nil {
+		t.Fatalf("failed to compress: %s", err)
+	}
+
+	if &(preAllocated[0]) != &(decompressed[0]) { // They should point to the same spot (no realloc)
+		t.Fatal("Compression buffer was changed")
+	}
+
+}
+
 func BenchmarkCompression(b *testing.B) {
 	if raw == nil {
 		b.Fatal(ErrNoPayloadEnv)

From 8eebfba5f0e71503643178d2b2c900bfc2405917 Mon Sep 17 00:00:00 2001
From: Vianney Tran <vianney.tran@datadoghq.com>
Date: Thu, 2 Jun 2022 10:57:22 -0400
Subject: [PATCH 2/2] [zstd] Always give a decompression hint of > 0

---
 zstd.go | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/zstd.go b/zstd.go
index 7f3d29d..b2f50b1 100644
--- a/zstd.go
+++ b/zstd.go
@@ -72,6 +72,9 @@ func decompressSizeHint(src []byte) int {
 		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
+		}
 	}
 
 	// Take the minimum of both