Fix a memory leak with command modules

crates/wasmtime/src/externals.rs

@@ -160,6 +160,17 @@ impl Extern {
  Extern::Module(_) => "module",
  }
  }
+
+ pub(crate) fn wasmtime_export(&self) -> wasmtime_runtime::Export {
+ match self {
+ Extern::Func(f) => f.wasmtime_export().clone().into(),
+ Extern::Global(f) => f.wasmtime_export().clone().into(),
+ Extern::Table(f) => f.wasmtime_export().clone().into(),
+ Extern::Memory(f) => f.wasmtime_export().clone().into(),
+ Extern::Instance(f) => wasmtime_runtime::Export::Instance(f.wasmtime_export().clone()),
+ Extern::Module(f) => wasmtime_runtime::Export::Module(Box::new(f.clone())),
+ }
+ }
 }
 
 impl From<Func> for Extern {

crates/wasmtime/src/instance.rs

@@ -121,6 +121,10 @@ impl Instance {
  ty
  }
 
+ pub(crate) fn wasmtime_export(&self) -> &RuntimeInstance {
+ &self.items
+ }
+
  /// Returns the associated [`Store`] that this `Instance` is compiled into.
  ///
  /// This is the [`Store`] that generally serves as a sort of global cache

crates/wasmtime/src/linker.rs

@@ -402,34 +402,53 @@ impl Linker {
  fn command(&mut self, module_name: &str, module: &Module) -> Result<&mut Self> {
  for export in module.exports() {
  if let Some(func_ty) = export.ty().func() {
- let imports = self.compute_imports(module)?;
- let store = self.store.clone();
+ let imports = self
+ .compute_imports(module)?
+ .into_iter()
+ .map(|e| e.wasmtime_export())
+ .collect::<Vec<_>>();
  let module = module.clone();
  let export_name = export.name().to_owned();
- let func = Func::new(&self.store, func_ty.clone(), move |_, params, results| {
- // Create a new instance for this command execution.
- let instance = Instance::new(&store, &module, &imports)?;
-
- // `unwrap()` everything here because we know the instance contains a
- // function export with the given name and signature because we're
- // iterating over the module it was instantiated from.
- let command_results = instance
- .get_export(&export_name)
- .unwrap()
- .into_func()
- .unwrap()
- .call(params)
- .map_err(|error| error.downcast::<Trap>().unwrap())?;
-
- // Copy the return values into the output slice.
- for (result, command_result) in
- results.iter_mut().zip(command_results.into_vec())
- {
- *result = command_result;
- }
-
- Ok(())
- });
+ let func = Func::new(
+ &self.store,
+ func_ty.clone(),
+ move |caller, params, results| {
+ let store = caller.store();
+
+ // Note that the unsafety here is due to the validity of
+ // `i` and the validity of `i` within `store`. For our
+ // case though these items all come from `imports` above
+ // so they're all valid. They're also all kept alive by
+ // the store itself used here so this should be safe.
+ let imports = imports
+ .iter()
+ .map(|i| unsafe { Extern::from_wasmtime_export(&i, &store) })
+ .collect::<Vec<_>>();
+
+ // Create a new instance for this command execution.
+ let instance = Instance::new(&store, &module, &imports)?;
+
+ // `unwrap()` everything here because we know the instance contains a
+ // function export with the given name and signature because we're
+ // iterating over the module it was instantiated from.
+ let command_results = instance
+ .get_export(&export_name)
+ .unwrap()
+ .into_func()
+ .unwrap()
+ .call(params)
+ .map_err(|error| error.downcast::<Trap>().unwrap())?;
+
+ // Copy the return values into the output slice.
+ for (result, command_result) in
+ results.iter_mut().zip(command_results.into_vec())
+ {
+ *result = command_result;
+ }
+
+ Ok(())
+ },
+ );
  self.insert(module_name, export.name(), func.into())?;
  } else if export.name() == "memory" && export.ty().memory().is_some() {
  // Allow an exported "memory" memory for now.

tests/all/linker.rs

@@ -1,4 +1,6 @@
 use anyhow::Result;
+use std::cell::Cell;
+use std::rc::Rc;
 use wasmtime::*;
 
 #[test]
@@ -160,3 +162,64 @@ fn module_interposition() -> Result<()> {
  assert_eq!(func()?, 112);
  Ok(())
 }
+
+#[test]
+fn no_leak() -> Result<()> {
+ struct DropMe(Rc<Cell<bool>>);
+
+ impl Drop for DropMe {
+ fn drop(&mut self) {
+ self.0.set(true);
+ }
+ }
+
+ let flag = Rc::new(Cell::new(false));
+ {
+ let store = Store::default();
+ let mut linker = Linker::new(&store);
+ let drop_me = DropMe(flag.clone());
+ linker.func("", "", move || drop(&drop_me))?;
+ let module = Module::new(
+ store.engine(),
+ r#"
+ (module
+ (func (export "_start"))
+ )
+ "#,
+ )?;
+ linker.module("a", &module)?;
+ }
+ assert!(flag.get(), "store was leaked");
+ Ok(())
+}
+
+#[test]
+fn no_leak_with_imports() -> Result<()> {
+ struct DropMe(Rc<Cell<bool>>);
+
+ impl Drop for DropMe {
+ fn drop(&mut self) {
+ self.0.set(true);
+ }
+ }
+
+ let flag = Rc::new(Cell::new(false));
+ {
+ let store = Store::default();
+ let mut linker = Linker::new(&store);
+ let drop_me = DropMe(flag.clone());
+ linker.func("", "", move || drop(&drop_me))?;
+ let module = Module::new(
+ store.engine(),
+ r#"
+ (module
+ (import "" "" (func))
+ (func (export "_start"))
+ )
+ "#,
+ )?;
+ linker.module("a", &module)?;
+ }
+ assert!(flag.get(), "store was leaked");
+ Ok(())
+}