From 91fdfcefc42381c63363418ad38d7d33233f79cd Mon Sep 17 00:00:00 2001
From: Spencer Nelson <spencer@twitch.tv>
Date: Thu, 5 Oct 2017 10:02:37 -0400
Subject: [PATCH] Check slice length when unmarshaling

With the help of github.com/dvyukov/go-fuzz, two potential panics were
discovered when unmarshaling a TDigest. Since we are setting a slice's
len from the bytes we unpack, we need to validate the len.

If that len is negative, we get a panic immediately when we try to
make the slice.

If that len is excessively huge, we can crash because allocating a
gigantic pile of memory takes excessively long (and we probably don't
want to allocate that much memory anyway!).
---
 fuzz.go      | 17 +++++++++++++++++
 fuzz_test.go | 28 ++++++++++++++++++++++++++++
 serde.go     |  6 ++++++
 3 files changed, 51 insertions(+)
 create mode 100644 fuzz.go
 create mode 100644 fuzz_test.go

diff --git a/fuzz.go b/fuzz.go
new file mode 100644
index 0000000..df96b8c
--- /dev/null
+++ b/fuzz.go
@@ -0,0 +1,17 @@
+// +build gofuzz
+
+package tdigest
+
+func Fuzz(data []byte) int {
+	v := new(TDigest)
+	err := v.UnmarshalBinary(data)
+	if err != nil {
+		return 0
+	}
+	_, err = v.MarshalBinary()
+	if err != nil {
+		panic(err)
+	}
+	return 1
+
+}
diff --git a/fuzz_test.go b/fuzz_test.go
new file mode 100644
index 0000000..6f9ce92
--- /dev/null
+++ b/fuzz_test.go
@@ -0,0 +1,28 @@
+package tdigest
+
+import "testing"
+
+func TestFuzzPanicRegressions(t *testing.T) {
+	// This test contains a list of byte sequences discovered by
+	// github.com/dvyukov/go-fuzz which, at one time, caused tdigest to panic. The
+	// test just makes sure that they no longer cause a panic.
+	testcase := func(crasher []byte) func(*testing.T) {
+		return func(*testing.T) {
+			v := new(TDigest)
+			err := v.UnmarshalBinary(crasher)
+			if err != nil {
+				return
+			}
+			_, err = v.MarshalBinary()
+			if err != nil {
+				panic(err)
+			}
+		}
+	}
+	t.Run("fuzz1", testcase([]byte(
+		"\x01\x00\x00\x000000000000000000"+
+			"00000000000\xfc")))
+	t.Run("fuzz2", testcase([]byte(
+		"\x01\x00\x00\x00\xdbF_\xbd\xdbF\x00\xbd\xe0\xdfʫ7172"+
+			"73747576777879(")))
+}
diff --git a/serde.go b/serde.go
index d8a56a3..c4dd87b 100644
--- a/serde.go
+++ b/serde.go
@@ -53,6 +53,12 @@ func unmarshalBinary(d *TDigest, p []byte) error {
 	if r.err != nil {
 		return r.err
 	}
+	if n < 0 {
+		return fmt.Errorf("invalid n, cannot be negative: %v", n)
+	}
+	if n > 1<<20 {
+		return fmt.Errorf("invalid n, cannot be greater than 2^20: %v", n)
+	}
 	d.centroids = make([]*centroid, int(n))
 	for i := 0; i < int(n); i++ {
 		c := new(centroid)