Support dirs/env vars in wasmtime serve

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

@@ -1,18 +1,12 @@
-pub mod bindings {
-    wit_bindgen::generate!({
-        path: "../wasi-http/wit",
-        world: "wasi:http/proxy",
-        default_bindings_module: "bindings",
-    });
-}
-
-use bindings::wasi::http::types::{IncomingRequest, ResponseOutparam};
+use test_programs::wasi::http::types::{
+    Headers, IncomingRequest, OutgoingBody, OutgoingResponse, 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 resp = bindings::wasi::http::types::OutgoingResponse::new(hdrs);
+        let hdrs = Headers::new();
+        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)
-            .expect("outgoing-body.finish");
+        OutgoingBody::finish(body, None).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()?);
-
-        // 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);
-        }
-
+        self.run.configure_wasip2(&mut builder)?;
         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}] :: "),