diff --git a/pkg/vm/stack.go b/pkg/vm/stack.go index 1a230b299..4f4db751e 100644 --- a/pkg/vm/stack.go +++ b/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 } diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go index 0043afa11..8f7db244b 100644 --- a/pkg/vm/vm.go +++ b/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") diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go index 599b930f6..2eedf0ffa 100644 --- a/pkg/vm/vm_test.go +++ b/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`