From 4000dd692c161cabba132cc52795045334fd8476 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Sun, 11 Jul 2021 19:37:06 +0300
Subject: [PATCH] vm: make cloning limits a bit more effective

Take actual reference counter value into account.
---
 pkg/vm/vm.go      | 10 +++++-----
 pkg/vm/vm_test.go |  2 ++
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index ad2fe2d1f..c91d5c881 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -1056,7 +1056,7 @@ func (v *VM) execute(ctx *Context, op opcode.Opcode, parameter []byte) (err erro
 		itemElem := v.estack.Pop()
 		arrElem := v.estack.Pop()
 
-		val, err := cloneIfStruct(itemElem.value)
+		val, err := cloneIfStruct(itemElem.value, MaxStackSize-v.refs.size)
 		if err != nil {
 			panic(err)
 		}
@@ -1373,7 +1373,7 @@ func (v *VM) execute(ctx *Context, op opcode.Opcode, parameter []byte) (err erro
 			src := t.Value().([]stackitem.Item)
 			arr = make([]stackitem.Item, len(src))
 			for i := range src {
-				arr[i], err = cloneIfStruct(src[i])
+				arr[i], err = cloneIfStruct(src[i], MaxStackSize-v.refs.size)
 				if err != nil {
 					panic(err)
 				}
@@ -1381,7 +1381,7 @@ func (v *VM) execute(ctx *Context, op opcode.Opcode, parameter []byte) (err erro
 		case *stackitem.Map:
 			arr = make([]stackitem.Item, 0, t.Len())
 			for k := range t.Value().([]stackitem.MapElement) {
-				elem, err := cloneIfStruct(t.Value().([]stackitem.MapElement)[k].Value)
+				elem, err := cloneIfStruct(t.Value().([]stackitem.MapElement)[k].Value, MaxStackSize-v.refs.size)
 				if err != nil {
 					panic(err)
 				}
@@ -1751,10 +1751,10 @@ func checkMultisig1(v *VM, curve elliptic.Curve, h []byte, pkeys [][]byte, sig [
 	return false
 }
 
-func cloneIfStruct(item stackitem.Item) (stackitem.Item, error) {
+func cloneIfStruct(item stackitem.Item, limit int) (stackitem.Item, error) {
 	switch it := item.(type) {
 	case *stackitem.Struct:
-		return it.Clone(MaxStackSize)
+		return it.Clone(limit)
 	default:
 		return it, nil
 	}
diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go
index 8a176a88f..484b3cda1 100644
--- a/pkg/vm/vm_test.go
+++ b/pkg/vm/vm_test.go
@@ -2440,6 +2440,8 @@ func TestNestedStructClone(t *testing.T) {
 		"5601c2c501fe0360589d604a12c0db415824f7cf45",
 		// VALUES for map with deeply nested struct.
 		"5601c84a11c501fe0060589d604a12c0db415824f7d0cd45",
+		// VALUES for a lot of not-so-deep nested structs.
+		"5601c5000a60589d604a12c0db415824f701fe03504a519d4a102afa01ff03c0cd45",
 	}
 	for _, h := range progs {
 		prog, err := hex.DecodeString(h)