From c7f0b7bd6849bde6b0c729f91e4e3bbd0a2546a0 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <kraunid@gmail.com>
Date: Tue, 29 Oct 2019 11:01:06 +0300
Subject: [PATCH] vm: restrict max invocation stack size

---
 pkg/vm/vm.go      | 23 +++++++++++++++++++++--
 pkg/vm/vm_test.go | 21 +++++++++++++++++++++
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index 122e4062b..f425a5a6f 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -41,6 +41,9 @@ const (
 	// MaxItemSize is the maximum item size allowed in the VM.
 	MaxItemSize = 1024 * 1024
 
+	// MaxInvocationStackSize is the maximum size of an invocation stack.
+	MaxInvocationStackSize = 1024
+
 	maxSHLArg = 256
 	minSHLArg = -256
 )
@@ -995,6 +998,8 @@ func (v *VM) execute(ctx *Context, op Instruction, parameter []byte) (err error)
 		}
 
 	case CALL:
+		v.checkInvocationStackSize()
+
 		newCtx := ctx.Copy()
 		newCtx.rvcount = -1
 		v.istack.PushVal(newCtx)
@@ -1017,6 +1022,10 @@ func (v *VM) execute(ctx *Context, op Instruction, parameter []byte) (err error)
 			panic("no getScript callback is set up")
 		}
 
+		if op == APPCALL {
+			v.checkInvocationStackSize()
+		}
+
 		hash, err := util.Uint160DecodeBytes(parameter)
 		if err != nil {
 			panic(err)
@@ -1227,8 +1236,12 @@ func (v *VM) execute(ctx *Context, op Instruction, parameter []byte) (err error)
 		if v.estack.Len() < pcount+addElement {
 			panic("missing some parameters")
 		}
-		if tailCall && ctx.rvcount != rvcount {
-			panic("context and parameter rvcount mismatch")
+		if tailCall {
+			if ctx.rvcount != rvcount {
+				panic("context and parameter rvcount mismatch")
+			}
+		} else {
+			v.checkInvocationStackSize()
 		}
 
 		if op == CALLI {
@@ -1313,3 +1326,9 @@ func validateMapKey(key *Element) {
 		panic("key can't be a collection")
 	}
 }
+
+func (v *VM) checkInvocationStackSize() {
+	if v.istack.len >= MaxInvocationStackSize {
+		panic("invocation stack is too big")
+	}
+}
diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go
index e3be5ce09..7eeb04a63 100644
--- a/pkg/vm/vm_test.go
+++ b/pkg/vm/vm_test.go
@@ -197,6 +197,27 @@ func TestPushData4Good(t *testing.T) {
 	assert.Equal(t, []byte{1, 2, 3}, vm.estack.Pop().Bytes())
 }
 
+func callNTimes(n uint16) []byte {
+	return makeProgram(
+		PUSHBYTES2, Instruction(n), Instruction(n>>8), // little-endian
+		TOALTSTACK, DUPFROMALTSTACK,
+		JMPIF, 0x4, 0, RET,
+		FROMALTSTACK, DEC,
+		CALL, 0xF8, 0xFF) // -8 -> JMP to TOALTSTACK)
+}
+
+func TestInvocationLimitGood(t *testing.T) {
+	prog := callNTimes(MaxInvocationStackSize - 1)
+	v := load(prog)
+	runVM(t, v)
+}
+
+func TestInvocationLimitBad(t *testing.T) {
+	prog := callNTimes(MaxInvocationStackSize)
+	v := load(prog)
+	checkVMFailed(t, v)
+}
+
 func TestNOTNoArgument(t *testing.T) {
 	prog := makeProgram(NOT)
 	vm := load(prog)