vm: restrict total stack item count

pkg/vm/stack.go

@@ -152,6 +152,9 @@ type Stack struct {
 	top  Element
 	name string
 	len  int
+
+	itemCount map[StackItem]int
+	size      *int
 }
 
 // NewStack returns a new stack name by the given name.
@@ -162,6 +165,8 @@ func NewStack(n string) *Stack {
 	s.top.next = &s.top
 	s.top.prev = &s.top
 	s.len = 0
+	s.itemCount = make(map[StackItem]int)
+	s.size = new(int)
 	return s
 }
 
@@ -192,9 +197,54 @@ func (s *Stack) insert(e, at *Element) *Element {
 	n.prev = e
 	e.stack = s
 	s.len++
+
+	s.updateSizeAdd(e.value)
+
 	return e
 }
 
+func (s *Stack) updateSizeAdd(item StackItem) {
+	*s.size++
+
+	s.itemCount[item]++
+	if s.itemCount[item] > 1 {
+		return
+	}
+
+	switch t := item.(type) {
+	case *ArrayItem, *StructItem:
+		for _, it := range item.Value().([]StackItem) {
+			s.updateSizeAdd(it)
+		}
+	case *MapItem:
+		for _, v := range t.value {
+			s.updateSizeAdd(v)
+		}
+	}
+}
+
+func (s *Stack) updateSizeRemove(item StackItem) {
+	*s.size--
+
+	if s.itemCount[item] > 1 {
+		s.itemCount[item]--
+		return
+	}
+
+	delete(s.itemCount, item)
+
+	switch t := item.(type) {
+	case *ArrayItem, *StructItem:
+		for _, it := range item.Value().([]StackItem) {
+			s.updateSizeRemove(it)
+		}
+	case *MapItem:
+		for _, v := range t.value {
+			s.updateSizeRemove(v)
+		}
+	}
+}
+
 // InsertAt inserts the given item (n) deep on the stack.
 // Be very careful using it and _always_ check both e and n before invocation
 // as it will silently do wrong things otherwise.
@@ -271,6 +321,9 @@ func (s *Stack) Remove(e *Element) *Element {
 	e.prev = nil // avoid memory leaks.
 	e.stack = nil
 	s.len--
+
+	s.updateSizeRemove(e.value)
+
 	return e
 }
 

pkg/vm/vm.go

@@ -44,6 +44,10 @@ const (
 	// MaxInvocationStackSize is the maximum size of an invocation stack.
 	MaxInvocationStackSize = 1024
 
+	// MaxStackSize is the maximum number of items allowed to be
+	// on all stacks at once.
+	MaxStackSize = 2 * 1024
+
 	maxSHLArg = 256
 	minSHLArg = -256
 )
@@ -64,6 +68,9 @@ type VM struct {
 
 	// Hash to verify in CHECKSIG/CHECKMULTISIG.
 	checkhash []byte
+
+	itemCount map[StackItem]int
+	size      int
 }
 
 // InteropFuncPrice represents an interop function with a price.
@@ -79,10 +86,13 @@ func New() *VM {
 		getScript: nil,
 		state:     haltState,
 		istack:    NewStack("invocation"),
-		estack:    NewStack("evaluation"),
-		astack:    NewStack("alt"),
+
+		itemCount: make(map[StackItem]int),
 	}
 
+	vm.estack = vm.newItemStack("evaluation")
+	vm.astack = vm.newItemStack("alt")
+
 	// Register native interop hooks.
 	vm.RegisterInteropFunc("Neo.Runtime.Log", runtimeLog, 1)
 	vm.RegisterInteropFunc("Neo.Runtime.Notify", runtimeNotify, 1)
@@ -94,6 +104,14 @@ func New() *VM {
 	return vm
 }
 
+func (v *VM) newItemStack(n string) *Stack {
+	s := NewStack(n)
+	s.size = &v.size
+	s.itemCount = v.itemCount
+
+	return s
+}
+
 // RegisterInteropFunc registers the given InteropFunc to the VM.
 func (v *VM) RegisterInteropFunc(name string, f InteropFunc, price int) {
 	v.interop[name] = InteropFuncPrice{f, price}
@@ -428,6 +446,9 @@ func (v *VM) execute(ctx *Context, op Instruction, parameter []byte) (err error)
 		if errRecover := recover(); errRecover != nil {
 			v.state = faultState
 			err = newError(ctx.ip, op, errRecover)
+		} else if v.size > MaxStackSize {
+			v.state = faultState
+			err = newError(ctx.ip, op, "stack is too big")
 		}
 	}()
 
@@ -855,6 +876,8 @@ func (v *VM) execute(ctx *Context, op Instruction, parameter []byte) (err error)
 			panic("APPEND: not of underlying type Array")
 		}
 
