From 0358d1d5271dc24d04adaef2b568733b0905372d Mon Sep 17 00:00:00 2001
From: John Jannotti <jannotti@gmail.com>
Date: Thu, 12 Oct 2023 14:50:51 -0400
Subject: [PATCH] AVM: Cleanly handle broken switch/match programs (#5782)

---
 data/transactions/logic/assembler.go |  1 +
 data/transactions/logic/eval.go      | 19 +++++++-
 data/transactions/logic/eval_test.go | 72 +++++++++++++++++++++++++++-
 3 files changed, 88 insertions(+), 4 deletions(-)

diff --git a/data/transactions/logic/assembler.go b/data/transactions/logic/assembler.go
index 6bc4381dd6..aa32c51eba 100644
--- a/data/transactions/logic/assembler.go
+++ b/data/transactions/logic/assembler.go
@@ -998,6 +998,7 @@ func asmBranch(ops *OpStream, spec *OpSpec, mnemonic token, args []token) *sourc
 	return nil
 }
 
+// asmSwitch assembles switch and match opcodes
 func asmSwitch(ops *OpStream, spec *OpSpec, mnemonic token, args []token) *sourceError {
 	numOffsets := len(args)
 	if numOffsets > math.MaxUint8 {
diff --git a/data/transactions/logic/eval.go b/data/transactions/logic/eval.go
index 871aa0399a..ad4acd69dd 100644
--- a/data/transactions/logic/eval.go
+++ b/data/transactions/logic/eval.go
@@ -2509,13 +2509,20 @@ func branchTarget(cx *EvalContext) (int, error) {
 }
 
 func switchTarget(cx *EvalContext, branchIdx uint64) (int, error) {
+	if cx.pc+1 >= len(cx.program) {
+		opcode := cx.program[cx.pc]
+		spec := &opsByOpcode[cx.version][opcode]
+		return 0, fmt.Errorf("bare %s opcode at end of program", spec.Name)
+	}
 	numOffsets := int(cx.program[cx.pc+1])
 
 	end := cx.pc + 2          // end of opcode + number of offsets, beginning of offset list
 	eoi := end + 2*numOffsets // end of instruction
 
 	if eoi > len(cx.program) { // eoi will equal len(p) if switch is last instruction
-		return 0, fmt.Errorf("switch claims to extend beyond program")
+		opcode := cx.program[cx.pc]
+		spec := &opsByOpcode[cx.version][opcode]
+		return 0, fmt.Errorf("%s opcode claims to extend beyond program", spec.Name)
 	}
 
 	offset := 0
@@ -2550,8 +2557,13 @@ func checkBranch(cx *EvalContext) error {
 	return nil
 }
 
-// checks switch is encoded properly (and calculates nextpc)
+// checks switch or match is encoded properly (and calculates nextpc)
 func checkSwitch(cx *EvalContext) error {
+	if cx.pc+1 >= len(cx.program) {
+		opcode := cx.program[cx.pc]
+		spec := &opsByOpcode[cx.version][opcode]
+		return fmt.Errorf("bare %s opcode at end of program", spec.Name)
+	}
 	numOffsets := int(cx.program[cx.pc+1])
 	eoi := cx.pc + 2 + 2*numOffsets
 
@@ -2628,6 +2640,9 @@ func opSwitch(cx *EvalContext) error {
 }
 
 func opMatch(cx *EvalContext) error {
+	if cx.pc+1 >= len(cx.program) {
+		return fmt.Errorf("bare match opcode at end of program")
+	}
 	n := int(cx.program[cx.pc+1])
 	// stack contains the n sized match list and the single match value
 	if n+1 > len(cx.Stack) {
diff --git a/data/transactions/logic/eval_test.go b/data/transactions/logic/eval_test.go
index e92bd9a486..230ad8c640 100644
--- a/data/transactions/logic/eval_test.go
+++ b/data/transactions/logic/eval_test.go
@@ -3156,8 +3156,7 @@ done:
 int 1
 `, v)
 			// cut two last bytes - intc_1 and last byte of bnz
-			ops.Program = ops.Program[:len(ops.Program)-2]
-			testLogicBytes(t, ops.Program, nil,
+			testLogicBytes(t, ops.Program[:len(ops.Program)-2], nil,
 				"bnz program ends short", "bnz program ends short")
 		})
 	}
@@ -5835,6 +5834,40 @@ switch done1 done2; done1: ; done2: ;
 `, 8)
 }
 
+// TestShortSwitch ensures a clean error, in Check and Eval, when a switch ends early
+func TestShortSwitch(t *testing.T) {
+	partitiontest.PartitionTest(t)
+	t.Parallel()
+
+	source := `
+	int 1
+	int 1
+	switch label1 label2
+	label1:
+	label2:
+	`
+	ops, err := AssembleStringWithVersion(source, AssemblerMaxVersion)
+	require.NoError(t, err)
+
+	// fine as is
+	testLogicBytes(t, ops.Program, nil)
+
+	beyond := "switch opcode claims to extend beyond program"
+
+	// bad if a label is gone
+	testLogicBytes(t, ops.Program[:len(ops.Program)-2], nil, beyond, beyond)
+
+	// chop off all the labels, but keep the label count
+	testLogicBytes(t, ops.Program[:len(ops.Program)-4], nil, beyond, beyond)
+
+	// chop off before the label count
+	testLogicBytes(t, ops.Program[:len(ops.Program)-5], nil,
+		"bare switch opcode at end of program", "bare switch opcode at end of program")
+
+	// chop off half of a label
+	testLogicBytes(t, ops.Program[:len(ops.Program)-1], nil, beyond, beyond)
+}
+
 func TestMatch(t *testing.T) {
 	partitiontest.PartitionTest(t)
 	t.Parallel()
@@ -5987,6 +6020,41 @@ one:  int 0;
 `, 8)
 }
 
+// TestShortMatch ensures a clean error when a match ends early
+func TestShortMatch(t *testing.T) {
+	partitiontest.PartitionTest(t)
+	t.Parallel()
+
+	source := `int 1
+    int 40
+    int 45
+    int 40
+	match label1 label2
+	label1:
+    label2:
+	`
+	ops, err := AssembleStringWithVersion(source, AssemblerMaxVersion)
+	require.NoError(t, err)
+
+	// fine as is
+	testLogicBytes(t, ops.Program, nil)
+
+	beyond := "match opcode claims to extend beyond program"
+
+	// bad if a label is gone
+	testLogicBytes(t, ops.Program[:len(ops.Program)-2], nil, beyond, beyond)
+
+	// chop off all the labels, but keep the label count
+	testLogicBytes(t, ops.Program[:len(ops.Program)-4], nil, beyond, beyond)
+
+	// chop off before the label count
+	testLogicBytes(t, ops.Program[:len(ops.Program)-5], nil,
+		"bare match opcode at end of program", "bare match opcode at end of program")
+
+	// chop off half of a label
+	testLogicBytes(t, ops.Program[:len(ops.Program)-1], nil, beyond, beyond)
+}
+
 func TestPushConsts(t *testing.T) {
 	partitiontest.PartitionTest(t)
 	t.Parallel()