Require a store in catch_unwind_and_record_trap (bytecodealliance#11441)

* Require a store in `catch_unwind_and_record_trap`

This commit does some preparatory refactoring for bytecodealliance#11326 to ensure that
a store is available when trap information is being processed. Currently
this doesn't leverage the new parameter but it should be leverage-able
in bytecodealliance#11326.

* Review comments

Author: alexcrichton

crates/wasmtime/src/runtime/component/func/host.rs

@@ -683,15 +683,12 @@ where
 {
     let cx = unsafe { VMComponentContext::from_opaque(cx) };
     unsafe {
-        ComponentInstance::from_vmctx(cx, |store, instance| {
+        ComponentInstance::enter_host_from_wasm(cx, |store, instance| {
             let mut store = store.unchecked_context_mut();
-
-            crate::runtime::vm::catch_unwind_and_record_trap(|| {
-                store.0.call_hook(CallHook::CallingHost)?;
-                let res = func(store.as_context_mut(), instance);
-                store.0.call_hook(CallHook::ReturningFromHost)?;
-                res
-            })
+            store.0.call_hook(CallHook::CallingHost)?;
+            let res = func(store.as_context_mut(), instance);
+            store.0.call_hook(CallHook::ReturningFromHost)?;
+            res
         })
     }
 }

crates/wasmtime/src/runtime/func.rs

@@ -1,10 +1,11 @@
 use crate::prelude::*;
 use crate::runtime::Uninhabited;
 use crate::runtime::vm::{
-    InterpreterRef, SendSyncPtr, StoreBox, VMArrayCallHostFuncContext, VMCommonStackInformation,
-    VMContext, VMFuncRef, VMFunctionImport, VMOpaqueContext, VMStoreContext,
+    self, InterpreterRef, SendSyncPtr, StoreBox, VMArrayCallHostFuncContext,
+    VMCommonStackInformation, VMContext, VMFuncRef, VMFunctionImport, VMOpaqueContext,
+    VMStoreContext,
 };
-use crate::store::{AutoAssertNoGc, StoreId, StoreOpaque};
+use crate::store::{AutoAssertNoGc, InstanceId, StoreId, StoreOpaque};
 use crate::type_registry::RegisteredType;
 use crate::{
     AsContext, AsContextMut, CallHook, Engine, Extern, FuncType, Instance, ModuleExport, Ref,
@@ -2041,46 +2042,29 @@ impl<T> Caller<'_, T> {
 
     /// Executes `f` with an appropriate `Caller`.
     ///
-    /// This is the entrypoint for host functions in core wasm and converts from
-    /// `VMContext` to `Caller`
-    ///
-    /// # Safety
-    ///
-    /// This requires that `caller` is safe to wrap up as a `Caller`,
-    /// effectively meaning that we just entered the host from wasm.
-    /// Additionally this `Caller`'s `T` parameter must match the actual `T` in
-    /// the store of the vmctx of `caller`.
-    unsafe fn with<F, R>(caller: NonNull<VMContext>, f: F) -> R
+    /// This is the entrypoint for host functions in core wasm. The `store` and
+    /// `instance` are created from `Instance::enter_host_from_wasm` and this
+    /// will further invoke the host function that `f` refers to.
+    fn with<F, R>(mut store: StoreContextMut<T>, caller: InstanceId, f: F) -> R
     where
         F: FnOnce(Caller<'_, T>) -> R,
     {
-        // SAFETY: it's a contract of this function itself that `from_vmctx` is
-        // safe to call. Additionally it's a contract of this function itself
-        // that the `T` of `Caller` matches the store.
-        unsafe {
-            crate::runtime::vm::InstanceAndStore::from_vmctx(caller, |pair| {
-                let (instance, store) = pair.unpack_mut();
-                let mut store = store.unchecked_context_mut::<T>();
-                let caller = Instance::from_wasmtime(instance.id(), store.0);
+        let caller = Instance::from_wasmtime(caller, store.0);
 
-                let (gc_lifo_scope, ret) = {
-                    let gc_lifo_scope = store.0.gc_roots().enter_lifo_scope();
+        let (gc_lifo_scope, ret) = {
+            let gc_lifo_scope = store.0.gc_roots().enter_lifo_scope();
 
-                    let ret = f(Caller {
-                        store: store.as_context_mut(),
-                        caller,
-                    });
+            let ret = f(Caller {
+                store: store.as_context_mut(),
+                caller,
+            });
 
-                    (gc_lifo_scope, ret)
-                };
+            (gc_lifo_scope, ret)
+        };
 
-                // Safe to recreate a mutable borrow of the store because `ret`
-                // cannot be borrowing from the store.
-                store.0.exit_gc_lifo_scope(gc_lifo_scope);
+        store.0.exit_gc_lifo_scope(gc_lifo_scope);
 
-                ret
-            })
-        }
+        ret
     }
 
     fn sub_caller(&mut self) -> Caller<'_, T> {
@@ -2409,10 +2393,14 @@ impl HostContext {
         // closure and then run it as part of `Caller::with`.
         //
         // SAFETY: this is an entrypoint of wasm which requires correct type
-        // ascription of `T` itself, meaning that this should be safe to call.
-        crate::runtime::vm::catch_unwind_and_record_trap(move || unsafe {
-            Caller::with(caller_vmctx, run)
-        })
+        // ascription of `T` itself, meaning that this should be safe to call
+        // both `enter_host_from_wasm` as well as `unchecked_context_mut`.
+        unsafe {
+            vm::Instance::enter_host_from_wasm(caller_vmctx, |store, instance| {
+                let store = store.unchecked_context_mut();
+                Caller::with(store, instance.id(), run)
+            })
+        }
     }
 }
 
@@ -2485,20 +2473,23 @@ impl HostFunc {
         T: 'static,
     {
         assert!(ty.comes_from_same_engine(engine));
-        // SAFETY: This is only only called in the raw entrypoint of wasm
-        // meaning that `caller_vmctx` is appropriate to read, and additionally
-        // the later usage of `{,in}to_func` will connect `T` to an actual
-        // store's `T` to ensure it's the same.
-        let func = move |caller_vmctx, values: &mut [ValRaw]| unsafe {
-            Caller::<T>::with(caller_vmctx, |mut caller| {
-                caller.store.0.call_hook(CallHook::CallingHost)?;
-                let result = func(caller.sub_caller(), values)?;
-                caller.store.0.call_hook(CallHook::ReturningFromHost)?;
-                Ok(result)
-            })
-        };
-        let ctx = crate::trampoline::create_array_call_function(&ty, func)
-            .expect("failed to create function");
+        let ctx = crate::trampoline::create_array_call_function(
+            &ty,
+            move |store, instance, values: &mut [ValRaw]| {
+                // SAFETY: the later usage of `{,in}to_func` will connect `T` to
+                // an actual store's `T` to ensure it's the same. This means
+                // that the store this is invoked with always has `T` as a type
+                // parameter which should make this cast safe.
+                let store = unsafe { store.unchecked_context_mut::<T>() };
+                Caller::with(store, instance, |mut caller| {
+                    caller.store.0.call_hook(CallHook::CallingHost)?;
+                    let result = func(caller.sub_caller(), values)?;
+                    caller.store.0.call_hook(CallHook::ReturningFromHost)?;
+                    Ok(result)
+                })
+            },
+        )
+        .expect("failed to create function");
         HostFunc::_new(engine, ctx.into())
     }
 

crates/wasmtime/src/runtime/trampoline/func.rs

@@ -1,7 +1,10 @@
 //! Support for a calling of an imported function.
 
 use crate::prelude::*;
-use crate::runtime::vm::{StoreBox, VMArrayCallHostFuncContext, VMContext, VMOpaqueContext};
+use crate::runtime::vm::{
+    Instance, StoreBox, VMArrayCallHostFuncContext, VMContext, VMOpaqueContext, VMStore,
+};
+use crate::store::InstanceId;
 use crate::type_registry::RegisteredType;
 use crate::{FuncType, ValRaw};
 use core::ptr::NonNull;
@@ -32,41 +35,45 @@ unsafe extern "C" fn array_call_shim<F>(
     values_vec_len: usize,
 ) -> bool
 where
-    F: Fn(NonNull<VMContext>, &mut [ValRaw]) -> Result<()> + 'static,
+    F: Fn(&mut dyn VMStore, InstanceId, &mut [ValRaw]) -> Result<()> + 'static,
 {
-    // Be sure to catch Rust panics to manually shepherd them across the wasm
-    // boundary, and then otherwise delegate as normal.
-    crate::runtime::vm::catch_unwind_and_record_trap(|| {
-        // SAFETY: this function itself requires that the `vmctx` is valid to
-        // use here.
-        let state = unsafe {
-            let vmctx = VMArrayCallHostFuncContext::from_opaque(vmctx);
-            vmctx.as_ref().host_state()
-        };
+    // SAFETY: this is an entrypoint of wasm calling a host and our parameters
+    // should reflect that making `enter_host_from_wasm` suitable. Further
+    // unsafe operations are commented below.
+    unsafe {
+        Instance::enter_host_from_wasm(caller_vmctx, |store, instance| {
+            let instance = instance.id();
+            // SAFETY: this function itself requires that the `vmctx` is valid to
+            // use here.
+            let state = {
+                let vmctx = VMArrayCallHostFuncContext::from_opaque(vmctx);
+                vmctx.as_ref().host_state()
+            };
 
-        // Double-check ourselves in debug mode, but we control
-        // the `Any` here so an unsafe downcast should also
-        // work.
-        //
-        // SAFETY: this function is only usable with `TrampolineState<F>`.
-        let state = unsafe {
-            debug_assert!(state.is::<TrampolineState<F>>());
-            &*(state as *const _ as *const TrampolineState<F>)
-        };
-        let mut values_vec = NonNull::slice_from_raw_parts(values_vec, values_vec_len);
-        // SAFETY: it's a contract of this function itself that the values
-        // provided are valid to view as a slice.
-        let values_vec = unsafe { values_vec.as_mut() };
-        (state.func)(caller_vmctx, values_vec)
-    })
+            // Double-check ourselves in debug mode, but we control
+            // the `Any` here so an unsafe downcast should also
+            // work.
+            //
+            // SAFETY: this function is only usable with `TrampolineState<F>`.
+            let state = {
+                debug_assert!(state.is::<TrampolineState<F>>());
+                &*(state as *const _ as *const TrampolineState<F>)
+            };
+            let mut values_vec = NonNull::slice_from_raw_parts(values_vec, values_vec_len);
+            // SAFETY: it's a contract of this function itself that the values
+            // provided are valid to view as a slice.
+            let values_vec = values_vec.as_mut();
+            (state.func)(store, instance, values_vec)
+        })
+    }
 }
 
 pub fn create_array_call_function<F>(
     ft: &FuncType,
     func: F,
 ) -> Result<StoreBox<VMArrayCallHostFuncContext>>
 where
-    F: Fn(NonNull<VMContext>, &mut [ValRaw]) -> Result<()> + Send + Sync + 'static,
+    F: Fn(&mut dyn VMStore, InstanceId, &mut [ValRaw]) -> Result<()> + Send + Sync + 'static,
 {
     let array_call = array_call_shim::<F>;
 

crates/wasmtime/src/runtime/vm.rs

@@ -97,8 +97,8 @@ pub use crate::runtime::vm::gc::*;
 pub use crate::runtime::vm::imports::Imports;
 pub use crate::runtime::vm::instance::{
     GcHeapAllocationIndex, Instance, InstanceAllocationRequest, InstanceAllocator,
-    InstanceAllocatorImpl, InstanceAndStore, InstanceHandle, MemoryAllocationIndex,
-    OnDemandInstanceAllocator, StorePtr, TableAllocationIndex, initialize_instance,
+    InstanceAllocatorImpl, InstanceHandle, MemoryAllocationIndex, OnDemandInstanceAllocator,
+    StorePtr, TableAllocationIndex, initialize_instance,
 };
 #[cfg(feature = "pooling-allocator")]
 pub use crate::runtime::vm::instance::{

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

@@ -11,9 +11,9 @@ use crate::runtime::component::ComponentInstanceId;
 use crate::runtime::vm::instance::{InstanceLayout, OwnedInstance, OwnedVMContext};
 use crate::runtime::vm::vmcontext::VMFunctionBody;
 use crate::runtime::vm::{
-    SendSyncPtr, VMArrayCallFunction, VMFuncRef, VMGlobalDefinition, VMMemoryDefinition,
-    VMOpaqueContext, VMStore, VMStoreRawPtr, VMTableImport, VMWasmCallFunction, ValRaw, VmPtr,
-    VmSafe,
+    HostResult, SendSyncPtr, VMArrayCallFunction, VMFuncRef, VMGlobalDefinition,
+    VMMemoryDefinition, VMOpaqueContext, VMStore, VMStoreRawPtr, VMTableImport, VMWasmCallFunction,
+    ValRaw, VmPtr, VmSafe, catch_unwind_and_record_trap,
 };
 use crate::store::InstanceId;
 use alloc::alloc::Layout;
@@ -203,17 +203,24 @@ impl ComponentInstance {
     /// Converts the `vmctx` provided into a `ComponentInstance` and runs the
     /// provided closure with that instance.
     ///
+    /// This function will also catch any failures that `f` produces and returns
+    /// an appropriate ABI value to return to wasm. This includes normal errors
+    /// such as traps as well as Rust-side panics which require wasm to unwind.
+    ///
     /// # Unsafety
     ///
     /// This is `unsafe` because `vmctx` cannot be guaranteed to be a valid
     /// pointer and it cannot be proven statically that it's safe to get a
     /// mutable reference at this time to the instance from `vmctx`. Note that
     /// it must be also safe to borrow the store mutably, meaning it can't
     /// already be in use elsewhere.
-    pub unsafe fn from_vmctx<R>(
+    pub unsafe fn enter_host_from_wasm<R>(
         vmctx: NonNull<VMComponentContext>,
         f: impl FnOnce(&mut dyn VMStore, Instance) -> R,
-    ) -> R {
+    ) -> R::Abi
+    where
+        R: HostResult,
+    {
         // SAFETY: it's a contract of this function that `vmctx` is a valid
         // allocation which can go backwards to a `ComponentInstance`.
         let mut ptr = unsafe {
@@ -230,7 +237,7 @@ impl ComponentInstance {
         let store = unsafe { &mut *reference.store.0.as_ptr() };
 
         let instance = Instance::from_wasmtime(store, reference.id);
-        f(store, instance)
+        catch_unwind_and_record_trap(store, |store| f(store, instance))
     }
 
     /// Returns the `InstanceId` associated with the `vmctx` provided.

crates/wasmtime/src/runtime/vm/component/libcalls.rs

@@ -99,11 +99,11 @@ mod trampolines {
                     {
                         $(shims!(@validate_param $pname $param);)*
 
-                        let ret = crate::runtime::vm::traphandlers::catch_unwind_and_record_trap(|| unsafe {
-                            ComponentInstance::from_vmctx(vmctx, |store, instance| {
+                        let ret = unsafe {
+                            ComponentInstance::enter_host_from_wasm(vmctx, |store, instance| {
                                 shims!(@invoke $name(store, instance,) $($pname)*)
                             })
-                        });
+                        };
                         shims!(@convert_ret ret $($pname: $param)*)
                     }
                     $(

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

@@ -1,7 +1,6 @@
 #![doc(hidden)]
 
-use crate::runtime::vm::instance::InstanceAndStore;
-use crate::runtime::vm::vmcontext::VMContext;
+use crate::runtime::vm::{Instance, VMContext};
 use core::ptr::NonNull;
 use wasmtime_environ::{EntityRef, MemoryIndex};
 use wasmtime_versioned_export_macros::versioned_export;
@@ -19,8 +18,7 @@ pub unsafe extern "C" fn resolve_vmctx_memory_ptr(p: *const u32) -> *const u8 {
             VMCTX_AND_MEMORY.0 != NonNull::dangling(),
             "must call `__vmctx->set()` before resolving Wasm pointers"
         );
-        InstanceAndStore::from_vmctx(VMCTX_AND_MEMORY.0, |handle| {
-            let (handle, _) = handle.unpack_mut();
+        Instance::enter_host_from_wasm(VMCTX_AND_MEMORY.0, |_store, handle| {
             assert!(
                 VMCTX_AND_MEMORY.1 < handle.env_module().memories.len(),
                 "memory index for debugger is out of bounds"