sausage/cli/src/runtime.rs

//! Wasmtime engine and process instances.
//!
//! An instance is one run of a `wasi:cli/command` component. The host writes
//! NDJSON lines to its stdin and reads NDJSON lines from its stdout.

use std::path::Path;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Result;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, DuplexStream};
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use wasmtime::component::{Component, Linker, ResourceTable};
use wasmtime::{Config, Engine, Store};
use wasmtime_wasi::cli::{AsyncStdinStream, AsyncStdoutStream};

use wasmtime_wasi::p2::bindings::Command;

use wasmtime_wasi::{I32Exit, WasiCtx, WasiCtxView, WasiView};
use wasmtime_wasi_http::{WasiHttpCtx, WasiHttpCtxView, WasiHttpHooks, WasiHttpView};

/// The HTTP transport hooks supplied to WASI processes.
pub struct HttpHooks {
    allow: bool,
}

impl WasiHttpHooks for HttpHooks {
    fn send_request(
        &mut self,
        request: http::Request<wasmtime_wasi_http::WasiBody>,
        options: Option<wasmtime_wasi_http::RequestOptions>,
        fut: Box<dyn Future<Output = Result<(), wasmtime_wasi_http::Error>> + Send>,
    ) -> Box<
        dyn Future<
                Output = wasmtime_wasi_http::Result<(
                    http::Response<wasmtime_wasi_http::WasiBody>,
                    Box<dyn Future<Output = Result<(), wasmtime_wasi_http::Error>> + Send>,
                )>,
            > + Send,
    > {
        _ = fut;
        let allow = self.allow;
        Box::new(async move {
            use http_body_util::BodyExt;
            if !allow {
                return Err(wasmtime_wasi_http::Error::HttpRequestDenied);
            }
            let (res, io) = wasmtime_wasi_http::default_send_request(request, options).await?;
            Ok((
                res.map(BodyExt::boxed_unsync),
                Box::new(io) as Box<dyn Future<Output = _> + Send>,
            ))
        })
    }
}

/// The WASI and HTTP resources owned by one store.
pub struct ProcessState {
    ctx: WasiCtx,
    table: ResourceTable,
    http: WasiHttpCtx,
    hooks: HttpHooks,
    fs: sausage_vfs::adapter::Adapter,
}

impl WasiView for ProcessState {
    fn ctx(&mut self) -> WasiCtxView<'_> {
        WasiCtxView {
            ctx: &mut self.ctx,
            table: &mut self.table,
        }
    }
}

impl WasiHttpView for ProcessState {
    fn http(&mut self) -> WasiHttpCtxView<'_> {
        WasiHttpCtxView {
            ctx: &mut self.http,
            table: &mut self.table,
            hooks: &mut self.hooks,
        }
    }
}

/// A Wasmtime engine and linker shared by the world.
pub struct Runtime {
    engine: Engine,
    linker: Linker<ProcessState>,
}

/// How an instance ended.
#[derive(Debug, Clone)]
pub enum Exit {
    /// `main` returned or the guest called `exit(0)`.
    Clean,
    /// The guest exited with a nonzero code.
    Code(i32),
    /// The guest trapped (panic, invalid memory access, ...).
    Trap(String),
}

impl std::fmt::Display for Exit {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Exit::Clean => write!(f, "exit 0"),
            Exit::Code(c) => write!(f, "exit {c}"),
            Exit::Trap(t) => write!(f, "trap: {t}"),
        }
    }
}

/// The filesystem, network, and environment supplied to one process.
pub struct InstanceConfig {
    /// The virtual namespace exposed through the single root preopen.
    pub filesystem: sausage_vfs::Mounts,
    /// The store completion barrier used by process cleanup.
    pub store_ended: tokio::sync::watch::Sender<bool>,
    /// Outgoing HTTP, and TCP, UDP, and DNS through `wasi:sockets`.
    pub network: bool,
    /// The environment variables supplied to the guest.
    pub env: Vec<(String, String)>,
}

/// A running guest. Dropping it kills the guest.
pub struct Instance {
    /// The host end of stdin, removed when the process stops.
    pub stdin: Option<DuplexStream>,
    /// Lines that the guest writes to stdout. Closed when the guest ends.
    pub stdout: mpsc::Receiver<String>,
    /// The stderr lines that the process writes.
    pub stderr: mpsc::Receiver<String>,
    task: Option<JoinHandle<Exit>>,
    store_ended: tokio::sync::watch::Sender<bool>,
}

impl Drop for Instance {
    fn drop(&mut self) {
        if let Some(task) = self.task.take() {
            task.abort();
            let ended = self.store_ended.clone();
            tokio::spawn(async move {
                let _ = task.await;
                ended.send_replace(true);
            });
        }
    }
}

impl Instance {
    /// Write one line to the guest's stdin, in one write. Meanwhile, pass the
    /// lines that the guest writes to `route`: a guest that blocks on a full
    /// stdout would never read its stdin.
    pub async fn send_line(
        &mut self,
        line: &str,
        mut route: impl FnMut(String),
    ) -> std::io::Result<()> {
        let stdin = self.stdin.as_mut().ok_or(std::io::ErrorKind::BrokenPipe)?;
        let bytes = format!("{line}\n");
        let write = async {
            stdin.write_all(bytes.as_bytes()).await?;
            stdin.flush().await
        };
        tokio::pin!(write);
        let mut open = true;
        loop {
            tokio::select! {
                done = &mut write => return done,
                out = self.stdout.recv(), if open => match out {
                    Some(out) => route(out),
                    None => open = false,
                },
            }
        }
    }

