From 18e36717817d8a83d9f2894e8ca6ef2871387086 Mon Sep 17 00:00:00 2001
From: Joonatan Saarhelo <joon.saar@gmail.com>
Date: Tue, 20 Aug 2024 19:46:02 +0200
Subject: [PATCH 1/2] prevent resuming after execution ends

---
 src/instruction_handlers/ret.rs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/instruction_handlers/ret.rs b/src/instruction_handlers/ret.rs
index df011f7d..30768c28 100644
--- a/src/instruction_handlers/ret.rs
+++ b/src/instruction_handlers/ret.rs
@@ -72,6 +72,10 @@ fn ret<T: Tracer, RT: TypeLevelReturnType, const TO_LABEL: bool>(
                 // the caller may take external snapshots while the VM is in the initial frame and
                 // these would break were the initial frame to be rolled back.
 
+                // But to continue execution would be nonsensical and can cause UB because there
+                // is no next instruction after a panic arising from some other instruction.
+                vm.state.current_frame.pc = invalid_instruction();
+
                 return if let Some(return_value) = return_value_or_panic {
                     let output = vm.state.heaps[return_value.memory_page]
                         .read_range_big_endian(

From 63b40ce503ddb2b134ce7117713b6b38a85783c9 Mon Sep 17 00:00:00 2001
From: Joonatan Saarhelo <joon.saar@gmail.com>
Date: Tue, 20 Aug 2024 20:35:13 +0200
Subject: [PATCH 2/2] save pc in snapshots

---
 src/callframe.rs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/callframe.rs b/src/callframe.rs
index 6409f588..b8af0ae5 100644
--- a/src/callframe.rs
+++ b/src/callframe.rs
@@ -171,6 +171,7 @@ impl<T> Callframe<T> {
 
             context_u128: self.context_u128,
             sp: self.sp,
+            pc: self.get_pc_as_u16(),
             gas: self.gas,
             near_calls: self.near_calls.clone(),
             heap_size: self.heap_size,
@@ -189,6 +190,7 @@ impl<T> Callframe<T> {
             stack,
             context_u128,
             sp,
+            pc,
             gas,
             near_calls,
             heap_size,
@@ -200,6 +202,7 @@ impl<T> Callframe<T> {
 
         self.context_u128 = context_u128;
         self.sp = sp;
+        self.set_pc_from_u16(pc);
         self.gas = gas;
         self.near_calls = near_calls;
         self.heap_size = heap_size;
@@ -221,6 +224,7 @@ pub(crate) struct CallframeSnapshot {
     stack: StackSnapshot,
     context_u128: u128,
     sp: u16,
+    pc: u16,
     gas: u32,
     near_calls: Vec<NearCallFrame>,
     heap_size: u32,