Skip to content

Commit

Permalink
Import sd listen
Browse files Browse the repository at this point in the history
Signed-off-by: Harald Hoyer <harald@profian.com>
  • Loading branch information
haraldh committed Jan 26, 2022
1 parent b36445b commit 598194d
Show file tree
Hide file tree
Showing 5 changed files with 121 additions and 6 deletions.
20 changes: 16 additions & 4 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
15 changes: 14 additions & 1 deletion crates/wasi-common/cap-std-sync/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down Expand Up @@ -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
}
Expand Down
19 changes: 19 additions & 0 deletions crates/wasi-common/tokio/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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
}
Expand Down
72 changes: 71 additions & 1 deletion src/commands/run.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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>,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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());
}

Expand Down Expand Up @@ -454,3 +492,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_listener((3 + i) as _, listener)?;
num_fd = 3 + i;
}
}

Ok((num_fd, builder))
}

0 comments on commit 598194d

Please sign in to comment.