From b4a402f41cf44b6094b5131286830ba9bb1eb290 Mon Sep 17 00:00:00 2001
From: Nicholas Wiersma <nick@wiersma.co.za>
Date: Sun, 16 Jul 2023 19:22:09 +0200
Subject: [PATCH] feat: add max byte slice size config (#273)

---
 README.md      |  7 +++++++
 config.go      | 14 ++++++++++++++
 reader.go      |  9 ++++++++-
 reader_test.go | 26 ++++++++++++++++++++++++++
 4 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 433683d6..34ba5ec8 100644
--- a/README.md
+++ b/README.md
@@ -125,11 +125,18 @@ encoding and decoding.
 Enums may also implement `TextMarshaler` and `TextUnmarshaler`, and must resolve to valid symbols in the given enum schema.
 
 ##### Identical Underlying Types
+
 One type can be [ConvertibleTo](https://go.dev/ref/spec#Conversions) another type if they have identical underlying types. 
 A non-native type is allowed be used if it can be convertible to *time.Time*, *big.Rat* or *avro.LogicalDuration* for the particular of *LogicalTypes*.
 
 Ex.: `type Timestamp time.Time`
 
+##### Untrusted Input With Bytes and Strings
+
+For security reasons, the configuration `Config.MaxByteSliceSize` restricts the maximum size of `bytes` and `string` types created
+by the `Reader`. The default maximum size is `1MiB` and is configurable. This is required to stop untrusted input from consuming all memory and
+crashing the application. Should this not be need, setting a negative number will disable the behaviour.
+
 ### Recursive Structs
 
 At this moment recursive structs are not supported. It is planned for the future.
diff --git a/config.go b/config.go
index b03a4c67..ad299a7d 100644
--- a/config.go
+++ b/config.go
@@ -8,6 +8,8 @@ import (
 	"github.com/modern-go/reflect2"
 )
 
+const maxByteSliceSize = 1024 * 1024
+
 // DefaultConfig is the default API.
 var DefaultConfig = Config{}.Freeze()
 
@@ -43,6 +45,10 @@ type Config struct {
 	// Disable caching layer for encoders and decoders, forcing them to get rebuilt on every
 	// call to Marshal() and Unmarshal()
 	DisableCaching bool
+
+	// MaxByteSliceSize is the maximum size of `bytes` or `string` types the Reader will create, defaulting to 1MiB.
+	// If this size is exceeded, the Reader returns an error. This can be disabled by setting a negative number.
+	MaxByteSliceSize int
 }
 
 // Freeze makes the configuration immutable.
@@ -252,3 +258,11 @@ func (c *frozenConfig) getBlockLength() int {
 	}
 	return blockSize
 }
+
+func (c *frozenConfig) getMaxByteSliceSize() int {
+	size := c.config.MaxByteSliceSize
+	if size == 0 {
+		return maxByteSliceSize
+	}
+	return size
+}
diff --git a/reader.go b/reader.go
index a5daabdf..96b80759 100644
--- a/reader.go
+++ b/reader.go
@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"strings"
 	"unsafe"
 )
 
@@ -216,12 +217,18 @@ func (r *Reader) ReadString() string {
 func (r *Reader) readBytes(op string) []byte {
 	size := int(r.ReadLong())
 	if size < 0 {
-		r.ReportError("ReadString", "invalid "+op+" length")
+		fnName := "Read" + strings.ToTitle(op)
+		r.ReportError(fnName, "invalid "+op+" length")
 		return nil
 	}
 	if size == 0 {
 		return []byte{}
 	}
+	if max := r.cfg.getMaxByteSliceSize(); max > 0 && size > max {
+		fnName := "Read" + strings.ToTitle(op)
+		r.ReportError(fnName, "size is greater than `Config.MaxByteSliceSize`")
+		return nil
+	}
 
 	// The bytes are entirely in the buffer and of a reasonable size.
 	// Use the byte slab.
diff --git a/reader_test.go b/reader_test.go
index 5488493f..f0835f20 100644
--- a/reader_test.go
+++ b/reader_test.go
@@ -507,6 +507,19 @@ func TestReader_ReadBytes(t *testing.T) {
 	}
 }
 
+func TestReader_ReadBytesLargerThanMaxByteSliceSize(t *testing.T) {
+	data := []byte{
+		246, 255, 255, 255, 255, 10, 255, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32,
+		32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32,
+		32, 32, 32, 32, 32, 32, 32,
+	}
+	r := avro.NewReader(bytes.NewReader(data), 4)
+
+	_ = r.ReadBytes()
+
+	assert.Error(t, r.Error)
+}
+
 func TestReader_ReadString(t *testing.T) {
 	tests := []struct {
 		data    []byte
@@ -583,6 +596,19 @@ func TestReader_ReadString(t *testing.T) {
 	}
 }
 
+func TestReader_ReadStringLargerThanMaxByteSliceSize(t *testing.T) {
+	data := []byte{
+		246, 255, 255, 255, 255, 10, 255, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32,
+		32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32,
+		32, 32, 32, 32, 32, 32, 32,
+	}
+	r := avro.NewReader(bytes.NewReader(data), 4)
+
+	_ = r.ReadString()
+
+	assert.Error(t, r.Error)
+}
+
 func TestReader_ReadStringFastPathIsntBoundToBuffer(t *testing.T) {
 	data := []byte{0x06, 0x66, 0x6F, 0x6F, 0x08, 0x61, 0x76, 0x72, 0x6F}
 	r := avro.NewReader(bytes.NewReader(data), 4)