Make vm::Instance::get_defined_memory safe (#11211)

This is a small step forward to making `vm::Instance` safe internally.
Notably the `get_defined_memory` helper now returns a safe mutable
reference instead of a raw pointer. This is then additionally coupled
with the canonicalization of always working with memories as "instance
plus defined memory index" instead of "instance plus memory index". This
enables removing some unsafe `VMContext` to `Instance` conversion as
well. This change, however, required updating libcalls to special-case
when an imported memory is operated on to load the vmcontext/index in
the libcall itself.

This change notably does not update the `memory_init` libcall just yet
due to the fact that the `DataIndex` is relative to the owning instance
even if the memory is owned by a different instance (e.g. it's
imported). Otherwise this chips away at some of the `unsafe` related to
memory/table management in `vm::Instance`.

Author: alexcrichton

crates/cranelift/src/func_environ.rs

@@ -282,21 +282,10 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
     }
 
     #[cfg(feature = "threads")]
-    fn get_memory_atomic_wait(
-        &mut self,
-        func: &mut Function,
-        memory_index: MemoryIndex,
-        ty: ir::Type,
-    ) -> (ir::FuncRef, usize) {
+    fn get_memory_atomic_wait(&mut self, func: &mut Function, ty: ir::Type) -> ir::FuncRef {
         match ty {
-            I32 => (
-                self.builtin_functions.memory_atomic_wait32(func),
-                memory_index.index(),
-            ),
-            I64 => (
-                self.builtin_functions.memory_atomic_wait64(func),
-                memory_index.index(),
-            ),
+            I32 => self.builtin_functions.memory_atomic_wait32(func),
+            I64 => self.builtin_functions.memory_atomic_wait64(func),
             x => panic!("get_memory_atomic_wait unsupported type: {x:?}"),
         }
     }
@@ -2744,6 +2733,45 @@ impl FuncEnvironment<'_> {
         Ok(())
     }
 
