wasmtime: add CLI options for pre-opened TCP listen sockets (#3729)

This patch  implements CLI options to insert pre-opened sockets.

`--listenfd` : Inherit environment variables and file descriptors following
               the systemd listen fd specification (UNIX only).

`--tcplisten <SOCKET ADDRESS>`: Grant access to the given TCP listen socket.

Signed-off-by: Harald Hoyer <harald@profian.com>

Cargo.lock

@@ -1585,6 +1585,17 @@ version = "0.0.37"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "95f5690fef754d905294c56f7ac815836f2513af966aa47f2e07ac79be07827f"
 
+[[package]]
+name = "listenfd"
+version = "0.3.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "809e514e2cb8a9624701346ea3e694c1766d76778e343e537d873c1c366e79a7"
+dependencies = [
+ "libc",
+ "uuid",
+ "winapi",
+]
+
 [[package]]
 name = "lock_api"
 version = "0.4.4"
@@ -3472,6 +3483,7 @@ dependencies = [
  "humantime 2.1.0",
  "lazy_static",
  "libc",
+ "listenfd",
  "memchr",
  "more-asserts",
  "num_cpus",

Cargo.toml

@@ -39,6 +39,7 @@ rayon = "1.5.0"
 humantime = "2.0.0"
 wasmparser = "0.82.0"
 lazy_static = "1.4.0"
+listenfd = "0.3.5"
 
 [target.'cfg(unix)'.dependencies]
 rustix = "0.33.0"

src/commands/run.rs

@@ -13,7 +13,7 @@ use std::{
 };
 use structopt::{clap::AppSettings, StructOpt};
 use wasmtime::{Engine, Func, Linker, Module, Store, Trap, Val, ValType};
-use wasmtime_wasi::sync::{ambient_authority, Dir, WasiCtxBuilder};
+use wasmtime_wasi::sync::{ambient_authority, Dir, TcpListener, WasiCtxBuilder};
 
 #[cfg(feature = "wasi-nn")]
 use wasmtime_wasi_nn::WasiNnCtx;
@@ -91,6 +91,19 @@ pub struct RunCommand {
     #[structopt(long = "allow-precompiled")]
     allow_precompiled: bool,
 
+    /// Inherit environment variables and file descriptors following the
+    /// systemd listen fd specification (UNIX only)
+    #[structopt(long = "listenfd")]
+    listenfd: bool,
+
+    /// Grant access to the given TCP listen socket
+    #[structopt(
+        long = "tcplisten",
+        number_of_values = 1,
+        value_name = "SOCKET ADDRESS"
+    )]
+    tcplisten: Vec<String>,
+
     /// Grant access to the given host directory
     #[structopt(long = "dir", number_of_values = 1, value_name = "DIRECTORY")]
     dirs: Vec<String>,
@@ -151,6 +164,8 @@ impl RunCommand {
         let engine = Engine::new(&config)?;
         let mut store = Store::new(&engine, Host::default());
 
+        let preopen_sockets = self.compute_preopen_sockets()?;
+
         // Make wasi available by default.
         let preopen_dirs = self.compute_preopen_dirs()?;
         let argv = self.compute_argv();
@@ -165,6 +180,8 @@ impl RunCommand {
             &argv,
             &self.vars,
             &self.common.wasi_modules.unwrap_or(WasiModules::default()),
+            self.listenfd,
+            preopen_sockets,
         )?;
 
         // Load the preload wasm modules.
@@ -243,6 +260,20 @@ impl RunCommand {
         Ok(preopen_dirs)
     }
 
+    fn compute_preopen_sockets(&self) -> Result<Vec<TcpListener>> {
+        let mut listeners = vec![];
+
+        for address in &self.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) -> Vec<String> {
         let mut result = Vec::new();
 
@@ -415,6 +446,8 @@ fn populate_with_wasi(
     argv: &[String],
     vars: &[(String, String)],
     wasi_modules: &WasiModules,
+    listenfd: bool,
+    mut tcplisten: Vec<TcpListener>,
 ) -> Result<()> {
     if wasi_modules.wasi_common {
         wasmtime_wasi::add_to_linker(linker, |host| host.wasi.as_mut().unwrap())?;
@@ -422,9 +455,23 @@ fn populate_with_wasi(
         let mut builder = WasiCtxBuilder::new();
         builder = builder.inherit_stdio().args(argv)?.envs(vars)?;
 
+        let mut num_fd: usize = 3;
+
+        if listenfd {
+            let (n, b) = ctx_set_listenfd(num_fd, builder)?;
+            num_fd = n;
+            builder = b;
+        }
+
+        for listener in tcplisten.drain(..) {
+            builder = builder.preopened_socket(num_fd as _, listener)?;
+            num_fd += 1;
+        }
+
         for (name, dir) in preopen_dirs.into_iter() {
             builder = builder.preopened_dir(dir, name)?;
         }
+
         store.data_mut().wasi = Some(builder.build());
     }
 
@@ -454,3 +501,35 @@ fn populate_with_wasi(
 
     Ok(())
 }
+
+#[cfg(not(unix))]
+fn ctx_set_listenfd(num_fd: usize, builder: WasiCtxBuilder) -> Result<(usize, WasiCtxBuilder)> {
+    Ok((num_fd, builder))
+}
+
+#[cfg(unix)]
+fn ctx_set_listenfd(num_fd: usize, builder: WasiCtxBuilder) -> Result<(usize, WasiCtxBuilder)> {
+    use listenfd::ListenFd;
+
+    let mut builder = builder;
+    let mut num_fd = num_fd;
+
+    for env in ["LISTEN_FDS", "LISTEN_FDNAMES"] {
+        if let Ok(val) = std::env::var(env) {
+            builder = builder.env(env, &val)?;
+        }
+    }
+
+    let mut listenfd = ListenFd::from_env();
+
+    for i in 0..listenfd.len() {
+        if let Some(stdlistener) = listenfd.take_tcp_listener(i)? {
+            let _ = stdlistener.set_nonblocking(true)?;
+            let listener = TcpListener::from_std(stdlistener);
+            builder = builder.preopened_socket((3 + i) as _, listener)?;
+            num_fd = 3 + i;
+        }
+    }
+
+    Ok((num_fd, builder))
+}