From cfa00331c6b23b4cd2dcb7f49b6a973c43bf46e2 Mon Sep 17 00:00:00 2001 From: Frank Emrich Date: Mon, 21 Apr 2025 23:09:11 +0100 Subject: [PATCH 1/5] Use EntryStoreContext to manage state for entering and exiting Wasm --- crates/wasmtime/src/runtime/func.rs | 221 +++++++++++------- .../wasmtime/src/runtime/vm/traphandlers.rs | 61 ++--- 2 files changed, 161 insertions(+), 121 deletions(-) diff --git a/crates/wasmtime/src/runtime/func.rs b/crates/wasmtime/src/runtime/func.rs index 3beb01362017..1f58eac71fe9 100644 --- a/crates/wasmtime/src/runtime/func.rs +++ b/crates/wasmtime/src/runtime/func.rs @@ -1,7 +1,7 @@ use crate::prelude::*; use crate::runtime::vm::{ ExportFunction, InterpreterRef, SendSyncPtr, StoreBox, VMArrayCallHostFuncContext, VMContext, - VMFuncRef, VMFunctionImport, VMOpaqueContext, + VMFuncRef, VMFunctionImport, VMOpaqueContext, VMStoreContext, }; use crate::runtime::Uninhabited; use crate::store::{AutoAssertNoGc, StoreData, StoreOpaque, Stored}; @@ -1604,103 +1604,154 @@ pub(crate) fn invoke_wasm_and_catch_traps( closure: impl FnMut(NonNull, Option>) -> bool, ) -> Result<()> { unsafe { - let exit = enter_wasm(store); + let previous_runtime_state = EntryStoreContext::enter_wasm(store); if let Err(trap) = store.0.call_hook(CallHook::CallingWasm) { - exit_wasm(store, exit); + // `previous_runtime_state` implicitly dropped here return Err(trap); } - let result = crate::runtime::vm::catch_traps(store, closure); - exit_wasm(store, exit); + let result = crate::runtime::vm::catch_traps(store, &previous_runtime_state, closure); + core::mem::drop(previous_runtime_state); store.0.call_hook(CallHook::ReturningFromWasm)?; result.map_err(|t| crate::trap::from_runtime_box(store.0, t)) } } -/// This function is called to register state within `Store` whenever -/// WebAssembly is entered within the `Store`. -/// -/// This function sets up various limits such as: -/// -/// * The stack limit. This is what ensures that we limit the stack space -/// allocated by WebAssembly code and it's relative to the initial stack -/// pointer that called into wasm. -/// -/// This function may fail if the stack limit can't be set because an -/// interrupt already happened. -fn enter_wasm(store: &mut StoreContextMut<'_, T>) -> Option { - // If this is a recursive call, e.g. our stack limit is already set, then - // we may be able to skip this function. - // - // For synchronous stores there's nothing else to do because all wasm calls - // happen synchronously and on the same stack. This means that the previous - // stack limit will suffice for the next recursive call. - // - // For asynchronous stores then each call happens on a separate native - // stack. This means that the previous stack limit is no longer relevant - // because we're on a separate stack. - if unsafe { *store.0.vm_store_context().stack_limit.get() } != usize::MAX - && !store.0.async_support() - { - return None; - } - - // Ignore this stack pointer business on miri since we can't execute wasm - // anyway and the concept of a stack pointer on miri is a bit nebulous - // regardless. - if cfg!(miri) { - return None; - } - - // When Cranelift has support for the host then we might be running native - // compiled code meaning we need to read the actual stack pointer. If - // Cranelift can't be used though then we're guaranteed to be running pulley - // in which case this stack pointer isn't actually used as Pulley has custom - // mechanisms for stack overflow. - #[cfg(has_host_compiler_backend)] - let stack_pointer = crate::runtime::vm::get_stack_pointer(); - #[cfg(not(has_host_compiler_backend))] - let stack_pointer = { - use wasmtime_environ::TripleExt; - debug_assert!(store.engine().target().is_pulley()); - usize::MAX - }; +/// This type helps managing the state of the runtime when entering and exiting +/// Wasm. To this end, it contains a subset of the data in `VMStoreContext`.. +/// Upon entering Wasm, it updates various runtime fields and their +/// original values saved in this struct. Upon exiting Wasm, the previous values +/// are restored. +// FIXME(frank-emrich) Do the fields in here need to be (Unsafe)Cells? +pub struct EntryStoreContext { + /// If set, contains value of `stack_limit` field to restore in + /// `VMRuntimeLimits` when exiting Wasm. + pub stack_limit: Option, + /// Contains value of `last_wasm_exit_pc` field to restore in + /// `VMRuntimeLimits` when exiting Wasm. + pub last_wasm_exit_pc: usize, + /// Contains value of `last_wasm_exit_fp` field to restore in + /// `VMRuntimeLimits` when exiting Wasm. + pub last_wasm_exit_fp: usize, + /// Contains value of `last_wasm_entry_fp` field to restore in + /// `VMRuntimeLimits` when exiting Wasm. + pub last_wasm_entry_fp: usize, + + /// We need a pointer to the runtime limits, so we can update them from + /// `drop`/`exit_wasm`. + vm_store_context: *const VMStoreContext, +} - // Determine the stack pointer where, after which, any wasm code will - // immediately trap. This is checked on the entry to all wasm functions. - // - // Note that this isn't 100% precise. We are requested to give wasm - // `max_wasm_stack` bytes, but what we're actually doing is giving wasm - // probably a little less than `max_wasm_stack` because we're - // calculating the limit relative to this function's approximate stack - // pointer. Wasm will be executed on a frame beneath this one (or next - // to it). In any case it's expected to be at most a few hundred bytes - // of slop one way or another. When wasm is typically given a MB or so - // (a million bytes) the slop shouldn't matter too much. - // - // After we've got the stack limit then we store it into the `stack_limit` - // variable. - let wasm_stack_limit = stack_pointer - store.engine().config().max_wasm_stack; - let prev_stack = unsafe { - mem::replace( - &mut *store.0.vm_store_context().stack_limit.get(), - wasm_stack_limit, - ) - }; +impl EntryStoreContext { + /// This function is called to update and save state when + /// WebAssembly is entered within the `Store`. + /// + /// This updates various fields such as: + /// + /// * The stack limit. This is what ensures that we limit the stack space + /// allocated by WebAssembly code and it's relative to the initial stack + /// pointer that called into wasm. + /// + /// It also saves the different last_wasm_* values in the `VMRuntimeLimits`. + pub fn enter_wasm(store: &mut StoreContextMut<'_, T>) -> Self { + let stack_limit; - Some(prev_stack) -} + // If this is a recursive call, e.g. our stack limit is already set, then + // we may be able to skip this function. + // + // For synchronous stores there's nothing else to do because all wasm calls + // happen synchronously and on the same stack. This means that the previous + // stack limit will suffice for the next recursive call. + // + // For asynchronous stores then each call happens on a separate native + // stack. This means that the previous stack limit is no longer relevant + // because we're on a separate stack. + if unsafe { *store.0.vm_store_context().stack_limit.get() } != usize::MAX + && !store.0.async_support() + { + stack_limit = None; + } + // Ignore this stack pointer business on miri since we can't execute wasm + // anyway and the concept of a stack pointer on miri is a bit nebulous + // regardless. + else if cfg!(miri) { + stack_limit = None; + } else { + // When Cranelift has support for the host then we might be running native + // compiled code meaning we need to read the actual stack pointer. If + // Cranelift can't be used though then we're guaranteed to be running pulley + // in which case this stack pointer isn't actually used as Pulley has custom + // mechanisms for stack overflow. + #[cfg(has_host_compiler_backend)] + let stack_pointer = crate::runtime::vm::get_stack_pointer(); + #[cfg(not(has_host_compiler_backend))] + let stack_pointer = { + use wasmtime_environ::TripleExt; + debug_assert!(store.engine().target().is_pulley()); + usize::MAX + }; -fn exit_wasm(store: &mut StoreContextMut<'_, T>, prev_stack: Option) { - // If we don't have a previous stack pointer to restore, then there's no - // cleanup we need to perform here. - let prev_stack = match prev_stack { - Some(stack) => stack, - None => return, - }; + // Determine the stack pointer where, after which, any wasm code will + // immediately trap. This is checked on the entry to all wasm functions. + // + // Note that this isn't 100% precise. We are requested to give wasm + // `max_wasm_stack` bytes, but what we're actually doing is giving wasm + // probably a little less than `max_wasm_stack` because we're + // calculating the limit relative to this function's approximate stack + // pointer. Wasm will be executed on a frame beneath this one (or next + // to it). In any case it's expected to be at most a few hundred bytes + // of slop one way or another. When wasm is typically given a MB or so + // (a million bytes) the slop shouldn't matter too much. + // + // After we've got the stack limit then we store it into the `stack_limit` + // variable. + let wasm_stack_limit = stack_pointer - store.engine().config().max_wasm_stack; + let prev_stack = unsafe { + mem::replace( + &mut *store.0.vm_store_context().stack_limit.get(), + wasm_stack_limit, + ) + }; + stack_limit = Some(prev_stack); + } - unsafe { - *store.0.vm_store_context().stack_limit.get() = prev_stack; + unsafe { + let last_wasm_exit_pc = *store.0.vm_store_context().last_wasm_exit_pc.get(); + let last_wasm_exit_fp = *store.0.vm_store_context().last_wasm_exit_fp.get(); + let last_wasm_entry_fp = *store.0.vm_store_context().last_wasm_entry_fp.get(); + + let vm_store_context = store.0.vm_store_context(); + + Self { + stack_limit, + last_wasm_exit_pc, + last_wasm_exit_fp, + last_wasm_entry_fp, + vm_store_context, + } + } + } + + /// This function restores the values stored in this struct. We invoke this + /// function through this type's `Drop` implementation. This ensures that we + /// even restore the values if we unwind the stack (e.g., because we are + /// panicing out of a Wasm execution). + fn exit_wasm(&mut self) { + unsafe { + self.stack_limit.inspect(|limit| { + *(&*self.vm_store_context).stack_limit.get() = *limit; + }); + + *(*self.vm_store_context).last_wasm_exit_fp.get() = self.last_wasm_exit_fp; + *(*self.vm_store_context).last_wasm_exit_pc.get() = self.last_wasm_exit_pc; + *(*self.vm_store_context).last_wasm_entry_fp.get() = self.last_wasm_entry_fp; + } + } +} + +impl Drop for EntryStoreContext { + fn drop(&mut self) { + self.exit_wasm(); } } diff --git a/crates/wasmtime/src/runtime/vm/traphandlers.rs b/crates/wasmtime/src/runtime/vm/traphandlers.rs index 95c9d9be2922..a405ff4af5a8 100644 --- a/crates/wasmtime/src/runtime/vm/traphandlers.rs +++ b/crates/wasmtime/src/runtime/vm/traphandlers.rs @@ -15,11 +15,11 @@ mod signals; #[cfg(all(has_native_signals))] pub use self::signals::*; -use crate::prelude::*; use crate::runtime::module::lookup_code; use crate::runtime::store::{ExecutorRef, StoreOpaque}; use crate::runtime::vm::sys::traphandlers; use crate::runtime::vm::{InterpreterRef, VMContext, VMStoreContext}; +use crate::{prelude::*, EntryStoreContext}; use crate::{StoreContextMut, WasmBacktrace}; use core::cell::Cell; use core::num::NonZeroU32; @@ -365,6 +365,7 @@ impl From for TrapReason { /// longjmp'd over and none of its destructors on the stack may be run. pub unsafe fn catch_traps( store: &mut StoreContextMut<'_, T>, + old_state: &EntryStoreContext, mut closure: F, ) -> Result<(), Box> where @@ -372,7 +373,7 @@ where { let caller = store.0.default_caller(); - let result = CallThreadState::new(store.0).with(|cx| match store.0.executor() { + let result = CallThreadState::new(store.0, old_state).with(|cx| match store.0.executor() { // In interpreted mode directly invoke the host closure since we won't // be using host-based `setjmp`/`longjmp` as that's not going to save // the context we want. @@ -424,6 +425,7 @@ where mod call_thread_state { use super::*; use crate::runtime::vm::Unwind; + use crate::EntryStoreContext; /// Temporary state stored on the stack which is registered in the `tls` /// module below for calls into wasm. @@ -462,17 +464,16 @@ mod call_thread_state { #[cfg(all(has_native_signals, unix))] pub(crate) async_guard_range: Range<*mut u8>, - // The values of `VMStoreContext::last_wasm_{exit_{pc,fp},entry_sp}` for - // the *previous* `CallThreadState` for this same store/limits. Our - // *current* last wasm PC/FP/SP are saved in `self.vm_store_context`. We - // save a copy of the old registers here because the `VMStoreContext` - // typically doesn't change across nested calls into Wasm (i.e. they are - // typically calls back into the same store and `self.vm_store_context - // == self.prev.vm_store_context`) and we must to maintain the list of - // contiguous-Wasm-frames stack regions for backtracing purposes. - old_last_wasm_exit_fp: Cell, - old_last_wasm_exit_pc: Cell, - old_last_wasm_entry_fp: Cell, + // The state of the runtime for the *previous* `CallThreadState` for + // this same store. Our *current* state is saved in `self.vm_store_context`, + // etc. We need access to the old values of these + // fields because the `VMStoreContext` typically doesn't change across + // nested calls into Wasm (i.e. they are typically calls back into the + // same store and `self.vm_store_context == self.prev.vm_store_context`) and we must to + // maintain the list of contiguous-Wasm-frames stack regions for + // backtracing purposes. + // FIXME(frank-emrich) Does this need to be an (Unsafe)Cell? + old_state: *const EntryStoreContext, } impl Drop for CallThreadState { @@ -480,13 +481,6 @@ mod call_thread_state { // Unwind information should not be present as it should have // already been processed. debug_assert!(self.unwind.replace(None).is_none()); - - unsafe { - let cx = self.vm_store_context.as_ref(); - *cx.last_wasm_exit_fp.get() = self.old_last_wasm_exit_fp.get(); - *cx.last_wasm_exit_pc.get() = self.old_last_wasm_exit_pc.get(); - *cx.last_wasm_entry_fp.get() = self.old_last_wasm_entry_fp.get(); - } } } @@ -494,7 +488,10 @@ mod call_thread_state { pub const JMP_BUF_INTERPRETER_SENTINEL: *mut u8 = 1 as *mut u8; #[inline] - pub(super) fn new(store: &mut StoreOpaque) -> CallThreadState { + pub(super) fn new( + store: &mut StoreOpaque, + old_state: *const EntryStoreContext, + ) -> CallThreadState { // Don't try to plumb #[cfg] everywhere for this field, just pretend // we're using it on miri/windows to silence compiler warnings. let _: Range<_> = store.async_guard_range(); @@ -512,31 +509,23 @@ mod call_thread_state { #[cfg(all(has_native_signals, unix))] async_guard_range: store.async_guard_range(), prev: Cell::new(ptr::null()), - old_last_wasm_exit_fp: Cell::new(unsafe { - *store.vm_store_context().last_wasm_exit_fp.get() - }), - old_last_wasm_exit_pc: Cell::new(unsafe { - *store.vm_store_context().last_wasm_exit_pc.get() - }), - old_last_wasm_entry_fp: Cell::new(unsafe { - *store.vm_store_context().last_wasm_entry_fp.get() - }), + old_state, } } /// Get the saved FP upon exit from Wasm for the previous `CallThreadState`. - pub fn old_last_wasm_exit_fp(&self) -> usize { - self.old_last_wasm_exit_fp.get() + pub unsafe fn old_last_wasm_exit_fp(&self) -> usize { + (&*self.old_state).last_wasm_exit_fp } /// Get the saved PC upon exit from Wasm for the previous `CallThreadState`. - pub fn old_last_wasm_exit_pc(&self) -> usize { - self.old_last_wasm_exit_pc.get() + pub unsafe fn old_last_wasm_exit_pc(&self) -> usize { + (&*self.old_state).last_wasm_exit_pc } /// Get the saved FP upon entry into Wasm for the previous `CallThreadState`. - pub fn old_last_wasm_entry_fp(&self) -> usize { - self.old_last_wasm_entry_fp.get() + pub unsafe fn old_last_wasm_entry_fp(&self) -> usize { + (&*self.old_state).last_wasm_entry_fp } /// Get the previous `CallThreadState`. From 70a2ddc109d6392255f0e8dcf751cdffbbaced16 Mon Sep 17 00:00:00 2001 From: Frank Emrich Date: Tue, 22 Apr 2025 20:45:33 +0100 Subject: [PATCH 2/5] use checked arithmetic for stack_pointer Co-authored-by: Nick Fitzgerald --- crates/wasmtime/src/runtime/func.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/crates/wasmtime/src/runtime/func.rs b/crates/wasmtime/src/runtime/func.rs index 1f58eac71fe9..2af2aed06d92 100644 --- a/crates/wasmtime/src/runtime/func.rs +++ b/crates/wasmtime/src/runtime/func.rs @@ -1705,7 +1705,9 @@ impl EntryStoreContext { // // After we've got the stack limit then we store it into the `stack_limit` // variable. - let wasm_stack_limit = stack_pointer - store.engine().config().max_wasm_stack; + let wasm_stack_limit = stack_pointer + .checked_sub(store.engine().config().max_wasm_stack) + .unwrap(); let prev_stack = unsafe { mem::replace( &mut *store.0.vm_store_context().stack_limit.get(), From bfb44d84398773abfb9670c0619739573112b9dd Mon Sep 17 00:00:00 2001 From: Frank Emrich Date: Tue, 22 Apr 2025 20:40:56 +0100 Subject: [PATCH 3/5] comment fixes --- crates/wasmtime/src/runtime/func.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/wasmtime/src/runtime/func.rs b/crates/wasmtime/src/runtime/func.rs index 2af2aed06d92..4127e938452d 100644 --- a/crates/wasmtime/src/runtime/func.rs +++ b/crates/wasmtime/src/runtime/func.rs @@ -1618,7 +1618,7 @@ pub(crate) fn invoke_wasm_and_catch_traps( } /// This type helps managing the state of the runtime when entering and exiting -/// Wasm. To this end, it contains a subset of the data in `VMStoreContext`.. +/// Wasm. To this end, it contains a subset of the data in `VMStoreContext`. /// Upon entering Wasm, it updates various runtime fields and their /// original values saved in this struct. Upon exiting Wasm, the previous values /// are restored. @@ -1628,13 +1628,13 @@ pub struct EntryStoreContext { /// `VMRuntimeLimits` when exiting Wasm. pub stack_limit: Option, /// Contains value of `last_wasm_exit_pc` field to restore in - /// `VMRuntimeLimits` when exiting Wasm. + /// `VMStoreContext` when exiting Wasm. pub last_wasm_exit_pc: usize, /// Contains value of `last_wasm_exit_fp` field to restore in - /// `VMRuntimeLimits` when exiting Wasm. + /// `VMStoreContext` when exiting Wasm. pub last_wasm_exit_fp: usize, /// Contains value of `last_wasm_entry_fp` field to restore in - /// `VMRuntimeLimits` when exiting Wasm. + /// `VMStoreContext` when exiting Wasm. pub last_wasm_entry_fp: usize, /// We need a pointer to the runtime limits, so we can update them from From b6d78d7bc442a6764d5a1c4c20bb323966a36dac Mon Sep 17 00:00:00 2001 From: Frank Emrich Date: Tue, 22 Apr 2025 20:46:29 +0100 Subject: [PATCH 4/5] address review comments --- crates/wasmtime/src/runtime/func.rs | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/crates/wasmtime/src/runtime/func.rs b/crates/wasmtime/src/runtime/func.rs index 4127e938452d..b6bf00325c63 100644 --- a/crates/wasmtime/src/runtime/func.rs +++ b/crates/wasmtime/src/runtime/func.rs @@ -1622,8 +1622,7 @@ pub(crate) fn invoke_wasm_and_catch_traps( /// Upon entering Wasm, it updates various runtime fields and their /// original values saved in this struct. Upon exiting Wasm, the previous values /// are restored. -// FIXME(frank-emrich) Do the fields in here need to be (Unsafe)Cells? -pub struct EntryStoreContext { +pub(crate) struct EntryStoreContext { /// If set, contains value of `stack_limit` field to restore in /// `VMRuntimeLimits` when exiting Wasm. pub stack_limit: Option, @@ -1740,9 +1739,9 @@ impl EntryStoreContext { /// panicing out of a Wasm execution). fn exit_wasm(&mut self) { unsafe { - self.stack_limit.inspect(|limit| { - *(&*self.vm_store_context).stack_limit.get() = *limit; - }); + if let Some(limit) = self.stack_limit { + *(&*self.vm_store_context).stack_limit.get() = limit; + } *(*self.vm_store_context).last_wasm_exit_fp.get() = self.last_wasm_exit_fp; *(*self.vm_store_context).last_wasm_exit_pc.get() = self.last_wasm_exit_pc; From f8a558ff6ee71bbe551baaa2f8b5d4742c7b1705 Mon Sep 17 00:00:00 2001 From: Frank Emrich Date: Tue, 22 Apr 2025 20:59:29 +0100 Subject: [PATCH 5/5] remove FIXME --- crates/wasmtime/src/runtime/vm/traphandlers.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/wasmtime/src/runtime/vm/traphandlers.rs b/crates/wasmtime/src/runtime/vm/traphandlers.rs index a405ff4af5a8..627b44d4629b 100644 --- a/crates/wasmtime/src/runtime/vm/traphandlers.rs +++ b/crates/wasmtime/src/runtime/vm/traphandlers.rs @@ -472,7 +472,6 @@ mod call_thread_state { // same store and `self.vm_store_context == self.prev.vm_store_context`) and we must to // maintain the list of contiguous-Wasm-frames stack regions for // backtracing purposes. - // FIXME(frank-emrich) Does this need to be an (Unsafe)Cell? old_state: *const EntryStoreContext, }