Allow to pass env. vars on CLI

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")?
@@ -586,6 +599,7 @@ pub async fn serve_stateless<C, S>(
     host_resources: Arc<HashMap<Box<str>, HashMap<Box<str>, (ResourceType, ResourceType)>>>,
     engine: &Engine,
     timeout: Duration,
+    vars: Vec<(String, String)>,
 ) -> anyhow::Result<()>
 where
     C: Invoke + Clone + 'static,
@@ -601,11 +615,19 @@ where
                 let clt = clt.clone();
                 let cx = cx.clone();
                 let engine = engine.clone();
+                let vars = vars.clone();
                 info!(?name, "serving root function");
                 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,
+                                vars.iter().map(|(k, v)| (k.as_str(), v.as_str())),
+                            )
                         },
                         pre.clone(),
                         Arc::clone(&host_resources),
@@ -655,6 +677,7 @@ where
                             let clt = clt.clone();
                             let engine = engine.clone();
                             let cx = cx.clone();
+                            let vars = vars.clone();
                             info!(?name, "serving instance function");
                             let invocations = srv
                                 .serve_function(
@@ -665,6 +688,7 @@ where
                                             cx.clone(),
                                             "reactor.wasm",
                                             timeout,
+                                            vars.iter().map(|(k, v)| (k.as_str(), v.as_str())),
                                         )
                                     },
                                     pre.clone(),
@@ -741,6 +765,7 @@ pub async fn handle_serve<C, S>(
     clt: C,
     cx: C::Context,
     timeout: Duration,
+    vars: Vec<(String, String)>,
     workload: &str,
 ) -> anyhow::Result<()>
 where
@@ -762,13 +787,22 @@ where
             host_resources,
             &engine,
             timeout,
+            vars,
         )
         .await?;
     } else {
+        let store = new_store(
+            &engine,
+            clt,
+            cx,
+            "reactor.wasm",
+            timeout,
+            vars.iter().map(|(k, v)| (k.as_str(), v.as_str())),
+        );
         serve_shared(
             &mut handlers,
             srv,
-            new_store(&engine, clt, cx, "reactor.wasm", timeout),
+            store,
             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 {
@@ -43,6 +58,14 @@ pub struct ServeArgs {
     #[arg(long, default_value = DEFAULT_ADDR)]
     export: 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,
 }
@@ -52,13 +75,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
@@ -70,6 +95,7 @@ pub async fn handle_serve(
         timeout,
         export,
         import,
+        vars,
         ref workload,
     }: ServeArgs,
 ) -> anyhow::Result<()> {
@@ -98,6 +124,7 @@ pub async fn handle_serve(
         wrpc_transport::tcp::Client::from(import),
         (),
         *timeout,
+        vars,
         workload,
     )
     .await;