From 3c53beca82cf9675148613627e9c69cf5b787158 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Wed, 18 Sep 2019 14:35:29 +0300
Subject: [PATCH] vm: restrict SHL/SHR arguments to -256..256

Also unify SHL/SHR implementation.
---
 pkg/vm/vm.go      | 25 ++++++++++++++-----------
 pkg/vm/vm_test.go | 18 ++++++++++++++++++
 2 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index 2bed4fb6c..0108b6a4e 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -22,6 +22,11 @@ var (
 	ModeMute Mode = 1 << 0
 )
 
+const (
+	maxSHLArg = 256
+	minSHLArg = -256
+)
+
 // VM represents the virtual machine.
 type VM struct {
 	state State
@@ -508,21 +513,19 @@ func (v *VM) execute(ctx *Context, op Instruction) {
 		a := v.estack.Pop().BigInt()
 		v.estack.PushVal(new(big.Int).Mod(a, b))
 
-	case SHL:
-		b := v.estack.Pop().BigInt()
-		if b.Int64() == 0 {
+	case SHL, SHR:
+		b := v.estack.Pop().BigInt().Int64()
+		if b == 0 {
 			return
+		} else if b < minSHLArg || b > maxSHLArg {
+			panic(fmt.Sprintf("operand must be between %d and %d", minSHLArg, maxSHLArg))
 		}
 		a := v.estack.Pop().BigInt()
-		v.estack.PushVal(new(big.Int).Lsh(a, uint(b.Int64())))
-
-	case SHR:
-		b := v.estack.Pop().BigInt()
-		if b.Int64() == 0 {
-			return
+		if op == SHL {
+			v.estack.PushVal(new(big.Int).Lsh(a, uint(b)))
+		} else {
+			v.estack.PushVal(new(big.Int).Rsh(a, uint(b)))
 		}
-		a := v.estack.Pop().BigInt()
-		v.estack.PushVal(new(big.Int).Rsh(a, uint(b.Int64())))
 
 	case BOOLAND:
 		b := v.estack.Pop().Bool()
diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go
index ac9dbe89e..80cceba9b 100644
--- a/pkg/vm/vm_test.go
+++ b/pkg/vm/vm_test.go
@@ -305,6 +305,15 @@ func TestSHRZero(t *testing.T) {
 	assert.Equal(t, makeStackItem([]byte{0, 1}), vm.estack.Pop().value)
 }
 
+func TestSHRSmallValue(t *testing.T) {
+	prog := makeProgram(SHR)
+	vm := load(prog)
+	vm.estack.PushVal(5)
+	vm.estack.PushVal(-257)
+	vm.Run()
+	assert.Equal(t, true, vm.state.HasFlag(faultState))
+}
+
 func TestSHLGood(t *testing.T) {
 	prog := makeProgram(SHL)
 	vm := load(prog)
@@ -327,6 +336,15 @@ func TestSHLZero(t *testing.T) {
 	assert.Equal(t, makeStackItem([]byte{0, 1}), vm.estack.Pop().value)
 }
 
+func TestSHLBigValue(t *testing.T) {
+	prog := makeProgram(SHL)
+	vm := load(prog)
+	vm.estack.PushVal(5)
+	vm.estack.PushVal(257)
+	vm.Run()
+	assert.Equal(t, true, vm.state.HasFlag(faultState))
+}
+
 func TestLT(t *testing.T) {
 	prog := makeProgram(LT)
 	vm := load(prog)