+    /// Returns two `ir::Value`s, the first of which is the vmctx for the memory
+    /// `index` and the second of which is the `DefinedMemoryIndex` for `index`.
+    ///
+    /// Handles internally whether `index` is an imported memory or not.
+    fn memory_vmctx_and_defined_index(
+        &mut self,
+        pos: &mut FuncCursor,
+        index: MemoryIndex,
+    ) -> (ir::Value, ir::Value) {
+        let cur_vmctx = self.vmctx_val(pos);
+        match self.module.defined_memory_index(index) {
+            // This is a defined memory, so the vmctx is our own and the defined
+            // index is `index` here.
+            Some(index) => (cur_vmctx, pos.ins().iconst(I32, i64::from(index.as_u32()))),
+
+            // This is an imported memory, so load the vmctx/defined index from
+            // the import definition itself.
+            None => {
+                let vmimport = self.offsets.vmctx_vmmemory_import(index);
+
+                let vmctx = pos.ins().load(
+                    self.isa.pointer_type(),
+                    ir::MemFlags::trusted(),
+                    cur_vmctx,
+                    i32::try_from(vmimport + u32::from(self.offsets.vmmemory_import_vmctx()))
+                        .unwrap(),
+                );
+                let index = pos.ins().load(
+                    ir::types::I32,
+                    ir::MemFlags::trusted(),
+                    cur_vmctx,
+                    i32::try_from(vmimport + u32::from(self.offsets.vmmemory_import_index()))
+                        .unwrap(),
+                );
+                (vmctx, index)
+            }
+        }
+    }
+
     pub fn translate_memory_grow(
         &mut self,
         builder: &mut FunctionBuilder<'_>,
@@ -2753,14 +2781,15 @@ impl FuncEnvironment<'_> {
     ) -> WasmResult<ir::Value> {
         let mut pos = builder.cursor();
         let memory_grow = self.builtin_functions.memory_grow(&mut pos.func);
-        let index_arg = index.index();
 
-        let memory_index = pos.ins().iconst(I32, index_arg as i64);
-        let vmctx = self.vmctx_val(&mut pos);
+        let (memory_vmctx, defined_memory_index) =
+            self.memory_vmctx_and_defined_index(&mut pos, index);
 
         let index_type = self.memory(index).idx_type;
         let val = self.cast_index_to_i64(&mut pos, val, index_type);
-        let call_inst = pos.ins().call(memory_grow, &[vmctx, val, memory_index]);
+        let call_inst = pos
+            .ins()
+            .call(memory_grow, &[memory_vmctx, val, defined_memory_index]);
         let result = *pos.func.dfg.inst_results(call_inst).first().unwrap();
         let single_byte_pages = match self.memory(index).page_size_log2 {
             16 => false,
@@ -2910,12 +2939,13 @@ impl FuncEnvironment<'_> {
         let memory_fill = self.builtin_functions.memory_fill(&mut pos.func);
         let dst = self.cast_index_to_i64(&mut pos, dst, self.memory(memory_index).idx_type);
         let len = self.cast_index_to_i64(&mut pos, len, self.memory(memory_index).idx_type);
-        let memory_index_arg = pos.ins().iconst(I32, i64::from(memory_index.as_u32()));
-
-        let vmctx = self.vmctx_val(&mut pos);
+        let (memory_vmctx, defined_memory_index) =
+            self.memory_vmctx_and_defined_index(&mut pos, memory_index);
 
-        pos.ins()
-            .call(memory_fill, &[vmctx, memory_index_arg, dst, val, len]);
+        pos.ins().call(
+            memory_fill,
+            &[memory_vmctx, defined_memory_index, dst, val, len],
+        );
 
         Ok(())
     }
@@ -3056,16 +3086,14 @@ impl FuncEnvironment<'_> {
             let mut pos = builder.cursor();
             let addr = self.cast_index_to_i64(&mut pos, addr, self.memory(memory_index).idx_type);
             let implied_ty = pos.func.dfg.value_type(expected);
-            let (wait_func, memory_index) =
-                self.get_memory_atomic_wait(&mut pos.func, memory_index, implied_ty);
-
-            let memory_index_arg = pos.ins().iconst(I32, memory_index as i64);
+            let wait_func = self.get_memory_atomic_wait(&mut pos.func, implied_ty);
 
-            let vmctx = self.vmctx_val(&mut pos);
+            let (memory_vmctx, defined_memory_index) =
+                self.memory_vmctx_and_defined_index(&mut pos, memory_index);
 
             let call_inst = pos.ins().call(
                 wait_func,
-                &[vmctx, memory_index_arg, addr, expected, timeout],
+                &[memory_vmctx, defined_memory_index, addr, expected, timeout],
             );
             let ret = pos.func.dfg.inst_results(call_inst)[0];
             Ok(builder.ins().ireduce(ir::types::I32, ret))
@@ -3093,11 +3121,12 @@ impl FuncEnvironment<'_> {
             let addr = self.cast_index_to_i64(&mut pos, addr, self.memory(memory_index).idx_type);
             let atomic_notify = self.builtin_functions.memory_atomic_notify(&mut pos.func);
 
-            let memory_index_arg = pos.ins().iconst(I32, memory_index.index() as i64);
-            let vmctx = self.vmctx_val(&mut pos);
-            let call_inst = pos
-                .ins()
-                .call(atomic_notify, &[vmctx, memory_index_arg, addr, count]);
+            let (memory_vmctx, defined_memory_index) =
+                self.memory_vmctx_and_defined_index(&mut pos, memory_index);
+            let call_inst = pos.ins().call(
+                atomic_notify,
+                &[memory_vmctx, defined_memory_index, addr, count],
+            );
             let ret = pos.func.dfg.inst_results(call_inst)[0];
             Ok(builder.ins().ireduce(ir::types::I32, ret))
         }

crates/environ/src/vmoffsets.rs

@@ -747,6 +747,18 @@ impl<P: PtrSize> VMOffsets<P> {
         0 * self.pointer_size()
     }
 
