Fix WASIX additional imports so it works with new threads as well

lib/api/src/imports.rs

@@ -187,6 +187,15 @@ impl<'a> Iterator for ImportsIterator<'a> {
     }
 }
 
+impl IntoIterator for Imports {
+    type IntoIter = std::collections::hash_map::IntoIter<(String, String), Extern>;
+    type Item = ((String, String), Extern);
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.map.into_iter()
+    }
+}
+
 impl IntoIterator for &Imports {
     type IntoIter = std::collections::hash_map::IntoIter<(String, String), Extern>;
     type Item = ((String, String), Extern);

lib/wasix/src/os/task/thread.rs

@@ -617,6 +617,8 @@ pub enum WasiThreadError {
     MemoryCreateFailed(MemoryError),
     #[error("{0}")]
     ExportError(ExportError),
+    #[error("Failed to create additional imports - {0}")]
+    AdditionalImportCreationFailed(Arc<anyhow::Error>),
     #[error("Failed to create the instance")]
     // Note: Boxed so we can keep the error size down
     InstanceCreateFailed(Box<InstantiationError>),
@@ -634,6 +636,7 @@ impl From<WasiThreadError> for Errno {
             WasiThreadError::MethodNotFound => Errno::Inval,
             WasiThreadError::MemoryCreateFailed(_) => Errno::Nomem,
             WasiThreadError::ExportError(_) => Errno::Noexec,
+            WasiThreadError::AdditionalImportCreationFailed(_) => Errno::Noexec,
             WasiThreadError::InstanceCreateFailed(_) => Errno::Noexec,
             WasiThreadError::InitFailed(_) => Errno::Noexec,
             WasiThreadError::InvalidWasmContext => Errno::Noexec,

lib/wasix/src/runners/wasi.rs

@@ -5,7 +5,7 @@ use std::{path::PathBuf, sync::Arc};
 use anyhow::{Context, Error};
 use tracing::Instrument;
 use virtual_fs::{ArcBoxFile, FileSystem, TmpFileSystem, VirtualFile};
-use wasmer::{Extern, Module};
+use wasmer::Module;
 use webc::metadata::{annotations::Wasi, Command};
 
 use crate::{
@@ -201,33 +201,6 @@ impl WasiRunner {
         self
     }
 
-    /// Add an item to the list of importable items provided to the instance.
-    pub fn with_import(
-        &mut self,
-        namespace: impl Into<String>,
-        name: impl Into<String>,
-        value: impl Into<Extern>,
-    ) -> &mut Self {
-        self.with_imports([((namespace, name), value)])
-    }
-
-    /// Add multiple import functions.
-    ///
-    /// This method will accept a [`&Imports`][wasmer::Imports] object.
-    pub fn with_imports<I, S1, S2, E>(&mut self, imports: I) -> &mut Self
-    where
-        I: IntoIterator<Item = ((S1, S2), E)>,
-        S1: Into<String>,
-        S2: Into<String>,
-        E: Into<Extern>,
-    {
-        let imports = imports
-            .into_iter()
-            .map(|((ns, n), e)| ((ns.into(), n.into()), e.into()));
-        self.wasi.additional_imports.extend(imports);
-        self
-    }
-
     #[tracing::instrument(level = "debug", skip_all)]
     pub fn prepare_webc_env(
         &self,

lib/wasix/src/runners/wasi_common.rs

@@ -9,7 +9,6 @@ use derivative::Derivative;
 use futures::future::BoxFuture;
 use tokio::runtime::Handle;
 use virtual_fs::{FileSystem, FsError, OverlayFileSystem, RootFileSystemBuilder, TmpFileSystem};
-use wasmer::Imports;
 use webc::metadata::annotations::Wasi as WasiAnnotation;
 
 use crate::{
@@ -46,7 +45,6 @@ pub(crate) struct CommonWasiOptions {
     pub(crate) snapshot_on: Vec<SnapshotTrigger>,
     pub(crate) snapshot_interval: Option<std::time::Duration>,
     pub(crate) current_dir: Option<PathBuf>,
-    pub(crate) additional_imports: Imports,
 }
 
 impl CommonWasiOptions {
@@ -88,8 +86,6 @@ impl CommonWasiOptions {
 
         *builder.capabilities_mut() = self.capabilities.clone();
 
-        builder.add_imports(&self.additional_imports);
-
         Ok(())
     }
 

lib/wasix/src/runtime/mod.rs

@@ -88,6 +88,11 @@ where
         }
     }
 
+    /// Create additional imports for a new WASIX thread in the provided [wasmer::Store].
+    fn additional_imports(&self, store: &mut wasmer::StoreMut) -> anyhow::Result<wasmer::Imports> {
+        Ok(wasmer::Imports::new())
+    }
+
     /// Get a custom HTTP client
     fn http_client(&self) -> Option<&DynHttpClient> {
         None
@@ -205,6 +210,9 @@ impl TtyBridge for DefaultTty {
     }
 }
 
+type MakeImportCallback =
+    dyn (Fn(&mut wasmer::StoreMut) -> anyhow::Result<wasmer::Imports>) + Send + Sync + 'static;
+
 #[derive(Clone, Derivative)]
 #[derivative(Debug)]
 pub struct PluggableRuntime {
@@ -220,6 +228,8 @@ pub struct PluggableRuntime {
     #[cfg(feature = "journal")]
     #[derivative(Debug = "ignore")]
     pub journals: Vec<Arc<DynJournal>>,
+    #[derivative(Debug = "ignore")]
+    pub additional_imports: Vec<Arc<MakeImportCallback>>,
 }
 
 impl PluggableRuntime {
@@ -256,6 +266,7 @@ impl PluggableRuntime {
             module_cache: Arc::new(module_cache::in_memory()),
             #[cfg(feature = "journal")]
             journals: Vec::new(),
+            additional_imports: Vec::new(),
         }
     }
 
@@ -311,6 +322,17 @@ impl PluggableRuntime {
         self.journals.push(journal);
         self
     }
+
+    pub fn with_additional_imports(
+        &mut self,
+        imports: impl (Fn(&mut wasmer::StoreMut) -> anyhow::Result<wasmer::Imports>)
+            + Send
+            + Sync
+            + 'static,
+    ) -> &mut Self {
+        self.additional_imports.push(Arc::new(imports));
+        self
+    }
 }
 
 impl Runtime for PluggableRuntime {
@@ -345,6 +367,14 @@ impl Runtime for PluggableRuntime {
             .unwrap_or_default()
     }
 
+    fn additional_imports(&self, store: &mut wasmer::StoreMut) -> anyhow::Result<wasmer::Imports> {
+        let mut imports = wasmer::Imports::new();
+        for cb in &self.additional_imports {
+            imports.extend((*cb)(store)?);
+        }
+        Ok(imports)
+    }
+
     fn task_manager(&self) -> &Arc<dyn VirtualTaskManager> {
         &self.rt
     }

lib/wasix/src/state/builder.rs

@@ -9,7 +9,7 @@ use std::{
 use rand::Rng;
 use thiserror::Error;
 use virtual_fs::{ArcFile, FileSystem, FsError, TmpFileSystem, VirtualFile};
-use wasmer::{AsStoreMut, Extern, Imports, Instance, Module, Store};
+use wasmer::{AsStoreMut, Instance, Module, Store};
 use wasmer_config::package::PackageId;
 
 #[cfg(feature = "journal")]
@@ -78,7 +78,6 @@ pub struct WasiEnvBuilder {
     pub(super) map_commands: HashMap<String, PathBuf>,
 
     pub(super) capabilites: Capabilities,
-    pub(super) additional_imports: Imports,
 
     #[cfg(feature = "journal")]
     pub(super) snapshot_on: Vec<SnapshotTrigger>,
@@ -673,49 +672,30 @@ impl WasiEnvBuilder {
         self.snapshot_interval.replace(interval);
     }
 
-    /// Add an item to the list of importable items provided to the instance.
-    pub fn import(
-        mut self,
-        namespace: impl Into<String>,
-        name: impl Into<String>,
-        value: impl Into<Extern>,
-    ) -> Self {
-        self.add_imports([((namespace, name), value)]);
-        self
-    }
-
-    /// Add an item to the list of importable items provided to the instance.
-    pub fn add_import(
-        &mut self,
-        namespace: impl Into<String>,
-        name: impl Into<String>,
-        value: impl Into<Extern>,
-    ) {
-        self.add_imports([((namespace, name), value)]);
-    }
+    pub fn get_runtime_or_create_default(&mut self) -> Arc<dyn Runtime + Send + Sync> {
+        if self.runtime.is_none() {
+            self.runtime = Some({
+                #[cfg(feature = "sys-thread")]
+                {
+                    #[allow(unused_mut)]
+                    let mut runtime = crate::runtime::PluggableRuntime::new(Arc::new(
+                        crate::runtime::task_manager::tokio::TokioTaskManager::default(),
+                    ));
+                    #[cfg(feature = "journal")]
+                    for journal in self.journals.clone() {
+                        runtime.add_journal(journal);
+                    }
+                    Arc::new(runtime)
+                }
 
-    pub fn add_imports<I, S1, S2, E>(&mut self, imports: I)
-    where
-        I: IntoIterator<Item = ((S1, S2), E)>,
-        S1: Into<String>,
-        S2: Into<String>,
-        E: Into<Extern>,
-    {
-        let imports = imports
-            .into_iter()
-            .map(|((ns, n), e)| ((ns.into(), n.into()), e.into()));
-        self.additional_imports.extend(imports);
-    }
+                #[cfg(not(feature = "sys-thread"))]
+                {
+                    panic!("this build does not support a default runtime - specify one with WasiEnvBuilder::runtime()");
+                }
+            });
+        }
 
-    pub fn imports<I, S1, S2, E>(mut self, imports: I) -> Self
-    where
-        I: IntoIterator<Item = ((S1, S2), E)>,
-        S1: Into<String>,
-        S2: Into<String>,
-        E: Into<Extern>,
-    {
-        self.add_imports(imports);
-        self
+        self.runtime.as_ref().unwrap().clone()
     }
 
     /// Consumes the [`WasiEnvBuilder`] and produces a [`WasiEnvInit`], which
@@ -863,6 +843,8 @@ impl WasiEnvBuilder {
             wasi_fs.has_unioned.lock().unwrap().insert(id.clone());
         }
 
+        let runtime = self.get_runtime_or_create_default();
+
         let state = WasiState {
             fs: wasi_fs,
             secret: rand::thread_rng().gen::<[u8; 32]>(),
@@ -874,24 +856,6 @@ impl WasiEnvBuilder {
             envs: std::sync::Mutex::new(conv_env_vars(self.envs)),
         };
 
-        let runtime = self.runtime.unwrap_or_else(|| {
-            #[cfg(feature = "sys-thread")]
-            {
-                #[allow(unused_mut)]
-                let mut runtime = crate::runtime::PluggableRuntime::new(Arc::new(crate::runtime::task_manager::tokio::TokioTaskManager::default()));
-                #[cfg(feature = "journal")]
-                for journal in self.journals.clone() {
-                    runtime.add_journal(journal);
-                }
-                Arc::new(runtime)
-            }
-
-            #[cfg(not(feature = "sys-thread"))]
-            {
-                panic!("this build does not support a default runtime - specify one with WasiEnvBuilder::runtime()");
-            }
-        });
-
         let uses = self.uses;
         let map_commands = self.map_commands;
 
@@ -925,7 +889,6 @@ impl WasiEnvBuilder {
             extra_tracing: true,
             #[cfg(feature = "journal")]
             snapshot_on: self.snapshot_on,
-            additional_imports: self.additional_imports,
         };
 
         Ok(init)
@@ -971,13 +934,17 @@ impl WasiEnvBuilder {
 
     #[allow(clippy::result_large_err)]
     pub fn instantiate_ext(
-        self,
+        mut self,
         module: Module,
         module_hash: ModuleHash,
         store: &mut impl AsStoreMut,
     ) -> Result<(Instance, WasiFunctionEnv), WasiRuntimeError> {
+        let runtime = self.get_runtime_or_create_default();
+        let additional_imports = runtime
+            .additional_imports(&mut store.as_store_mut())
+            .map_err(Arc::new)?;
         let init = self.build_init()?;
-        WasiEnv::instantiate(init, module, module_hash, store)
+        WasiEnv::instantiate(init, module, module_hash, store, additional_imports)
     }
 
     #[allow(clippy::result_large_err)]