From ee4a64793411a8e9aa441814ec9168d6b632a48e Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Mon, 19 Jul 2021 10:07:05 +0300
Subject: [PATCH] stackitem: limit EQUAL for deeply nested structs

See neo-project/neo-vm#428.
---
 pkg/vm/stackitem/item.go      | 30 +++++++++++++++++++++-----
 pkg/vm/stackitem/item_test.go | 40 +++++++++++++++++++++++++++++++++++
 pkg/vm/vm_test.go             |  8 +++++++
 3 files changed, 73 insertions(+), 5 deletions(-)

diff --git a/pkg/vm/stackitem/item.go b/pkg/vm/stackitem/item.go
index 9d23dccf8..3f554f0b1 100644
--- a/pkg/vm/stackitem/item.go
+++ b/pkg/vm/stackitem/item.go
@@ -24,6 +24,8 @@ const (
 	MaxArraySize = 1024
 	// MaxSize is the maximum item size allowed in the VM.
 	MaxSize = 1024 * 1024
+	// MaxComparableNumOfItems is the maximum number of items that can be compared for structs.
+	MaxComparableNumOfItems = 2048
 	// MaxByteArrayComparableSize is the maximum allowed length of ByteArray for Equals method.
 	// It is set to be the maximum uint16 value.
 	MaxByteArrayComparableSize = math.MaxUint16
@@ -268,17 +270,35 @@ func (i *Struct) TryInteger() (*big.Int, error) {
 
 // Equals implements Item interface.
 func (i *Struct) Equals(s Item) bool {
-	if i == s {
-		return true
-	} else if s == nil {
+	if s == nil {
 		return false
 	}
 	val, ok := s.(*Struct)
-	if !ok || len(i.value) != len(val.value) {
+	if !ok {
+		return false
+	}
+	var limit = MaxComparableNumOfItems - 1 // 1 for current element.
+	return i.equalStruct(val, &limit)
+}
+
+func (i *Struct) equalStruct(s *Struct, limit *int) bool {
+	if i == s {
+		return true
+	} else if len(i.value) != len(s.value) {
 		return false
 	}
 	for j := range i.value {
-		if !i.value[j].Equals(val.value[j]) {
+		*limit--
+		if *limit == 0 {
+			panic(errTooBigElements)
+		}
+		sa, oka := i.value[j].(*Struct)
+		sb, okb := s.value[j].(*Struct)
+		if oka && okb {
+			if !sa.equalStruct(sb, limit) {
+				return false
+			}
+		} else if !i.value[j].Equals(s.value[j]) {
 			return false
 		}
 	}
diff --git a/pkg/vm/stackitem/item_test.go b/pkg/vm/stackitem/item_test.go
index 1244074a0..3f62c5b4d 100644
--- a/pkg/vm/stackitem/item_test.go
+++ b/pkg/vm/stackitem/item_test.go
@@ -172,6 +172,11 @@ var equalsTestCases = map[string][]struct {
 			item2:  NewStruct([]Item{NewBigInteger(big.NewInt(1))}),
 			result: true,
 		},
+		{
+			item1:  NewStruct([]Item{NewBigInteger(big.NewInt(1)), NewStruct([]Item{})}),
+			item2:  NewStruct([]Item{NewBigInteger(big.NewInt(1)), NewStruct([]Item{})}),
+			result: true,
+		},
 	},
 	"bigint": {
 		{
@@ -381,6 +386,41 @@ func TestEquals(t *testing.T) {
 	}
 }
 
+func TestEqualsDeepStructure(t *testing.T) {
+	const perStruct = 4
+	var items = []Item{}
+	var num int
+	for i := 0; i < perStruct; i++ {
+		items = append(items, Make(0))
+		num++
+	}
+	var layerUp = func(sa *Struct, num int) (*Struct, int) {
+		items := []Item{}
+		for i := 0; i < perStruct; i++ {
+			clon, err := sa.Clone(100500)
+			require.NoError(t, err)
+			items = append(items, clon)
+		}
+		num *= perStruct
+		num++
+		return NewStruct(items), num
+	}
+	var sa = NewStruct(items)
+	for i := 0; i < 4; i++ {
+		sa, num = layerUp(sa, num)
+	}
+	require.Less(t, num, MaxComparableNumOfItems)
+	sb, err := sa.Clone(num)
+	require.NoError(t, err)
+	require.True(t, sa.Equals(sb))
+	sa, num = layerUp(sa, num)
+
+	require.Less(t, MaxComparableNumOfItems, num)
+	sb, err = sa.Clone(num)
+	require.NoError(t, err)
+	require.Panics(t, func() { sa.Equals(sb) })
+}
+
 var marshalJSONTestCases = []struct {
 	input  Item
 	result []byte
diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go
index b497ffc9b..dffe8e4d1 100644
--- a/pkg/vm/vm_test.go
+++ b/pkg/vm/vm_test.go
@@ -2450,6 +2450,14 @@ func TestNestedStructClone(t *testing.T) {
 	}
 }
 
+func TestNestedStructEquals(t *testing.T) {
+	h := "560112c501fe0160589d604a12c0db415824f7509d4a102aec4597" // See neo-project/neo-vm#426.
+	prog, err := hex.DecodeString(h)
+	require.NoError(t, err)
+	vm := load(prog)
+	checkVMFailed(t, vm)
+}
+
 func makeProgram(opcodes ...opcode.Opcode) []byte {
 	prog := make([]byte, len(opcodes)+1) // RET
 	for i := 0; i < len(opcodes); i++ {