From f65545023dde18e5910d7a72f4ad3dd264455087 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Fri, 15 Nov 2019 22:33:28 +0300
Subject: [PATCH] io: restrict ReadArray max array size

---
 pkg/io/binaryReader.go  | 20 ++++++++++++++++++--
 pkg/io/binaryrw_test.go | 11 +++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/pkg/io/binaryReader.go b/pkg/io/binaryReader.go
index 126593480..645aee5c1 100644
--- a/pkg/io/binaryReader.go
+++ b/pkg/io/binaryReader.go
@@ -3,10 +3,15 @@ package io
 import (
 	"bytes"
 	"encoding/binary"
+	"fmt"
 	"io"
 	"reflect"
 )
 
+// maxArraySize is a maximums size of an array which can be decoded.
+// It is taken from https://github.com/neo-project/neo/blob/master/neo/IO/Helper.cs#L130
+const maxArraySize = 0x1000000
+
 // BinReader is a convenient wrapper around a io.Reader and err object.
 // Used to simplify error handling when reading into a struct with many fields.
 type BinReader struct {
@@ -36,7 +41,7 @@ func (r *BinReader) ReadLE(v interface{}) {
 
 // ReadArray reads array into value which must be
 // a pointer to a slice.
-func (r *BinReader) ReadArray(t interface{}) {
+func (r *BinReader) ReadArray(t interface{}, maxSize ...int) {
 	value := reflect.ValueOf(t)
 	if value.Kind() != reflect.Ptr || value.Elem().Kind() != reflect.Slice {
 		panic(value.Type().String() + " is not a pointer to a slice")
@@ -55,7 +60,18 @@ func (r *BinReader) ReadArray(t interface{}) {
 		return
 	}
 
-	l := int(r.ReadVarUint())
+	ms := maxArraySize
+	if len(maxSize) != 0 {
+		ms = maxSize[0]
+	}
+
+	lu := r.ReadVarUint()
+	if lu > uint64(ms) {
+		r.Err = fmt.Errorf("array is too big (%d)", lu)
+		return
+	}
+
+	l := int(lu)
 	arr := reflect.MakeSlice(sliceType, l, l)
 
 	for i := 0; i < l; i++ {
diff --git a/pkg/io/binaryrw_test.go b/pkg/io/binaryrw_test.go
index 8bc9e1578..672e14ffd 100644
--- a/pkg/io/binaryrw_test.go
+++ b/pkg/io/binaryrw_test.go
@@ -266,6 +266,17 @@ func TestBinReader_ReadArray(t *testing.T) {
 	require.NoError(t, r.Err)
 	require.Equal(t, elems, arrVal)
 
+	r = NewBinReaderFromBuf(data)
+	arrVal = []testSerializable{}
+	r.ReadArray(&arrVal, 3)
+	require.NoError(t, r.Err)
+	require.Equal(t, elems, arrVal)
+
+	r = NewBinReaderFromBuf(data)
+	arrVal = []testSerializable{}
+	r.ReadArray(&arrVal, 2)
+	require.Error(t, r.Err)
+
 	r = NewBinReaderFromBuf([]byte{0})
 	r.ReadArray(&arrVal)
 	require.NoError(t, r.Err)