From 57a8a0936bb9b6095061fc5784a786d3e5daf415 Mon Sep 17 00:00:00 2001
From: Amanieu d'Antras <amanieu@gmail.com>
Date: Mon, 29 Nov 2021 16:42:09 +0000
Subject: [PATCH] Provide WASI imports when invoking an explicit export from
 the CLI

---
 lib/cli/src/commands/run.rs      | 64 +++++++++++++++++++-------------
 lib/cli/src/commands/run/wasi.rs | 21 +++++++----
 2 files changed, 51 insertions(+), 34 deletions(-)

diff --git a/lib/cli/src/commands/run.rs b/lib/cli/src/commands/run.rs
index 96e15fbc958..2ac40f45234 100644
--- a/lib/cli/src/commands/run.rs
+++ b/lib/cli/src/commands/run.rs
@@ -94,21 +94,6 @@ impl Run {
 
     fn inner_execute(&self) -> Result<()> {
         let module = self.get_module()?;
-        // Do we want to invoke a function?
-        if let Some(ref invoke) = self.invoke {
-            let imports = imports! {};
-            let instance = Instance::new(&module, &imports)?;
-            let result = self.invoke_function(&instance, &invoke, &self.args)?;
-            println!(
-                "{}",
-                result
-                    .iter()
-                    .map(|val| val.to_string())
-                    .collect::<Vec<String>>()
-                    .join(" ")
-            );
-            return Ok(());
-        }
         #[cfg(feature = "emscripten")]
         {
             use wasmer_emscripten::{
@@ -117,6 +102,9 @@ impl Run {
             };
             // TODO: refactor this
             if is_emscripten_module(&module) {
+                if self.invoke.is_some() {
+                    bail!("--invoke is not supported with emscripten modules");
+                }
                 let mut emscripten_globals = EmscriptenGlobals::new(module.store(), &module)
                     .map_err(|e| anyhow!("{}", e))?;
                 let mut em_env = EmEnv::new(&emscripten_globals.data, Default::default());
@@ -154,7 +142,7 @@ impl Run {
 
         // If WASI is enabled, try to execute it with it
         #[cfg(feature = "wasi")]
-        {
+        let instance = {
             use std::collections::BTreeSet;
             use wasmer_wasi::WasiVersion;
 
@@ -187,21 +175,45 @@ impl Run {
                                 .map(|f| f.to_string_lossy().to_string())
                         })
                         .unwrap_or_default();
-                    return self
-                        .wasi
-                        .execute(module, program_name, self.args.clone())
-                        .with_context(|| "WASI execution failed");
+                    self.wasi
+                        .instantiate(&module, program_name, self.args.clone())
+                        .with_context(|| "failed to instantiate WASI module")?
                 }
                 // not WASI
-                _ => (),
+                _ => Instance::new(&module, &imports! {})?,
             }
+        };
+        #[cfg(not(feature = "wasi"))]
+        let instance = Instance::new(&module, &imports! {})?;
+
+        // If this module exports an _initialize function, run that first.
+        if let Ok(initialize) = instance.exports.get_function("_initialize") {
+            initialize
+                .call(&[])
+                .with_context(|| "failed to run _initialize function")?;
         }
 
-        // Try to instantiate the wasm file, with no provided imports
-        let imports = imports! {};
-        let instance = Instance::new(&module, &imports)?;
-        let start: Function = self.try_find_function(&instance, "_start", &[])?;
-        start.call(&[])?;
+        // Do we want to invoke a function?
+        if let Some(ref invoke) = self.invoke {
+            let imports = imports! {};
+            let instance = Instance::new(&module, &imports)?;
+            let result = self.invoke_function(&instance, &invoke, &self.args)?;
+            println!(
+                "{}",
+                result
+                    .iter()
+                    .map(|val| val.to_string())
+                    .collect::<Vec<String>>()
+                    .join(" ")
+            );
+        } else {
+            let start: Function = self.try_find_function(&instance, "_start", &[])?;
+            let result = start.call(&[]);
+            #[cfg(feature = "wasi")]
+            self.wasi.handle_result(result)?;
+            #[cfg(not(feature = "wasi"))]
+            result?;
+        }
 
         Ok(())
     }
diff --git a/lib/cli/src/commands/run/wasi.rs b/lib/cli/src/commands/run/wasi.rs
index 318f599863a..8f408fd443e 100644
--- a/lib/cli/src/commands/run/wasi.rs
+++ b/lib/cli/src/commands/run/wasi.rs
@@ -1,8 +1,8 @@
 use crate::utils::{parse_envvar, parse_mapdir};
-use anyhow::{Context, Result};
+use anyhow::Result;
 use std::collections::BTreeSet;
 use std::path::PathBuf;
-use wasmer::{Instance, Module};
+use wasmer::{Instance, Module, RuntimeError, Val};
 use wasmer_wasi::{get_wasi_versions, WasiError, WasiState, WasiVersion};
 
 use structopt::StructOpt;
@@ -69,8 +69,13 @@ impl Wasi {
         get_wasi_versions(&module, false).is_some()
     }
 
-    /// Helper function for executing Wasi from the `Run` command.
-    pub fn execute(&self, module: Module, program_name: String, args: Vec<String>) -> Result<()> {
+    /// Helper function for instantiating a module with Wasi imports for the `Run` command.
+    pub fn instantiate(
+        &self,
+        module: &Module,
+        program_name: String,
+        args: Vec<String>,
+    ) -> Result<Instance> {
         let args = args.iter().cloned().map(|arg| arg.into_bytes());
 
         let mut wasi_state_builder = WasiState::new(program_name);
@@ -91,10 +96,11 @@ impl Wasi {
         let mut wasi_env = wasi_state_builder.finalize()?;
         let resolver = wasi_env.import_object_for_all_wasi_versions(&module)?;
         let instance = Instance::new(&module, &resolver)?;
+        Ok(instance)
+    }
 
-        let start = instance.exports.get_function("_start")?;
-        let result = start.call(&[]);
-
+    /// Helper function for handling the result of a Wasi _start function.
+    pub fn handle_result(&self, result: Result<Box<[Val]>, RuntimeError>) -> Result<()> {
         match result {
             Ok(_) => Ok(()),
             Err(err) => {
@@ -109,7 +115,6 @@ impl Wasi {
                 Err(err)
             }
         }
-        .with_context(|| "failed to run WASI `_start` function")
     }
 
     pub fn for_binfmt_interpreter() -> Result<Self> {