Make limiter growth callbacks unwind-safe and kind-aware

crates/fuzzing/src/oracles.rs

@@ -162,6 +162,7 @@ impl ResourceLimiter for StoreLimits {
         current: usize,
         desired: usize,
         _maximum: Option<usize>,
+        _kind: MemoryKind,
     ) -> Result<bool> {
         Ok(self.alloc(desired - current))
     }

crates/wasmtime/src/runtime/memory.rs

@@ -1118,20 +1118,21 @@ mod tests {
     use alloc::vec::Vec;
 
     #[derive(Default)]
-    struct SuccessfulGrowths(Vec<(usize, usize)>);
+    struct SuccessfulGrowths(Vec<(usize, usize, MemoryKind)>);
 
     impl ResourceLimiter for SuccessfulGrowths {
         fn memory_growing(
             &mut self,
             _current: usize,
             _desired: usize,
             _maximum: Option<usize>,
+            _kind: MemoryKind,
         ) -> Result<bool> {
             Ok(true)
         }
 
-        fn memory_grown(&mut self, current: usize, desired: usize) {
-            self.0.push((current, desired));
+        fn memory_grown(&mut self, current: usize, desired: usize, kind: MemoryKind) {
+            self.0.push((current, desired, kind));
         }
 
         fn table_growing(
@@ -1155,10 +1156,16 @@ mod tests {
 
         let memory = instance.get_memory(&mut store, "m").unwrap();
         assert_eq!(memory.grow(&mut store, 1)?, 1);
-        assert_eq!(store.data().0, [(65536, 2 * 65536)]);
+        assert_eq!(
+            store.data().0,
+            [(65536, 2 * 65536, MemoryKind::LinearMemory)]
+        );
 
         assert!(memory.grow(&mut store, 1).is_err());
-        assert_eq!(store.data().0, [(65536, 2 * 65536)]);
+        assert_eq!(
+            store.data().0,
+            [(65536, 2 * 65536, MemoryKind::LinearMemory)]
+        );
         Ok(())
     }
 

crates/wasmtime/src/runtime/store.rs

@@ -330,27 +330,32 @@ impl StoreResourceLimiter<'_> {
         current: usize,
         desired: usize,
         maximum: Option<usize>,
+        kind: crate::MemoryKind,
     ) -> Result<bool, Error> {
         match self {
-            Self::Sync(s) => s.memory_growing(current, desired, maximum),
+            Self::Sync(s) => s.memory_growing(current, desired, maximum, kind),
             #[cfg(feature = "async")]
-            Self::Async(s) => s.memory_growing(current, desired, maximum).await,
+            Self::Async(s) => s.memory_growing(current, desired, maximum, kind).await,
         }
     }
 
-    pub(crate) fn memory_grow_failed(&mut self, error: crate::Error) -> Result<()> {
+    pub(crate) fn memory_grow_failed(
+        &mut self,
+        error: crate::Error,
+        kind: crate::MemoryKind,
+    ) -> Result<()> {
         match self {
-            Self::Sync(s) => s.memory_grow_failed(error),
+            Self::Sync(s) => s.memory_grow_failed(error, kind),
             #[cfg(feature = "async")]
-            Self::Async(s) => s.memory_grow_failed(error),
+            Self::Async(s) => s.memory_grow_failed(error, kind),
         }
     }
 
-    pub(crate) fn memory_grown(&mut self, current: usize, desired: usize) {
+    pub(crate) fn memory_grown(&mut self, current: usize, desired: usize, kind: crate::MemoryKind) {
         match self {
-            Self::Sync(s) => s.memory_grown(current, desired),
+            Self::Sync(s) => s.memory_grown(current, desired, kind),
             #[cfg(feature = "async")]
-            Self::Async(s) => s.memory_grown(current, desired),
+            Self::Async(s) => s.memory_grown(current, desired, kind),
         }
     }
 

crates/wasmtime/src/runtime/store/gc.rs

@@ -224,7 +224,7 @@ impl StoreOpaque {
     /// Returns an error if growing the GC heap fails.
     pub(crate) async fn grow_gc_heap(
         &mut self,
-        limiter: Option<&mut StoreResourceLimiter<'_>>,
+        mut limiter: Option<&mut StoreResourceLimiter<'_>>,
         bytes_needed: u64,
         asyncness: Asyncness,
     ) -> Result<()> {
@@ -289,13 +289,13 @@ impl StoreOpaque {
         // Safety: we pair growing the GC heap with updating its associated
         // `VMMemoryDefinition` in the `VMStoreContext` immediately
         // afterwards.
-        unsafe {
+        let (old_byte_size, new_byte_size) = unsafe {
             heap.memory
-                .grow(delta_pages_for_alloc, limiter)
+                .grow(delta_pages_for_alloc, limiter.as_deref_mut())
                 .await
                 .context(GcHeapGrowthFailed)?
-                .ok_or(GcHeapGrowthFailed)?;
-        }
+                .ok_or(GcHeapGrowthFailed)?
+        };
         *heap.store.vm_store_context.gc_heap.get_mut() = heap.memory.vmmemory();
 
         let new_size_in_bytes = u64::try_from(heap.memory.byte_size())?;
@@ -311,6 +311,25 @@ impl StoreOpaque {
             "  -> grew GC heap by {:#x} bytes: new size is {new_size_in_bytes:#x} bytes",
             heap.delta_bytes_grown
         );
+
+        // Put the grown memory, and the size delta the collector needs in order
+        // to use the new capacity, back into the `GcStore` before notifying the
+        // embedder. Running the notification while `heap` is still alive would
+        // hand control to embedder code at a point where the store's GC heap has
+        // no memory at all and its free list has not been extended.
+        drop(heap);
+
+        // Resolve the `memory_growing` that permitted this growth. See
+        // `ResourceLimiter::memory_grown` for the contract; the ordering above
+        // is what upholds its unwind-safety guarantee. This reports the kind
+        // directly rather than going through `Memory::notify_grown` because the
+        // memory now lives inside the `GcStore`. Growth here is always nonzero:
+        // a `bytes_needed` of zero returned early and the delta was asserted
+        // above.
+        if let Some(limiter) = limiter {
+            limiter.memory_grown(old_byte_size, new_byte_size, crate::MemoryKind::GcHeap);
+        }
+
         return Ok(());
 
         struct TakenGcHeap<'a> {

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

@@ -787,7 +787,7 @@ impl Instance {
     /// successful.
     pub(crate) async fn memory_grow(
         mut self: Pin<&mut Self>,
-        limiter: Option<&mut StoreResourceLimiter<'_>>,
+        mut limiter: Option<&mut StoreResourceLimiter<'_>>,
         idx: DefinedMemoryIndex,
         delta: u64,
     ) -> Result<Option<usize>, Error> {
@@ -796,16 +796,28 @@ impl Instance {
         // SAFETY: this is the safe wrapper around `Memory::grow` because it
         // automatically updates the `VMMemoryDefinition` in this instance after
         // a growth operation below.
-        let result = unsafe { memory.grow(delta, limiter).await };
+        let result = unsafe { memory.grow(delta, limiter.as_deref_mut()).await };
 
         // Update the state used by a non-shared Wasm memory in case the base
         // pointer and/or the length changed.
         if memory.as_shared_memory().is_none() {
             let vmmemory = memory.vmmemory();
-            self.set_memory(idx, vmmemory);
+            self.as_mut().set_memory(idx, vmmemory);
         }
 
-        result
+        // Only now that the `VMContext` holds the committed base pointer and
+        // length is it safe to hand control to the embedder's post-growth
+        // notification. It may panic, and if that panic is caught this instance
+        // must still be consistent and reusable.
+        match result {
+            Ok(Some((old, new))) => {
+                let memory = &self.as_mut().memories_mut()[idx].1;
+                memory.notify_grown(limiter, old, new);
+                Ok(Some(old))
+            }
+            Ok(None) => Ok(None),
+            Err(e) => Err(e),
+        }
     }
 
     /// Performs a grow operation on the `table_index` specified using `grow`.

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

@@ -246,7 +246,7 @@ impl Memory {
         limiter: Option<&mut StoreResourceLimiter<'_>>,
         kind: MemoryKind,
     ) -> Result<Self> {
-        let (minimum, maximum) = Self::limit_new(ty, limiter).await?;
+        let (minimum, maximum) = Self::limit_new(ty, kind, limiter).await?;
         let tunables = engine.tunables();
         let memory_tunables = MemoryTunables::new(tunables, kind);
         let allocation = creator.new_memory(ty, &memory_tunables, minimum, maximum)?;
@@ -270,7 +270,7 @@ impl Memory {
         memory_image: MemoryImageSlot,
         limiter: Option<&mut StoreResourceLimiter<'_>>,
     ) -> Result<Self> {
-        let (minimum, maximum) = Self::limit_new(ty, limiter).await?;
+        let (minimum, maximum) = Self::limit_new(ty, kind, limiter).await?;
         let pooled_memory = StaticMemory::new(base, base_capacity, minimum, maximum)?;
         let allocation = try_new::<Box<_>>(pooled_memory)?;
 
@@ -299,6 +299,7 @@ impl Memory {
     /// size) of the memory, all in bytes.
     pub(crate) async fn limit_new(
         ty: &wasmtime_environ::Memory,
+        kind: MemoryKind,
         limiter: Option<&mut StoreResourceLimiter<'_>>,
     ) -> Result<(usize, Option<usize>)> {
         let page_size = usize::try_from(ty.page_size()).unwrap();
@@ -344,7 +345,7 @@ impl Memory {
         // now the expected uses of limiter means that's ok.
         if let Some(limiter) = limiter {
             if !limiter
-                .memory_growing(0, minimum.unwrap_or(absolute_max), maximum)
+                .memory_growing(0, minimum.unwrap_or(absolute_max), maximum, kind.into())
                 .await?
             {
                 bail!(
@@ -399,8 +400,14 @@ impl Memory {
     /// Grow memory by the specified amount of wasm pages.
     ///
     /// Returns `None` if memory can't be grown by the specified amount
-    /// of wasm pages. Returns `Some` with the old size of memory, in bytes, on
-    /// successful growth.
+    /// of wasm pages. Returns `Some` with the old and new sizes of memory, in
+    /// bytes, on successful growth.
+    ///
+    /// Note that this does *not* emit the limiter's `memory_grown`
+    /// notification: at the point this returns the backing allocation has
+    /// committed but the owning `VMContext` has not yet been refreshed. Callers
+    /// are responsible for publishing the new [`VMMemoryDefinition`] and then
+    /// calling [`Self::notify_grown`].
     ///
     /// # Safety
     ///
@@ -417,14 +424,35 @@ impl Memory {
         &mut self,
         delta_pages: u64,
         limiter: Option<&mut StoreResourceLimiter<'_>>,
-    ) -> Result<Option<usize>, Error> {
-        let result = match self {
-            Memory::Local(mem) => mem.grow(delta_pages, limiter).await?,
-            Memory::Shared(mem) => mem.grow(delta_pages)?,
-        };
-        match result {
-            Some((old, _new)) => Ok(Some(old)),
-            None => Ok(None),
+    ) -> Result<Option<(usize, usize)>, Error> {
+        match self {
+            Memory::Local(mem) => mem.grow(delta_pages, limiter).await,
+            Memory::Shared(mem) => mem.grow(delta_pages),
+        }
+    }
+
+    /// Emits the limiter's `memory_grown` notification for a growth that
+    /// [`Self::grow`] just reported as successful.
+    ///
+    /// This must be called only *after* the refreshed [`VMMemoryDefinition`]
+    /// has been published to whatever owns this memory, so that embedder code
+    /// which panics cannot leave a `VMContext` pointing at a stale base or
+    /// length.
+    ///
+    /// The notification is suppressed for shared memories, which have no
+    /// associated limiter, and for growth of zero bytes.
+    pub fn notify_grown(
+        &self,
+        limiter: Option<&mut StoreResourceLimiter<'_>>,
+        old_byte_size: usize,
+        new_byte_size: usize,
+    ) {
+        let Memory::Local(mem) = self else { return };
+        if new_byte_size <= old_byte_size {
+            return;
+        }
+        if let Some(limiter) = limiter {
+            limiter.memory_grown(old_byte_size, new_byte_size, mem.limiter_kind());
         }
     }
 
@@ -600,6 +628,14 @@ impl LocalMemory {
         &self.ty
     }
 
+    /// This memory's kind as the embedder-facing [`crate::MemoryKind`].
+    ///
+    /// Note that `MemoryKind` in this module refers to the `wasmtime_environ`
+    /// type; this converts to the public one the limiter callbacks take.
+    fn limiter_kind(&self) -> crate::MemoryKind {
+        self.kind.into()
+    }
+
     /// Grows a memory by `delta_pages`.
     ///
     /// This performs the necessary checks on the growth before delegating to
@@ -645,15 +681,15 @@ impl LocalMemory {
         if !self.ty().allow_growth_to(new_byte_size) {
             if let Some(limiter) = limiter {
                 let err = crate::format_err!("memory growth exceeds memory type's limits");
-                limiter.memory_grow_failed(err)?;
+                limiter.memory_grow_failed(err, self.limiter_kind())?;
             }
             return Ok(None);
         }
 
         // Store limiter gets first chance to reject memory_growing.
         if let Some(limiter) = &mut limiter {
             if !limiter
-                .memory_growing(old_byte_size, new_byte_size, maximum)
+                .memory_growing(old_byte_size, new_byte_size, maximum, self.limiter_kind())
                 .await?
             {
                 return Ok(None);
@@ -714,12 +750,13 @@ impl LocalMemory {
                     assert_eq!(base_ptr_before, self.alloc.base().as_mut_ptr());
                 }
 
-                if matches!(self.kind, MemoryKind::LinearMemory)
-                    && let Some(limiter) = limiter
-                {
-                    limiter.memory_grown(old_byte_size, new_byte_size);
-                }
-
+                // NB: the `memory_grown` notification is deliberately *not*
+                // emitted here. The backing allocation has committed but the
+                // owning `VMContext` still holds the pre-growth base pointer
+                // and length, so running embedder code at this point would let
+                // a panic escape with compiled wasm observing stale bounds or a
+                // freed base. Callers emit the notification once they have
+                // published the refreshed `VMMemoryDefinition`.
                 Ok(Some((old_byte_size, new_byte_size)))
             }
             Err(e) => {
@@ -728,7 +765,7 @@ impl LocalMemory {
                 // dropped
                 // (https://github.com/bytecodealliance/wasmtime/issues/4240).
                 if let Some(limiter) = limiter {
-                    limiter.memory_grow_failed(e)?;
+                    limiter.memory_grow_failed(e, self.limiter_kind())?;
                 }
                 Ok(None)
             }

crates/wasmtime/src/runtime/vm/memory/shared_memory.rs

@@ -43,7 +43,11 @@ impl SharedMemory {
         );
         // Note that without a limiter being passed to `limit_new` this
         // `assert_ready` should never panic.
-        let (minimum_bytes, maximum_bytes) = vm::assert_ready(Memory::limit_new(ty, None))?;
+        let (minimum_bytes, maximum_bytes) = vm::assert_ready(Memory::limit_new(
+            ty,
+            wasmtime_environ::MemoryKind::LinearMemory,
+            None,
+        ))?;
         let mmap_memory = MmapMemory::new(ty, &memory_tunables, minimum_bytes, maximum_bytes)?;
         let boxed: Box<dyn crate::runtime::vm::RuntimeLinearMemory> =
             try_new::<Box<_>>(mmap_memory)?;