Fix cyclic ref in WasmerEnv

CHANGELOG.md

@@ -35,6 +35,7 @@ Looking for changes that affect our C API? See the [C API Changelog](lib/c-api/C
 - [#2149](https://github.com/wasmerio/wasmer/pull/2144) `wasmer-engine-native` looks for clang-11 instead of clang-10.
 
 ### Fixed
+- [#2327](https://github.com/wasmerio/wasmer/pull/2327) Fix memory leak preventing internal instance memory from being freed when a WasmerEnv contained an exported extern (e.g. Memory, etc.).
 - [#2247](https://github.com/wasmerio/wasmer/pull/2247) Internal WasiFS logic updated to be closer to what WASI libc does when finding a preopened fd for a path.
 - [#2208](https://github.com/wasmerio/wasmer/pull/2208) Fix ownership in Wasm C API of `wasm_extern_as_func`, `wasm_extern_as_memory`, `wasm_extern_as_table`, `wasm_extern_as_global`, `wasm_func_as_extern`, `wasm_memory_as_extern`, `wasm_table_as_extern`, and `wasm_global_as_extern`. These functions no longer allocate memory and thus their results should not be freed. This is a breaking change to align more closely with the Wasm C API's stated ownership.
 - [#2210](https://github.com/wasmerio/wasmer/pull/2210) Fix a memory leak in the Wasm C API in the strings used to identify imports and exports coming from user code.

lib/api/src/env.rs

@@ -72,12 +72,16 @@ impl From<ExportError> for HostEnvInitError {
 ///
 /// impl WasmerEnv for MyEnv {
 ///     fn init_with_instance(&mut self, instance: &Instance) -> Result<(), HostEnvInitError> {
-///         let memory = instance.exports.get_memory("memory").unwrap();
+///         let memory: Memory = instance.exports.get_with_generics_weak("memory").unwrap();
 ///         self.memory.initialize(memory.clone());
 ///         Ok(())
 ///     }
 /// }
 /// ```
+///
+/// When implementing the trait manually, it's important to get a "weak" export to
+/// prevent a cyclic reference leaking memory. You can access a "weak" export with
+/// a method like `get_with_generics_weak`.
 pub trait WasmerEnv: Clone + Send + Sync {
     /// The function that Wasmer will call on your type to let it finish
     /// setting up the environment with data from the `Instance`.

lib/api/src/exports.rs

@@ -166,6 +166,19 @@ impl Exports {
         }
     }
 
+    /// Like `get_with_generics` but with a WeakReference to the `InstanceRef` internally.
+    /// This is useful for passing data into `WasmerEnv`, for example.
+    pub fn get_with_generics_weak<'a, T, Args, Rets>(&'a self, name: &str) -> Result<T, ExportError>
+    where
+        Args: WasmTypeList,
+        Rets: WasmTypeList,
+        T: ExportableWithGenerics<'a, Args, Rets>,
+    {
+        let mut out: T = self.get_with_generics(name)?;
+        out.into_weak_instance_ref();
+        Ok(out)
+    }
+
     /// Get an export as an `Extern`.
     pub fn get_extern(&self, name: &str) -> Option<&Extern> {
         self.map.get(name)
@@ -294,6 +307,11 @@ pub trait Exportable<'a>: Sized {
     ///
     /// [`Instance`]: crate::Instance
     fn get_self_from_extern(_extern: &'a Extern) -> Result<&'a Self, ExportError>;
+
+    /// Convert the extern internally to hold a weak reference to the `InstanceRef`.
+    /// This is useful for preventing cycles, for example for data stored in a
+    /// type implementing `WasmerEnv`.
+    fn into_weak_instance_ref(&mut self);
 }
 
 /// A trait for accessing exports (like [`Exportable`]) but it takes generic
@@ -302,6 +320,10 @@ pub trait Exportable<'a>: Sized {
 pub trait ExportableWithGenerics<'a, Args: WasmTypeList, Rets: WasmTypeList>: Sized {
     /// Get an export with the given generics.
     fn get_self_from_extern_with_generics(_extern: &'a Extern) -> Result<Self, ExportError>;
+    /// Convert the extern internally to hold a weak reference to the `InstanceRef`.
+    /// This is useful for preventing cycles, for example for data stored in a
+    /// type implementing `WasmerEnv`.
+    fn into_weak_instance_ref(&mut self);
 }
 
 /// We implement it for all concrete [`Exportable`] types (that are `Clone`)
@@ -310,4 +332,8 @@ impl<'a, T: Exportable<'a> + Clone + 'static> ExportableWithGenerics<'a, (), ()>
     fn get_self_from_extern_with_generics(_extern: &'a Extern) -> Result<Self, ExportError> {
         T::get_self_from_extern(_extern).map(|i| i.clone())
     }
+
+    fn into_weak_instance_ref(&mut self) {
+        <Self as Exportable>::into_weak_instance_ref(self);
+    }
 }

lib/api/src/externals/function.rs

@@ -39,7 +39,7 @@ use wasmer_vm::{
 ///   with native functions. Attempting to create a native `Function` with one will
 ///   result in a panic.
 ///   [Closures as host functions tracking issue](https://github.com/wasmerio/wasmer/issues/1840)
-#[derive(Clone, PartialEq, MemoryUsage)]
+#[derive(PartialEq, MemoryUsage)]
 pub struct Function {
     pub(crate) store: Store,
     pub(crate) exported: ExportFunction,
@@ -712,6 +712,18 @@ impl Function {
     fn closures_unsupported_panic() -> ! {
         unimplemented!("Closures (functions with captured environments) are currently unsupported with native functions. See: https://github.com/wasmerio/wasmer/issues/1840")
     }
+
+    /// Get access to the backing VM value for this extern. This function is for
+    /// tests it should not be called by users of the Wasmer API.
+    ///
+    /// # Safety
+    /// This function is unsafe to call outside of tests for the wasmer crate
+    /// because there is no stability guarantee for the returned type and we may
+    /// make breaking changes to it at any time or remove this method.
+    #[doc(hidden)]
+    pub unsafe fn get_vm_function(&self) -> &VMFunction {
+        &self.exported.vm_function
+    }
 }
 
 impl<'a> Exportable<'a> for Function {
@@ -725,6 +737,26 @@ impl<'a> Exportable<'a> for Function {
             _ => Err(ExportError::IncompatibleType),
         }
     }
+
+    fn into_weak_instance_ref(&mut self) {
+        self.exported
+            .vm_function
+            .instance_ref
+            .as_mut()
+            .map(|v| *v = v.downgrade());
+    }
+}
+
+impl Clone for Function {
+    fn clone(&self) -> Self {
+        let mut exported = self.exported.clone();
+        exported.vm_function.upgrade_instance_ref().unwrap();
+
+        Self {
+            store: self.store.clone(),
+            exported,
+        }
+    }
 }
 
 impl fmt::Debug for Function {

lib/api/src/externals/global.rs

@@ -17,7 +17,7 @@ use wasmer_vm::{Global as RuntimeGlobal, VMGlobal};
 /// It consists of an individual value and a flag indicating whether it is mutable.
 ///
 /// Spec: <https://webassembly.github.io/spec/core/exec/runtime.html#global-instances>
-#[derive(Clone, MemoryUsage)]
+#[derive(MemoryUsage)]
 pub struct Global {
     store: Store,
     vm_global: VMGlobal,
@@ -208,6 +208,30 @@ impl Global {
     pub fn same(&self, other: &Self) -> bool {
         Arc::ptr_eq(&self.vm_global.from, &other.vm_global.from)
     }
+
+    /// Get access to the backing VM value for this extern. This function is for
+    /// tests it should not be called by users of the Wasmer API.
+    ///
+    /// # Safety
+    /// This function is unsafe to call outside of tests for the wasmer crate
+    /// because there is no stability guarantee for the returned type and we may
+    /// make breaking changes to it at any time or remove this method.
+    #[doc(hidden)]
+    pub unsafe fn get_vm_global(&self) -> &VMGlobal {
+        &self.vm_global
+    }
+}
+
+impl Clone for Global {
+    fn clone(&self) -> Self {
+        let mut vm_global = self.vm_global.clone();
+        vm_global.upgrade_instance_ref().unwrap();
+
+        Self {
+            store: self.store.clone(),
+            vm_global,
+        }
+    }
 }
 
 impl fmt::Debug for Global {
@@ -231,4 +255,11 @@ impl<'a> Exportable<'a> for Global {
             _ => Err(ExportError::IncompatibleType),
         }
     }
+
+    fn into_weak_instance_ref(&mut self) {
+        self.vm_global
+            .instance_ref
+            .as_mut()
+            .map(|v| *v = v.downgrade());
+    }
 }

lib/api/src/externals/memory.rs

@@ -24,7 +24,7 @@ use wasmer_vm::{MemoryError, VMMemory};
 /// mutable from both host and WebAssembly.
 ///
 /// Spec: <https://webassembly.github.io/spec/core/exec/runtime.html#memory-instances>
-#[derive(Debug, Clone, MemoryUsage)]
+#[derive(Debug, MemoryUsage)]
 pub struct Memory {
     store: Store,
     vm_memory: VMMemory,
@@ -248,6 +248,30 @@ impl Memory {
     pub fn same(&self, other: &Self) -> bool {
         Arc::ptr_eq(&self.vm_memory.from, &other.vm_memory.from)
     }
+
+    /// Get access to the backing VM value for this extern. This function is for
+    /// tests it should not be called by users of the Wasmer API.
+    ///
+    /// # Safety
+    /// This function is unsafe to call outside of tests for the wasmer crate
+    /// because there is no stability guarantee for the returned type and we may
+    /// make breaking changes to it at any time or remove this method.
+    #[doc(hidden)]
+    pub unsafe fn get_vm_memory(&self) -> &VMMemory {
+        &self.vm_memory
+    }
+}
+
+impl Clone for Memory {
+    fn clone(&self) -> Self {
+        let mut vm_memory = self.vm_memory.clone();
+        vm_memory.upgrade_instance_ref().unwrap();
+
+        Self {
+            store: self.store.clone(),
+            vm_memory,
+        }
+    }
 }
 
 impl<'a> Exportable<'a> for Memory {
@@ -261,4 +285,11 @@ impl<'a> Exportable<'a> for Memory {
             _ => Err(ExportError::IncompatibleType),
         }
     }
+
+    fn into_weak_instance_ref(&mut self) {
+        self.vm_memory
+            .instance_ref
+            .as_mut()
+            .map(|v| *v = v.downgrade());
+    }
 }

lib/api/src/externals/mod.rs

@@ -72,6 +72,15 @@ impl<'a> Exportable<'a> for Extern {
         // Since this is already an extern, we can just return it.
         Ok(_extern)
     }
+
+    fn into_weak_instance_ref(&mut self) {
+        match self {
+            Self::Function(f) => f.into_weak_instance_ref(),
+            Self::Global(g) => g.into_weak_instance_ref(),
+            Self::Memory(m) => m.into_weak_instance_ref(),
+            Self::Table(t) => t.into_weak_instance_ref(),
+        }
+    }
 }
 
 impl StoreObject for Extern {

lib/api/src/externals/table.rs

@@ -18,7 +18,7 @@ use wasmer_vm::{Table as RuntimeTable, TableElement, VMTable};
 /// mutable from both host and WebAssembly.
 ///
 /// Spec: <https://webassembly.github.io/spec/core/exec/runtime.html#table-instances>
-#[derive(Clone, MemoryUsage)]
+#[derive(MemoryUsage)]
 pub struct Table {
     store: Store,
     vm_table: VMTable,
@@ -146,6 +146,30 @@ impl Table {
     pub fn same(&self, other: &Self) -> bool {
         Arc::ptr_eq(&self.vm_table.from, &other.vm_table.from)
     }
+
+    /// Get access to the backing VM value for this extern. This function is for
+    /// tests it should not be called by users of the Wasmer API.
+    ///
+    /// # Safety
+    /// This function is unsafe to call outside of tests for the wasmer crate
+    /// because there is no stability guarantee for the returned type and we may
+    /// make breaking changes to it at any time or remove this method.
+    #[doc(hidden)]
+    pub unsafe fn get_vm_table(&self) -> &VMTable {
+        &self.vm_table
+    }
+}
+
+impl Clone for Table {
+    fn clone(&self) -> Self {
+        let mut vm_table = self.vm_table.clone();
+        vm_table.upgrade_instance_ref().unwrap();
+
+        Self {
+            store: self.store.clone(),
+            vm_table,
+        }
+    }
 }
 
 impl<'a> Exportable<'a> for Table {
@@ -159,4 +183,11 @@ impl<'a> Exportable<'a> for Table {
             _ => Err(ExportError::IncompatibleType),
         }
     }
+
+    fn into_weak_instance_ref(&mut self) {
+        self.vm_table
+            .instance_ref
+            .as_mut()
+            .map(|v| *v = v.downgrade());
+    }
 }

lib/api/src/native.rs

@@ -18,7 +18,6 @@ use wasmer_vm::{VMDynamicFunctionContext, VMFunctionBody, VMFunctionEnvironment,
 
 /// A WebAssembly function that can be called natively
 /// (using the Native ABI).
-#[derive(Clone)]
 pub struct NativeFunc<Args = (), Rets = ()> {
     store: Store,
     exported: ExportFunction,
@@ -55,6 +54,18 @@ where
     pub(crate) fn arg_kind(&self) -> VMFunctionKind {
         self.exported.vm_function.kind
     }
+
+    /// Get access to the backing VM value for this extern. This function is for
+    /// tests it should not be called by users of the Wasmer API.
+    ///
+    /// # Safety
+    /// This function is unsafe to call outside of tests for the wasmer crate
+    /// because there is no stability guarantee for the returned type and we may
+    /// make breaking changes to it at any time or remove this method.
+    #[doc(hidden)]
+    pub unsafe fn get_vm_function(&self) -> &wasmer_vm::VMFunction {
+        &self.exported.vm_function
+    }
 }
 
 /*
@@ -76,6 +87,19 @@ where
     }
 }*/
 
+impl<Args: WasmTypeList, Rets: WasmTypeList> Clone for NativeFunc<Args, Rets> {
+    fn clone(&self) -> Self {
+        let mut exported = self.exported.clone();
+        exported.vm_function.upgrade_instance_ref().unwrap();
+
+        Self {
+            store: self.store.clone(),
+            exported,
+            _phantom: PhantomData,
+        }
+    }
+}
+
 impl<Args, Rets> From<&NativeFunc<Args, Rets>> for ExportFunction
 where
     Args: WasmTypeList,
@@ -194,8 +218,8 @@ macro_rules! impl_native_traits {
                         }
                     }
                 }
-
             }
+
         }
 
         #[allow(unused_parens)]
@@ -208,6 +232,10 @@ macro_rules! impl_native_traits {
                 use crate::exports::Exportable;
                 crate::Function::get_self_from_extern(_extern)?.native().map_err(|_| crate::exports::ExportError::IncompatibleType)
             }
+
+            fn into_weak_instance_ref(&mut self) {
+                self.exported.vm_function.instance_ref.as_mut().map(|v| *v = v.downgrade());
+            }
         }
     };
 }