diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go index 46e60a5f7..584b7c14e 100644 --- a/pkg/vm/vm.go +++ b/pkg/vm/vm.go @@ -705,11 +705,11 @@ func (v *VM) execute(ctx *Context, op Instruction) { v.estack.PushVal(l) case PICKITEM: - var ( - key = v.estack.Pop() - obj = v.estack.Pop() - index = int(key.BigInt().Int64()) - ) + key := v.estack.Pop() + validateMapKey(key) + + obj := v.estack.Pop() + index := int(key.BigInt().Int64()) switch t := obj.value.(type) { // Struct and Array items have their underlying value as []StackItem. @@ -720,6 +720,12 @@ func (v *VM) execute(ctx *Context, op Instruction) { } item := arr[index] v.estack.PushVal(item) + case *MapItem: + if !t.Has(key.value) { + panic("invalid key") + } + k := toMapKey(key.value) + v.estack.Push(&Element{value: t.value[k]}) default: arr := obj.Bytes() if index < 0 || index >= len(arr) { @@ -730,21 +736,24 @@ func (v *VM) execute(ctx *Context, op Instruction) { } case SETITEM: - var ( - item = v.estack.Pop().value - key = v.estack.Pop() - obj = v.estack.Pop() - index = int(key.BigInt().Int64()) - ) + item := v.estack.Pop().value + key := v.estack.Pop() + validateMapKey(key) + + obj := v.estack.Pop() switch t := obj.value.(type) { // Struct and Array items have their underlying value as []StackItem. case *ArrayItem, *StructItem: arr := t.Value().([]StackItem) + index := int(key.BigInt().Int64()) if index < 0 || index >= len(arr) { panic("SETITEM: invalid index") } arr[index] = item + case *MapItem: + t.Add(key.value, item) + default: panic(fmt.Sprintf("SETITEM: invalid item type %s", t)) } diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go index 2c0842bae..ad5aa475d 100644 --- a/pkg/vm/vm_test.go +++ b/pkg/vm/vm_test.go @@ -646,6 +646,39 @@ func TestPICKITEMByteArray(t *testing.T) { assert.Equal(t, makeStackItem(2), vm.estack.Pop().value) } +func TestPICKITEMMap(t *testing.T) { + prog := makeProgram(PICKITEM) + vm := load(prog) + + m := NewMapItem() + m.Add(makeStackItem(5), makeStackItem(3)) + vm.estack.Push(&Element{value: m}) + vm.estack.PushVal(makeStackItem(5)) + + vm.Run() + assert.Equal(t, false, vm.HasFailed()) + assert.Equal(t, 1, vm.estack.Len()) + assert.Equal(t, makeStackItem(3), vm.estack.Pop().value) +} + +func TestSETITEMMap(t *testing.T) { + prog := makeProgram(SETITEM, PICKITEM) + vm := load(prog) + + m := NewMapItem() + m.Add(makeStackItem(5), makeStackItem(3)) + vm.estack.Push(&Element{value: m}) + vm.estack.PushVal(5) + vm.estack.Push(&Element{value: m}) + vm.estack.PushVal(5) + vm.estack.PushVal([]byte{0, 1}) + + vm.Run() + assert.Equal(t, false, vm.HasFailed()) + assert.Equal(t, 1, vm.estack.Len()) + assert.Equal(t, makeStackItem([]byte{0, 1}), vm.estack.Pop().value) +} + func TestSIZENoArgument(t *testing.T) { prog := makeProgram(SIZE) vm := load(prog)