Merge pull request #462 from nspcc-dev/feature/stack_size

Restrict total stack item count in the VM. Refs. #373.
This commit is contained in:
Roman Khimov 2019-11-06 11:09:13 +03:00 committed by GitHub
commit 5544ff1768
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 182 additions and 3 deletions

View file

@ -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
}

View file

@ -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")

View file

@ -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`