From 59b3d4da3d3ffa908c295f946351ca53ba37f18a Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Wed, 6 Dec 2023 18:32:11 +0100
Subject: [PATCH] huff0: Improve fuzz base set

---
 huff0/fuzz_test.go       | 21 ++++++++++++---
 internal/fuzz/helpers.go | 58 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 4 deletions(-)

diff --git a/huff0/fuzz_test.go b/huff0/fuzz_test.go
index fabb12e3fe..1a81439634 100644
--- a/huff0/fuzz_test.go
+++ b/huff0/fuzz_test.go
@@ -10,6 +10,7 @@ import (
 
 func FuzzCompress(f *testing.F) {
 	fuzz.AddFromZip(f, "testdata/fse_compress.zip", fuzz.TypeRaw, false)
+	fuzz.AddFromZip(f, "testdata/regression.zip", fuzz.TypeRaw, testing.Short())
 	f.Fuzz(func(t *testing.T, buf0 []byte) {
 		//use of Compress1X
 		var s Scratch
@@ -76,15 +77,27 @@ func FuzzCompress(f *testing.F) {
 
 func FuzzDecompress1x(f *testing.F) {
 	fuzz.AddFromZip(f, "testdata/huff0_decompress1x.zip", fuzz.TypeRaw, false)
+	var s Scratch
+	addCompressed := func(b []byte) {
+		s.Reuse = ReusePolicyNone
+		if b2, _, err := Compress1X(b, &s); err == nil {
+			f.Add(b2)
+		}
+		s.Reuse = ReusePolicyNone
+		if b2, _, err := Compress4X(b, &s); err == nil {
+			f.Add(b2)
+		}
+	}
+	fuzz.ReturnFromZip(f, "testdata/regression.zip", fuzz.TypeRaw, addCompressed)
+	fuzz.ReturnFromZip(f, "testdata/fse_compress.zip", fuzz.TypeRaw, addCompressed)
+
 	f.Fuzz(func(t *testing.T, buf0 []byte) {
 		var s Scratch
 		_, remain, err := ReadTable(buf0, &s)
 		if err != nil {
 			return
 		}
-		out, err := s.Decompress1X(remain)
-		if err != nil || out == nil {
-			return
-		}
+		s.Decompress1X(remain)
+		s.Decompress4X(remain, len(buf0))
 	})
 }
diff --git a/internal/fuzz/helpers.go b/internal/fuzz/helpers.go
index eb49d66bf6..7f1cc1db53 100644
--- a/internal/fuzz/helpers.go
+++ b/internal/fuzz/helpers.go
@@ -89,6 +89,64 @@ func AddFromZip(f *testing.F, filename string, t InputType, short bool) {
 	}
 }
 
+// ReturnFromZip will read the supplied zip and add all as corpus for f.
+// Byte slices only.
+func ReturnFromZip(tb testing.TB, filename string, t InputType, fn func([]byte)) {
+	file, err := os.Open(filename)
+	if err != nil {
+		tb.Fatal(err)
+	}
+	fi, err := file.Stat()
+	if err != nil {
+		tb.Fatal(err)
+	}
+	zr, err := zip.NewReader(file, fi.Size())
+	if err != nil {
+		tb.Fatal(err)
+	}
+	for _, file := range zr.File {
+		rc, err := file.Open()
+		if err != nil {
+			tb.Fatal(err)
+		}
+
+		b, err := io.ReadAll(rc)
+		if err != nil {
+			tb.Fatal(err)
+		}
+		rc.Close()
+		t := t
+		if t == TypeOSSFuzz {
+			t = TypeRaw // Fallback
+			if len(b) >= 4 {
+				sz := binary.BigEndian.Uint32(b)
+				if sz <= uint32(len(b))-4 {
+					fn(b[4 : 4+sz])
+					continue
+				}
+			}
+		}
+
+		if bytes.HasPrefix(b, []byte("go test fuzz")) {
+			t = TypeGoFuzz
+		} else {
+			t = TypeRaw
+		}
+
+		if t == TypeRaw {
+			fn(b)
+			continue
+		}
+		vals, err := unmarshalCorpusFile(b)
+		if err != nil {
+			tb.Fatal(err)
+		}
+		for _, v := range vals {
+			fn(v)
+		}
+	}
+}
+
 // unmarshalCorpusFile decodes corpus bytes into their respective values.
 func unmarshalCorpusFile(b []byte) ([][]byte, error) {
 	if len(b) == 0 {