From e670df9041014067c6fcc5cd9bdcbf9daefd5173 Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Tue, 26 Sep 2023 04:08:01 -0700
Subject: [PATCH] s2: Fix EstimateBlockSize on 6&7 length input (#867)

Fixes `panic: runtime error: index out of range [7] with length 5` when providing short input to EstimateBlockSize.

Found via #866
---
 s2/encode.go    | 2 +-
 s2/encode_go.go | 2 ++
 s2/s2_test.go   | 8 ++++++++
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/s2/encode.go b/s2/encode.go
index e6c2310212..0c9088adfe 100644
--- a/s2/encode.go
+++ b/s2/encode.go
@@ -57,7 +57,7 @@ func Encode(dst, src []byte) []byte {
 // The function returns -1 if no improvement could be achieved.
 // Using actual compression will most often produce better compression than the estimate.
 func EstimateBlockSize(src []byte) (d int) {
-	if len(src) < 6 || int64(len(src)) > 0xffffffff {
+	if len(src) <= inputMargin || int64(len(src)) > 0xffffffff {
 		return -1
 	}
 	if len(src) <= 1024 {
diff --git a/s2/encode_go.go b/s2/encode_go.go
index 0d39c7b0e0..6b393c34d3 100644
--- a/s2/encode_go.go
+++ b/s2/encode_go.go
@@ -316,6 +316,7 @@ func matchLen(a []byte, b []byte) int {
 	return len(a) + checked
 }
 
+// input must be > inputMargin
 func calcBlockSize(src []byte) (d int) {
 	// Initialize the hash table.
 	const (
@@ -501,6 +502,7 @@ emitRemainder:
 	return d
 }
 
+// length must be > inputMargin.
 func calcBlockSizeSmall(src []byte) (d int) {
 	// Initialize the hash table.
 	const (
diff --git a/s2/s2_test.go b/s2/s2_test.go
index 7135d11087..9fd954ce2f 100644
--- a/s2/s2_test.go
+++ b/s2/s2_test.go
@@ -1652,6 +1652,14 @@ func downloadBenchmarkFiles(b testing.TB, basename string) (errRet error) {
 	return nil
 }
 
+func TestEstimateBlockSize(t *testing.T) {
+	var input []byte
+	for i := 0; i < 100; i++ {
+		EstimateBlockSize(input)
+		input = append(input, 0)
+	}
+}
+
 func benchFile(b *testing.B, i int, decode bool) {
 	if err := downloadBenchmarkFiles(b, testFiles[i].filename); err != nil {
 		b.Fatalf("failed to download testdata: %s", err)