+    /// The offset of the `vmctx` field.
+    #[inline]
+    pub fn vmmemory_import_vmctx(&self) -> u8 {
+        1 * self.pointer_size()
+    }
+
+    /// The offset of the `index` field.
+    #[inline]
+    pub fn vmmemory_import_index(&self) -> u8 {
+        2 * self.pointer_size()
+    }
+
     /// Return the size of `VMMemoryImport`.
     #[inline]
     pub fn size_of_vmmemory_import(&self) -> u8 {

crates/wasmtime/src/runtime/memory.rs

@@ -603,7 +603,9 @@ impl Memory {
     /// ```
     pub fn grow(&self, mut store: impl AsContextMut, delta: u64) -> Result<u64> {
         let store = store.as_context_mut().0;
-        let mem = self.wasmtime_memory(store);
+        // FIXME(#11179) shouldn't use a raw pointer to work around the borrow
+        // checker here.
+        let mem: *mut _ = self.wasmtime_memory(store);
         unsafe {
             match (*mem).grow(delta, Some(store))? {
                 Some(size) => {
@@ -638,7 +640,10 @@ impl Memory {
         store.on_fiber(|store| self.grow(store, delta)).await?
     }
 
-    fn wasmtime_memory(&self, store: &mut StoreOpaque) -> *mut crate::runtime::vm::Memory {
+    fn wasmtime_memory<'a>(
+        &self,
+        store: &'a mut StoreOpaque,
+    ) -> &'a mut crate::runtime::vm::Memory {
         self.instance.get_mut(store).get_defined_memory(self.index)
     }
 
@@ -1040,15 +1045,11 @@ impl SharedMemory {
             )
         )]
         crate::runtime::vm::Instance::from_vmctx(wasmtime_export.vmctx, |handle| {
-            let memory_index = handle.env_module().memory_index(wasmtime_export.index);
-            let page_size = handle.memory_page_size(memory_index);
-            debug_assert!(page_size.is_power_of_two());
-            let page_size_log2 = u8::try_from(page_size.ilog2()).unwrap();
+            let module = handle.env_module();
+            let memory_index = module.memory_index(wasmtime_export.index);
+            let page_size_log2 = module.memories[memory_index].page_size_log2;
 
-            let memory = handle
-                .get_defined_memory(wasmtime_export.index)
-                .as_mut()
-                .unwrap();
+            let memory = handle.get_defined_memory(wasmtime_export.index);
             match memory.as_shared_memory() {
                 Some(mem) => Self {
                     vm: mem.clone(),

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

@@ -423,22 +423,6 @@ impl Instance {
         }
     }
 
-    /// Get a locally defined or imported memory.
-    #[cfg(feature = "threads")]
-    pub(crate) fn get_runtime_memory(self: Pin<&mut Self>, index: MemoryIndex) -> &mut Memory {
-        if let Some(defined_index) = self.env_module().defined_memory_index(index) {
-            unsafe { &mut *self.get_defined_memory(defined_index) }
-        } else {
-            let import = self.imported_memory(index);
-            unsafe {
-                let ptr = Instance::from_vmctx(import.vmctx.as_non_null(), |i| {
-                    i.get_defined_memory(import.index)
-                });
-                &mut *ptr
-            }
-        }
-    }
-
     /// Return the indexed `VMMemoryDefinition`, loaded from vmctx memory
     /// already.
     pub fn memory(&self, index: DefinedMemoryIndex) -> VMMemoryDefinition {
@@ -698,36 +682,12 @@ impl Instance {
         index
     }
 
-    /// Get the given memory's page size, in bytes.
-    pub(crate) fn memory_page_size(&self, index: MemoryIndex) -> usize {
-        usize::try_from(self.env_module().memories[index].page_size()).unwrap()
-    }
-
     /// Grow memory by the specified amount of pages.
     ///
     /// Returns `None` if memory can't be grown by the specified amount
     /// of pages. Returns `Some` with the old size in bytes if growth was
     /// successful.
     pub(crate) fn memory_grow(
-        self: Pin<&mut Self>,
-        store: &mut dyn VMStore,
-        index: MemoryIndex,
-        delta: u64,
-    ) -> Result<Option<usize>, Error> {
-        match self.env_module().defined_memory_index(index) {
-            Some(idx) => self.defined_memory_grow(store, idx, delta),
-            None => {
-                let import = self.imported_memory(index);
-                unsafe {
-                    Instance::from_vmctx(import.vmctx.as_non_null(), |i| {
-                        i.defined_memory_grow(store, import.index, delta)
-                    })
-                }
-            }
-        }
-    }
-
-    fn defined_memory_grow(
         mut self: Pin<&mut Self>,
         store: &mut dyn VMStore,
         idx: DefinedMemoryIndex,
@@ -1064,10 +1024,8 @@ impl Instance {
     }
 
     /// Get a locally-defined memory.
-    pub fn get_defined_memory(self: Pin<&mut Self>, index: DefinedMemoryIndex) -> *mut Memory {
-        // SAFETY: the `unsafe` here is projecting from `*mut (A, B)` to
-        // `*mut A`, which should be a safe operation to do.
-        unsafe { &raw mut (*self.memories_mut().get_raw_mut(index).unwrap()).1 }
+    pub fn get_defined_memory(self: Pin<&mut Self>, index: DefinedMemoryIndex) -> &mut Memory {
+        &mut self.memories_mut()[index].1
     }
 
     /// Do a `memory.copy`
@@ -1126,11 +1084,12 @@ impl Instance {
     /// Returns a `Trap` error if the memory range is out of bounds.
     pub(crate) fn memory_fill(
         self: Pin<&mut Self>,
-        memory_index: MemoryIndex,
+        memory_index: DefinedMemoryIndex,
         dst: u64,
         val: u8,
         len: u64,
     ) -> Result<(), Trap> {
+        let memory_index = self.env_module().memory_index(memory_index);
         let memory = self.get_memory(memory_index);
         let dst = self.validate_inbounds(memory.current_length(), dst, len)?;
         let len = usize::try_from(len).unwrap();

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

@@ -69,7 +69,9 @@ use core::pin::Pin;
 use core::ptr::NonNull;
 #[cfg(feature = "threads")]
 use core::time::Duration;
-use wasmtime_environ::{DataIndex, ElemIndex, FuncIndex, MemoryIndex, TableIndex, Trap};
+use wasmtime_environ::{
+    DataIndex, DefinedMemoryIndex, ElemIndex, FuncIndex, MemoryIndex, TableIndex, Trap,
+};
 #[cfg(feature = "wmemcheck")]
 use wasmtime_wmemcheck::AccessError::{
     DoubleMalloc, InvalidFree, InvalidRead, InvalidWrite, OutOfBounds,
@@ -168,13 +170,14 @@ fn memory_grow(
     delta: u64,
     memory_index: u32,
 ) -> Result<Option<AllocationSize>, TrapReason> {
-    let memory_index = MemoryIndex::from_u32(memory_index);
+    let memory_index = DefinedMemoryIndex::from_u32(memory_index);
+    let module = instance.env_module();
+    let page_size_log2 = module.memories[module.memory_index(memory_index)].page_size_log2;
+
     let result = instance
         .as_mut()
         .memory_grow(store, memory_index, delta)?
-        .map(|size_in_bytes| {
-            AllocationSize(size_in_bytes / instance.memory_page_size(memory_index))
-        });
+        .map(|size_in_bytes| AllocationSize(size_in_bytes >> page_size_log2));
 
     Ok(result)
 }
@@ -437,7 +440,7 @@ fn memory_fill(
     val: u32,
     len: u64,
 ) -> Result<(), Trap> {
-    let memory_index = MemoryIndex::from_u32(memory_index);
+    let memory_index = DefinedMemoryIndex::from_u32(memory_index);
     #[expect(clippy::cast_possible_truncation, reason = "known to truncate here")]
     instance.memory_fill(memory_index, dst, val as u8, len)
 }
@@ -1056,9 +1059,9 @@ fn memory_atomic_notify(
     addr_index: u64,
     count: u32,
 ) -> Result<u32, Trap> {
-    let memory = MemoryIndex::from_u32(memory_index);
+    let memory = DefinedMemoryIndex::from_u32(memory_index);
     instance
-        .get_runtime_memory(memory)
+        .get_defined_memory(memory)
         .atomic_notify(addr_index, count)
 }
 
@@ -1073,9 +1076,9 @@ fn memory_atomic_wait32(
     timeout: u64,
 ) -> Result<u32, Trap> {
     let timeout = (timeout as i64 >= 0).then(|| Duration::from_nanos(timeout));
-    let memory = MemoryIndex::from_u32(memory_index);
+    let memory = DefinedMemoryIndex::from_u32(memory_index);
     Ok(instance
-        .get_runtime_memory(memory)
+        .get_defined_memory(memory)
         .atomic_wait32(addr_index, expected, timeout)? as u32)
 }
 
@@ -1090,9 +1093,9 @@ fn memory_atomic_wait64(
     timeout: u64,
 ) -> Result<u32, Trap> {
     let timeout = (timeout as i64 >= 0).then(|| Duration::from_nanos(timeout));
-    let memory = MemoryIndex::from_u32(memory_index);
+    let memory = DefinedMemoryIndex::from_u32(memory_index);
     Ok(instance
-        .get_runtime_memory(memory)
+        .get_defined_memory(memory)
         .atomic_wait64(addr_index, expected, timeout)? as u32)
 }
 

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

@@ -234,6 +234,14 @@ mod test_vmmemory_import {
             offset_of!(VMMemoryImport, from),
             usize::from(offsets.vmmemory_import_from())
         );
+        assert_eq!(
+            offset_of!(VMMemoryImport, vmctx),
+            usize::from(offsets.vmmemory_import_vmctx())
+        );
+        assert_eq!(
+            offset_of!(VMMemoryImport, index),
+            usize::from(offsets.vmmemory_import_index())
+        );
     }
 }