Save exit Wasm FP and PC in component-to-host trampolines

Fixes bytecodealliance#4535

cranelift/codegen/src/ir/immediates.rs

@@ -426,6 +426,12 @@ impl From<i32> for Offset32 {
     }
 }
 
+impl From<u8> for Offset32 {
+    fn from(val: u8) -> Offset32 {
+        Self(val.into())
+    }
+}
+
 impl Display for Offset32 {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
         // 0 displays as an empty offset.

crates/cranelift/src/compiler.rs

@@ -30,8 +30,8 @@ use std::mem;
 use std::sync::Mutex;
 use wasmtime_environ::{
     AddressMapSection, CompileError, FilePos, FlagValue, FunctionBodyData, FunctionInfo,
-    InstructionAddressMap, Module, ModuleTranslation, ModuleTypes, StackMapInformation, Trampoline,
-    TrapCode, TrapEncodingBuilder, TrapInformation, Tunables, VMOffsets,
+    InstructionAddressMap, Module, ModuleTranslation, ModuleTypes, PtrSize, StackMapInformation,
+    Trampoline, TrapCode, TrapEncodingBuilder, TrapInformation, Tunables, VMOffsets,
 };
 
 #[cfg(feature = "component-model")]
@@ -196,7 +196,7 @@ impl wasmtime_environ::Compiler for Compiler {
         });
         let stack_limit = context.func.create_global_value(ir::GlobalValueData::Load {
             base: interrupts_ptr,
-            offset: i32::try_from(func_env.offsets.vmruntime_limits_stack_limit())
+            offset: i32::try_from(func_env.offsets.ptr.vmruntime_limits_stack_limit())
                 .unwrap()
                 .into(),
             global_type: isa.pointer_type(),

crates/cranelift/src/compiler/component.rs

@@ -13,7 +13,7 @@ use wasmtime_environ::component::{
     AlwaysTrapInfo, CanonicalOptions, Component, ComponentCompiler, ComponentTypes, FunctionInfo,
     LowerImport, LoweredIndex, RuntimeAlwaysTrapIndex, VMComponentOffsets,
 };
-use wasmtime_environ::{PrimaryMap, SignatureIndex, Trampoline, TrapCode, WasmFuncType};
+use wasmtime_environ::{PrimaryMap, PtrSize, SignatureIndex, Trampoline, TrapCode, WasmFuncType};
 
 impl ComponentCompiler for Compiler {
     fn compile_lowered_trampoline(
@@ -46,6 +46,41 @@ impl ComponentCompiler for Compiler {
             self.wasm_to_host_spill_args(ty, &mut builder, block0);
         let vmctx = builder.func.dfg.block_params(block0)[0];
 
+        // Save the exit FP and return address for stack walking purposes.
+        //
+        // First we need to get the `VMRuntimeLimits`.
+        let limits = builder.ins().load(
+            pointer_type,
+            MemFlags::trusted(),
+            vmctx,
+            i32::try_from(offsets.limits()).unwrap(),
+        );
+        // Then save the exit Wasm FP to the limits.
+        let trampoline_fp = builder.ins().get_frame_pointer(pointer_type);
+        let wasm_fp = builder.ins().load(
+            pointer_type,
+            MemFlags::trusted(),
+            trampoline_fp,
+            // The FP always points to the next older FP for all supported
+            // targets. See assertion in
+            // `crates/runtime/src/traphandlers/backtrace.rs`.
+            0,
+        );
+        builder.ins().store(
+            MemFlags::trusted(),
+            wasm_fp,
+            limits,
+            offsets.ptr.vmruntime_limits_last_wasm_exit_fp(),
+        );
+        // Finally save the Wasm return address to the limits.
+        let wasm_pc = builder.ins().get_return_address(pointer_type);
+        builder.ins().store(
+            MemFlags::trusted(),
+            wasm_pc,
+            limits,
+            offsets.ptr.vmruntime_limits_last_wasm_exit_pc(),
+        );
+
         // Below this will incrementally build both the signature of the host
         // function we're calling as well as the list of arguments since the
         // list is somewhat long.

crates/cranelift/src/func_environ.rs

@@ -528,7 +528,7 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
     ) -> (ir::Value, ir::immediates::Offset32) {
         (
             builder.use_var(self.vmruntime_limits_ptr),
-            i32::from(self.offsets.vmruntime_limits_fuel_consumed()).into(),
+            i32::from(self.offsets.ptr.vmruntime_limits_fuel_consumed()).into(),
         )
     }
 
@@ -628,12 +628,15 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
 
     fn epoch_load_deadline_into_var(&mut self, builder: &mut FunctionBuilder<'_>) {
         let interrupts = builder.use_var(self.vmruntime_limits_ptr);
-        let deadline = builder.ins().load(
-            ir::types::I64,
-            ir::MemFlags::trusted(),
-            interrupts,
-            ir::immediates::Offset32::new(self.offsets.vmruntime_limits_epoch_deadline() as i32),
-        );
+        let deadline =
+            builder.ins().load(
+                ir::types::I64,
+                ir::MemFlags::trusted(),
+                interrupts,
+                ir::immediates::Offset32::new(
+                    self.offsets.ptr.vmruntime_limits_epoch_deadline() as i32
+                ),
+            );
         builder.def_var(self.epoch_deadline_var, deadline);
     }
 

crates/environ/src/vmoffsets.rs

@@ -135,6 +135,39 @@ pub trait PtrSize {
     fn size_of_vmglobal_definition(&self) -> u8 {
         16
     }
+
+    /// Return the offset of the `stack_limit` field of `VMRuntimeLimits`
+    #[inline]
+    fn vmruntime_limits_stack_limit(&self) -> u8 {
+        0
+    }
+
+    /// Return the offset of the `fuel_consumed` field of `VMRuntimeLimits`
+    #[inline]
+    fn vmruntime_limits_fuel_consumed(&self) -> u8 {
+        self.size()
+    }
+
+    /// Return the offset of the `epoch_deadline` field of `VMRuntimeLimits`
+    #[inline]
+    fn vmruntime_limits_epoch_deadline(&self) -> u8 {
+        self.vmruntime_limits_fuel_consumed() + 8 // `stack_limit` is a pointer; `fuel_consumed` is an `i64`
+    }
+
+    /// Return the offset of the `last_wasm_exit_fp` field of `VMRuntimeLimits`.
+    fn vmruntime_limits_last_wasm_exit_fp(&self) -> u8 {
+        self.vmruntime_limits_epoch_deadline() + 8
+    }
+
+    /// Return the offset of the `last_wasm_exit_pc` field of `VMRuntimeLimits`.
+    fn vmruntime_limits_last_wasm_exit_pc(&self) -> u8 {
+        self.vmruntime_limits_last_wasm_exit_fp() + self.size()
+    }
+
+    /// Return the offset of the `last_enty_sp` field of `VMRuntimeLimits`.
+    fn vmruntime_limits_last_wasm_entry_sp(&self) -> u8 {
+        self.vmruntime_limits_last_wasm_exit_pc() + self.size()
+    }
 }
 
 /// Type representing the size of a pointer for the current compilation host
@@ -545,42 +578,6 @@ impl<P: PtrSize> VMOffsets<P> {
     }
 }
 
-/// Offsets for `VMRuntimeLimits`.
-impl<P: PtrSize> VMOffsets<P> {
-    /// Return the offset of the `stack_limit` field of `VMRuntimeLimits`
-    #[inline]
-    pub fn vmruntime_limits_stack_limit(&self) -> u8 {
-        0
-    }
-
-    /// Return the offset of the `fuel_consumed` field of `VMRuntimeLimits`
-    #[inline]
-    pub fn vmruntime_limits_fuel_consumed(&self) -> u8 {
-        self.pointer_size()
-    }
-
-    /// Return the offset of the `epoch_deadline` field of `VMRuntimeLimits`
-    #[inline]
-    pub fn vmruntime_limits_epoch_deadline(&self) -> u8 {
-        self.vmruntime_limits_fuel_consumed() + 8 // `stack_limit` is a pointer; `fuel_consumed` is an `i64`
-    }
-
-    /// Return the offset of the `last_wasm_exit_fp` field of `VMRuntimeLimits`.
-    pub fn vmruntime_limits_last_wasm_exit_fp(&self) -> u8 {
-        self.vmruntime_limits_epoch_deadline() + 8
-    }
-
-    /// Return the offset of the `last_wasm_exit_pc` field of `VMRuntimeLimits`.
-    pub fn vmruntime_limits_last_wasm_exit_pc(&self) -> u8 {
-        self.vmruntime_limits_last_wasm_exit_fp() + self.pointer_size()
-    }
-
-    /// Return the offset of the `last_enty_sp` field of `VMRuntimeLimits`.
-    pub fn vmruntime_limits_last_wasm_entry_sp(&self) -> u8 {
-        self.vmruntime_limits_last_wasm_exit_pc() + self.pointer_size()
-    }
-}
-
 /// Offsets for `VMContext`.
 impl<P: PtrSize> VMOffsets<P> {
     /// Return the offset to the `magic` value in this `VMContext`.

crates/runtime/src/traphandlers/backtrace.rs

@@ -114,6 +114,7 @@ impl Backtrace {
         trap_pc_and_fp: Option<(usize, usize)>,
         mut f: impl FnMut(Frame) -> ControlFlow<()>,
     ) {
+        log::trace!("====== Capturing Backtrace ======");
         let (last_wasm_exit_pc, last_wasm_exit_fp) = match trap_pc_and_fp {
             // If we exited Wasm by catching a trap, then the Wasm-to-host
             // trampoline did not get a chance to save the last Wasm PC and FP,
@@ -137,6 +138,7 @@ impl Backtrace {
             *(*state.limits).last_wasm_entry_sp.get(),
             &mut f,
         ) {
+            log::trace!("====== Done Capturing Backtrace ======");
             return;
         }
 
@@ -151,6 +153,7 @@ impl Backtrace {
                 debug_assert_eq!(state.old_last_wasm_exit_pc, 0);
                 debug_assert_eq!(state.old_last_wasm_exit_fp, 0);
                 debug_assert_eq!(state.old_last_wasm_entry_sp, 0);
+                log::trace!("====== Done Capturing Backtrace ======");
                 return;
             }
 
@@ -160,9 +163,12 @@ impl Backtrace {
                 state.old_last_wasm_entry_sp,
                 &mut f,
             ) {
+                log::trace!("====== Done Capturing Backtrace ======");
                 return;
             }
         }
+
+        unreachable!()
     }
 
     /// Walk through a contiguous sequence of Wasm frames starting with the
@@ -245,7 +251,13 @@ impl Backtrace {
 
             pc = arch::get_next_older_pc_from_fp(fp);
 
-            let next_older_fp = arch::get_next_older_fp_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);
+
+            let next_older_fp = *(fp as *mut usize).add(arch::NEXT_OLDER_FP_FROM_FP_OFFSET);
             // Because the stack always grows down, the older FP must be greater
             // than the current FP.
             assert!(next_older_fp > fp, "{next_older_fp:#x} > {fp:#x}");

crates/runtime/src/traphandlers/backtrace/aarch64.rs

@@ -30,9 +30,9 @@ pub unsafe fn get_next_older_pc_from_fp(fp: usize) -> usize {
 
     pc
 }
-pub unsafe fn get_next_older_fp_from_fp(fp: usize) -> usize {
-    *(fp as *mut usize)
-}
+
+// And the current frame pointer points to the next older frame pointer.
+pub const NEXT_OLDER_FP_FROM_FP_OFFSET: usize = 0;
 
 pub fn reached_entry_sp(fp: usize, first_wasm_sp: usize) -> bool {
     // Calls in aarch64 push two i64s (old FP and return PC) so our entry SP is

crates/runtime/src/traphandlers/backtrace/s390x.rs

@@ -5,11 +5,9 @@ pub unsafe fn get_next_older_pc_from_fp(fp: usize) -> usize {
     *(fp as *mut usize).offset(14)
 }
 
-pub unsafe fn get_next_older_fp_from_fp(fp: usize) -> usize {
-    // The next older "FP" (backchain pointer) was saved in the slot pointed to
-    // by the current "FP".
-    *(fp as *mut usize)
-}
+// The next older "FP" (backchain pointer) was saved in the slot pointed to
+// by the current "FP".
+pub const NEXT_OLDER_FP_FROM_FP_OFFSET: usize = 0;
 
 pub fn reached_entry_sp(fp: usize, first_wasm_sp: usize) -> bool {
     // The "FP" (backchain pointer) holds the value of the stack pointer at

crates/runtime/src/traphandlers/backtrace/x86_64.rs

@@ -4,10 +4,8 @@ pub unsafe fn get_next_older_pc_from_fp(fp: usize) -> usize {
     *(fp as *mut usize).offset(1)
 }
 
-pub unsafe fn get_next_older_fp_from_fp(fp: usize) -> usize {
-    // And the current frame pointer points to the next older frame pointer.
-    *(fp as *mut usize)
-}
+// And the current frame pointer points to the next older frame pointer.
+pub const NEXT_OLDER_FP_FROM_FP_OFFSET: usize = 0;
 
 pub fn reached_entry_sp(fp: usize, first_wasm_sp: usize) -> bool {
     // When the FP is just below the SP (because we are in a function prologue