generic-sandbox: Fix invalid jump detection on macOS / Rosetta 2

On Rosetta 2, indirect jump to a non-canonical address doesn't trigger
a page fault (SIGBUS) right away. Instead, it jumps to the address, and then
triggers a page fault (SIGBUS) when trying to fetch the instruction.

This means that previous program counter is now lost, and now we have no
way of detecting if the guest program triggered an invalid jump or not.

We fix this program by updating the vmctx before making the jump and in
the signal handler we check if the current program counter is non-canonical.

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

Author: aman4150

crates/polkavm/src/compiler/amd64.rs

@@ -920,6 +920,22 @@ where
                 asm.assert_reserved_exactly_as_needed();
             }
             SandboxKind::Generic => {
+                // On Rosetta 2, indirect jump to a non-canonical address doesn't trigger a SIGBUS right away.
+                // Instead, it jumps to the address, and then triggers a SIGBUS when trying to fetch the instruction.
+                // This means that previous program counter is now lost. There are other ways we can design it (by adding a
+                // dedicated jump table subroutine) but that would be less efficient.
+                // Therefore, let's store the executing address to the program counter before jumping.
+                #[cfg(target_os = "macos")]
+                {
+                    let label_start = self.asm.create_label();
+                    self.asm.push(lea_rip_label(TMP_REG, label_start));
+                    self.push(store(
+                        RegSize::R64,
+                        Self::vmctx_field(S::offset_table().next_native_program_counter),
+                        TMP_REG,
+                    ));
+                }
+
                 // TODO: This also could be more efficient.
                 self.push(lea_rip_label(TMP_REG, self.jump_table_label));
                 self.push(push(conv_reg(base)));

crates/polkavm/src/sandbox/generic.rs

@@ -415,6 +415,21 @@ unsafe extern "C" fn signal_handler(signal: c_int, info: &sys::siginfo_t, contex
 
         let rip = fetch_reg!(rip);
         let vmctx = &mut *vmctx;
+
+        // On Rosetta 2, the JMP emulation logic doesn't work same as on x64.
+        // Instead of triggering a GPF immdiately, it jumps to that address and then trigger a PF.
+        // Therefore the original program counter is lost.
+        // We fix this problem by storing the program counter in vmctx before jumping.
+        // See jump_indirect_impl for more details.
+        #[cfg(target_os = "macos")]
+        {
+            let is_invalid_rip = (rip >> 48) != 0;
+            if is_invalid_rip {
+                log::trace!("Jump table invalid address hit, returning to host");
+                trigger_exit(vmctx, ExitReason::Signal);
+            }
+        }
+
         if vmctx.program_range.contains(&rip) {
             use polkavm_common::regmap::NativeReg;
             for reg in polkavm_common::program::Reg::ALL {