Properly manage ownership of GC refs in passive element segments (#13179) (#13183)

* Properly manage ownership of GC refs in passive element segments

We need to `clone_gc_ref` when pushing into the segment and `drop_gc_ref` when
clearing the segment, similar to what we do for e.g. setting
`VMGlobalDefinition`s.

Fixes #13066

* fix cfg warnings

* fix cfg_attr

* remove cfg_attr

Co-authored-by: Nick Fitzgerald <fitzgen@gmail.com>

Author: alexcrichton

crates/wasmtime/src/runtime/func.rs

@@ -1064,10 +1064,6 @@ impl Func {
         self.vm_func_ref(store).as_ptr().cast()
     }
 
-    pub(crate) fn to_val_raw(&self, store: &mut StoreOpaque) -> ValRaw {
-        ValRaw::funcref(self.to_raw_(store))
-    }
-
     /// Invokes this function with the `params` given, returning the results
     /// asynchronously.
     ///

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

@@ -13,7 +13,7 @@ use crate::runtime::vm::vmcontext::{
     VMTableDefinition, VMTableImport, VMTagDefinition, VMTagImport,
 };
 use crate::runtime::vm::{
-    GcStore, HostResult, Imports, ModuleRuntimeInfo, SendSyncPtr, VMGlobalKind, VMStore,
+    GcStore, HostResult, Imports, ModuleRuntimeInfo, SendSyncPtr, VMGcRef, VMGlobalKind, VMStore,
     VMStoreRawPtr, VmPtr, VmSafe, WasmFault, catch_unwind_and_record_trap,
 };
 use crate::store::{
@@ -133,7 +133,7 @@ pub struct Instance {
     //
     // TODO(#12621): This should be a `TrySecondaryMap<PassiveElemIndex, _>`
     // but that type is currently footgun-y / isn't actually OOM-safe yet.
-    passive_elements: TryVec<Option<(NeedsGcRooting, TryVec<ValRaw>)>>,
+    passive_elements: TryVec<PassiveElementSegment>,
 
     /// Stores the dropped passive data segments in this instantiation by index.
     /// If the index is present in the set, the segment has been dropped.
@@ -225,8 +225,8 @@ impl Instance {
         gc_roots: &mut crate::vm::GcRootsList,
     ) {
         for segment in self.passive_elements_mut().iter_mut() {
-            if let Some((wasmtime_environ::NeedsGcRooting::Yes, elems)) = segment {
-                for e in elems {
+            if segment.needs_gc_rooting() == NeedsGcRooting::Yes {
+                for e in segment.elements() {
                     let Some(root) = e.as_vmgc_ref_ptr() else {
                         continue;
                     };
@@ -912,16 +912,12 @@ impl Instance {
             return &[];
         };
 
-        let Some((_, seg)) = &self.passive_elements[passive.index()] else {
-            return &[];
-        };
-
-        &**seg
+        self.passive_elements[passive.index()].elements()
     }
 
     pub(crate) fn passive_elements_mut(
         self: Pin<&mut Self>,
-    ) -> Pin<&mut TryVec<Option<(NeedsGcRooting, TryVec<ValRaw>)>>> {
+    ) -> Pin<&mut TryVec<PassiveElementSegment>> {
         // SAFETY: Not moving data out of `self`.
         Pin::new(&mut unsafe { self.get_unchecked_mut() }.passive_elements)
     }
@@ -996,6 +992,7 @@ impl Instance {
     /// Drop an element.
     pub(crate) fn elem_drop(
         self: Pin<&mut Self>,
+        gc_store: Option<&mut GcStore>,
         elem_index: ElemIndex,
     ) -> Result<(), OutOfMemory> {
         // https://webassembly.github.io/reference-types/core/exec/instructions.html#exec-elem-drop
@@ -1010,7 +1007,7 @@ impl Instance {
             return Ok(());
         };
 
-        self.passive_elements_mut()[passive_index.index()] = None;
+        self.passive_elements_mut()[passive_index.index()].clear(gc_store);
         Ok(())
     }
 
@@ -1922,3 +1919,77 @@ impl<T: InstanceLayout> Drop for OwnedInstance<T> {
         }
     }
 }
+
+#[derive(Debug)]
+pub(crate) struct PassiveElementSegment {
+    needs_gc_rooting: NeedsGcRooting,
+    elements: TryVec<ValRaw>,
+}
+
+impl PassiveElementSegment {
+    /// Create a new passive element segment with the given capacity.
+    pub(crate) fn new(
+        needs_gc_rooting: NeedsGcRooting,
+        capacity: usize,
+    ) -> Result<Self, OutOfMemory> {
+        Ok(Self {
+            needs_gc_rooting,
+            elements: TryVec::with_capacity(capacity)?,
+        })
+    }
+
+    /// Push a value onto this passive element segment.
+    ///
+    /// NB: Does not type check the value, relies on callers to ensure the value
+    /// is of the correct type (generally, due to validation).
+    pub(crate) fn push(&mut self, store: &mut StoreOpaque, val: Val) -> Result<()> {
+        let val = {
+            let mut store = AutoAssertNoGc::new(store);
+            val.to_raw_(&mut store)?
+        };
+        let val = self.clone_gc_ref(store, val);
+        self.elements.push(val)?;
+        Ok(())
+    }
+
+    fn clone_gc_ref(&mut self, store: &mut StoreOpaque, val: ValRaw) -> ValRaw {
+        if let NeedsGcRooting::Yes = self.needs_gc_rooting {
+            let gc_ref = val.get_anyref();
+            if let Some(gc_ref) = VMGcRef::from_raw_u32(gc_ref) {
+                if let Some(gc_store) = store.optional_gc_store_mut() {
+                    return ValRaw::anyref(gc_store.clone_gc_ref(&gc_ref).as_raw_u32());
+                }
+            }
+        }
+        val
+    }
+
+    /// Clear this segment's elements.
+    pub(crate) fn clear(&mut self, mut gc_store: Option<&mut GcStore>) {
+        for val in mem::take(&mut self.elements) {
+            self.drop_gc_ref(&mut gc_store, val);
+        }
+    }
+
+    fn drop_gc_ref(&mut self, gc_store: &mut Option<&mut GcStore>, val: ValRaw) {
+        if let NeedsGcRooting::Yes = self.needs_gc_rooting {
+            let gc_ref = val.get_anyref();
+            if let Some(gc_ref) = VMGcRef::from_raw_u32(gc_ref) {
+                if let Some(gc_store) = gc_store.as_deref_mut() {
+                    let _ = gc_store.drop_gc_ref(gc_ref);
+                }
+            }
+        }
+    }
+
+    /// Whether this segment needs GC rooting and tracing.
+    #[cfg(feature = "gc")]
+    pub(crate) fn needs_gc_rooting(&self) -> NeedsGcRooting {
+        self.needs_gc_rooting
+    }
+
+    /// The elements of this segment.
+    pub(crate) fn elements(&self) -> &[ValRaw] {
+        &self.elements
+    }
+}

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

@@ -6,7 +6,8 @@ use crate::runtime::vm::memory::Memory;
 use crate::runtime::vm::mpk::ProtectionKey;
 use crate::runtime::vm::table::Table;
 use crate::runtime::vm::{CompiledModuleId, ModuleRuntimeInfo};
-use crate::store::{Asyncness, AutoAssertNoGc, InstanceId, StoreOpaque, StoreResourceLimiter};
+use crate::store::{Asyncness, InstanceId, StoreOpaque, StoreResourceLimiter};
+use crate::vm::instance::PassiveElementSegment;
 use crate::{OpaqueRootScope, Val};
 use core::future::Future;
 use core::pin::Pin;
@@ -859,40 +860,34 @@ async fn initialize_passive_elements(
     for (idx, segment) in &module.passive_elements {
         match segment {
             TableSegmentElements::Functions(func_indices) => {
-                let mut vals = TryVec::with_capacity(func_indices.len())?;
+                let mut segment =
+                    PassiveElementSegment::new(NeedsGcRooting::No, func_indices.len())?;
                 for func_idx in func_indices {
                     let (instance, registry) =
                         store.instance_and_module_registry_mut(context.instance);
                     // SAFETY: `store_id` is for the store that owns this instance.
                     let func = unsafe { instance.get_exported_func(registry, store_id, *func_idx) };
-                    vals.push(func.to_val_raw(store))?;
+                    segment.push(store, func.into())?;
                 }
                 let instance = store.instance_mut(context.instance);
                 debug_assert_eq!(instance.passive_elements.len(), idx.index());
-                instance
-                    .passive_elements_mut()
-                    .push(Some((NeedsGcRooting::No, vals)))?;
+                instance.passive_elements_mut().push(segment)?;
             }
             TableSegmentElements::Expressions {
                 needs_gc_rooting,
                 exprs,
             } => {
-                let mut vals = TryVec::with_capacity(exprs.len())?;
+                let mut segment = PassiveElementSegment::new(*needs_gc_rooting, exprs.len())?;
                 for expr in exprs {
                     let mut store = OpaqueRootScope::new(&mut *store);
-
                     let val = const_evaluator
                         .eval(&mut store, limiter.as_deref_mut(), context, expr)
                         .await?;
-
-                    let mut store = AutoAssertNoGc::new(&mut store);
-                    vals.push(val.to_raw_(&mut store)?)?;
+                    segment.push(&mut store, *val)?;
                 }
                 let instance = store.instance_mut(context.instance);
                 debug_assert_eq!(instance.passive_elements.len(), idx.index());
-                instance
-                    .passive_elements_mut()
-                    .push(Some((*needs_gc_rooting, vals)))?;
+                instance.passive_elements_mut().push(segment)?;
             }
         }
     }

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

@@ -528,7 +528,8 @@ fn table_init(
 // Implementation of `elem.drop`.
 fn elem_drop(store: &mut dyn VMStore, instance: InstanceId, elem_index: u32) -> Result<()> {
     let elem_index = ElemIndex::from_u32(elem_index);
-    store.instance_mut(instance).elem_drop(elem_index)?;
+    let (gc_store, instance) = store.optional_gc_store_and_instance_mut(instance);
+    instance.elem_drop(gc_store, elem_index)?;
     Ok(())
 }
 

tests/all/gc.rs

@@ -1705,6 +1705,67 @@ fn select_gc_ref_stack_map() -> Result<()> {
     Ok(())
 }
 
+#[test]
+#[cfg_attr(miri, ignore)]
+fn dropped_passive_elems_do_not_leak() -> Result<()> {
+    let _ = env_logger::try_init();
+
+    let mut config = Config::new();
+    config.wasm_function_references(true);
+    config.wasm_gc(true);
+    config.collector(Collector::DeferredReferenceCounting);
+
+    let engine = Engine::new(&config)?;
+
+    let module = Module::new(
+        &engine,
+        r#"
+            (module
+                (type $a (array (mut externref)))
+                (type $aa (array (ref $a)))
+
+                (elem $e (ref $a) (array.new_default $a (i32.const 1)))
+
+                (func (export "run") (param $x externref)
+                    (local $l (ref null $aa))
+
+                    ;; Create an array from the passive element segment.
+                    (local.set $l (array.new_elem $aa $e (i32.const 0) (i32.const 1)))
+
+                    ;; Put our externref into the passive element segment's array:
+                    ;; `l[0][0] = x`.
+                    (array.set $a
+                               (array.get $aa (local.get $l) (i32.const 0))
+                               (i32.const 0)
+                               (local.get $x))
+
+                    ;; Drop the passive element segment. This should make `x`
+                    ;; unreachable upon return, since we aren't keeping our
+                    ;; array live nor the passive element segment.
+                    (elem.drop $e)
+                )
+            )
+        "#,
+    )?;
+
+    let mut store = Store::new(&engine, ());
+    let dropped = Arc::new(AtomicBool::new(false));
+
+    {
+        let mut store = RootScope::new(&mut store);
+        let e = ExternRef::new(&mut store, SetFlagOnDrop(dropped.clone()))?;
+
+        let instance = Instance::new(&mut store, &module, &[])?;
+        let run = instance.get_typed_func::<Rooted<ExternRef>, ()>(&mut store, "run")?;
+        run.call(&mut store, e)?;
+    }
+
+    store.gc(None)?;
+    assert!(dropped.load(SeqCst));
+
+    Ok(())
+}
+
 #[test]
 #[cfg_attr(miri, ignore)]
 fn gc_heap_does_not_grow_unboundedly() -> Result<()> {

tests/misc_testsuite/gc/issue-13066.wast

@@ -0,0 +1,31 @@
+;;! gc = true
+;;! bulk_memory = true
+
+(module
+  (import "wasmtime" "gc" (func $gc))
+  (type $s (struct (field i32)))
+  (table $t 10 anyref)
+
+  ;; Passive element segment with a GC struct
+  (elem $e anyref (struct.new $s (i32.const 42)))
+
+  ;; Copy passive element into table
+  (func (export "init")
+      (table.init $t $e (i32.const 0) (i32.const 0) (i32.const 1))
+  )
+
+  ;; Read the struct from the table and return its field value
+  (func (export "get_field") (result i32)
+      (struct.get $s 0
+          (ref.cast (ref $s)
+              (table.get $t (i32.const 0))
+          )
+      )
+  )
+
+  (export "gc" (func $gc))
+)
+
+(assert_return (invoke "gc"))
+(assert_return (invoke "init"))
+(assert_return (invoke "get_field") (i32.const 42))