wasmtime serve: Allow inheriting sockets from system manager

Cargo.toml

@@ -106,6 +106,7 @@ pin-project-lite = { workspace = true, optional = true }
 
 [target.'cfg(unix)'.dependencies]
 rustix = { workspace = true, features = ["mm", "process"] }
+libc = { workspace = true, optional = true }
 
 [dev-dependencies]
 env_logger = { workspace = true }
@@ -617,6 +618,7 @@ serve = [
   "dep:http-body-util",
   "dep:http",
   "dep:pin-project-lite",
+  "dep:libc",
   "wasmtime-cli-flags/async",
   "wasmtime-wasi-http?/p2",
 ]

RELEASES.md

@@ -4,8 +4,13 @@ Unreleased.
 
 ### Added
 
+- Add `--listenfd` option to `wasmtime serve`, which allows launching wasmtime
+  with sockets inherited from a service manager (e.g. systemd socket units).
+
 ### Changed
 
+- Remove non-functional `listenfd` WASI CLI option.
+
 --------------------------------------------------------------------------------
 
 Release notes for previous releases of Wasmtime can be found on the respective

crates/cli-flags/src/lib.rs

@@ -516,10 +516,6 @@ wasmtime_option_group! {
         pub config: Option<bool>,
         /// Enable support for WASI key-value imports (experimental)
         pub keyvalue: Option<bool>,
-        /// Inherit environment variables and file descriptors following the
-        /// systemd listen fd specification (UNIX only) (legacy wasip1
-        /// implementation only)
-        pub listenfd: Option<bool>,
         /// Grant access to the given TCP listen socket (experimental, legacy
         /// wasip1 implementation only)
         #[serde(default)]

src/bin/wasmtime.rs

@@ -176,9 +176,22 @@ impl CompletionCommand {
 
 #[allow(unreachable_code, reason = "empty enum with all features disabled")]
 fn main() -> Result<()> {
+    setup();
+
     return Wasmtime::parse().execute();
 }
 
+#[cfg(all(unix, feature = "serve"))]
+fn setup() {
+    unsafe {
+        // Safety: This is called first in main
+        wasmtime_cli::init_inherited_fds()
+    }
+}
+
+#[cfg(not(all(unix, feature = "serve")))]
+fn setup() {}
+
 #[test]
 fn verify_cli() {
     use clap::CommandFactory;

src/commands/serve.rs

@@ -18,6 +18,7 @@ use std::{
     time::{Duration, Instant},
 };
 use tokio::io::{self, AsyncWrite};
+use tokio::net::TcpListener;
 use tokio::sync::{Notify, Semaphore};
 use wasmtime::component::{Component, GuestTaskId, Linker};
 use wasmtime::error::Context as _;
@@ -121,6 +122,11 @@ pub struct ServeCommand {
     #[arg(long)]
     no_logging_prefix: bool,
 
+    /// Use sockets passed via the 'LISTEN_FDS' environment variable (set e.g. by systemd when
+    /// launching a service from socket units). Not available on Windows.
+    #[arg(long)]
+    listenfd: bool,
+
     /// The WebAssembly component to run.
     #[arg(value_name = "WASM", required = true)]
     component: PathBuf,
@@ -615,25 +621,41 @@ impl ServeCommand {
             });
         }
 
-        let socket = match &self.addr {
-            SocketAddr::V4(_) => tokio::net::TcpSocket::new_v4()?,
-            SocketAddr::V6(_) => tokio::net::TcpSocket::new_v6()?,
+        let inherited_socket = if self.listenfd {
+            Self::inherit_socket().with_context(|| "Resolve inherited socket")?
+        } else {
+            None
+        };
+
+        let listener = match inherited_socket {
+            Some(listener) => {
+                eprintln!("Serving HTTP on inherited socket");
+                log::info!("Listening on inherited socket");
+
+                listener
+            }
+            None => {
+                let socket = match &self.addr {
+                    SocketAddr::V4(_) => tokio::net::TcpSocket::new_v4()?,
+                    SocketAddr::V6(_) => tokio::net::TcpSocket::new_v6()?,
+                };
+                // Conditionally enable `SO_REUSEADDR` depending on the current
+                // platform. On Unix we want this to be able to rebind an address in
+                // the `TIME_WAIT` state which can happen then a server is killed with
+                // active TCP connections and then restarted. On Windows though if
+                // `SO_REUSEADDR` is specified then it enables multiple applications to
+                // bind the port at the same time which is not something we want. Hence
+                // this is conditionally set based on the platform (and deviates from
+                // Tokio's default from always-on).
+                socket.set_reuseaddr(!cfg!(windows))?;
+                socket.bind(self.addr)?;
+                let listener = socket.listen(100)?;
+
+                eprintln!("Serving HTTP on http://{}/", listener.local_addr()?);
+                log::info!("Listening on {}", self.addr);
+                listener
+            }
         };
-        // Conditionally enable `SO_REUSEADDR` depending on the current
-        // platform. On Unix we want this to be able to rebind an address in
-        // the `TIME_WAIT` state which can happen then a server is killed with
-        // active TCP connections and then restarted. On Windows though if
-        // `SO_REUSEADDR` is specified then it enables multiple applications to
-        // bind the port at the same time which is not something we want. Hence
-        // this is conditionally set based on the platform (and deviates from
-        // Tokio's default from always-on).
-        socket.set_reuseaddr(!cfg!(windows))?;
-        socket.bind(self.addr)?;
-        let listener = socket.listen(100)?;
-
-        eprintln!("Serving HTTP on http://{}/", listener.local_addr()?);
-
-        log::info!("Listening on {}", self.addr);
 
         let epoch_interval = if let Some(Profile::Guest { interval, .. }) = self.run.profile {
             Some(interval)
@@ -753,6 +775,28 @@ impl ServeCommand {
 
         Ok(())
     }
+
+    #[cfg(unix)]
+    fn inherit_socket() -> Result<Option<TcpListener>> {
+        use crate::inherited_fd::{InheritedFileDescriptor, take_inherited_fds};
+
+        let sockets = take_inherited_fds();
+        let Some(inherited_socket) = ({
+            sockets.into_iter().find_map(|fd| match fd {
+                InheritedFileDescriptor::TcpSocket(listener) => Some(listener),
+            })
+        }) else {
+            return Ok(None);
+        };
+
+        inherited_socket.set_nonblocking(true)?;
+        Ok(Some(TcpListener::from_std(inherited_socket)?))
+    }
+
+    #[cfg(not(unix))]
+    fn inherit_socket() -> Result<Option<TcpListener>> {
+        bail!("The --listenfd option is not available on Windows")
+    }
 }
 
 pin_project! {

src/common.rs

@@ -356,9 +356,6 @@ impl RunCommon {
             builder.initial_cwd(cwd);
         }
 
-        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");
         }

src/inherited_fd.rs

@@ -0,0 +1,121 @@
+use std::env;
+use std::mem::{self, MaybeUninit};
+use std::os::fd::AsRawFd;
+use std::os::fd::FromRawFd;
+use std::os::fd::RawFd;
+use std::process;
+use std::{
+    net::TcpListener,
+    os::fd::OwnedFd,
+    sync::{Mutex, OnceLock},
+};
+
+/// A file descriptor the current process has inherited from a parent process.
+pub enum InheritedFileDescriptor {
+    /// An inherited socket verified to be a TCP socket.
+    TcpSocket(TcpListener),
+}
+
+static INHERITED_FDS: OnceLock<Mutex<Vec<InheritedFileDescriptor>>> = OnceLock::new();
+
+/// Takes ownership of file descriptors this process has inherited from a parent process like a
+/// service manager.
+///
+/// These are looked up with the [protocol from systemd](https://www.freedesktop.org/software/systemd/man/latest/sd_listen_fds.html#Notes).
+/// This is used to implement socket activation for `wasmtime serve`.
+///
+/// # Safety
+///
+/// This must be called first in `main()`, before the file descriptors can be used by anything
+/// else.
+pub unsafe fn init_inherited_fds() {
+    // The logic here is taken from https://github.com/systemd/systemd/blob/main/src/libsystemd/sd-daemon/sd-daemon.c,
+    // the protocol is described in the "Notes" section of https://www.freedesktop.org/software/systemd/man/latest/sd_listen_fds.html#Notes.
+
+    if !env::var("LISTEN_PID")
+        .ok()
+        .and_then(|pid| pid.parse().ok())
+        .is_some_and(|pid: u32| pid == process::id())
+    {
+        // Not meant for this process, ignore.
+        return;
+    }
+
+    let Some(num_fds) = env::var("LISTEN_FDS").ok().and_then(|fds| fds.parse().ok()) else {
+        return;
+    };
+
+    let first_fd: RawFd = 3;
+    let Some(last_fd) = first_fd.checked_add(num_fds) else {
+        return;
+    };
+
+    let mut descriptors = Vec::with_capacity(num_fds as usize);
+    for fd in first_fd..last_fd {
+        let fd = unsafe {
+            // Safety: This is called first in main and we checked the PID, so we have exclusive
+            // access to this fd.
+            libc::fcntl(fd, libc::F_SETFD, libc::FD_CLOEXEC);
+            OwnedFd::from_raw_fd(fd)
+        };
+
+        if is_tcp_socket(&fd) {
+            descriptors.push(InheritedFileDescriptor::TcpSocket(fd.into()));
+        }
+    }
+
+    let _ = INHERITED_FDS.set(Mutex::new(descriptors));
+}
+
+/// Consumes inherited file descriptors obtained through [init_inherited_fds].
+pub fn take_inherited_fds() -> Vec<InheritedFileDescriptor> {
+    let Some(fds) = INHERITED_FDS.get() else {
+        return Default::default();
+    };
+
+    let mut guard = fds.lock().unwrap();
+    return mem::take(&mut *guard);
+}
+
+fn is_tcp_socket(fd: &OwnedFd) -> bool {
+    let mut stat: MaybeUninit<libc::stat> = MaybeUninit::uninit();
+    if unsafe { libc::fstat(fd.as_raw_fd(), stat.as_mut_ptr()) } != 0 {
+        return false;
+    }
+
+    if (unsafe { stat.assume_init() }.st_mode & libc::S_IFMT) != libc::S_IFSOCK {
+        return false;
+    }
+
+    let sa_family = unsafe {
+        let mut sockaddr: MaybeUninit<libc::sockaddr> = MaybeUninit::uninit();
+        let mut len = mem::size_of::<libc::sockaddr>() as libc::c_uint;
+
+        if libc::getsockname(fd.as_raw_fd(), sockaddr.as_mut_ptr(), &mut len) != 0 {
+            return false;
+        }
+        sockaddr.assume_init().sa_family
+    } as libc::c_int;
+
+    if sa_family != libc::AF_INET && sa_family != libc::AF_INET6 {
+        return false;
+    }
+
+    let mut socket_type: libc::c_int = 0;
+    unsafe {
+        let mut type_len = mem::size_of_val(&socket_type) as libc::c_uint;
+
+        if libc::getsockopt(
+            fd.as_raw_fd(),
+            libc::SOL_SOCKET,
+            libc::SO_TYPE,
+            std::ptr::from_mut(&mut socket_type).cast(),
+            &mut type_len,
+        ) != 0
+        {
+            return false;
+        }
+    }
+
+    socket_type == libc::SOCK_STREAM
+}

src/lib.rs

@@ -11,3 +11,9 @@ pub(crate) mod common;
 
 #[cfg(any(feature = "objdump", all(feature = "hot-blocks", target_os = "linux")))]
 pub(crate) mod disas;
+
+#[cfg(all(unix, feature = "serve"))]
+pub(crate) mod inherited_fd;
+
+#[cfg(all(unix, feature = "serve"))]
+pub use inherited_fd::init_inherited_fds;