Fix a use-after-free of trampoline code

This commit fixes an issue with wasmtime where it was possible for a
trampoline from one module to get used for another module after it was
freed. This issue arises because we register a module's native
trampolines *before* it's fully instantiated, which is a fallible
process. Some fallibility is predictable, such as import type
mismatches, but other fallibility is less predictable, such as failure
to allocate a linear memory.

The problem happened when a module was registered with a `Store`,
retaining information about its trampolines, but then instantiation
failed and the module's code was never persisted within the `Store`.
Unlike as documented in #2374 the `Module` inside an `Instance` is not
the primary way to hold on to a module's code, but rather the
`Arc<ModuleCode>` is persisted within the global frame information off
on the side. This persistence only made its way into the store through
the `Box<Any>` field of `InstanceHandle`, but that's never made if
instantiation fails during import matching.

The fix here is to build on the refactoring of #2407 to not store module
code in frame information but rather explicitly in the `Store`.
Registration is now deferred until just-before an instance handle is
created, and during module registration we insert the `Arc<ModuleCode>`
into a set stored within the `Store`.

Author: alexcrichton

crates/wasmtime/src/instance.rs

@@ -16,6 +16,11 @@ fn instantiate(
     imports: Imports<'_>,
     host: Box<dyn Any>,
 ) -> Result<StoreInstanceHandle, Error> {
+    // Register the module just before instantiation to ensure we have a
+    // trampoline registered for every signature and to preserve the module's
+    // compiled JIT code within the `Store`.
+    store.register_module(compiled_module);
+
     let config = store.engine().config();
     let instance = unsafe {
         let instance = compiled_module.instantiate(
@@ -98,11 +103,6 @@ fn instantiate(
 #[derive(Clone)]
 pub struct Instance {
     pub(crate) handle: StoreInstanceHandle,
-    // Note that this is required to keep the module's code memory alive while
-    // we have a handle to this `Instance`. We may eventually want to shrink
-    // this to only hold onto the bare minimum each instance needs to allow
-    // deallocating some `Module` resources early, but until then we just hold
-    // on to everything.
     module: Module,
 }
 
@@ -165,7 +165,6 @@ impl Instance {
             bail!("cross-`Engine` instantiation is not currently supported");
         }
 
-        store.register_module(module.compiled_module());
         let handle = with_imports(store, module.compiled_module(), imports, |imports| {
             instantiate(store, module.compiled_module(), imports, Box::new(()))
         })?;

crates/wasmtime/src/store.rs

@@ -5,6 +5,7 @@ use crate::Engine;
 use anyhow::{bail, Result};
 use std::any::Any;
 use std::cell::RefCell;
+use std::collections::HashSet;
 use std::fmt;
 use std::hash::{Hash, Hasher};
 use std::rc::{Rc, Weak};
@@ -63,9 +64,9 @@ pub(crate) struct StoreInner {
     /// Information about JIT code which allows us to test if a program counter
     /// is in JIT code, lookup trap information, etc.
     frame_info: RefCell<StoreFrameInfo>,
-    /// List of all compiled modules that we're holding a strong reference to
+    /// Set of all compiled modules that we're holding a strong reference to
     /// the module's code for. This includes JIT functions, trampolines, etc.
-    modules: RefCell<Vec<Arc<ModuleCode>>>,
+    modules: RefCell<HashSet<ArcModuleCode>>,
 }
 
 struct HostInfoKey(VMExternRef);
@@ -169,6 +170,15 @@ impl Store {
         // a `Func` wrapper for any function in the module, which requires that
         // we know about the signature and trampoline for all instances.
         self.register_signatures(module);
+
+        // And finally with a module being instantiated into this `Store` we
+        // need to preserve its jit-code. References to this module's code and
+        // trampolines are not owning-references so it's our responsibility to
+        // keep it all alive within the `Store`.
+        self.inner
+            .modules
+            .borrow_mut()
+            .insert(ArcModuleCode(module.code().clone()));
     }
 
     fn register_jit_code(&self, module: &CompiledModule) {
@@ -180,7 +190,6 @@ impl Store {
         // Only register this module if it hasn't already been registered.
         if !self.is_wasm_code(first_pc) {
             self.inner.frame_info.borrow_mut().register(module);
-            self.inner.modules.borrow_mut().push(module.code().clone());
         }
     }
 
@@ -435,3 +444,21 @@ impl InterruptHandle {
         self.interrupts.interrupt()
     }
 }
+
+// Wrapper struct to implement hash/equality based on the pointer value of the
+// `Arc` in question.
+struct ArcModuleCode(Arc<ModuleCode>);
+
+impl PartialEq for ArcModuleCode {
+    fn eq(&self, other: &ArcModuleCode) -> bool {
+        Arc::ptr_eq(&self.0, &other.0)
+    }
+}
+
+impl Eq for ArcModuleCode {}
+
+impl Hash for ArcModuleCode {
+    fn hash<H: Hasher>(&self, hasher: &mut H) {
+        Arc::as_ptr(&self.0).hash(hasher)
+    }
+}

tests/all/func.rs

@@ -522,3 +522,30 @@ fn externref_signature_no_reference_types() -> anyhow::Result<()> {
     );
     Ok(())
 }
+
+#[test]
+fn trampolines_always_valid() -> anyhow::Result<()> {
+    let func = {
+        // Compile two modules up front
+        let store = Store::default();
+        let module1 = Module::new(store.engine(), "(module (import \"\" \"\" (func)))")?;
+        let module2 = Module::new(store.engine(), "(module (func (export \"\")))")?;
+        // Start instantiating the first module, but this will fail.
+        // Historically this registered the module's trampolines with `Store`
+        // before the failure, but then after the failure the `Store` didn't
+        // hold onto the trampoline.
+        drop(Instance::new(&store, &module1, &[]));
+        drop(module1);
+
+        // Then instantiate another module which has the same function type (no
+        // parameters or results) which tries to use the trampoline defined in
+        // the previous module. Then we extract the function and, in another
+        // scope where everything is dropped, we call the func.
+        let i = Instance::new(&store, &module2, &[])?;
+        i.get_func("").unwrap()
+    };
+
+    // ... and no segfaults! right? right? ...
+    func.call(&[])?;
+    Ok(())
+}