Allow CLI to exec WASI w/ multiple namespaces

Added a warning when it's done and a flag to opt out of the warning: `--allow-multiple-wasi-versions`

lib/cli/src/commands/run.rs

@@ -152,8 +152,12 @@ impl Run {
         // If WASI is enabled, try to execute it with it
         #[cfg(feature = "wasi")]
         {
-            let wasi_version = Wasi::get_version(&module);
-            if wasi_version.is_some() {
+            let wasi_versions = Wasi::get_versions(&module);
+            if let Some(wasi_versions) = wasi_versions {
+                if wasi_versions.len() >= 2 && !self.wasi.allow_multiple_wasi_versions {
+                    warning!("Found more than 1 WASI version in this module. If this is intentional, pass `--allow-multiple-wasi-versions` to suppress this warning.");
+                }
+
                 let program_name = self
                     .command_name
                     .clone()

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

@@ -1,8 +1,9 @@
 use crate::utils::{parse_envvar, parse_mapdir};
 use anyhow::{Context, Result};
+use std::collections::BTreeSet;
 use std::path::PathBuf;
 use wasmer::{Instance, Module};
-use wasmer_wasi::{get_wasi_version, WasiError, WasiState, WasiVersion};
+use wasmer_wasi::{get_wasi_versions, WasiError, WasiState, WasiVersion};
 
 use clap::Clap;
 
@@ -25,22 +26,25 @@ pub struct Wasi {
     #[cfg(feature = "experimental-io-devices")]
     #[clap(long = "enable-experimental-io-devices")]
     enable_experimental_io_devices: bool,
+
+    #[clap(long = "allow-multiple-wasi-versions")]
+    pub allow_multiple_wasi_versions: bool,
 }
 
 #[allow(dead_code)]
 impl Wasi {
     /// Gets the WASI version (if any) for the provided module
-    pub fn get_version(module: &Module) -> Option<WasiVersion> {
+    pub fn get_versions(module: &Module) -> Option<BTreeSet<WasiVersion>> {
         // Get the wasi version in strict mode, so no other imports are
         // allowed.
-        get_wasi_version(&module, true)
+        get_wasi_versions(&module, true)
     }
 
     /// Checks if a given module has any WASI imports at all.
     pub fn has_wasi_imports(module: &Module) -> bool {
         // Get the wasi version in non-strict mode, so no other imports
         // are allowed
-        get_wasi_version(&module, false).is_some()
+        get_wasi_versions(&module, false).is_some()
     }
 
     /// Helper function for executing Wasi from the `Run` command.
@@ -63,8 +67,9 @@ impl Wasi {
         }
 
         let mut wasi_env = wasi_state_builder.finalize()?;
-        let import_object = wasi_env.import_object(&module)?;
-        let instance = Instance::new(&module, &import_object)?;
+        //let import_object = wasi_env.import_object(&module)?;
+        let resolver = wasi_env.multiple_wasi_resolver(&module)?;
+        let instance = Instance::new(&module, &resolver)?;
 
         let start = instance.exports.get_function("_start")?;
         let result = start.call(&[]);

lib/wasi/src/lib.rs

@@ -26,10 +26,13 @@ pub use crate::state::{
     WasiStateCreationError, ALL_RIGHTS, VIRTUAL_ROOT_FD,
 };
 pub use crate::syscalls::types;
-pub use crate::utils::{get_wasi_version, is_wasi_module, WasiVersion};
+pub use crate::utils::{get_wasi_version, get_wasi_versions, is_wasi_module, WasiVersion};
 
 use thiserror::Error;
-use wasmer::{imports, Function, ImportObject, LazyInit, Memory, Module, Store, WasmerEnv};
+use wasmer::{
+    imports, ChainableNamedResolver, Function, ImportObject, LazyInit, Memory, Module,
+    NamedResolver, Store, WasmerEnv,
+};
 #[cfg(all(target_os = "macos", target_arch = "aarch64",))]
 use wasmer::{FunctionType, ValType};
 
@@ -67,6 +70,7 @@ impl WasiEnv {
         }
     }
 
+    /// Get an `ImportObject` for a specific version of WASI detected in the module.
     pub fn import_object(&mut self, module: &Module) -> Result<ImportObject, WasiError> {
         let wasi_version = get_wasi_version(module, false).ok_or(WasiError::UnknownWasiVersion)?;
         Ok(generate_import_object_from_env(
@@ -76,6 +80,34 @@ impl WasiEnv {
         ))
     }
 
+    /// Like `import_object` but containing all the WASI versions detected in
+    /// the module.
+    pub fn multiple_wasi_resolver(
+        &mut self,
+        module: &Module,
+    ) -> Result<Box<dyn NamedResolver>, WasiError> {
+        let wasi_versions =
+            get_wasi_versions(module, false).ok_or(WasiError::UnknownWasiVersion)?;
+        let mut version_iter = wasi_versions.iter();
+        let mut resolver: Box<dyn NamedResolver>  = {
+            let version = version_iter.next().ok_or(WasiError::UnknownWasiVersion)?;
+            Box::new(generate_import_object_from_env(
+                module.store(),
+                self.clone(),
+                *version,
+            ))
+        };
+        for version in version_iter {
+            let new_import_object = generate_import_object_from_env(
+                module.store(),
+                self.clone(),
+                *version,
+            );
+            resolver = Box::new(new_import_object.chain_front(resolver));
+        }
+        Ok(resolver)
+    }
+
     /// Get the WASI state
     ///
     /// Be careful when using this in host functions that call into Wasm:

lib/wasi/src/utils.rs

@@ -1,4 +1,5 @@
-use wasmer::{ExternType, Module};
+use std::collections::BTreeSet;
+use wasmer::Module;
 
 #[allow(dead_code)]
 /// Check if a provided module is compiled for some version of WASI.
@@ -9,7 +10,7 @@ pub fn is_wasi_module(module: &Module) -> bool {
 
 /// The version of WASI. This is determined by the imports namespace
 /// string.
-#[derive(Debug, Clone, Copy, PartialEq)]
+#[derive(Debug, Clone, Copy, Eq)]
 pub enum WasiVersion {
     /// `wasi_unstable`.
     Snapshot0,
@@ -30,6 +31,41 @@ pub enum WasiVersion {
     Latest,
 }
 
+impl PartialEq<WasiVersion> for WasiVersion {
+    fn eq(&self, other: &Self) -> bool {
+        match (*self, *other) {
+            (Self::Snapshot1, Self::Latest)
+            | (Self::Latest, Self::Snapshot1)
+            | (Self::Latest, Self::Latest)
+            | (Self::Snapshot0, Self::Snapshot0)
+            | (Self::Snapshot1, Self::Snapshot1) => true,
+            _ => false,
+        }
+    }
+}
+
+impl PartialOrd for WasiVersion {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for WasiVersion {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        if self == other {
+            return std::cmp::Ordering::Equal;
+        }
+        match (*self, *other) {
+            // if snapshot0 is not equal, it must be less
+            (Self::Snapshot0, _) => std::cmp::Ordering::Less,
+            (Self::Snapshot1, Self::Snapshot0) | (Self::Latest, Self::Snapshot0) => {
+                std::cmp::Ordering::Greater
+            }
+            _ => unreachable!("Missing info about ordering of WasiVerison"),
+        }
+    }
+}
+
 /// Namespace for the `Snapshot0` version.
 const SNAPSHOT0_NAMESPACE: &str = "wasi_unstable";
 
@@ -44,10 +80,7 @@ const SNAPSHOT1_NAMESPACE: &str = "wasi_snapshot_preview1";
 /// namespace exits to detect the version. Note that the strict
 /// detection is faster than the non-strict one.
 pub fn get_wasi_version(module: &Module, strict: bool) -> Option<WasiVersion> {
-    let mut imports = module.imports().filter_map(|extern_| match extern_.ty() {
-        ExternType::Function(_f) => Some(extern_.module().to_owned()),
-        _ => None,
-    });
+    let mut imports = module.imports().functions().map(|f| f.module().to_owned());
 
     if strict {
         let first_module = imports.next()?;
@@ -70,3 +103,68 @@ pub fn get_wasi_version(module: &Module, strict: bool) -> Option<WasiVersion> {
         })
     }
 }
+
+/// Like [`get_wasi_version`] but detects multiple WASI versions in a single module.
+/// Thus `strict` behaves differently in this function as multiple versions are
+/// always supported. `strict` indicates whether non-WASI imports should trigger a
+/// failure or be ignored.
+pub fn get_wasi_versions(module: &Module, strict: bool) -> Option<BTreeSet<WasiVersion>> {
+    let mut out = BTreeSet::new();
+    let imports = module.imports().functions().map(|f| f.module().to_owned());
+
+    let mut non_wasi_seen = false;
+    for ns in imports {
+        match ns.as_str() {
+            SNAPSHOT0_NAMESPACE => {
+                out.insert(WasiVersion::Snapshot0);
+            }
+            SNAPSHOT1_NAMESPACE => {
+                out.insert(WasiVersion::Snapshot1);
+            }
+            _ => {
+                non_wasi_seen = true;
+            }
+        }
+    }
+    if strict && non_wasi_seen {
+        None
+    } else {
+        Some(out)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn wasi_version_equality() {
+        assert_eq!(WasiVersion::Snapshot0, WasiVersion::Snapshot0);
+        assert_eq!(WasiVersion::Snapshot1, WasiVersion::Snapshot1);
+        assert_eq!(WasiVersion::Snapshot1, WasiVersion::Latest);
+        assert_eq!(WasiVersion::Latest, WasiVersion::Snapshot1);
+        assert_eq!(WasiVersion::Latest, WasiVersion::Latest);
+        assert!(WasiVersion::Snapshot0 != WasiVersion::Snapshot1);
+        assert!(WasiVersion::Snapshot1 != WasiVersion::Snapshot0);
+        assert!(WasiVersion::Snapshot0 != WasiVersion::Latest);
+        assert!(WasiVersion::Latest != WasiVersion::Snapshot0);
+    }
+
+    #[test]
+    fn wasi_version_ordering() {
+        assert!(WasiVersion::Snapshot0 <= WasiVersion::Snapshot0);
+        assert!(WasiVersion::Snapshot1 <= WasiVersion::Snapshot1);
+        assert!(WasiVersion::Latest <= WasiVersion::Latest);
+        assert!(WasiVersion::Snapshot0 >= WasiVersion::Snapshot0);
+        assert!(WasiVersion::Snapshot1 >= WasiVersion::Snapshot1);
+        assert!(WasiVersion::Latest >= WasiVersion::Latest);
+
+        assert!(WasiVersion::Snapshot0 < WasiVersion::Snapshot1);
+        assert!(WasiVersion::Snapshot1 > WasiVersion::Snapshot0);
+        assert!(WasiVersion::Snapshot0 < WasiVersion::Latest);
+        assert!(WasiVersion::Latest > WasiVersion::Snapshot0);
+
+        assert!(!(WasiVersion::Snapshot1 < WasiVersion::Latest));
+        assert!(!(WasiVersion::Latest > WasiVersion::Snapshot1));
+    }
+}