Implement custom stack unwinding for Pulley (#9758)

* Implement custom stack unwinding for Pulley

The pulley calling convention may not match the native calling
convention, which is the case for s390x, so the unwinding that Wasmtime
does needs to be parameterized over what's being unwound. This commit
adds a new `Unwind` trait with various methods behind it that the
backtrace implementation is then updated to use.

* Fix alignment check for 32-bit

Author: alexcrichton

crates/wasmtime/src/runtime/profiling.rs

@@ -137,7 +137,7 @@ impl GuestProfiler {
         let now = Timestamp::from_nanos_since_reference(
             self.start.elapsed().as_nanos().try_into().unwrap(),
         );
-        let backtrace = Backtrace::new(store.as_context().0.vmruntime_limits());
+        let backtrace = Backtrace::new(store.as_context().0);
         let frames = lookup_frames(&self.modules, &backtrace);
         self.profile
             .add_sample(self.thread, now, frames, delta.into(), 1);
@@ -154,7 +154,7 @@ impl GuestProfiler {
         match kind {
             CallHook::CallingWasm | CallHook::ReturningFromWasm => {}
             CallHook::CallingHost => {
-                let backtrace = Backtrace::new(store.as_context().0.vmruntime_limits());
+                let backtrace = Backtrace::new(store.as_context().0);
                 let frames = lookup_frames(&self.modules, &backtrace);
                 self.profile.add_marker_with_stack(
                     self.thread,

crates/wasmtime/src/runtime/store.rs

@@ -85,7 +85,8 @@ use crate::runtime::vm::mpk::{self, ProtectionKey, ProtectionMask};
 use crate::runtime::vm::{
     Backtrace, ExportGlobal, GcRootsList, GcStore, InstanceAllocationRequest, InstanceAllocator,
     InstanceHandle, Interpreter, InterpreterRef, ModuleRuntimeInfo, OnDemandInstanceAllocator,
-    SignalHandler, StoreBox, StorePtr, VMContext, VMFuncRef, VMGcRef, VMRuntimeLimits,
+    SignalHandler, StoreBox, StorePtr, Unwind, UnwindHost, UnwindPulley, VMContext, VMFuncRef,
+    VMGcRef, VMRuntimeLimits,
 };
 use crate::trampoline::VMHostGlobalContext;
 use crate::type_registry::RegisteredType;
@@ -1739,7 +1740,7 @@ impl StoreOpaque {
 
         log::trace!("Begin trace GC roots :: Wasm stack");
 
-        Backtrace::trace(self.vmruntime_limits().cast_const(), |frame| {
+        Backtrace::trace(self, |frame| {
             let pc = frame.pc();
             debug_assert!(pc != 0, "we should always get a valid PC for Wasm frames");
 
@@ -2148,6 +2149,14 @@ at https://bytecodealliance.org/security.
         let i = self.interpreter.as_mut()?;
         Some(i.as_interpreter_ref())
     }
+
+    pub(crate) fn unwinder(&self) -> &'static dyn Unwind {
+        if self.interpreter.is_some() {
+            &UnwindPulley
+        } else {
+            &UnwindHost
+        }
+    }
 }
 
 impl<T> StoreContextMut<'_, T> {

crates/wasmtime/src/runtime/trap.rs

@@ -257,11 +257,7 @@ impl WasmBacktrace {
     /// always captures a backtrace.
     pub fn force_capture(store: impl AsContext) -> WasmBacktrace {
         let store = store.as_context();
-        Self::from_captured(
-            store.0,
-            crate::runtime::vm::Backtrace::new(store.0.runtime_limits()),
-            None,
-        )
+        Self::from_captured(store.0, crate::runtime::vm::Backtrace::new(store.0), None)
     }
 
     fn from_captured(

crates/wasmtime/src/runtime/vm.rs

@@ -35,6 +35,7 @@ mod store_box;
 mod sys;
 mod table;
 mod traphandlers;
+mod unwind;
 mod vmcontext;
 
 #[cfg(feature = "threads")]
@@ -82,6 +83,7 @@ pub use crate::runtime::vm::sys::mmap::open_file_for_mmap;
 pub use crate::runtime::vm::sys::unwind::UnwindRegistration;
 pub use crate::runtime::vm::table::{Table, TableElement};
 pub use crate::runtime::vm::traphandlers::*;
+pub use crate::runtime::vm::unwind::*;
 pub use crate::runtime::vm::vmcontext::{
     VMArrayCallFunction, VMArrayCallHostFuncContext, VMContext, VMFuncRef, VMFunctionBody,
     VMFunctionImport, VMGlobalDefinition, VMGlobalImport, VMMemoryDefinition, VMMemoryImport,

crates/wasmtime/src/runtime/vm/traphandlers.rs

@@ -407,6 +407,7 @@ where
 // usage of its accessor methods.
 mod call_thread_state {
     use super::*;
+    use crate::runtime::vm::Unwind;
 
     /// Temporary state stored on the stack which is registered in the `tls` module
     /// below for calls into wasm.
@@ -420,6 +421,7 @@ mod call_thread_state {
         pub(super) capture_coredump: bool,
 
         pub(crate) limits: *const VMRuntimeLimits,
+        pub(crate) unwinder: &'static dyn Unwind,
 
         pub(super) prev: Cell<tls::Ptr>,
         #[cfg(all(feature = "signals-based-traps", unix, not(miri)))]
@@ -465,6 +467,7 @@ mod call_thread_state {
 
             CallThreadState {
                 unwind: Cell::new(None),
+                unwinder: store.unwinder(),
                 jmp_buf: Cell::new(ptr::null()),
                 #[cfg(all(feature = "signals-based-traps", not(miri)))]
                 signal_handler: store.signal_handler(),
@@ -608,7 +611,7 @@ impl CallThreadState {
             return None;
         }
 
-        Some(unsafe { Backtrace::new_with_trap_state(limits, self, trap_pc_and_fp) })
+        Some(unsafe { Backtrace::new_with_trap_state(limits, self.unwinder, self, trap_pc_and_fp) })
     }
 
     pub(crate) fn iter<'a>(&'a self) -> impl Iterator<Item = &'a Self> + 'a {

crates/wasmtime/src/runtime/vm/traphandlers/backtrace.rs

@@ -22,10 +22,10 @@
 //! frame is a host frame).
 
 use crate::prelude::*;
-use crate::runtime::vm::arch;
+use crate::runtime::store::StoreOpaque;
 use crate::runtime::vm::{
     traphandlers::{tls, CallThreadState},
-    VMRuntimeLimits,
+    Unwind, VMRuntimeLimits,
 };
 use core::ops::ControlFlow;
 
@@ -59,9 +59,11 @@ impl Backtrace {
     }
 
     /// Capture the current Wasm stack in a backtrace.
-    pub fn new(limits: *const VMRuntimeLimits) -> Backtrace {
+    pub fn new(store: &StoreOpaque) -> Backtrace {
+        let limits = store.runtime_limits();
+        let unwind = store.unwinder();
         tls::with(|state| match state {
-            Some(state) => unsafe { Self::new_with_trap_state(limits, state, None) },
+            Some(state) => unsafe { Self::new_with_trap_state(limits, unwind, state, None) },
             None => Backtrace(vec![]),
         })
     }
@@ -73,21 +75,24 @@ impl Backtrace {
     /// instead of looking them up in `VMRuntimeLimits`.
     pub(crate) unsafe fn new_with_trap_state(
         limits: *const VMRuntimeLimits,
+        unwind: &dyn Unwind,
         state: &CallThreadState,
         trap_pc_and_fp: Option<(usize, usize)>,
     ) -> Backtrace {
         let mut frames = vec![];
-        Self::trace_with_trap_state(limits, state, trap_pc_and_fp, |frame| {
+        Self::trace_with_trap_state(limits, unwind, state, trap_pc_and_fp, |frame| {
             frames.push(frame);
             ControlFlow::Continue(())
         });
         Backtrace(frames)
     }
 
     /// Walk the current Wasm stack, calling `f` for each frame we walk.
-    pub fn trace(limits: *const VMRuntimeLimits, f: impl FnMut(Frame) -> ControlFlow<()>) {
+    pub fn trace(store: &StoreOpaque, f: impl FnMut(Frame) -> ControlFlow<()>) {
+        let limits = store.runtime_limits();
+        let unwind = store.unwinder();
         tls::with(|state| match state {
-            Some(state) => unsafe { Self::trace_with_trap_state(limits, state, None, f) },
+            Some(state) => unsafe { Self::trace_with_trap_state(limits, unwind, state, None, f) },
             None => {}
         });
     }
@@ -99,6 +104,7 @@ impl Backtrace {
     /// instead of looking them up in `VMRuntimeLimits`.
     pub(crate) unsafe fn trace_with_trap_state(
         limits: *const VMRuntimeLimits,
+        unwind: &dyn Unwind,
         state: &CallThreadState,
         trap_pc_and_fp: Option<(usize, usize)>,
         mut f: impl FnMut(Frame) -> ControlFlow<()>,
@@ -148,7 +154,7 @@ impl Backtrace {
         });
 
         for (pc, fp, sp) in activations {
-            if let ControlFlow::Break(()) = Self::trace_through_wasm(pc, fp, sp, &mut f) {
+            if let ControlFlow::Break(()) = Self::trace_through_wasm(unwind, pc, fp, sp, &mut f) {
                 log::trace!("====== Done Capturing Backtrace (closure break) ======");
                 return;
             }
@@ -160,6 +166,7 @@ impl Backtrace {
     /// Walk through a contiguous sequence of Wasm frames starting with the
     /// frame at the given PC and FP and ending at `trampoline_sp`.
     unsafe fn trace_through_wasm(
+        unwind: &dyn Unwind,
         mut pc: usize,
         mut fp: usize,
         trampoline_fp: usize,
@@ -228,25 +235,25 @@ impl Backtrace {
             // Wasm. Finally also assert that it's aligned correctly as an
             // additional sanity check.
             assert!(trampoline_fp > fp, "{trampoline_fp:#x} > {fp:#x}");
-            arch::assert_fp_is_aligned(fp);
+            unwind.assert_fp_is_aligned(fp);
 
             log::trace!("--- Tracing through one Wasm frame ---");
             log::trace!("pc = {:p}", pc as *const ());
             log::trace!("fp = {:p}", fp as *const ());
 
             f(Frame { pc, fp })?;
 
-            pc = arch::get_next_older_pc_from_fp(fp);
+            pc = unwind.get_next_older_pc_from_fp(fp);
 
             // We rely on this offset being zero for all supported architectures
             // in `crates/cranelift/src/component/compiler.rs` when we set the
             // Wasm exit FP. If this ever changes, we will need to update that
             // code as well!
-            assert_eq!(arch::NEXT_OLDER_FP_FROM_FP_OFFSET, 0);
+            assert_eq!(unwind.next_older_fp_from_fp_offset(), 0);
 
             // Get the next older frame pointer from the current Wasm frame
             // pointer.
-            let next_older_fp = *(fp as *mut usize).add(arch::NEXT_OLDER_FP_FROM_FP_OFFSET);
+            let next_older_fp = *(fp as *mut usize).add(unwind.next_older_fp_from_fp_offset());
 
             // Because the stack always grows down, the older FP must be greater
             // than the current FP.

crates/wasmtime/src/runtime/vm/traphandlers/coredump_enabled.rs

@@ -31,7 +31,8 @@ impl CallThreadState {
         if !self.capture_coredump {
             return None;
         }
-        let bt = unsafe { Backtrace::new_with_trap_state(limits, self, trap_pc_and_fp) };
+        let bt =
+            unsafe { Backtrace::new_with_trap_state(limits, self.unwinder, self, trap_pc_and_fp) };
 
         Some(CoreDumpStack {
             bt,

crates/wasmtime/src/runtime/vm/unwind.rs

@@ -0,0 +1,56 @@
+//! Support for low-level primitives of unwinding the stack.
+
+use crate::runtime::vm::arch;
+
+/// Implementation necessary to unwind the stack, used by `Backtrace`.
+pub unsafe trait Unwind {
+    /// Returns the offset, from the current frame pointer, of where to get to
+    /// the previous frame pointer on the stack.
+    fn next_older_fp_from_fp_offset(&self) -> usize;
+
+    /// Load the return address of a frame given the frame pointer for that
+    /// frame.
+    unsafe fn get_next_older_pc_from_fp(&self, fp: usize) -> usize;
+
+    /// Debug assertion that the frame pointer is aligned.
+    fn assert_fp_is_aligned(&self, fp: usize);
+}
+
+/// A host-backed implementation of unwinding, using the native platform ABI
+/// that Cranelift has.
+pub struct UnwindHost;
+
+unsafe impl Unwind for UnwindHost {
+    fn next_older_fp_from_fp_offset(&self) -> usize {
+        arch::NEXT_OLDER_FP_FROM_FP_OFFSET
+    }
+    unsafe fn get_next_older_pc_from_fp(&self, fp: usize) -> usize {
+        arch::get_next_older_pc_from_fp(fp)
+    }
+    fn assert_fp_is_aligned(&self, fp: usize) {
+        arch::assert_fp_is_aligned(fp)
+    }
+}
+
+/// An implementation specifically designed for unwinding Pulley's runtime stack
+/// (which might not match the native host).
+pub struct UnwindPulley;
+
+unsafe impl Unwind for UnwindPulley {
+    fn next_older_fp_from_fp_offset(&self) -> usize {
+        0
+    }
+    unsafe fn get_next_older_pc_from_fp(&self, fp: usize) -> usize {
+        // The calling convention always pushes the return pointer (aka the PC
+        // of the next older frame) just before this frame.
+        *(fp as *mut usize).offset(1)
+    }
+    fn assert_fp_is_aligned(&self, fp: usize) {
+        let expected = if cfg!(target_pointer_width = "32") {
+            8
+        } else {
+            16
+        };
+        assert_eq!(fp % expected, 0, "stack should always be aligned");
+    }
+}