Make module information lookup from runtime safe.

This commit uses a two-phase lookup of stack map information from modules
rather than giving back raw pointers to stack maps.

First the runtime looks up information about a module from a pc value, which
returns an `Arc` it keeps a reference on while completing the stack map lookup.

Second it then queries the module information for the stack map from a pc
value, getting a reference to the stack map (which is now safe because of the
`Arc` held by the runtime).

crates/environ/src/vmoffsets.rs

@@ -6,7 +6,7 @@
 // struct VMContext {
 //      interrupts: *const VMInterrupts,
 //      externref_activations_table: *mut VMExternRefActivationsTable,
-//      stack_map_lookup: *const dyn StackMapLookup,
+//      module_info_lookup: *const dyn ModuleInfoLookup,
 //      signature_ids: [VMSharedSignatureIndex; module.num_signature_ids],
 //      imported_functions: [VMFunctionImport; module.num_imported_functions],
 //      imported_tables: [VMTableImport; module.num_imported_tables],
@@ -77,7 +77,7 @@ pub struct VMOffsets {
     // precalculated offsets of various member fields
     interrupts: u32,
     externref_activations_table: u32,
-    stack_map_lookup: u32,
+    module_info_lookup: u32,
     signature_ids: u32,
     imported_functions: u32,
     imported_tables: u32,
@@ -149,7 +149,7 @@ impl From<VMOffsetsFields> for VMOffsets {
             num_defined_globals: fields.num_defined_globals,
             interrupts: 0,
             externref_activations_table: 0,
-            stack_map_lookup: 0,
+            module_info_lookup: 0,
             signature_ids: 0,
             imported_functions: 0,
             imported_tables: 0,
@@ -168,12 +168,12 @@ impl From<VMOffsetsFields> for VMOffsets {
             .interrupts
             .checked_add(u32::from(fields.pointer_size))
             .unwrap();
-        ret.stack_map_lookup = ret
+        ret.module_info_lookup = ret
             .externref_activations_table
             .checked_add(u32::from(fields.pointer_size))
             .unwrap();
         ret.signature_ids = ret
-            .stack_map_lookup
+            .module_info_lookup
             .checked_add(u32::from(fields.pointer_size * 2))
             .unwrap();
         ret.imported_functions = ret
@@ -507,10 +507,10 @@ impl VMOffsets {
         self.externref_activations_table
     }
 
-    /// The offset of the `*const dyn StackMapLookup` member.
+    /// The offset of the `*const dyn ModuleInfoLookup` member.
     #[inline]
-    pub fn vmctx_stack_map_lookup(&self) -> u32 {
+    pub fn vmctx_module_info_lookup(&self) -> u32 {
-        self.stack_map_lookup
+        self.module_info_lookup
     }
 
     /// The offset of the `signature_ids` array.

crates/runtime/src/externref.rs

@@ -99,7 +99,6 @@
 //! Examination of Deferred Reference Counting and Cycle Detection* by Quinane:
 //! <https://openresearch-repository.anu.edu.au/bitstream/1885/42030/2/hon-thesis.pdf>
 
-use std::alloc::Layout;
 use std::any::Any;
 use std::cell::{Cell, RefCell, UnsafeCell};
 use std::cmp::Ordering;
@@ -108,6 +107,7 @@ use std::hash::{Hash, Hasher};
 use std::mem;
 use std::ops::Deref;
 use std::ptr::{self, NonNull};
+use std::{alloc::Layout, sync::Arc};
 use wasmtime_environ::ir::StackMap;
 
 /// An external reference to some opaque data.
@@ -594,10 +594,10 @@ impl VMExternRefActivationsTable {
     pub unsafe fn insert_with_gc(
         &self,
         externref: VMExternRef,
-        stack_map_lookup: &dyn StackMapLookup,
+        module_info_lookup: &dyn ModuleInfoLookup,
     ) {
         if let Err(externref) = self.try_insert(externref) {
-            self.gc_and_insert_slow(externref, stack_map_lookup);
+            self.gc_and_insert_slow(externref, module_info_lookup);
         }
     }
 
@@ -605,9 +605,9 @@ impl VMExternRefActivationsTable {
     unsafe fn gc_and_insert_slow(
         &self,
         externref: VMExternRef,
-        stack_map_lookup: &dyn StackMapLookup,
+        module_info_lookup: &dyn ModuleInfoLookup,
     ) {
-        gc(stack_map_lookup, self);
+        gc(module_info_lookup, self);
 
         // Might as well insert right into the hash set, rather than the bump
         // chunk, since we are already on a slow path and we get de-duplication
@@ -741,29 +741,28 @@ impl VMExternRefActivationsTable {
     }
 }
 
-/// Used by the runtime to lookup a stack map from a PC value.
+/// Used by the runtime to lookup information about a module given a
-///
+/// program counter value.
-/// # Safety
+pub trait ModuleInfoLookup: 'static {
-///
+    /// Lookup the module information from a program counter value.
-/// This trait is unsafe as it returns pointers to a stack map without
+    fn lookup(&self, pc: usize) -> Option<Arc<dyn ModuleInfo>>;
-/// any clear ownership.
-///
-/// It is the responsibility of the caller to not have the pointer outlive
-/// the stack map lookup trait object.
-pub unsafe trait StackMapLookup: 'static {
-    /// Lookup the stack map at a program counter (PC) value.
-    fn lookup(&self, pc: usize) -> Option<*const StackMap>;
 }
 
-pub(crate) struct EmptyStackMapLookup;
+/// Used by the runtime to query module information.
+pub trait ModuleInfo {
+    /// Lookup the stack map at a program counter value.
+    fn lookup_stack_map(&self, pc: usize) -> Option<&StackMap>;
+}
 
-unsafe impl StackMapLookup for EmptyStackMapLookup {
+pub(crate) struct EmptyModuleInfoLookup;
-    fn lookup(&self, _pc: usize) -> Option<*const StackMap> {
+
+impl ModuleInfoLookup for EmptyModuleInfoLookup {
+    fn lookup(&self, _pc: usize) -> Option<Arc<dyn ModuleInfo>> {
         None
     }
 }
 
-pub(crate) const EMPTY_STACK_MAP_LOOKUP: EmptyStackMapLookup = EmptyStackMapLookup;
+pub(crate) const EMPTY_MODULE_LOOKUP: EmptyModuleInfoLookup = EmptyModuleInfoLookup;
 
 #[derive(Debug, Default)]
 struct DebugOnly<T> {
@@ -810,7 +809,7 @@ impl<T> std::ops::DerefMut for DebugOnly<T> {
 /// Additionally, you must have registered the stack maps for every Wasm module
 /// that has frames on the stack with the given `stack_maps_registry`.
 pub unsafe fn gc(
-    stack_map_lookup: &dyn StackMapLookup,
+    module_info_lookup: &dyn ModuleInfoLookup,
     externref_activations_table: &VMExternRefActivationsTable,
 ) {
     // We borrow the precise stack roots `RefCell` for the whole duration of
@@ -848,8 +847,7 @@ pub unsafe fn gc(
             if cfg!(debug_assertions) {
                 // Assert that there aren't any Wasm frames on the stack.
                 backtrace::trace(|frame| {
-                    let stack_map = stack_map_lookup.lookup(frame.ip() as usize);
+                    assert!(module_info_lookup.lookup(frame.ip() as usize).is_none());
-                    assert!(stack_map.is_none());
                     true
                 });
             }
@@ -893,28 +891,30 @@ pub unsafe fn gc(
         let pc = frame.ip() as usize;
         let sp = frame.sp() as usize;
 
-        if let Some(stack_map) = stack_map_lookup.lookup(pc) {
+        if let Some(module_info) = module_info_lookup.lookup(pc) {
-            debug_assert!(sp != 0, "we should always get a valid SP for Wasm frames");
+            if let Some(stack_map) = module_info.lookup_stack_map(pc) {
+                debug_assert!(sp != 0, "we should always get a valid SP for Wasm frames");
 
-            for i in 0..((*stack_map).mapped_words() as usize) {
+                for i in 0..(stack_map.mapped_words() as usize) {
-                if (*stack_map).get_bit(i) {
+                    if stack_map.get_bit(i) {
-                    // Stack maps have one bit per word in the frame, and the
+                        // Stack maps have one bit per word in the frame, and the
-                    // zero^th bit is the *lowest* addressed word in the frame,
+                        // zero^th bit is the *lowest* addressed word in the frame,
-                    // i.e. the closest to the SP. So to get the `i`^th word in
+                        // i.e. the closest to the SP. So to get the `i`^th word in
-                    // this frame, we add `i * sizeof(word)` to the SP.
+                        // this frame, we add `i * sizeof(word)` to the SP.
-                    let ptr_to_ref = sp + i * mem::size_of::<usize>();
+                        let ptr_to_ref = sp + i * mem::size_of::<usize>();
 
-                    let r = std::ptr::read(ptr_to_ref as *const *mut VMExternData);
+                        let r = std::ptr::read(ptr_to_ref as *const *mut VMExternData);
-                    debug_assert!(
+                        debug_assert!(
-                        r.is_null() || activations_table_set.contains(&r),
+                            r.is_null() || activations_table_set.contains(&r),
-                        "every on-stack externref inside a Wasm frame should \
+                            "every on-stack externref inside a Wasm frame should \
-                         have an entry in the VMExternRefActivationsTable"
+                            have an entry in the VMExternRefActivationsTable"
-                    );
-                    if let Some(r) = NonNull::new(r) {
-                        VMExternRefActivationsTable::insert_precise_stack_root(
-                            &mut precise_stack_roots,
-                            r,
                         );
+                        if let Some(r) = NonNull::new(r) {
+                            VMExternRefActivationsTable::insert_precise_stack_root(
+                                &mut precise_stack_roots,
+                                r,
+                            );
+                        }
                     }
                 }
             }

crates/runtime/src/instance.rs

@@ -3,7 +3,7 @@
 //! `InstanceHandle` is a reference-counting handle for an `Instance`.
 
 use crate::export::Export;
-use crate::externref::{StackMapLookup, VMExternRefActivationsTable};
+use crate::externref::{ModuleInfoLookup, VMExternRefActivationsTable};
 use crate::memory::{Memory, RuntimeMemoryCreator};
 use crate::table::{Table, TableElement};
 use crate::traphandlers::Trap;
@@ -249,9 +249,9 @@ impl Instance {
         unsafe { self.vmctx_plus_offset(self.offsets.vmctx_externref_activations_table()) }
     }
 
-    /// Return a pointer to the `StackMapLookup`.
+    /// Return a pointer to the `ModuleInfoLookup`.
-    pub fn stack_map_lookup(&self) -> *mut *const dyn StackMapLookup {
+    pub fn module_info_lookup(&self) -> *mut *const dyn ModuleInfoLookup {
-        unsafe { self.vmctx_plus_offset(self.offsets.vmctx_stack_map_lookup()) }
+        unsafe { self.vmctx_plus_offset(self.offsets.vmctx_module_info_lookup()) }
     }
 
     /// Return a reference to the vmctx used by compiled wasm code.

crates/runtime/src/instance/allocator.rs

@@ -1,4 +1,4 @@
-use crate::externref::{StackMapLookup, VMExternRefActivationsTable, EMPTY_STACK_MAP_LOOKUP};
+use crate::externref::{ModuleInfoLookup, VMExternRefActivationsTable, EMPTY_MODULE_LOOKUP};
 use crate::imports::Imports;
 use crate::instance::{Instance, InstanceHandle, RuntimeMemoryCreator};
 use crate::memory::{DefaultMemoryCreator, Memory};
@@ -57,8 +57,8 @@ pub struct InstanceAllocationRequest<'a> {
     /// The pointer to the reference activations table to use for the instance.
     pub externref_activations_table: *mut VMExternRefActivationsTable,
 
-    /// The pointer to the stack map lookup to use for the instance.
+    /// The pointer to the module info lookup to use for the instance.
-    pub stack_map_lookup: Option<*const dyn StackMapLookup>,
+    pub module_info_lookup: Option<*const dyn ModuleInfoLookup>,
 }
 
 /// An link error while instantiating a module.
@@ -447,7 +447,7 @@ unsafe fn initialize_vmcontext(instance: &Instance, req: InstanceAllocationReque
 
     *instance.interrupts() = req.interrupts;
     *instance.externref_activations_table() = req.externref_activations_table;
-    *instance.stack_map_lookup() = req.stack_map_lookup.unwrap_or(&EMPTY_STACK_MAP_LOOKUP);
+    *instance.module_info_lookup() = req.module_info_lookup.unwrap_or(&EMPTY_MODULE_LOOKUP);
 
     // Initialize shared signatures
     let mut ptr = instance.signature_ids_ptr();

crates/runtime/src/instance/allocator/pooling.rs

@@ -1370,7 +1370,7 @@ mod test {
                             host_state: Box::new(()),
                             interrupts: std::ptr::null(),
                             externref_activations_table: std::ptr::null_mut(),
-                            stack_map_lookup: None,
+                            module_info_lookup: None,
                         },
                     )
                     .expect("allocation should succeed"),
@@ -1394,7 +1394,7 @@ mod test {
                 host_state: Box::new(()),
                 interrupts: std::ptr::null(),
                 externref_activations_table: std::ptr::null_mut(),
-                stack_map_lookup: None,
+                module_info_lookup: None,
             },
         ) {
             Err(InstantiationError::Limit(3)) => {}

crates/runtime/src/instance/allocator/pooling/uffd.rs

@@ -523,7 +523,7 @@ mod test {
                                 host_state: Box::new(()),
                                 interrupts: ptr::null(),
                                 externref_activations_table: ptr::null_mut(),
-                                stack_map_lookup: None,
+                                module_info_lookup: None,
                             },
                         )
                         .expect("instance should allocate"),

crates/runtime/src/libcalls.rs

@@ -449,8 +449,8 @@ pub unsafe extern "C" fn wasmtime_activations_table_insert_with_gc(
     let externref = VMExternRef::clone_from_raw(externref);
     let instance = (&mut *vmctx).instance();
     let activations_table = &**instance.externref_activations_table();
-    let stack_map_lookup = &**instance.stack_map_lookup();
+    let module_info_lookup = &**instance.module_info_lookup();
-    activations_table.insert_with_gc(externref, stack_map_lookup);
+    activations_table.insert_with_gc(externref, module_info_lookup);
 }
 
 /// Perform a Wasm `global.get` for `externref` globals.
@@ -466,8 +466,8 @@ pub unsafe extern "C" fn wasmtime_externref_global_get(
         Some(externref) => {
             let raw = externref.as_raw();
             let activations_table = &**instance.externref_activations_table();
-            let stack_map_lookup = &**instance.stack_map_lookup();
+            let module_info_lookup = &**instance.module_info_lookup();
-            activations_table.insert_with_gc(externref, stack_map_lookup);
+            activations_table.insert_with_gc(externref, module_info_lookup);
             raw
         }
     }

crates/wasmtime/src/func/typed.rs

@@ -207,7 +207,7 @@ unsafe impl WasmTy for Option<ExternRef> {
             unsafe {
                 store
                     .externref_activations_table()
-                    .insert_with_gc(x.inner, store.stack_map_lookup());
+                    .insert_with_gc(x.inner, store.module_info_lookup());
             }
             abi
         } else {

crates/wasmtime/src/instance.rs

@@ -524,7 +524,7 @@ impl<'a> Instantiator<'a> {
                 externref_activations_table: self.store.externref_activations_table()
                     as *const VMExternRefActivationsTable
                     as *mut _,
-                stack_map_lookup: Some(self.store.stack_map_lookup()),
+                module_info_lookup: Some(self.store.module_info_lookup()),
             })?;
 
             // After we've created the `InstanceHandle` we still need to run

crates/wasmtime/src/module/registry.rs

@@ -6,10 +6,13 @@ use std::{
     sync::{Arc, Mutex},
 };
 use wasmtime_environ::{
-    entity::EntityRef, ir, wasm::DefinedFuncIndex, FunctionAddressMap, TrapInformation,
+    entity::EntityRef,
+    ir::{self, StackMap},
+    wasm::DefinedFuncIndex,
+    FunctionAddressMap, TrapInformation,
 };
 use wasmtime_jit::CompiledModule;
-use wasmtime_runtime::{VMCallerCheckedAnyfunc, VMTrampoline};
+use wasmtime_runtime::{ModuleInfo, VMCallerCheckedAnyfunc, VMTrampoline};
 
 lazy_static::lazy_static! {
     static ref GLOBAL_MODULES: Mutex<GlobalModuleRegistry> = Default::default();
@@ -27,7 +30,7 @@ fn func_by_pc(module: &CompiledModule, pc: usize) -> Option<(DefinedFuncIndex, u
 ///
 /// The `BTreeMap` is used to quickly locate a module based on a program counter value.
 #[derive(Default)]
-pub struct ModuleRegistry(BTreeMap<usize, RegisteredModule>);
+pub struct ModuleRegistry(BTreeMap<usize, Arc<RegisteredModule>>);
 
 impl ModuleRegistry {
     /// Fetches frame information about a program counter in a backtrace.
@@ -48,12 +51,13 @@ impl ModuleRegistry {
         self.module(pc)?.lookup_trap_info(pc)
     }
 
-    /// Looks up a stack map from a program counter.
+    /// Fetches information about a registered module given a program counter value.
-    pub fn lookup_stack_map<'a>(&'a self, pc: usize) -> Option<&'a ir::StackMap> {
+    pub fn lookup_module(&self, pc: usize) -> Option<Arc<dyn ModuleInfo>> {
-        self.module(pc)?.lookup_stack_map(pc)
+        self.module(pc)
+            .map(|m| -> Arc<dyn ModuleInfo> { m.clone() })
     }
 
-    fn module(&self, pc: usize) -> Option<&RegisteredModule> {
+    fn module(&self, pc: usize) -> Option<&Arc<RegisteredModule>> {
         let (end, info) = self.0.range(pc..).next()?;
         if pc < info.start || *end < pc {
             return None;
@@ -94,11 +98,11 @@ impl ModuleRegistry {
 
         let prev = self.0.insert(
             end,
-            RegisteredModule {
+            Arc::new(RegisteredModule {
                 start,
                 module: compiled_module.clone(),
                 signatures: module.signatures().clone(),
-            },
+            }),
         );
         assert!(prev.is_none());
 
@@ -209,8 +213,29 @@ impl RegisteredModule {
         Some(&info.traps[idx])
     }
 
-    /// Looks up a stack map from a program counter
+    fn instr_pos(offset: u32, addr_map: &FunctionAddressMap) -> Option<usize> {
-    pub fn lookup_stack_map(&self, pc: usize) -> Option<&ir::StackMap> {
+        // Use our relative position from the start of the function to find the
+        // machine instruction that corresponds to `pc`, which then allows us to
+        // map that to a wasm original source location.
+        match addr_map
+            .instructions
+            .binary_search_by_key(&offset, |map| map.code_offset)
+        {
+            // Exact hit!
+            Ok(pos) => Some(pos),
+
+            // This *would* be at the first slot in the array, so no
+            // instructions cover `pc`.
+            Err(0) => None,
+
+            // This would be at the `nth` slot, so we're at the `n-1`th slot.
+            Err(n) => Some(n - 1),
+        }
+    }
+}
+
+impl ModuleInfo for RegisteredModule {
+    fn lookup_stack_map(&self, pc: usize) -> Option<&StackMap> {
         let (index, offset) = func_by_pc(&self.module, pc)?;
         let info = self.module.func_info(index);
 
@@ -275,26 +300,6 @@ impl RegisteredModule {
 
         Some(&info.stack_maps[index].stack_map)
     }
-
-    fn instr_pos(offset: u32, addr_map: &FunctionAddressMap) -> Option<usize> {
-        // Use our relative position from the start of the function to find the
-        // machine instruction that corresponds to `pc`, which then allows us to
-        // map that to a wasm original source location.
-        match addr_map
-            .instructions
-            .binary_search_by_key(&offset, |map| map.code_offset)
-        {
-            // Exact hit!
-            Ok(pos) => Some(pos),
-
-            // This *would* be at the first slot in the array, so no
-            // instructions cover `pc`.
-            Err(0) => None,
-
-            // This would be at the `nth` slot, so we're at the `n-1`th slot.
-            Err(n) => Some(n - 1),
-        }
-    }
 }
 
 // Counterpart to `RegisteredModule`, but stored in the global registry.

crates/wasmtime/src/store.rs

@@ -16,9 +16,9 @@ use std::rc::Rc;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 use wasmtime_runtime::{
-    InstanceAllocator, InstanceHandle, OnDemandInstanceAllocator, SignalHandler, TrapInfo,
+    InstanceAllocator, InstanceHandle, ModuleInfo, OnDemandInstanceAllocator, SignalHandler,
-    VMCallerCheckedAnyfunc, VMContext, VMExternRef, VMExternRefActivationsTable, VMInterrupts,
+    TrapInfo, VMCallerCheckedAnyfunc, VMContext, VMExternRef, VMExternRefActivationsTable,
-    VMTrampoline,
+    VMInterrupts, VMTrampoline,
 };
 
 /// Used to associate instances with the store.
@@ -440,7 +440,7 @@ impl Store {
     }
 
     #[inline]
-    pub(crate) fn stack_map_lookup(&self) -> &dyn wasmtime_runtime::StackMapLookup {
+    pub(crate) fn module_info_lookup(&self) -> &dyn wasmtime_runtime::ModuleInfoLookup {
         self.inner.as_ref()
     }
 
@@ -910,10 +910,9 @@ impl Drop for StoreInner {
     }
 }
 
-unsafe impl wasmtime_runtime::StackMapLookup for StoreInner {
+impl wasmtime_runtime::ModuleInfoLookup for StoreInner {
-    fn lookup(&self, pc: usize) -> Option<*const wasmtime_environ::ir::StackMap> {
+    fn lookup(&self, pc: usize) -> Option<Arc<dyn ModuleInfo>> {
-        // The address of the stack map is stable for the lifetime of the store
+        self.modules.borrow().lookup_module(pc)
-        self.modules.borrow().lookup_stack_map(pc).map(|m| m as _)
     }
 }
 

crates/wasmtime/src/trampoline.rs

@@ -77,7 +77,7 @@ fn create_handle(
                 externref_activations_table: store.externref_activations_table()
                     as *const VMExternRefActivationsTable
                     as *mut _,
-                stack_map_lookup: Some(store.stack_map_lookup()),
+                module_info_lookup: Some(store.module_info_lookup()),
             },
         )?;
 

crates/wasmtime/src/trampoline/create_handle.rs

@@ -43,7 +43,7 @@ pub(crate) fn create_handle(
                 externref_activations_table: store.externref_activations_table()
                     as *const VMExternRefActivationsTable
                     as *mut _,
-                stack_map_lookup: &store,
+                module_info_lookup: &store,
             })?;
 
         Ok(store.add_instance(handle, true))

crates/wasmtime/src/trampoline/func.rs

@@ -287,7 +287,7 @@ pub fn create_function(
                 host_state: Box::new(trampoline_state),
                 interrupts: std::ptr::null(),
                 externref_activations_table: std::ptr::null_mut(),
-                stack_map_lookup: None,
+                module_info_lookup: None,
             })?,
             trampoline,
         ))
@@ -319,7 +319,7 @@ pub unsafe fn create_raw_function(
             host_state,
             interrupts: std::ptr::null(),
             externref_activations_table: std::ptr::null_mut(),
-            stack_map_lookup: None,
+            module_info_lookup: None,
         })?,
     )
 }

crates/wasmtime/src/values.rs

@@ -98,7 +98,7 @@ impl Val {
                 let externref_ptr = x.inner.as_raw();
                 store
                     .externref_activations_table()
-                    .insert_with_gc(x.inner, store.stack_map_lookup());
+                    .insert_with_gc(x.inner, store.module_info_lookup());
                 ptr::write(p as *mut *mut u8, externref_ptr)
             }
             Val::FuncRef(f) => ptr::write(