use crate::cli::{StdinStream, StdoutStream, WasiCliCtx};
use crate::clocks::{HostMonotonicClock, HostWallClock, WasiClocksCtx};
use crate::filesystem::{Dir, WasiFilesystemCtx};
use crate::random::WasiRandomCtx;
use crate::sockets::{SocketAddrCheck, SocketAddrUse, WasiSocketsCtx};
use crate::{DirPerms, FilePerms, OpenMode};
use cap_rand::RngCore;
use cap_std::ambient_authority;
use std::future::Future;
use std::mem;
use std::net::SocketAddr;
use std::path::Path;
use std::pin::Pin;
use tokio::io::{stderr, stdin, stdout};
use wasmtime::Result;
#[derive(Default)]
pub struct WasiCtxBuilder {
cli: WasiCliCtx,
clocks: WasiClocksCtx,
filesystem: WasiFilesystemCtx,
random: WasiRandomCtx,
sockets: WasiSocketsCtx,
built: bool,
}
impl WasiCtxBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn stdin(&mut self, stdin: impl StdinStream + 'static) -> &mut Self {
self.cli.stdin = Box::new(stdin);
self
}
pub fn stdout(&mut self, stdout: impl StdoutStream + 'static) -> &mut Self {
self.cli.stdout = Box::new(stdout);
self
}
pub fn stderr(&mut self, stderr: impl StdoutStream + 'static) -> &mut Self {
self.cli.stderr = Box::new(stderr);
self
}
pub fn inherit_stdin(&mut self) -> &mut Self {
self.stdin(stdin())
}
pub fn inherit_stdout(&mut self) -> &mut Self {
self.stdout(stdout())
}
pub fn inherit_stderr(&mut self) -> &mut Self {
self.stderr(stderr())
}
pub fn inherit_stdio(&mut self) -> &mut Self {
self.inherit_stdin().inherit_stdout().inherit_stderr()
}
pub fn allow_blocking_current_thread(&mut self, enable: bool) -> &mut Self {
self.filesystem.allow_blocking_current_thread = enable;
self
}
pub fn envs(&mut self, env: &[(impl AsRef<str>, impl AsRef<str>)]) -> &mut Self {
self.cli.environment.extend(
env.iter()
.map(|(k, v)| (k.as_ref().to_owned(), v.as_ref().to_owned())),
);
self
}
pub fn env(&mut self, k: impl AsRef<str>, v: impl AsRef<str>) -> &mut Self {
self.cli
.environment
.push((k.as_ref().to_owned(), v.as_ref().to_owned()));
self
}
pub fn inherit_env(&mut self) -> &mut Self {
self.cli.environment.extend(std::env::vars());
self
}
pub fn args(&mut self, args: &[impl AsRef<str>]) -> &mut Self {
self.cli
.arguments
.extend(args.iter().map(|a| a.as_ref().to_owned()));
self
}
pub fn arg(&mut self, arg: impl AsRef<str>) -> &mut Self {
self.cli.arguments.push(arg.as_ref().to_owned());
self
}
pub fn inherit_args(&mut self) -> &mut Self {
self.cli.arguments.extend(std::env::args());
self
}
pub fn preopened_dir(
&mut self,
host_path: impl AsRef<Path>,
guest_path: impl AsRef<str>,
dir_perms: DirPerms,
file_perms: FilePerms,
) -> Result<&mut Self> {
let dir = cap_std::fs::Dir::open_ambient_dir(host_path.as_ref(), ambient_authority())?;
let mut open_mode = OpenMode::empty();
if dir_perms.contains(DirPerms::READ) {
open_mode |= OpenMode::READ;
}
if dir_perms.contains(DirPerms::MUTATE) {
open_mode |= OpenMode::WRITE;
}
self.filesystem.preopens.push((
Dir::new(
dir,
dir_perms,
file_perms,
open_mode,
self.filesystem.allow_blocking_current_thread,
),
guest_path.as_ref().to_owned(),
));
Ok(self)
}
pub fn secure_random(&mut self, random: impl RngCore + Send + 'static) -> &mut Self {
self.random.random = Box::new(random);
self
}
pub fn insecure_random(&mut self, insecure_random: impl RngCore + Send + 'static) -> &mut Self {
self.random.insecure_random = Box::new(insecure_random);
self
}
pub fn insecure_random_seed(&mut self, insecure_random_seed: u128) -> &mut Self {
self.random.insecure_random_seed = insecure_random_seed;
self
}
pub fn max_random_size(&mut self, max_size: u64) -> &mut Self {
self.random.max_size = max_size;
self
}
pub fn wall_clock(&mut self, clock: impl HostWallClock + 'static) -> &mut Self {
self.clocks.wall_clock = Box::new(clock);
self
}
pub fn monotonic_clock(&mut self, clock: impl HostMonotonicClock + 'static) -> &mut Self {
self.clocks.monotonic_clock = Box::new(clock);
self
}
pub fn inherit_network(&mut self) -> &mut Self {
self.socket_addr_check(|_, _| Box::pin(async { true }))
}
pub fn socket_addr_check<F>(&mut self, check: F) -> &mut Self
where
F: Fn(SocketAddr, SocketAddrUse) -> Pin<Box<dyn Future<Output = bool> + Send + Sync>>
+ Send
+ Sync
+ 'static,
{
self.sockets.socket_addr_check = SocketAddrCheck::new(check);
self
}
pub fn allow_ip_name_lookup(&mut self, enable: bool) -> &mut Self {
self.sockets.allowed_network_uses.ip_name_lookup = enable;
self
}
pub fn allow_udp(&mut self, enable: bool) -> &mut Self {
self.sockets.allowed_network_uses.udp = enable;
self
}
pub fn allow_tcp(&mut self, enable: bool) -> &mut Self {
self.sockets.allowed_network_uses.tcp = enable;
self
}
pub fn build(&mut self) -> WasiCtx {
assert!(!self.built);
let Self {
cli,
clocks,
filesystem,
random,
sockets,
built: _,
} = mem::replace(self, Self::new());
self.built = true;
WasiCtx {
cli,
clocks,
filesystem,
random,
sockets,
}
}
#[cfg(feature = "p1")]
pub fn build_p1(&mut self) -> crate::p1::WasiP1Ctx {
let wasi = self.build();
crate::p1::WasiP1Ctx::new(wasi)
}
}
#[derive(Default)]
pub struct WasiCtx {
pub(crate) cli: WasiCliCtx,
pub(crate) clocks: WasiClocksCtx,
pub(crate) filesystem: WasiFilesystemCtx,
pub(crate) random: WasiRandomCtx,
pub(crate) sockets: WasiSocketsCtx,
}
impl WasiCtx {
pub fn builder() -> WasiCtxBuilder {
WasiCtxBuilder::new()
}
pub fn random(&mut self) -> &mut WasiRandomCtx {
&mut self.random
}
pub fn clocks(&mut self) -> &mut WasiClocksCtx {
&mut self.clocks
}
pub fn filesystem(&mut self) -> &mut WasiFilesystemCtx {
&mut self.filesystem
}
pub fn cli(&mut self) -> &mut WasiCliCtx {
&mut self.cli
}
pub fn sockets(&mut self) -> &mut WasiSocketsCtx {
&mut self.sockets
}
}