+		v.estack.updateSizeAdd(val)
+
 	case PACK:
 		n := int(v.estack.Pop().BigInt().Int64())
 		if n < 0 || n > v.estack.Len() || n > MaxArraySize {
@@ -922,12 +945,17 @@ func (v *VM) execute(ctx *Context, op Instruction, parameter []byte) (err error)
 			if index < 0 || index >= len(arr) {
 				panic("SETITEM: invalid index")
 			}
+			v.estack.updateSizeRemove(arr[index])
 			arr[index] = item
+			v.estack.updateSizeAdd(arr[index])
 		case *MapItem:
-			if !t.Has(key.value) && len(t.value) >= MaxArraySize {
+			if t.Has(key.value) {
+				v.estack.updateSizeRemove(item)
+			} else if len(t.value) >= MaxArraySize {
 				panic("too big map")
 			}
 			t.Add(key.value, item)
+			v.estack.updateSizeAdd(item)
 
 		default:
 			panic(fmt.Sprintf("SETITEM: invalid item type %s", t))
@@ -952,6 +980,7 @@ func (v *VM) execute(ctx *Context, op Instruction, parameter []byte) (err error)
 			if k < 0 || k >= len(a) {
 				panic("REMOVE: invalid index")
 			}
+			v.estack.updateSizeRemove(a[k])
 			a = append(a[:k], a[k+1:]...)
 			t.value = a
 		case *StructItem:
@@ -960,11 +989,13 @@ func (v *VM) execute(ctx *Context, op Instruction, parameter []byte) (err error)
 			if k < 0 || k >= len(a) {
 				panic("REMOVE: invalid index")
 			}
+			v.estack.updateSizeRemove(a[k])
 			a = append(a[:k], a[k+1:]...)
 			t.value = a
 		case *MapItem:
 			m := t.value
 			k := toMapKey(key.value)
+			v.estack.updateSizeRemove(m[k])
 			delete(m, k)
 		default:
 			panic("REMOVE: invalid type")

pkg/vm/vm_test.go

@@ -84,6 +84,101 @@ func checkVMFailed(t *testing.T, vm *VM) {
 	assert.Equal(t, true, vm.HasFailed())
 }
 
+func TestStackLimitPUSH1Good(t *testing.T) {
+	prog := make([]byte, MaxStackSize*2)
+	for i := 0; i < MaxStackSize; i++ {
+		prog[i] = byte(PUSH1)
+	}
+	for i := MaxStackSize; i < MaxStackSize*2; i++ {
+		prog[i] = byte(DROP)
+	}
+
+	v := load(prog)
+	runVM(t, v)
+}
+
+func TestStackLimitPUSH1Bad(t *testing.T) {
+	prog := make([]byte, MaxStackSize+1)
+	for i := range prog {
+		prog[i] = byte(PUSH1)
+	}
+	v := load(prog)
+	checkVMFailed(t, v)
+}
+
+// appendBigStruct returns a program which:
+// 1. pushes size Structs on stack
+// 2. packs them into a new struct
+// 3. appends them to a zero-length array
+// Resulting stack size consists of:
+// - struct (size+1)
+// - array (1) of struct (size+1)
+// which equals to size*2+3 elements in total.
+func appendBigStruct(size uint16) []Instruction {
+	prog := make([]Instruction, size*2)
+	for i := uint16(0); i < size; i++ {
+		prog[i*2] = PUSH0
+		prog[i*2+1] = NEWSTRUCT
+	}
+
+	return append(prog,
+		PUSHBYTES2, Instruction(size), Instruction(size>>8), // LE
+		PACK, NEWSTRUCT,
+		DUP,
+		PUSH0, NEWARRAY, TOALTSTACK, DUPFROMALTSTACK,
+		SWAP,
+		APPEND, RET)
+}
+
+func TestStackLimitAPPENDStructGood(t *testing.T) {
+	prog := makeProgram(appendBigStruct(MaxStackSize/2 - 2)...)
+	v := load(prog)
+	runVM(t, v) // size = 2047 = (Max/2-2)*2+3 = Max-1
+}
+
+func TestStackLimitAPPENDStructBad(t *testing.T) {
+	prog := makeProgram(appendBigStruct(MaxStackSize/2 - 1)...)
+	v := load(prog)
+	checkVMFailed(t, v) // size = 2049 = (Max/2-1)*2+3 = Max+1
+}
+
+func TestStackLimit(t *testing.T) {
+	expected := []struct {
+		inst Instruction
+		size int
+	}{
+		{PUSH2, 1},
+		{NEWARRAY, 3}, // array + 2 items
+		{TOALTSTACK, 3},
+		{DUPFROMALTSTACK, 4},
+		{NEWSTRUCT, 6}, // all items are copied
+		{NEWMAP, 7},
+		{DUP, 8},
+		{PUSH2, 9},
+		{DUPFROMALTSTACK, 10},
+		{SETITEM, 8}, // -3 items and 1 new element in map
+		{DUP, 9},
+		{PUSH2, 10},
+		{DUPFROMALTSTACK, 11},
+		{SETITEM, 8}, // -3 items and no new elements in map
+		{DUP, 9},
+		{PUSH2, 10},
+		{REMOVE, 7}, // as we have right after NEWMAP
+		{DROP, 6},   // DROP map with no elements
+	}
+
+	prog := make([]Instruction, len(expected))
+	for i := range expected {
+		prog[i] = expected[i].inst
+	}
+
+	vm := load(makeProgram(prog...))
+	for i := range expected {
+		require.NoError(t, vm.Step())
+		require.Equal(t, expected[i].size, vm.size)
+	}
+}
+
 func TestPushBytesShort(t *testing.T) {
 	prog := make([]byte, 10)
 	prog[0] = byte(PUSHBYTES10) // but only 9 left in the `prog`