Import sd listen

Signed-off-by: Harald Hoyer <[email protected]>

Author: haraldh

Cargo.lock

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

Cargo.toml

@@ -43,12 +43,13 @@ pub mod tcpstream;
 
 pub use cap_std::ambient_authority;
 pub use cap_std::fs::Dir;
+pub use cap_std::net::TcpListener;
 pub use clocks::clocks_ctx;
 pub use sched::sched_ctx;
 
 use cap_rand::RngCore;
 use std::path::Path;
-use wasi_common::{table::Table, Error, WasiCtx, WasiFile};
+use wasi_common::{file::FileCaps, table::Table, Error, WasiCtx, WasiFile};
 
 pub struct WasiCtxBuilder(WasiCtx);
 
@@ -122,6 +123,18 @@ impl WasiCtxBuilder {
         self.0.push_preopened_dir(dir, guest_path)?;
         Ok(self)
     }
+    pub fn preopened_listener(mut self, fd: u32, listener: TcpListener) -> Result<Self, Error> {
+        let file: Box<dyn WasiFile> =
+            Box::new(crate::tcplisten::TcpListener::from_cap_std(listener));
+
+        let caps = FileCaps::FDSTAT_SET_FLAGS
+            | FileCaps::FILESTAT_GET
+            | FileCaps::READ
+            | FileCaps::POLL_READWRITE;
+
+        self.0.insert_file(fd, file, caps);
+        Ok(self)
+    }
     pub fn build(self) -> WasiCtx {
         self.0
     }

crates/wasi-common/cap-std-sync/src/lib.rs

@@ -14,6 +14,9 @@ use wasi_common::{Error, Table, WasiCtx, WasiFile};
 
 pub use dir::Dir;
 pub use file::File;
+pub use tcplisten::TcpListener;
+pub use tcpstream::TcpStream;
+use wasi_common::file::FileCaps;
 
 use crate::sched::sched_ctx;
 
@@ -93,6 +96,22 @@ impl WasiCtxBuilder {
         self.0.push_preopened_dir(dir, guest_path)?;
         Ok(self)
     }
+    pub fn preopened_listener(
+        mut self,
+        fd: u32,
+        listener: cap_std::net::TcpListener,
+    ) -> Result<Self, Error> {
+        let file: Box<dyn WasiFile> = Box::new(TcpListener::from_cap_std(listener));
+
+        let caps = FileCaps::FDSTAT_SET_FLAGS
+            | FileCaps::FILESTAT_GET
+            | FileCaps::READ
+            | FileCaps::POLL_READWRITE;
+
+        self.0.insert_file(fd, file, caps);
+        Ok(self)
+    }
+
     pub fn build(self) -> WasiCtx {
         self.0
     }

crates/wasi-common/tokio/src/lib.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>,
@@ -165,6 +178,8 @@ impl RunCommand {
             &argv,
             &self.vars,
             &self.common.wasi_modules.unwrap_or(WasiModules::default()),
+            self.listenfd,
+            &self.tcplisten,
         )?;
 
         // Load the preload wasm modules.
@@ -415,16 +430,39 @@ fn populate_with_wasi(
     argv: &[String],
     vars: &[(String, String)],
     wasi_modules: &WasiModules,
+    listenfd: bool,
+    tcplisten: &Vec<String>,
 ) -> Result<()> {
     if wasi_modules.wasi_common {
         wasmtime_wasi::add_to_linker(linker, |host| host.wasi.as_mut().unwrap())?;
 
         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 address in tcplisten.iter() {
+            let stdlistener = std::net::TcpListener::bind(address)
+                .with_context(|| format!("failed to bind to address '{}'", address))?;
+
+            let _ = stdlistener.set_nonblocking(true)?;
+
+            let listener = TcpListener::from_std(stdlistener);
+
+            builder = builder.preopened_listener(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 +492,55 @@ fn populate_with_wasi(
 
     Ok(())
 }
+
+#[cfg(not(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;
+
+    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_listener((3 + i) as _, listener)?;
+            num_fd = 3 + i;
+        }
+    }
+
+    if num_fd > 3 {
+        builder = builder.env("LISTEN_FDS", &(num_fd - 3).to_string())?;
+    }
+
+    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_listener((3 + i) as _, listener)?;
+            num_fd = 3 + i;
+        }
+    }
+
+    Ok((num_fd, builder))
+}