interpreter: Fix out-of-range execution ProgramCounter

Currently, the program counter always returns 0 when the last instruction
is branch instruction and it evaluates to false. The correct behavior
should be returning the faulting address.

This is also weird because the program counter returns correct value if
the last instruction is something else.

Also, on re-compiler we do return correct faulting address.

This patch fixes interpreter behavior and makes it consistent with
re-compiler.

Signed-off-by: Aman <[email protected]>

Author: aman4150

.github/workflows/rust.yml

@@ -46,6 +46,8 @@ jobs:
     - uses: actions/checkout@v4
     - name: Install clippy (base toolchain)
       run: rustup component add clippy
+    - name: Install clippy (zygote toolchain)
+      run: cd crates/polkavm-zygote && rustup component add clippy
     - name: Install clippy (benchtool toolchain)
       run: cd tools/benchtool && rustup component add clippy
     - name: Install clippy (guest-programs toolchain)

ci/jobs/fuzz.sh

@@ -11,3 +11,9 @@ cargo fuzz run fuzz_generic_allocator -- -runs=1000000
 
 echo ">> cargo fuzz run (fuzz_shm_allocator)"
 cargo fuzz run fuzz_shm_allocator -- -runs=1000000
+
+echo ">> cargo fuzz run (fuzz_linker)"
+cargo fuzz run fuzz_linker -- -runs=10000
+
+echo ">> cargo fuzz run (fuzz_polkavm)"
+cargo fuzz run fuzz_polkavm -- -runs=10000

crates/polkavm/src/interpreter.rs

@@ -411,6 +411,7 @@ pub(crate) struct InterpretedInstance {
     compiled_offset: u32,
     interrupt: InterruptKind,
     step_tracing: bool,
+    unresolved_program_counter: Option<ProgramCounter>,
 }
 
 impl InterpretedInstance {
@@ -433,6 +434,7 @@ impl InterpretedInstance {
             compiled_offset: 0,
             interrupt: InterruptKind::Finished,
             step_tracing,
+            unresolved_program_counter: None,
         };
 
         instance.initialize_module();
@@ -715,9 +717,15 @@ impl InterpretedInstance {
             };
 
             self.program_counter = program_counter;
-            self.compiled_offset = self.resolve_arbitrary_jump::<DEBUG>(program_counter).unwrap_or(TARGET_OUT_OF_RANGE);
             self.next_program_counter_changed = false;
 
+            if let Some(offset) = self.resolve_arbitrary_jump::<DEBUG>(program_counter) {
+                self.compiled_offset = offset;
+            } else {
+                self.compiled_offset = TARGET_OUT_OF_RANGE;
+                self.unresolved_program_counter = Some(program_counter);
+            }
+
             if DEBUG {
                 log::debug!("Starting execution at: {} [{}]", program_counter, self.compiled_offset);
             }
@@ -1035,10 +1043,24 @@ impl<'a> Visitor<'a> {
     ) -> Option<Target> {
         let s1 = self.get64(s1);
         let s2 = self.get64(s2);
-        if callback(s1, s2) {
-            Some(target_true)
+
+        let (packed_target, branch_taken) = if callback(s1, s2) {
+            (target_true, true)
+        } else {
+            (target_false, false)
+        };
+
+        let (is_jump_target_valid, target) = InterpretedInstance::unpack_target(NonZeroU32::new(packed_target).unwrap());
+
+        if !is_jump_target_valid {
+            self.inner.unresolved_program_counter = Some(ProgramCounter(target));
+            if branch_taken {
+                Some(TARGET_INVALID_BRANCH)
+            } else {
+                Some(TARGET_OUT_OF_RANGE)
+            }
         } else {
-            Some(target_false)
+            Some(target)
         }
     }
 
@@ -1857,6 +1879,7 @@ fn trap_impl<const DEBUG: bool>(visitor: &mut Visitor, program_counter: ProgramC
     visitor.inner.program_counter_valid = true;
     visitor.inner.next_program_counter = None;
     visitor.inner.next_program_counter_changed = true;
+    visitor.inner.unresolved_program_counter = None;
     visitor.inner.interrupt = InterruptKind::Trap;
     None
 }
@@ -1891,10 +1914,21 @@ macro_rules! handle_unresolved_branch {
         }
 
         let offset = $visitor.inner.compiled_offset;
-        let target_false = $visitor.inner.resolve_jump::<DEBUG>($tf).unwrap_or(TARGET_OUT_OF_RANGE);
-        let target_true = $visitor.inner.resolve_jump::<DEBUG>($tt).unwrap_or(TARGET_INVALID_BRANCH);
+
+        let packed_tt = if let Some(target_true) = $visitor.inner.resolve_jump::<DEBUG>($tt) {
+            InterpretedInstance::pack_target(cast(target_true).to_usize(), true).into()
+        } else {
+            InterpretedInstance::pack_target(cast($tt.0).to_usize(), false).into()
+        };
+
+        let packed_tf = if let Some(target_false) = $visitor.inner.resolve_jump::<DEBUG>($tf) {
+            InterpretedInstance::pack_target(cast(target_false).to_usize(), true).into()
+        } else {
+            InterpretedInstance::pack_target(cast($tf.0).to_usize(), false).into()
+        };
+
         $visitor.inner.compiled_handlers[cast(offset).to_usize()] = cast_handler!(raw_handlers::$name::<DEBUG>);
-        $visitor.inner.compiled_args[cast(offset).to_usize()] = Args::$name($s1, $s2, target_true, target_false);
+        $visitor.inner.compiled_args[cast(offset).to_usize()] = Args::$name($s1, $s2, packed_tt, packed_tf);
         Some(offset)
     }};
 }
