fix(wasix): Linker must retain the Engine

Author: theduke

examples/wasi.rs

@@ -43,8 +43,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     println!("Compiling module...");
     // Let's compile the Wasm module.
-    let data =
-        wasmer_wasix::runtime::module_cache::HashedModuleData::new_sha256(wasm_bytes.clone());
+    let data = wasmer_wasix::runtime::module_cache::HashedModuleData::new(wasm_bytes.clone());
     let module = runtime.load_hashed_module_sync(data)?;
 
     // Create a pipe for the module's stdout.

lib/cli/src/commands/run/mod.rs

@@ -50,7 +50,9 @@ use wasmer_wasix::{
         MappedCommand, MappedDirectory, Runner,
     },
     runtime::{
-        module_cache::CacheError, package_loader::PackageLoader, resolver::QueryError,
+        module_cache::{CacheError, HashedModuleData},
+        package_loader::PackageLoader,
+        resolver::QueryError,
         task_manager::VirtualTaskManagerExt,
     },
     Runtime, WasiError,
@@ -843,15 +845,17 @@ impl ExecutableTarget {
         match TargetOnDisk::from_file(path)? {
             TargetOnDisk::WebAssemblyBinary | TargetOnDisk::Wat => {
                 let wasm = std::fs::read(path)?;
+                let module_data = HashedModuleData::new(wasm);
+                let module_hash = module_data.hash().clone();
 
                 pb.set_message("Compiling to WebAssembly");
                 let module = runtime
-                    .load_module_sync(&wasm)
+                    .load_hashed_module_sync(module_data, None)
                     .with_context(|| format!("Unable to compile \"{}\"", path.display()))?;
 
                 Ok(ExecutableTarget::WebAssembly {
                     module,
-                    module_hash: ModuleHash::xxhash(&wasm),
+                    module_hash,
                     path: path.to_path_buf(),
                 })
             }

lib/wasix/src/bin_factory/exec.rs

@@ -89,7 +89,7 @@ pub async fn spawn_load_module(
     wasm: HashedModuleData,
     runtime: &Arc<dyn Runtime + Send + Sync + 'static>,
 ) -> Result<Module, SpawnError> {
-    match runtime.load_hashed_module(wasm).await {
+    match runtime.load_hashed_module(wasm, None).await {
         Ok(module) => Ok(module),
         Err(err) => {
             tracing::error!(

lib/wasix/src/bin_factory/mod.rs

@@ -86,7 +86,7 @@ impl BinFactory {
             // Execute
             match executable {
                 Executable::Wasm(bytes) => {
-                    let data = HashedModuleData::new_xxhash(bytes.clone());
+                    let data = HashedModuleData::new(bytes.clone());
                     spawn_exec_wasm(data, name.as_str(), env, &self.runtime).await
                 }
                 Executable::BinaryPackage(pkg) => {

lib/wasix/src/os/command/builtins/cmd_wasmer.rs

@@ -139,7 +139,7 @@ impl CmdWasmer {
                     spawn_exec(binary, name, env, &self.runtime).await
                 }
                 Executable::Wasm(bytes) => {
-                    let data = HashedModuleData::new_xxhash(bytes);
+                    let data = HashedModuleData::new(bytes);
                     spawn_exec_wasm(data, name, env, &self.runtime).await
                 }
             }

lib/wasix/src/runtime/mod.rs

@@ -17,7 +17,7 @@ use std::{
 
 use futures::future::BoxFuture;
 use virtual_net::{DynVirtualNetworking, VirtualNetworking};
-use wasmer::{CompileError, Module, RuntimeError};
+use wasmer::{CompileError, Engine, Module, RuntimeError};
 use wasmer_wasix_types::wasi::ExitCode;
 
 #[cfg(feature = "journal")]
@@ -76,14 +76,14 @@ where
     fn source(&self) -> Arc<dyn Source + Send + Sync>;
 
     /// Get a [`wasmer::Engine`] for module compilation.
-    fn engine(&self) -> wasmer::Engine {
-        wasmer::Engine::default()
+    fn engine(&self) -> Engine {
+        Engine::default()
     }
 
     fn engine_with_suggested_opts(
         &self,
         suggested_opts: &SuggestedCompilerOptimizations,
-    ) -> Result<wasmer::Engine, CompileError> {
+    ) -> Result<Engine, CompileError> {
         let mut engine = self.engine();
         engine.with_opts(&WasmerSuggestedCompilerOptimizations {
             pass_params: suggested_opts.pass_params,
@@ -157,7 +157,7 @@ where
     fn load_module<'a>(&'a self, wasm: &'a [u8]) -> BoxFuture<'a, Result<Module, SpawnError>> {
         let engine = self.engine();
         let module_cache = self.module_cache();
-        let data = HashedModuleData::new_xxhash(wasm.to_vec());
+        let data = HashedModuleData::new(wasm.to_vec());
         let task = async move { load_module(&engine, &module_cache, &data).await };
         Box::pin(task)
     }
@@ -178,16 +178,21 @@ where
     fn load_hashed_module(
         &self,
         module: HashedModuleData,
+        engine: Option<&Engine>,
     ) -> BoxFuture<'_, Result<Module, SpawnError>> {
-        let engine = self.engine();
+        let engine = engine.cloned().unwrap_or_else(|| self.engine());
         let module_cache = self.module_cache();
         let task = async move { load_module(&engine, &module_cache, &module).await };
         Box::pin(task)
     }
 
     /// Synchronous version of [`Self::load_hashed_module`].
-    fn load_hashed_module_sync(&self, wasm: HashedModuleData) -> Result<Module, SpawnError> {
-        InlineWaker::block_on(self.load_hashed_module(wasm))
+    fn load_hashed_module_sync(
+        &self,
+        wasm: HashedModuleData,
+        engine: Option<&Engine>,
+    ) -> Result<Module, SpawnError> {
+        InlineWaker::block_on(self.load_hashed_module(wasm, engine))
     }
 
     /// Callback thats invokes whenever the instance is tainted, tainting can occur
@@ -223,7 +228,7 @@ pub type DynRuntime = dyn Runtime + Send + Sync;
 // implementing [`Runtime::load_module`], so custom logic can be added on top.
 #[tracing::instrument(level = "debug", skip_all)]
 pub async fn load_module(
-    engine: &wasmer::Engine,
+    engine: &Engine,
     module_cache: &(dyn ModuleCache + Send + Sync),
     module: &HashedModuleData,
 ) -> Result<Module, crate::SpawnError> {
@@ -291,7 +296,7 @@ pub struct PluggableRuntime {
     pub http_client: Option<DynHttpClient>,
     pub package_loader: Arc<dyn PackageLoader + Send + Sync>,
     pub source: Arc<dyn Source + Send + Sync>,
-    pub engine: wasmer::Engine,
+    pub engine: Engine,
     pub module_cache: Arc<dyn ModuleCache + Send + Sync>,
     pub tty: Option<Arc<dyn TtyBridge + Send + Sync>>,
     #[cfg(feature = "journal")]
@@ -347,7 +352,7 @@ impl PluggableRuntime {
         self
     }
 
-    pub fn set_engine(&mut self, engine: wasmer::Engine) -> &mut Self {
+    pub fn set_engine(&mut self, engine: Engine) -> &mut Self {
         self.engine = engine;
         self
     }
@@ -416,7 +421,7 @@ impl Runtime for PluggableRuntime {
         Arc::clone(&self.source)
     }
 
-    fn engine(&self) -> wasmer::Engine {
+    fn engine(&self) -> Engine {
         self.engine.clone()
     }
 
@@ -462,7 +467,7 @@ pub struct OverriddenRuntime {
     http_client: Option<DynHttpClient>,
     package_loader: Option<Arc<dyn PackageLoader + Send + Sync>>,
     source: Option<Arc<dyn Source + Send + Sync>>,
-    engine: Option<wasmer::Engine>,
+    engine: Option<Engine>,
     module_cache: Option<Arc<dyn ModuleCache + Send + Sync>>,
     tty: Option<Arc<dyn TtyBridge + Send + Sync>>,
     #[cfg(feature = "journal")]
@@ -518,7 +523,7 @@ impl OverriddenRuntime {
         self
     }
 
-    pub fn with_engine(mut self, engine: wasmer::Engine) -> Self {
+    pub fn with_engine(mut self, engine: Engine) -> Self {
         self.engine.replace(engine);
         self
     }
@@ -587,7 +592,7 @@ impl Runtime for OverriddenRuntime {
         }
     }
 
-    fn engine(&self) -> wasmer::Engine {
+    fn engine(&self) -> Engine {
         if let Some(engine) = self.engine.clone() {
             engine
         } else {
@@ -646,17 +651,16 @@ impl Runtime for OverriddenRuntime {
         }
     }
 
-    fn load_module<'a>(&'a self, wasm: &'a [u8]) -> BoxFuture<'a, Result<Module, SpawnError>> {
-        if self.engine.is_some() || self.module_cache.is_some() {
-            let engine = self.engine();
-            let module_cache = self.module_cache();
-
-            let data = HashedModuleData::new_xxhash(wasm.to_vec());
-            let task = async move { load_module(&engine, &module_cache, &data).await };
+    fn load_module<'a>(&'a self, data: &[u8]) -> BoxFuture<'a, Result<Module, SpawnError>> {
+        #[allow(deprecated)]
+        if let (Some(engine), Some(module_cache)) = (&self.engine, self.module_cache.clone()) {
+            let engine = engine.clone();
+            let hashed_data = HashedModuleData::new(data.to_vec());
+            let task = async move { load_module(&engine, &module_cache, &hashed_data).await };
             Box::pin(task)
         } else {
-            #[allow(deprecated)]
-            self.inner.load_module(wasm)
+            let data = data.to_vec();
+            Box::pin(async move { self.inner.load_module(&data).await })
         }
     }
 
@@ -672,22 +676,27 @@ impl Runtime for OverriddenRuntime {
     fn load_hashed_module(
         &self,
         data: HashedModuleData,
+        engine: Option<&Engine>,
     ) -> BoxFuture<'_, Result<Module, SpawnError>> {
         if self.engine.is_some() || self.module_cache.is_some() {
             let engine = self.engine();
             let module_cache = self.module_cache();
             let task = async move { load_module(&engine, &module_cache, &data).await };
             Box::pin(task)
         } else {
-            self.inner.load_hashed_module(data)
+            self.inner.load_hashed_module(data, engine)
         }
     }
 
-    fn load_hashed_module_sync(&self, wasm: HashedModuleData) -> Result<Module, SpawnError> {
+    fn load_hashed_module_sync(
+        &self,
+        wasm: HashedModuleData,
+        engine: Option<&Engine>,
+    ) -> Result<Module, SpawnError> {
         if self.engine.is_some() || self.module_cache.is_some() {
-            InlineWaker::block_on(self.load_hashed_module(wasm))
+            InlineWaker::block_on(self.load_hashed_module(wasm, engine))
         } else {
-            self.inner.load_hashed_module_sync(wasm)
+            self.inner.load_hashed_module_sync(wasm, engine)
         }
     }
 }

lib/wasix/src/runtime/module_cache/hashed_module.rs

@@ -20,6 +20,10 @@ pub struct HashedModuleData {
 }
 
 impl HashedModuleData {
+    pub fn new(bytes: impl Into<OwnedBuffer>) -> Self {
+        Self::new_sha256(bytes)
+    }
+
     pub fn new_sha256(bytes: impl Into<OwnedBuffer>) -> Self {
         let wasm = bytes.into();
         let hash = ModuleHash::sha256(&wasm);

lib/wasix/src/state/env.rs

@@ -421,6 +421,7 @@ impl WasiEnv {
         let pid = self.process.pid();
 
         let mut store = store.as_store_mut();
+        let engine = self.runtime().engine();
         let mut func_env = WasiFunctionEnv::new(&mut store, self);
 
         let is_dl = super::linker::is_dynamically_linked(&module);
@@ -451,6 +452,7 @@ impl WasiEnv {
 
                     // TODO: make stack size configurable
                     Linker::new(
+                        engine,
                         &module,
                         &mut store,
                         memory,

lib/wasix/src/state/linker.rs

@@ -277,8 +277,8 @@ use tracing::trace;
 use virtual_fs::{AsyncReadExt, FileSystem, FsError};
 use virtual_mio::InlineWaker;
 use wasmer::{
-    AsStoreMut, AsStoreRef, ExportError, Exportable, Extern, ExternType, Function, FunctionEnv,
-    FunctionEnvMut, FunctionType, Global, GlobalType, ImportType, Imports, Instance,
+    AsStoreMut, AsStoreRef, Engine, ExportError, Exportable, Extern, ExternType, Function,
+    FunctionEnv, FunctionEnvMut, FunctionType, Global, GlobalType, ImportType, Imports, Instance,
     InstantiationError, Memory, MemoryError, Module, RuntimeError, StoreMut, Table, Tag, Type,
     Value, WasmTypeList, WASM_PAGE_SIZE,
 };
@@ -817,6 +817,8 @@ struct InstanceGroupState {
 
 // There is only one LinkerState for all instance groups
 struct LinkerState {
+    engine: Engine,
+
     main_module: Module,
     main_module_dylink_info: DylinkInfo,
 
@@ -947,6 +949,7 @@ impl Linker {
     /// running, and a Linker instance is returned which can then be used for the
     /// loading/linking of further side modules.
     pub fn new(
+        engine: Engine,
         main_module: &Module,
         store: &mut StoreMut<'_>,
         memory: Option<Memory>,
@@ -1082,6 +1085,7 @@ impl Linker {
         };
 
         let mut linker_state = LinkerState {
+            engine,
             main_module: main_module.clone(),
             main_module_dylink_info: dylink_section,
             side_modules: BTreeMap::new(),
@@ -2310,14 +2314,14 @@ impl LinkerState {
                 }
 
                 (
-                    HashedModuleData::new_sha256(bytes),
+                    HashedModuleData::new(bytes),
                     Some((full_path, ld_library_path)),
                 )
             }
-            DlModuleSpec::Memory { bytes, .. } => (HashedModuleData::new_sha256(bytes), None),
+            DlModuleSpec::Memory { bytes, .. } => (HashedModuleData::new(bytes), None),
         };
 
-        let module = runtime.load_hashed_module_sync(module_data)?;
+        let module = runtime.load_hashed_module_sync(module_data, Some(&self.engine))?;
 
         let dylink_info = parse_dylink0_section(&module)?;