From 7ddd6e864687b7ce380f04b881c2f6cf8407311d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20Mei=C3=9Fner?= <AlexanderMeissner@gmx.net>
Date: Sat, 23 Sep 2023 19:44:34 +0200
Subject: [PATCH] Restricts the final instruction in each function to be
 diverted.

---
 src/verifier.rs    |  2 +-
 tests/execution.rs | 22 +++++++++++++---------
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/verifier.rs b/src/verifier.rs
index 535592f7..b6502582 100644
--- a/src/verifier.rs
+++ b/src/verifier.rs
@@ -243,7 +243,7 @@ impl Verifier for RequisiteVerifier {
                 function_range.end = *function_iter.peek().unwrap_or(&program_range.end);
                 let insn = ebpf::get_insn(prog, function_range.end.saturating_sub(1));
                 match insn.opc {
-                    ebpf::JA | ebpf::CALL_IMM | ebpf::CALL_REG | ebpf::EXIT => {},
+                    ebpf::JA | ebpf::EXIT => {},
                     _ => return Err(VerifierError::InvalidFunction(
                         function_range.end.saturating_sub(1),
                     )),
diff --git a/tests/execution.rs b/tests/execution.rs
index bb2f9a95..1fc9e0ff 100644
--- a/tests/execution.rs
+++ b/tests/execution.rs
@@ -2630,8 +2630,9 @@ fn test_tight_infinite_recursion_callx() {
         "
         mov64 r8, 0x1
         lsh64 r8, 0x20
-        or64 r8, 0x20
+        or64 r8, 0x28
         call function_foo
+        exit
         function_foo:
         mov64 r3, 0x41414141
         callx r8
@@ -2639,7 +2640,7 @@ fn test_tight_infinite_recursion_callx() {
         [],
         (),
         TestContextObject::new(8),
-        ProgramResult::Err(Box::new(EbpfError::ExceededMaxInstructions(35))),
+        ProgramResult::Err(Box::new(EbpfError::ExceededMaxInstructions(36))),
     );
 }
 
@@ -2665,7 +2666,7 @@ fn test_err_instruction_count_syscall_capped() {
     test_interpreter_and_jit_asm!(
         "
         mov64 r2, 0x5
-        call 0
+        syscall bpf_syscall_string
         mov64 r0, 0x0
         exit",
         [72, 101, 108, 108, 111],
@@ -2778,22 +2779,24 @@ fn test_err_exit_capped() {
         "
         mov64 r1, 0x1
         lsh64 r1, 0x20
-        or64 r1, 0x20
+        or64 r1, 0x28
         callx r1
+        exit
         function_foo:
         exit
         ",
         [],
         (),
         TestContextObject::new(5),
-        ProgramResult::Err(Box::new(EbpfError::ExceededMaxInstructions(34))),
+        ProgramResult::Err(Box::new(EbpfError::ExceededMaxInstructions(35))),
     );
     test_interpreter_and_jit_asm!(
         "
         mov64 r1, 0x1
         lsh64 r1, 0x20
-        or64 r1, 0x20
+        or64 r1, 0x28
         callx r1
+        exit
         function_foo:
         mov r0, r0
         exit
@@ -2801,18 +2804,19 @@ fn test_err_exit_capped() {
         [],
         (),
         TestContextObject::new(6),
-        ProgramResult::Err(Box::new(EbpfError::ExceededMaxInstructions(35))),
+        ProgramResult::Err(Box::new(EbpfError::ExceededMaxInstructions(36))),
     );
     test_interpreter_and_jit_asm!(
         "
-        call 0
+        call 1
+        exit
         mov r0, r0
         exit
         ",
         [],
         (),
         TestContextObject::new(3),
-        ProgramResult::Err(Box::new(EbpfError::ExceededMaxInstructions(32))),
+        ProgramResult::Err(Box::new(EbpfError::ExceededMaxInstructions(33))),
     );
 }