@@ -1929,7 +1963,8 @@ define_interpreter! {
             log::trace!("[{}]: trap (out of range)", visitor.inner.compiled_offset);
         }
 
-        let program_counter = visitor.inner.program_counter;
+        let program_counter = visitor.inner.unresolved_program_counter.unwrap_or(visitor.inner.program_counter);
+
         let new_gas = visitor.inner.gas - i64::from(gas);
         if new_gas < 0 {
             not_enough_gas_impl::<DEBUG>(visitor, program_counter, new_gas)
@@ -1960,8 +1995,11 @@ define_interpreter! {
             log::trace!("[{}]: step (out of range)", visitor.inner.compiled_offset);
         }
 
+        let program_counter = visitor.inner.unresolved_program_counter.unwrap_or(visitor.inner.program_counter);
+
+        visitor.inner.program_counter = program_counter;
         visitor.inner.program_counter_valid = true;
-        visitor.inner.next_program_counter = Some(visitor.inner.program_counter);
+        visitor.inner.next_program_counter = Some(program_counter);
         visitor.inner.next_program_counter_changed = false;
         visitor.inner.interrupt = InterruptKind::Step;
         visitor.inner.compiled_offset += 1;
@@ -3640,6 +3678,7 @@ define_interpreter! {
         } else {
             visitor.inner.compiled_handlers[cast(offset).to_usize()] = cast_handler!(raw_handlers::jump::<DEBUG>);
             visitor.inner.compiled_args[cast(offset).to_usize()] = Args::jump(TARGET_OUT_OF_RANGE);
+            visitor.inner.unresolved_program_counter = Some(jump_to);
             Some(TARGET_OUT_OF_RANGE)
         }
     }

crates/polkavm/src/tests.rs

@@ -692,6 +692,28 @@ fn simple_test(engine_config: Config) {
     assert_eq!(instance.reg(Reg::A1), 100);
 }
 
+fn out_of_range_execution(engine_config: Config) {
+    let _ = env_logger::try_init();
+    let engine = Engine::new(&engine_config).unwrap();
+    let mut builder = ProgramBlobBuilder::new();
+    builder.add_export_by_basic_block(0, b"main");
+    builder.set_code(&[asm::load_imm(A0, 1), asm::load_imm(A0, 2), asm::branch_eq_imm(RA, 0, 0)], &[]);
+
+    let blob = ProgramBlob::parse(builder.into_vec().unwrap().into()).unwrap();
+    let module = Module::from_blob(&engine, &ModuleConfig::new(), blob).unwrap();
+    let next_offsets: Vec<_> = module
+        .blob()
+        .instructions(DefaultInstructionSet::default())
+        .map(|inst| inst.next_offset)
+        .collect();
+
+    let mut instance = module.instantiate().unwrap();
+    instance.set_reg(Reg::RA, crate::RETURN_TO_HOST);
+    instance.set_next_program_counter(ProgramCounter(0));
+    match_interrupt!(instance.run().unwrap(), InterruptKind::Trap);
+    assert_eq!(instance.program_counter(), Some(next_offsets[2]));
+}
+
 fn jump_into_middle_of_basic_block_from_outside(engine_config: Config) {
     let _ = env_logger::try_init();
     let engine = Engine::new(&engine_config).unwrap();
@@ -3861,6 +3883,8 @@ run_tests! {
     trapping_preserves_all_registers_normal_trap
     trapping_preserves_all_registers_segfault
 
+    out_of_range_execution
+
     memset_basic
     memset_with_dynamic_paging
 

fuzz/fuzz_targets/fuzz_polkavm.rs

@@ -605,7 +605,7 @@ fn correctness_fuzzer_harness(data: Vec<OperationKind>) {
         //
 
         match interrupt_interpreter.clone() {
-            InterruptKind::NotEnoughGas | InterruptKind::Trap => {
+            InterruptKind::Finished | InterruptKind::NotEnoughGas | InterruptKind::Trap => {
                 break;
             }
             polkavm::InterruptKind::Ecalli(_) => {