    /// Wait until the guest ends.
    pub async fn wait(&mut self) -> Exit {
        let result = self.task.as_mut().unwrap().await;
        self.task.take();
        self.store_ended.send_replace(true);
        match result {
            Ok(exit) => exit,
            Err(e) if e.is_cancelled() => Exit::Trap("killed".into()),
            Err(e) => Exit::Trap(format!("host task failed: {e}")),
        }
    }
}

/// The most useful line of a trap: the `wasm trap: ...` reason.
fn short_trap(full: &str) -> String {
    let line = full
        .lines()
        .find_map(|l| l.find("wasm trap:").map(|i| l[i..].to_string()))
        .unwrap_or_else(|| full.lines().next().unwrap_or("trap").to_string());
    let mut line = line.trim().to_string();
    line.truncate(300);
    line
}

impl Runtime {
    /// Create an engine with disk caching and a ten-millisecond epoch tick.
    pub fn new(cache: &Path) -> Result<Arc<Self>> {
        let mut config = Config::new();
        config.epoch_interruption(true);
        std::fs::create_dir_all(cache)?;
        let mut cache_config = wasmtime::CacheConfig::new();
        cache_config.with_directory(cache.canonicalize()?);
        config.cache(Some(wasmtime::Cache::new(cache_config)?));
        let engine = Engine::new(&config)?;

        let mut linker: Linker<ProcessState> = Linker::new(&engine);
        wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
        sausage_vfs::adapter::add_to_linker(&mut linker, |state| sausage_vfs::adapter::View {
            adapter: &mut state.fs,
            table: &mut state.table,
        })?;
        wasmtime_wasi_http::p2::add_only_http_to_linker_async(&mut linker)?;
        // Time slices: every 10 ms each running guest yields to the scheduler.
        let runtime = Arc::new(Self { engine, linker });
        let ticker = Arc::downgrade(&runtime);
        std::thread::spawn(move || {
            loop {
                std::thread::sleep(Duration::from_millis(10));
                let Some(runtime) = ticker.upgrade() else {
                    break;
                };
                runtime.engine.increment_epoch();
            }
        });

        Ok(runtime)
    }

    /// Compile component bytes through Wasmtime's disk cache.
    pub fn compile(&self, bytes: &[u8]) -> Result<Component> {
        Ok(Component::new(&self.engine, bytes)?)
    }

    /// Start a WASI task with the supplied virtual namespace.
    pub async fn start(
        self: &Arc<Self>,
        component: &Component,
        cfg: InstanceConfig,
    ) -> Result<Instance> {
        self.start_with_args(component, cfg, &[]).await
    }

    /// Start a WASI task with explicit CLI arguments, including its program name.
    pub async fn start_with_args(
        self: &Arc<Self>,
        component: &Component,
        cfg: InstanceConfig,
        args: &[String],
    ) -> Result<Instance> {
        let (host_in, guest_in) = tokio::io::duplex(64 * 1024);
        let (guest_out, host_out) = tokio::io::duplex(64 * 1024);

        let (guest_err, host_err) = tokio::io::duplex(64 * 1024);
        let mut b = WasiCtx::builder();
        b.args(args);
        if cfg.network {
            b.inherit_network()
                .allow_tcp(true)
                .allow_udp(true)
                .allow_ip_name_lookup(true);
        }
        b.stdin(AsyncStdinStream::new(guest_in));
        b.stdout(AsyncStdoutStream::new(8192, guest_out));
        b.stderr(AsyncStdoutStream::new(8192, guest_err));
        for (k, v) in &cfg.env {
            b.env(k, v);
        }

        let state = ProcessState {
            ctx: b.build(),
            table: ResourceTable::new(),
            http: WasiHttpCtx::new(),
            hooks: HttpHooks { allow: cfg.network },
            fs: sausage_vfs::adapter::Adapter::new(cfg.filesystem),
        };

        let (line_tx, line_rx) = mpsc::channel(1024);
        tokio::spawn(async move {
            let mut reader = BufReader::new(host_out);
            let mut bytes = Vec::new();
            while let Ok(n) = reader.read_until(b'\n', &mut bytes).await {
                if n == 0 {
                    break;
                }
                // Preserve a non-blank invalid line so routing emits its warning.
                let line = String::from_utf8(std::mem::take(&mut bytes))
                    .unwrap_or_else(|_| "\u{fffd}".into());
                if line_tx.send(line).await.is_err() {
                    break;
                }
            }
        });

        let (err_tx, err_rx) = mpsc::channel(1024);
        tokio::spawn(async move {
            let mut lines = BufReader::new(host_err).lines();
            while let Ok(Some(line)) = lines.next_line().await {
                if err_tx.send(line).await.is_err() {
                    break;
                }
            }
        });
        let rt = self.clone();
        let component = component.clone();
        cfg.store_ended.send_replace(false);
        let task = tokio::spawn(async move {
            let mut store = Store::new(&rt.engine, state);
            store.epoch_deadline_async_yield_and_update(1);
            let result = async {
                let cmd = Command::instantiate_async(&mut store, &component, &rt.linker).await?;
                cmd.wasi_cli_run().call_run(&mut store).await
            }
            .await;
            match result {
                Ok(Ok(())) => Exit::Clean,
                Ok(Err(())) => Exit::Code(1),
                Err(e) => match e.downcast_ref::<I32Exit>() {
                    Some(I32Exit(0)) => Exit::Clean,
                    Some(I32Exit(c)) => Exit::Code(*c),
                    None => {
                        let full = format!("{e:?}");
                        Exit::Trap(short_trap(&full))
                    }
                },
            }
        });

        Ok(Instance {
            stdin: Some(host_in),
            stdout: line_rx,
            stderr: err_rx,
            task: Some(task),
            store_ended: cfg.store_ended,
        })
    }
}