Allow to pass env. vars on CLI

Author: psiegl

crates/wasmtime-cli/src/lib.rs

@@ -359,25 +359,30 @@ where
     Ok((pre, engine, guest_resources, host_resources))
 }
 
-fn new_store<C: Invoke>(
+fn new_store<'a, C: Invoke>(
     engine: &Engine,
     wrpc: C,
     cx: C::Context,
     arg0: &str,
     timeout: Duration,
+    vars: impl IntoIterator<Item = (&'a str, &'a str)>,
 ) -> wasmtime::Store<Ctx<C>> {
+    let mut builder = WasiCtxBuilder::new();
+    builder
+        .inherit_env()
+        .inherit_stdio()
+        .inherit_network()
+        .allow_ip_name_lookup(true)
+        .allow_tcp(true)
+        .allow_udp(true)
+        .args(&[arg0]);
+    for (k, v) in vars {
+        builder.env(k, v);
+    }
     Store::new(
         engine,
         Ctx {
-            wasi: WasiCtxBuilder::new()
-                .inherit_env()
-                .inherit_stdio()
-                .inherit_network()
-                .allow_ip_name_lookup(true)
-                .allow_tcp(true)
-                .allow_udp(true)
-                .args(&[arg0])
-                .build(),
+            wasi: builder.build(),
             http: WasiHttpCtx::new(),
             table: ResourceTable::new(),
             wrpc: WrpcCtx {
@@ -395,6 +400,7 @@ pub async fn handle_run<C>(
     clt: C,
     cx: C::Context,
     timeout: Duration,
+    vars: Vec<(String, String)>,
     workload: &str,
 ) -> anyhow::Result<()>
 where
@@ -403,7 +409,14 @@ where
 {
     let (pre, engine, _, _) =
         instantiate_pre(WASI_SNAPSHOT_PREVIEW1_COMMAND_ADAPTER, workload).await?;
-    let mut store = new_store(&engine, clt, cx, "command.wasm", timeout);
+    let mut store = new_store(
+        &engine,
+        clt,
+        cx,
+        "command.wasm",
+        timeout,
+        vars.iter().map(|(k, v)| (k.as_str(), v.as_str())),
+    );
     let cmd = wasmtime_wasi::p2::bindings::CommandPre::new(pre)
         .map_err(anyhow::Error::from)
         .context("failed to construct `command` instance")?
@@ -605,7 +618,14 @@ where
                 let invocations = srv
                     .serve_function(
                         move || {
-                            new_store(&engine, clt.clone(), cx.clone(), "reactor.wasm", timeout)
+                            new_store(
+                                &engine,
+                                clt.clone(),
+                                cx.clone(),
+                                "reactor.wasm",
+                                timeout,
+                                std::iter::empty(),
+                            )
                         },
                         pre.clone(),
                         Arc::clone(&host_resources),
@@ -665,6 +685,7 @@ where
                                             cx.clone(),
                                             "reactor.wasm",
                                             timeout,
+                                            std::iter::empty(),
                                         )
                                     },
                                     pre.clone(),
@@ -768,7 +789,14 @@ where
         serve_shared(
             &mut handlers,
             srv,
-            new_store(&engine, clt, cx, "reactor.wasm", timeout),
+            new_store(
+                &engine,
+                clt,
+                cx,
+                "reactor.wasm",
+                timeout,
+                std::iter::empty(),
+            ),
             pre,
             guest_resources,
             host_resources,

crates/wasmtime-cli/src/tcp.rs

@@ -24,10 +24,25 @@ pub struct RunArgs {
     #[arg(long, default_value = DEFAULT_ADDR)]
     import: String,
 
+    /// Pass an environment variable to the program.
+    ///
+    /// `--env NAME=VALUE` sets the environment variable `NAME` to `VALUE`
+    /// for the guest. Host environment variables are already inherited by
+    /// default, so only the `NAME=VALUE` form is accepted.
+    #[arg(long = "env", number_of_values = 1, value_name = "NAME=VALUE", value_parser = parse_env_var)]
+    vars: Vec<(String, String)>,
+
     /// Path or URL to Wasm command component
     workload: String,
 }
 
+fn parse_env_var(s: &str) -> Result<(String, String), String> {
+    let (key, val) = s
+        .split_once('=')
+        .ok_or_else(|| format!("invalid `--env` value `{s}`: expected `NAME=VALUE`"))?;
+    Ok((key.to_string(), val.to_string()))
+}
+
 /// Serve a reactor component
 #[derive(Parser, Debug)]
 pub struct ServeArgs {
@@ -52,13 +67,15 @@ pub async fn handle_run(
     RunArgs {
         timeout,
         import,
+        vars,
         ref workload,
     }: RunArgs,
 ) -> anyhow::Result<()> {
     crate::handle_run(
         wrpc_transport::tcp::Client::from(import),
         (),
         *timeout,
+        vars,
         workload,
     )
     .await