Add new --invoke arg for choosing entry function for single WASM modules + fix --invoke not working for WASI(X) modules

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

@@ -76,9 +76,12 @@ pub struct Run {
     /// Set the default stack size (default is 1048576)
     #[clap(long = "stack-size")]
     stack_size: Option<usize>,
-    /// The function or command to invoke.
-    #[clap(short, long, aliases = &["command", "invoke", "command-name"])]
+    /// The entrypoint module for webc packages.
+    #[clap(short, long, aliases = &["command", "command-name"])]
     entrypoint: Option<String>,
+    /// The function to invoke.
+    #[clap(short, long)]
+    invoke: Option<String>,
     /// Generate a coredump at this path if a WebAssembly trap occurs
     #[clap(name = "COREDUMP_PATH", long)]
     coredump_on_trap: Option<PathBuf>,
@@ -377,19 +380,19 @@ impl Run {
         let instance = Instance::new(store, module, &imports)
             .context("Unable to instantiate the WebAssembly module")?;
 
-        let entrypoint  = match &self.entrypoint {
+        let entry_function  = match &self.invoke {
             Some(entry) => {
                 instance.exports
                     .get_function(entry)
-                    .with_context(|| format!("The module doesn't contain a \"{entry}\" function"))?
+                    .with_context(|| format!("The module doesn't export a function named \"{entry}\""))?
             },
             None => {
                 instance.exports.get_function("_start")
-                    .context("The module doesn't contain a \"_start\" function. Either implement it or specify an entrypoint function.")?
+                    .context("The module doesn't export a \"_start\" function. Either implement it or specify an entry function with --invoke")?
             }
         };
 
-        let return_values = invoke_function(&instance, store, entrypoint, &self.args)?;
+        let return_values = invoke_function(&instance, store, entry_function, &self.args)?;
 
         println!(
             "{}",
@@ -425,6 +428,10 @@ impl Run {
             .with_forward_host_env(self.wasi.forward_host_env)
             .with_capabilities(self.wasi.capabilities());
 
+        if let Some(ref entry_function) = self.invoke {
+            runner.with_entry_function(entry_function);
+        }
+
         #[cfg(feature = "journal")]
         {
             for trigger in self.wasi.snapshot_on.iter().cloned() {
@@ -519,6 +526,7 @@ impl Run {
             wcgi: WcgiOptions::default(),
             stack_size: None,
             entrypoint: Some(original_executable.to_string()),
+            invoke: None,
             coredump_on_trap: None,
             input: PackageSource::infer(executable)?,
             args: args.to_vec(),

lib/wasix/src/runners/wasi.rs

@@ -34,6 +34,20 @@ impl WasiRunner {
         WasiRunner::default()
     }
 
+    /// Returns the current entry function for this `WasiRunner`
+    pub fn entry_function(&self) -> Option<String> {
+        self.wasi.entry_function.clone()
+    }
+
+    /// Builder method to set the name of the entry function for this `WasiRunner`
+    pub fn with_entry_function<S>(&mut self, entry_function: S) -> &mut Self
+    where
+        S: Into<String>,
+    {
+        self.wasi.entry_function = Some(entry_function.into());
+        self
+    }
+
     /// Returns the current arguments for this `WasiRunner`
     pub fn get_args(&self) -> Vec<String> {
         self.wasi.args.clone()

lib/wasix/src/runners/wasi_common.rs

@@ -32,6 +32,7 @@ pub struct MappedCommand {
 #[derive(Derivative, Default, Clone)]
 #[derivative(Debug)]
 pub(crate) struct CommonWasiOptions {
+    pub(crate) entry_function: Option<String>,
     pub(crate) args: Vec<String>,
     pub(crate) env: HashMap<String, String>,
     pub(crate) forward_host_env: bool,
@@ -57,6 +58,10 @@ impl CommonWasiOptions {
         wasi: &WasiAnnotation,
         root_fs: Option<TmpFileSystem>,
     ) -> Result<(), anyhow::Error> {
+        if let Some(ref entry_function) = self.entry_function {
+            builder.set_entry_function(entry_function);
+        }
+
         let root_fs = root_fs.unwrap_or_else(|| {
             RootFileSystemBuilder::default()
                 .with_tmp(!self.is_tmp_mapped)

lib/wasix/src/state/builder.rs

@@ -49,6 +49,8 @@ use super::env::WasiEnvInit;
 /// ```
 #[derive(Default)]
 pub struct WasiEnvBuilder {
+    /// Name of entry function. Defaults to running `_start` if not specified.
+    pub(super) entry_function: Option<String>,
     /// Command line arguments.
     pub(super) args: Vec<String>,
     /// Environment variables.
@@ -97,6 +99,7 @@ impl std::fmt::Debug for WasiEnvBuilder {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         // TODO: update this when stable
         f.debug_struct("WasiEnvBuilder")
+            .field("entry_function", &self.entry_function)
             .field("args", &self.args)
             .field("envs", &self.envs)
             .field("preopens", &self.preopens)
@@ -240,6 +243,21 @@ impl WasiEnvBuilder {
         &mut self.envs
     }
 
+    pub fn entry_function<S>(mut self, entry_function: S) -> Self
+    where
+        S: AsRef<str>,
+    {
+        self.set_entry_function(entry_function);
+        self
+    }
+
+    pub fn set_entry_function<S>(&mut self, entry_function: S)
+    where
+        S: AsRef<str>,
+    {
+        self.entry_function = Some(entry_function.as_ref().to_owned());
+    }
+
     /// Add an argument.
     ///
     /// Arguments must not contain the nul (0x0) byte
@@ -1024,6 +1042,8 @@ impl WasiEnvBuilder {
             );
         }
 
+        let entry_function = self.entry_function.clone();
+
         let (instance, env) = self.instantiate_ext(module, module_hash, store)?;
 
         // Bootstrap the process
@@ -1036,7 +1056,9 @@ impl WasiEnvBuilder {
                 .map_err(|exit| WasiRuntimeError::Wasi(WasiError::Exit(exit)))?;
         }
 
-        let start = instance.exports.get_function("_start")?;
+        let start = instance
+            .exports
+            .get_function(entry_function.as_deref().unwrap_or("_start"))?;
         env.data(&store).thread.set_status_running();
 
         let result = crate::run_wasi_func_start(start, store);