Use after free in Wasmtime

@@ -17,6 +17,7 @@
 //! See the main module comment in `mod.rs` for more details on the VCode-based
 //! backend pipeline.
 
+use crate::fx::FxHashMap;
 use crate::ir::{self, types, Constant, ConstantData, SourceLoc};
 use crate::machinst::*;
 use crate::settings;
@@ -478,6 +479,19 @@ impl<I: VCodeInst> VCode<I> {
         let mut inst_end_offsets = vec![0; self.insts.len()];
         let mut label_inst_indices = vec![0; self.num_blocks()];
 
+        // Map from instruction index to index in
+        // `safepoint_slots`. We need this because we emit
+        // instructions out-of-order, while the safepoint_insns /
+        // safepoint_slots data structures are sorted in instruction
+        // order.
+        let mut safepoint_indices: FxHashMap<u32, usize> = FxHashMap::default();
+        for (safepoint_idx, iix) in self.safepoint_insns.iter().enumerate() {
+            // Disregard safepoints that ended up having no live refs.
+            if self.safepoint_slots[safepoint_idx].len() > 0 {
+                safepoint_indices.insert(*iix, safepoint_idx);
+            }
+        }
+
         // Construct the final order we emit code in: cold blocks at the end.
         let mut final_order: SmallVec<[BlockIndex; 16]> = smallvec![];
         let mut cold_blocks: SmallVec<[BlockIndex; 16]> = smallvec![];
@@ -493,7 +507,6 @@ impl<I: VCodeInst> VCode<I> {
         final_order.extend(cold_blocks.clone());
 
         // Emit blocks.
-        let mut safepoint_idx = 0;
         let mut cur_srcloc = None;
         let mut last_offset = None;
         let mut start_of_cold_code = None;
@@ -541,17 +554,11 @@ impl<I: VCodeInst> VCode<I> {
                 }
                 state.pre_sourceloc(cur_srcloc.unwrap_or(SourceLoc::default()));
 
-                if safepoint_idx < self.safepoint_insns.len()
-                    && self.safepoint_insns[safepoint_idx] == iix
-                {
-                    if self.safepoint_slots[safepoint_idx].len() > 0 {
-                        let stack_map = self.abi.spillslots_to_stack_map(
-                            &self.safepoint_slots[safepoint_idx][..],
-                            &state,
-                        );
-                        state.pre_safepoint(stack_map);
-                    }
-                    safepoint_idx += 1;
+                if let Some(safepoint_idx) = safepoint_indices.get(&iix) {
+                    let stack_map = self
+                        .abi
+                        .spillslots_to_stack_map(&self.safepoint_slots[*safepoint_idx][..], &state);
+                    state.pre_safepoint(stack_map);
                 }
 
                 self.insts[iix as usize].emit(&mut buffer, &self.emit_info, &mut state);

@@ -6,6 +6,7 @@ use wasmtime::*;
 #[test]
 fn pass_funcref_in_and_out_of_wasm() -> anyhow::Result<()> {
     let (mut store, module) = ref_types_module(
+        false,
         r#"
             (module
                 (func (export "func") (param funcref) (result funcref)
@@ -60,7 +61,8 @@ fn pass_funcref_in_and_out_of_wasm() -> anyhow::Result<()> {
 
     // Passing in a `funcref` from another store fails.
     {
-        let (mut other_store, other_module) = ref_types_module(r#"(module (func (export "f")))"#)?;
+        let (mut other_store, other_module) =
+            ref_types_module(false, r#"(module (func (export "f")))"#)?;
         let other_store_instance = Instance::new(&mut other_store, &other_module, &[])?;
         let f = other_store_instance
             .get_func(&mut other_store, "f")
@@ -77,6 +79,7 @@ fn pass_funcref_in_and_out_of_wasm() -> anyhow::Result<()> {
 #[test]
 fn receive_null_funcref_from_wasm() -> anyhow::Result<()> {
     let (mut store, module) = ref_types_module(
+        false,
         r#"
             (module
                 (func (export "get-null") (result funcref)

@@ -14,7 +14,17 @@ impl Drop for SetFlagOnDrop {
 
 #[test]
 fn smoke_test_gc() -> anyhow::Result<()> {
+    smoke_test_gc_impl(false)
+}
+
+#[test]
+fn smoke_test_gc_epochs() -> anyhow::Result<()> {
+    smoke_test_gc_impl(true)
+}
+
+fn smoke_test_gc_impl(use_epochs: bool) -> anyhow::Result<()> {
     let (mut store, module) = ref_types_module(
+        use_epochs,
         r#"
             (module
                 (import "" "" (func $do_gc))
@@ -69,6 +79,7 @@ fn smoke_test_gc() -> anyhow::Result<()> {
 #[test]
 fn wasm_dropping_refs() -> anyhow::Result<()> {
     let (mut store, module) = ref_types_module(
+        false,
         r#"
             (module
                 (func (export "drop_ref") (param externref)
@@ -145,7 +156,7 @@ fn many_live_refs() -> anyhow::Result<()> {
         ",
     );
 
-    let (mut store, module) = ref_types_module(&wat)?;
+    let (mut store, module) = ref_types_module(false, &wat)?;
 
     let live_refs = Arc::new(AtomicUsize::new(0));
 
@@ -191,6 +202,7 @@ fn many_live_refs() -> anyhow::Result<()> {
 #[test]
 fn drop_externref_via_table_set() -> anyhow::Result<()> {
     let (mut store, module) = ref_types_module(
+        false,
         r#"
             (module
                 (table $t 1 externref)
@@ -400,6 +412,7 @@ fn gee_i_sure_hope_refcounting_is_atomic() -> anyhow::Result<()> {
 #[test]
 fn global_init_no_leak() -> anyhow::Result<()> {
     let (mut store, module) = ref_types_module(
+        false,
         r#"
             (module
                 (import "" "" (global externref))
@@ -424,6 +437,7 @@ fn global_init_no_leak() -> anyhow::Result<()> {
 #[test]
 fn no_gc_middle_of_args() -> anyhow::Result<()> {
     let (mut store, module) = ref_types_module(
+        false,
         r#"
             (module
                 (import "" "return_some" (func $return (result externref externref externref)))

@@ -33,6 +33,7 @@ mod wast;
 
 /// A helper to compile a module in a new store with reference types enabled.
 pub(crate) fn ref_types_module(
+    use_epochs: bool,
     source: &str,
 ) -> anyhow::Result<(wasmtime::Store<()>, wasmtime::Module)> {
     use wasmtime::*;
@@ -41,9 +42,15 @@ pub(crate) fn ref_types_module(
 
     let mut config = Config::new();
     config.wasm_reference_types(true);
+    if use_epochs {
+        config.epoch_interruption(true);
+    }
 
     let engine = Engine::new(&config)?;
-    let store = Store::new(&engine, ());
+    let mut store = Store::new(&engine, ());
+    if use_epochs {
+        store.set_epoch_deadline(1);
+    }
 
     let module = Module::new(&engine, source)?;