From c22d09adad4d83b0d8aeec628d18bd3ff0408524 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Mon, 20 Jan 2020 14:49:07 +0300
Subject: [PATCH] vm: set GAS limit

Make it possible to set maximum amount of GAS which can be spent
during execution.
---
 pkg/vm/vm.go      | 10 ++++++++++
 pkg/vm/vm_test.go | 16 +++++++++++++++-
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index 77c69c528..b91565aa2 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -80,6 +80,7 @@ type VM struct {
 	size      int
 
 	gasConsumed util.Fixed8
+	gasLimit    util.Fixed8
 
 	// Public keys cache.
 	keys map[string]*keys.PublicKey
@@ -130,6 +131,12 @@ func (v *VM) GasConsumed() util.Fixed8 {
 	return v.gasConsumed
 }
 
+// SetGasLimit sets maximum amount of gas which v can spent.
+// If max <= 0, no limit is imposed.
+func (v *VM) SetGasLimit(max util.Fixed8) {
+	v.gasLimit = max
+}
+
 // Estack returns the evaluation stack so interop hooks can utilize this.
 func (v *VM) Estack() *Stack {
 	return v.estack
@@ -483,6 +490,9 @@ func (v *VM) execute(ctx *Context, op opcode.Opcode, parameter []byte) (err erro
 
 	if v.getPrice != nil && ctx.ip < len(ctx.prog) {
 		v.gasConsumed += v.getPrice(v, op, parameter)
+		if v.gasLimit > 0 && v.gasConsumed > v.gasLimit {
+			panic("gas limit is exceeded")
+		}
 	}
 
 	if op >= opcode.PUSHBYTES1 && op <= opcode.PUSHBYTES75 {
diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go
index 012515124..cab01864a 100644
--- a/pkg/vm/vm_test.go
+++ b/pkg/vm/vm_test.go
@@ -92,7 +92,21 @@ func TestVM_SetPriceGetter(t *testing.T) {
 		v.Load(prog)
 		runVM(t, v)
 
-		require.EqualValues(t, 9, v.gasConsumed)
+		require.EqualValues(t, 9, v.GasConsumed())
+	})
+
+	t.Run("with sufficient gas limit", func(t *testing.T) {
+		v.Load(prog)
+		v.SetGasLimit(9)
+		runVM(t, v)
+
+		require.EqualValues(t, 9, v.GasConsumed())
+	})
+
+	t.Run("with small gas limit", func(t *testing.T) {
+		v.Load(prog)
+		v.SetGasLimit(8)
+		checkVMFailed(t, v)
 	})
 }