From 3917b4b4e49fa1ace5999e0aac57d23b5807bbd1 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Thu, 21 May 2020 14:57:16 +0300
Subject: [PATCH] vm: update stack size in SETITEM properly

---
 pkg/vm/vm.go      |  4 ++--
 pkg/vm/vm_test.go | 20 ++++++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index 515008ea1..3a084d55a 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -1123,8 +1123,8 @@ func (v *VM) execute(ctx *Context, op opcode.Opcode, parameter []byte) (err erro
 			arr[index] = item
 			v.refs.Add(arr[index])
 		case *MapItem:
-			if t.Has(key.value) {
-				v.refs.Remove(item)
+			if i := t.Index(key.value); i >= 0 {
+				v.refs.Remove(t.value[i].Value)
 			} else if len(t.value) >= MaxArraySize {
 				panic("too big map")
 			}
diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go
index 2a906edab..da504ae45 100644
--- a/pkg/vm/vm_test.go
+++ b/pkg/vm/vm_test.go
@@ -1354,6 +1354,26 @@ func TestSETITEMBigMapBad(t *testing.T) {
 	runWithArgs(t, prog, nil, m, MaxArraySize, 0)
 }
 
+// This test checks is SETITEM properly updates reference counter.
+// 1. Create 2 arrays of size MaxArraySize - 3. (MaxStackSize = 2 * MaxArraySize)
+// 2. SETITEM each of them to a map.
+// 3. Replace each of them with a scalar value.
+func TestSETITEMMapStackLimit(t *testing.T) {
+	size := MaxArraySize - 3
+	m := NewMapItem()
+	m.Add(NewBigIntegerItem(big.NewInt(1)), NewArrayItem(makeArrayOfType(size, BooleanT)))
+	m.Add(NewBigIntegerItem(big.NewInt(2)), NewArrayItem(makeArrayOfType(size, BooleanT)))
+
+	prog := makeProgram(
+		opcode.DUP, opcode.PUSH1, opcode.PUSH1, opcode.SETITEM,
+		opcode.DUP, opcode.PUSH2, opcode.PUSH2, opcode.SETITEM,
+		opcode.DUP, opcode.PUSH3, opcode.PUSH3, opcode.SETITEM,
+		opcode.DUP, opcode.PUSH4, opcode.PUSH4, opcode.SETITEM)
+	v := load(prog)
+	v.estack.PushVal(m)
+	runVM(t, v)
+}
+
 func TestSETITEMBigMapGood(t *testing.T) {
 	prog := makeProgram(opcode.SETITEM)
 	vm := load(prog)