Support dirs/env vars in `wasmtime serve` (#8279)

* More flags like `--dir` and `--env` are moved into `RunCommon` to be
  shared between `wasmtime serve` and `wasmtime run`, meaning that the
  `serve` command can now configure environment variables.

* A small test has been added as well as infrastructure for running
  tests with `wasmtime serve` itself. Previously there were no tests
  that executed `wasmtime serve`.

* The `test_programs` crate had a small refactoring to avoid
  double-generation of http bindings.

crates/test-programs/src/bin/api_proxy.rs

@@ -1,18 +1,12 @@
-pub mod bindings {
+use test_programs::wasi::http::types::{
-    wit_bindgen::generate!({
+    Headers, IncomingRequest, OutgoingBody, OutgoingResponse, ResponseOutparam,
-        path: "../wasi-http/wit",
+};
-        world: "wasi:http/proxy",
-        default_bindings_module: "bindings",
-    });
-}
-
-use bindings::wasi::http::types::{IncomingRequest, ResponseOutparam};
 
 struct T;
 
-bindings::export!(T);
+test_programs::proxy::export!(T);
 
-impl bindings::exports::wasi::http::incoming_handler::Guest for T {
+impl test_programs::proxy::exports::wasi::http::incoming_handler::Guest for T {
     fn handle(request: IncomingRequest, outparam: ResponseOutparam) {
         assert!(request.scheme().is_some());
         assert!(request.authority().is_some());
@@ -41,19 +35,18 @@ impl bindings::exports::wasi::http::incoming_handler::Guest for T {
             "forbidden host header present in incoming request"
         );
 
-        let hdrs = bindings::wasi::http::types::Headers::new();
+        let hdrs = Headers::new();
-        let resp = bindings::wasi::http::types::OutgoingResponse::new(hdrs);
+        let resp = OutgoingResponse::new(hdrs);
         let body = resp.body().expect("outgoing response");
 
-        bindings::wasi::http::types::ResponseOutparam::set(outparam, Ok(resp));
+        ResponseOutparam::set(outparam, Ok(resp));
 
         let out = body.write().expect("outgoing stream");
         out.blocking_write_and_flush(b"hello, world!")
             .expect("writing response");
 
         drop(out);
-        bindings::wasi::http::types::OutgoingBody::finish(body, None)
+        OutgoingBody::finish(body, None).expect("outgoing-body.finish");
-            .expect("outgoing-body.finish");
     }
 }
 

crates/test-programs/src/bin/api_proxy_streaming.rs

@@ -1,26 +1,18 @@
 use anyhow::{anyhow, bail, Result};
-use bindings::wasi::http::types::{
+use futures::{future, stream, Future, SinkExt, StreamExt, TryStreamExt};
+use test_programs::wasi::http::types::{
     Fields, IncomingRequest, IncomingResponse, Method, OutgoingBody, OutgoingRequest,
     OutgoingResponse, ResponseOutparam, Scheme,
 };
-use futures::{future, stream, Future, SinkExt, StreamExt, TryStreamExt};
 use url::Url;
 
-mod bindings {
-    wit_bindgen::generate!({
-        path: "../wasi-http/wit",
-        world: "wasi:http/proxy",
-        default_bindings_module: "bindings",
-    });
-}
-
 const MAX_CONCURRENCY: usize = 16;
 
 struct Handler;
 
-bindings::export!(Handler);
+test_programs::proxy::export!(Handler);
 
-impl bindings::exports::wasi::http::incoming_handler::Guest for Handler {
+impl test_programs::proxy::exports::wasi::http::incoming_handler::Guest for Handler {
     fn handle(request: IncomingRequest, response_out: ResponseOutparam) {
         executor::run(async move {
             handle_request(request, response_out).await;
@@ -312,16 +304,6 @@ async fn hash(url: &Url) -> Result<String> {
 fn main() {}
 
 mod executor {
-    use super::bindings::wasi::{
-        http::{
-            outgoing_handler,
-            types::{
-                self, IncomingBody, IncomingResponse, InputStream, OutgoingBody, OutgoingRequest,
-                OutputStream,
-            },
-        },
-        io::{self, streams::StreamError},
-    };
     use anyhow::{anyhow, Error, Result};
     use futures::{future, sink, stream, Sink, Stream};
     use std::{
@@ -332,6 +314,16 @@ mod executor {
         sync::{Arc, Mutex},
         task::{Context, Poll, Wake, Waker},
     };
+    use test_programs::wasi::{
+        http::{
+            outgoing_handler,
+            types::{
+                self, IncomingBody, IncomingResponse, InputStream, OutgoingBody, OutgoingRequest,
+                OutputStream,
+            },
+        },
+        io::{self, streams::StreamError},
+    };
 
     const READ_SIZE: u64 = 16 * 1024;
 

crates/test-programs/src/bin/cli_serve_echo_env.rs

@@ -0,0 +1,26 @@
+use test_programs::proxy;
+use test_programs::wasi::http::types::{
+    Fields, IncomingRequest, OutgoingResponse, ResponseOutparam,
+};
+
+struct T;
+
+proxy::export!(T);
+
+impl proxy::exports::wasi::http::incoming_handler::Guest for T {
+    fn handle(request: IncomingRequest, outparam: ResponseOutparam) {
+        let headers = request.headers();
+        let header_key = "env".to_string();
+        let env_var = headers.get(&header_key);
+        assert!(env_var.len() == 1, "should have exactly one `env` header");
+        let key = std::str::from_utf8(&env_var[0]).unwrap();
+        let fields = Fields::new();
+        if let Ok(val) = std::env::var(key) {
+            fields.set(&header_key, &[val.into_bytes()]).unwrap();
+        }
+        let resp = OutgoingResponse::new(fields);
+        ResponseOutparam::set(outparam, Ok(resp));
+    }
+}
+
+fn main() {}

crates/test-programs/src/lib.rs

@@ -3,3 +3,16 @@ pub mod preview1;
 pub mod sockets;
 
 wit_bindgen::generate!("test-command" in "../wasi/wit");
+
+pub mod proxy {
+    wit_bindgen::generate!({
+        path: "../wasi-http/wit",
+        world: "wasi:http/proxy",
+        default_bindings_module: "test_programs::proxy",
+        pub_export_macro: true,
+        with: {
+            "wasi:http/types@0.2.0": crate::wasi::http::types,
+            "wasi:http/outgoing-handler@0.2.0": crate::wasi::http::outgoing_handler,
+        },
+    });
+}

src/commands/run.rs

@@ -26,24 +26,6 @@ use wasmtime_wasi_threads::WasiThreadsCtx;
 #[cfg(feature = "wasi-http")]
 use wasmtime_wasi_http::WasiHttpCtx;
 
-fn parse_env_var(s: &str) -> Result<(String, Option<String>)> {
-    let mut parts = s.splitn(2, '=');
-    Ok((
-        parts.next().unwrap().to_string(),
-        parts.next().map(|s| s.to_string()),
-    ))
-}
-
-fn parse_dirs(s: &str) -> Result<(String, String)> {
-    let mut parts = s.split("::");
-    let host = parts.next().unwrap();
-    let guest = match parts.next() {
-        Some(guest) => guest,
-        None => host,
-    };
-    Ok((host.into(), guest.into()))
-}
-
 fn parse_preloads(s: &str) -> Result<(String, PathBuf)> {
     let parts: Vec<&str> = s.splitn(2, '=').collect();
     if parts.len() != 2 {
@@ -59,25 +41,6 @@ pub struct RunCommand {
     #[allow(missing_docs)]
     pub run: RunCommon,
 
-    /// Grant access of a host directory to a guest.
-    ///
-    /// If specified as just `HOST_DIR` then the same directory name on the
-    /// host is made available within the guest. If specified as `HOST::GUEST`
-    /// then the `HOST` directory is opened and made available as the name
-    /// `GUEST` in the guest.
-    #[arg(long = "dir", value_name = "HOST_DIR[::GUEST_DIR]", value_parser = parse_dirs)]
-    pub dirs: Vec<(String, String)>,
-
-    /// Pass an environment variable to the program.
-    ///
-    /// The `--env FOO=BAR` form will set the environment variable named `FOO`
-    /// to the value `BAR` for the guest program using WASI. The `--env FOO`
-    /// form will set the environment variable named `FOO` to the same value it
-    /// has in the calling process for the guest, or in other words it will
-    /// cause the environment variable `FOO` to be inherited.
-    #[arg(long = "env", number_of_values = 1, value_name = "NAME[=VAL]", value_parser = parse_env_var)]
-    pub vars: Vec<(String, Option<String>)>,
-
     /// The name of the function to run
     #[arg(long, value_name = "FUNCTION")]
     pub invoke: Option<String>,
@@ -263,20 +226,6 @@ impl RunCommand {
         Ok(())
     }
 
-    fn compute_preopen_sockets(&self) -> Result<Vec<TcpListener>> {
-        let mut listeners = vec![];
-
-        for address in &self.run.common.wasi.tcplisten {
-            let stdlistener = std::net::TcpListener::bind(address)
-                .with_context(|| format!("failed to bind to address '{}'", address))?;
-
-            let _ = stdlistener.set_nonblocking(true)?;
-
-            listeners.push(TcpListener::from_std(stdlistener))
-        }
-        Ok(listeners)
-    }
-
     fn compute_argv(&self) -> Result<Vec<String>> {
         let mut result = Vec::new();
 
@@ -762,7 +711,7 @@ impl RunCommand {
                 builder.env(&k, &v)?;
             }
         }
-        for (key, value) in self.vars.iter() {
+        for (key, value) in self.run.vars.iter() {
             let value = match value {
                 Some(value) => value.clone(),
                 None => match std::env::var_os(key) {
@@ -784,12 +733,13 @@ impl RunCommand {
             num_fd = ctx_set_listenfd(num_fd, &mut builder)?;
         }
 
-        for listener in self.compute_preopen_sockets()? {
+        for listener in self.run.compute_preopen_sockets()? {
+            let listener = TcpListener::from_std(listener);
             builder.preopened_socket(num_fd as _, listener)?;
             num_fd += 1;
         }
 
-        for (host, guest) in self.dirs.iter() {
+        for (host, guest) in self.run.dirs.iter() {
             let dir = Dir::open_ambient_dir(host, ambient_authority())
                 .with_context(|| format!("failed to open directory '{}'", host))?;
             builder.preopened_dir(dir, guest)?;
@@ -802,63 +752,7 @@ impl RunCommand {
     fn set_preview2_ctx(&self, store: &mut Store<Host>) -> Result<()> {
         let mut builder = wasmtime_wasi::WasiCtxBuilder::new();
         builder.inherit_stdio().args(&self.compute_argv()?);
-
+        self.run.configure_wasip2(&mut builder)?;
-        // It's ok to block the current thread since we're the only thread in
-        // the program as the CLI. This helps improve the performance of some
-        // blocking operations in WASI, for example, by skipping the
-        // back-and-forth between sync and async.
-        builder.allow_blocking_current_thread(true);
-
-        if self.run.common.wasi.inherit_env == Some(true) {
-            for (k, v) in std::env::vars() {
-                builder.env(&k, &v);
-            }
-        }
-        for (key, value) in self.vars.iter() {
-            let value = match value {
-                Some(value) => value.clone(),
-                None => match std::env::var_os(key) {
-                    Some(val) => val
-                        .into_string()
-                        .map_err(|_| anyhow!("environment variable `{key}` not valid utf-8"))?,
-                    None => {
-                        // leave the env var un-set in the guest
-                        continue;
-                    }
-                },
-            };
-            builder.env(key, &value);
-        }
-
-        if self.run.common.wasi.listenfd == Some(true) {
-            bail!("components do not support --listenfd");
-        }
-        for _ in self.compute_preopen_sockets()? {
-            bail!("components do not support --tcplisten");
-        }
-
-        for (host, guest) in self.dirs.iter() {
-            builder.preopened_dir(
-                host,
-                guest,
-                wasmtime_wasi::DirPerms::all(),
-                wasmtime_wasi::FilePerms::all(),
-            )?;
-        }
-
-        if self.run.common.wasi.inherit_network == Some(true) {
-            builder.inherit_network();
-        }
-        if let Some(enable) = self.run.common.wasi.allow_ip_name_lookup {
-            builder.allow_ip_name_lookup(enable);
-        }
-        if let Some(enable) = self.run.common.wasi.tcp {
-            builder.allow_tcp(enable);
-        }
-        if let Some(enable) = self.run.common.wasi.udp {
-            builder.allow_udp(enable);
-        }
-
         let ctx = builder.build_p1();
         store.data_mut().preview2_ctx = Some(Arc::new(Mutex::new(ctx)));
         Ok(())

src/commands/serve.rs

@@ -132,8 +132,9 @@ impl ServeCommand {
 
     fn new_store(&self, engine: &Engine, req_id: u64) -> Result<Store<Host>> {
         let mut builder = WasiCtxBuilder::new();
+        self.run.configure_wasip2(&mut builder)?;
 
-        builder.envs(&[("REQUEST_ID", req_id.to_string())]);
+        builder.env("REQUEST_ID", req_id.to_string());
 
         builder.stdout(LogStream {
             prefix: format!("stdout [{req_id}] :: "),

src/common.rs

@@ -1,10 +1,12 @@
 //! Common functionality shared between command implementations.
 
-use anyhow::{bail, Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 use clap::Parser;
+use std::net::TcpListener;
 use std::{path::Path, time::Duration};
 use wasmtime::{Engine, Module, Precompiled, StoreLimits, StoreLimitsBuilder};
 use wasmtime_cli_flags::{opt::WasmtimeOptionValue, CommonOptions};
+use wasmtime_wasi::WasiCtxBuilder;
 
 #[cfg(feature = "component-model")]
 use wasmtime::component::Component;
@@ -70,6 +72,43 @@ pub struct RunCommon {
         value_parser = Profile::parse,
     )]
     pub profile: Option<Profile>,
+
+    /// Grant access of a host directory to a guest.
+    ///
+    /// If specified as just `HOST_DIR` then the same directory name on the
+    /// host is made available within the guest. If specified as `HOST::GUEST`
+    /// then the `HOST` directory is opened and made available as the name
+    /// `GUEST` in the guest.
+    #[arg(long = "dir", value_name = "HOST_DIR[::GUEST_DIR]", value_parser = parse_dirs)]
+    pub dirs: Vec<(String, String)>,
+
+    /// Pass an environment variable to the program.
+    ///
+    /// The `--env FOO=BAR` form will set the environment variable named `FOO`
+    /// to the value `BAR` for the guest program using WASI. The `--env FOO`
+    /// form will set the environment variable named `FOO` to the same value it
+    /// has in the calling process for the guest, or in other words it will
+    /// cause the environment variable `FOO` to be inherited.
+    #[arg(long = "env", number_of_values = 1, value_name = "NAME[=VAL]", value_parser = parse_env_var)]
+    pub vars: Vec<(String, Option<String>)>,
+}
+
+fn parse_env_var(s: &str) -> Result<(String, Option<String>)> {
+    let mut parts = s.splitn(2, '=');
+    Ok((
+        parts.next().unwrap().to_string(),
+        parts.next().map(|s| s.to_string()),
+    ))
+}
+
+fn parse_dirs(s: &str) -> Result<(String, String)> {
+    let mut parts = s.split("::");
+    let host = parts.next().unwrap();
+    let guest = match parts.next() {
+        Some(guest) => guest,
+        None => host,
+    };
+    Ok((host.into(), guest.into()))
 }
 
 impl RunCommon {
@@ -222,6 +261,80 @@ impl RunCommon {
             }
         })
     }
+
+    pub fn configure_wasip2(&self, builder: &mut WasiCtxBuilder) -> Result<()> {
+        // It's ok to block the current thread since we're the only thread in
+        // the program as the CLI. This helps improve the performance of some
+        // blocking operations in WASI, for example, by skipping the
+        // back-and-forth between sync and async.
+        builder.allow_blocking_current_thread(true);
+
+        if self.common.wasi.inherit_env == Some(true) {
+            for (k, v) in std::env::vars() {
+                builder.env(&k, &v);
+            }
+        }
+        for (key, value) in self.vars.iter() {
+            let value = match value {
+                Some(value) => value.clone(),
+                None => match std::env::var_os(key) {
+                    Some(val) => val
+                        .into_string()
+                        .map_err(|_| anyhow!("environment variable `{key}` not valid utf-8"))?,
+                    None => {
+                        // leave the env var un-set in the guest
+                        continue;
+                    }
+                },
+            };
+            builder.env(key, &value);
+        }
+
+        for (host, guest) in self.dirs.iter() {
+            builder.preopened_dir(
+                host,
+                guest,
+                wasmtime_wasi::DirPerms::all(),
+                wasmtime_wasi::FilePerms::all(),
+            )?;
+        }
+
+        if self.common.wasi.listenfd == Some(true) {
+            bail!("components do not support --listenfd");
+        }
+        for _ in self.compute_preopen_sockets()? {
+            bail!("components do not support --tcplisten");
+        }
+
+        if self.common.wasi.inherit_network == Some(true) {
+            builder.inherit_network();
+        }
+        if let Some(enable) = self.common.wasi.allow_ip_name_lookup {
+            builder.allow_ip_name_lookup(enable);
+        }
+        if let Some(enable) = self.common.wasi.tcp {
+            builder.allow_tcp(enable);
+        }
+        if let Some(enable) = self.common.wasi.udp {
+            builder.allow_udp(enable);
+        }
+
+        Ok(())
+    }
+
+    pub fn compute_preopen_sockets(&self) -> Result<Vec<TcpListener>> {
+        let mut listeners = vec![];
+
+        for address in &self.common.wasi.tcplisten {
+            let stdlistener = std::net::TcpListener::bind(address)
+                .with_context(|| format!("failed to bind to address '{}'", address))?;
+
+            let _ = stdlistener.set_nonblocking(true)?;
+
+            listeners.push(stdlistener)
+        }
+        Ok(listeners)
+    }
 }
 
 #[derive(Clone, PartialEq)]

src/old_cli.rs

@@ -816,14 +816,14 @@ impl RunCommand {
             common,
             allow_precompiled,
             profile: profile.map(|p| p.convert()),
+            dirs,
+            vars,
         };
 
         let mut module_and_args = vec![module.into()];
         module_and_args.extend(module_args.into_iter().map(|s| s.into()));
         crate::commands::RunCommand {
             run,
-            dirs,
-            vars,
             invoke,
             preloads,
             module_and_args,

tests/all/cli_tests.rs

@@ -1260,10 +1260,14 @@ fn mpk_without_pooling() -> Result<()> {
 
 mod test_programs {
     use super::{get_wasmtime_command, run_wasmtime};
-    use anyhow::Result;
+    use anyhow::{bail, Context, Result};
-    use std::io::{Read, Write};
+    use http_body_util::BodyExt;
-    use std::process::Stdio;
+    use hyper::header::HeaderValue;
+    use std::io::{BufRead, BufReader, Read, Write};
+    use std::net::SocketAddr;
+    use std::process::{Child, Command, Stdio};
     use test_programs_artifacts::*;
+    use tokio::net::TcpStream;
 
     macro_rules! assert_test_exists {
         ($name:ident) => {
@@ -1603,6 +1607,174 @@ mod test_programs {
         run_wasmtime(&["run", CLI_SLEEP_COMPONENT])?;
         Ok(())
     }
+
+    /// Helper structure to manage an invocation of `wasmtime serve`
+    struct WasmtimeServe {
+        child: Option<Child>,
+        addr: SocketAddr,
+    }
+
+    impl WasmtimeServe {
+        /// Creates a new server which will serve the wasm component pointed to
+        /// by `wasm`.
+        ///
+        /// A `configure` callback is provided to specify how `wasmtime serve`
+        /// will be invoked and configure arguments such as headers.
+        fn new(wasm: &str, configure: impl FnOnce(&mut Command)) -> Result<WasmtimeServe> {
+            // Spawn `wasmtime serve` on port 0 which will randomly assign it a
+            // port.
+            let mut cmd = super::get_wasmtime_command()?;
+            cmd.stdin(Stdio::null());
+            cmd.stdout(Stdio::piped());
+            cmd.stderr(Stdio::piped());
+            cmd.arg("serve").arg("--addr=127.0.0.1:0").arg(wasm);
+            configure(&mut cmd);
+            let mut child = cmd.spawn()?;
+
+            // Read the first line of stderr which will say which address it's
+            // listening on.
+            //
+            // NB: this intentionally discards any extra buffered data in the
+            // `BufReader` once the newline is found. The server shouldn't print
+            // anything interesting other than the address so once we get a line
+            // all remaining output is left to be captured by future requests
+            // send to the server.
+            let mut line = String::new();
+            BufReader::new(child.stderr.as_mut().unwrap()).read_line(&mut line)?;
+            let addr_start = line.find("127.0.0.1").unwrap();
+            let addr = &line[addr_start..];
+            let addr_end = addr.find("/").unwrap();
+            let addr = &addr[..addr_end];
+            Ok(WasmtimeServe {
+                child: Some(child),
+                addr: addr.parse().unwrap(),
+            })
+        }
+
+        /// Completes this server gracefully by printing the output on failure.
+        fn finish(mut self) -> Result<()> {
+            let mut child = self.child.take().unwrap();
+
+            // If the child process has already exited then collect the output
+            // and test if it succeeded. Otherwise it's still running so kill it
+            // and then reap it. Assume that if it's still running then the test
+            // has otherwise passed so no need to print the output.
+            if child.try_wait()?.is_some() {
+                let output = child.wait_with_output()?;
+                if output.status.success() {
+                    return Ok(());
+                }
+                bail!("child failed {output:?}");
+            } else {
+                child.kill()?;
+                child.wait_with_output()?;
+            }
+            Ok(())
+        }
+
+        /// Send a request to this server and wait for the response.
+        async fn send_request(&self, req: http::Request<String>) -> Result<http::Response<String>> {
+            let tcp = TcpStream::connect(&self.addr)
+                .await
+                .context("failed to connect")?;
+            let tcp = wasmtime_wasi_http::io::TokioIo::new(tcp);
+            let (mut send, conn) = hyper::client::conn::http1::handshake(tcp)
+                .await
+                .context("failed http handshake")?;
+
+            let conn_task = tokio::task::spawn(conn);
+
+            let response = send
+                .send_request(req)
+                .await
+                .context("error sending request")?;
+            drop(send);
+            let (parts, body) = response.into_parts();
+
+            let body = body.collect().await.context("failed to read body")?;
+            assert!(body.trailers().is_none());
+            let body = std::str::from_utf8(&body.to_bytes())?.to_string();
+
+            conn_task.await??;
+
+            Ok(http::Response::from_parts(parts, body))
+        }
+    }
+
+    // Don't leave child processes running by accident so kill the child process
+    // if our server goes away.
+    impl Drop for WasmtimeServe {
+        fn drop(&mut self) {
+            let mut child = match self.child.take() {
+                Some(child) => child,
+                None => return,
+            };
+            if child.kill().is_err() {
+                return;
+            }
+            let output = match child.wait_with_output() {
+                Ok(output) => output,
+                Err(_) => return,
+            };
+
+            println!("server status: {}", output.status);
+            if !output.stdout.is_empty() {
+                println!(
+                    "server stdout:\n{}",
+                    String::from_utf8_lossy(&output.stdout)
+                );
+            }
+            if !output.stderr.is_empty() {
+                println!(
+                    "server stderr:\n{}",
+                    String::from_utf8_lossy(&output.stderr)
+                );
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn cli_serve_echo_env() -> Result<()> {
+        let server = WasmtimeServe::new(CLI_SERVE_ECHO_ENV_COMPONENT, |cmd| {
+            cmd.arg("--env=FOO=bar");
+            cmd.arg("--env=BAR");
+            cmd.arg("-Scli");
+            cmd.env_remove("BAR");
+        })?;
+
+        let foo_env = server
+            .send_request(
+                hyper::Request::builder()
+                    .uri("http://localhost/")
+                    .header("env", "FOO")
+                    .body(String::new())
+                    .context("failed to make request")?,
+            )
+            .await?;
+
+        assert!(foo_env.status().is_success());
+        assert!(foo_env.body().is_empty());
+        let headers = foo_env.headers();
+        assert_eq!(headers.get("env"), Some(&HeaderValue::from_static("bar")));
+
+        let bar_env = server
+            .send_request(
+                hyper::Request::builder()
+                    .uri("http://localhost/")
+                    .header("env", "BAR")
+                    .body(String::new())
+                    .context("failed to make request")?,
+            )
+            .await?;
+
+        assert!(bar_env.status().is_success());
+        assert!(bar_env.body().is_empty());
+        let headers = bar_env.headers();
+        assert_eq!(headers.get("env"), None);
+
+        server.finish()?;
+        Ok(())
+    }
 }
 
 #[test]