diff --git a/crates/wasmtime/src/externals.rs b/crates/wasmtime/src/externals.rs index 00bcaa499b7a..e71fb7d75c4b 100644 --- a/crates/wasmtime/src/externals.rs +++ b/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 for Extern { diff --git a/crates/wasmtime/src/instance.rs b/crates/wasmtime/src/instance.rs index 47115d90ae90..8bc6085c11cc 100644 --- a/crates/wasmtime/src/instance.rs +++ b/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 diff --git a/crates/wasmtime/src/linker.rs b/crates/wasmtime/src/linker.rs index 396327b5d57f..5875cf8891f8 100644 --- a/crates/wasmtime/src/linker.rs +++ b/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::>(); 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::().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::>(); + + // 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::().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. diff --git a/tests/all/linker.rs b/tests/all/linker.rs index 675c109cf18c..5312af31f524 100644 --- a/tests/all/linker.rs +++ b/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>); + + 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>); + + 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(()) +}