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::CommandPre;

/// A component with its imports resolved once for this runtime.
pub(crate) struct Prepared(Result<CommandPre<ProcessState>, String>);

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>,
    stderr_task: Option<JoinHandle<()>>,
    task: Option<JoinHandle<Exit>>,
    store_ended: tokio::sync::watch::Sender<bool>,
}

impl Drop for Instance {
    fn drop(&mut self) {
        let task = self.task.take();
        let stderr = self.stderr_task.take();
        if let Some(task) = &task {
            task.abort();
        }
        if task.is_some() || stderr.is_some() {
            let ended = self.store_ended.clone();
            tokio::spawn(async move {
                if let Some(task) = task {
                    let _ = task.await;
                }
                if let Some(stderr) = stderr {
                    let _ = stderr.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();
        if let Some(stderr) = self.stderr_task.as_mut() {
            let _ = stderr.await;
        }
        self.stderr_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();
    let mut end = line.len().min(300);
    while !line.is_char_boundary(end) {
        end -= 1;
    }
    line.truncate(end);
    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);
        let disk_cache = (|| -> Result<_> {
            let cache = cache.join("wasmtime");
            std::fs::create_dir_all(&cache)?;
            let mut cache_config = wasmtime::CacheConfig::new();
            cache_config.with_directory(cache.canonicalize()?);
            Ok(wasmtime::Cache::new(cache_config)?)
        })();
        match disk_cache {
            Ok(cache) => {
                config.cache(Some(cache));
            }
            Err(error) => eprintln!(
                "compile cache {} unavailable: {error}; running without disk cache",
                cache.display()
            ),
        }
        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)?)
    }

    pub(crate) fn prepare(&self, component: &Component) -> Prepared {
        // Preserve link failures as process exits, rather than load errors.
        Prepared(
            self.linker
                .instantiate_pre(component)
                .and_then(CommandPre::new)
                .map_err(|error| error.to_string()),
        )
    }

    pub(crate) async fn start_logged(
        self: &Arc<Self>,
        prepared: &Prepared,
        cfg: InstanceConfig,
        journal: Arc<crate::journal::Journal>,
        pid: String,
        plugin: String,
    ) -> Result<Instance> {
        self.start_inner(prepared.0.clone(), cfg, &[], Some((journal, pid, plugin)))
            .await
    }

    /// 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> {
        self.start_inner(self.prepare(component).0, cfg, args, None)
            .await
    }

    async fn start_inner(
        self: &Arc<Self>,
        prepared: Result<CommandPre<ProcessState>, String>,
        cfg: InstanceConfig,
        args: &[String],
        log: Option<(Arc<crate::journal::Journal>, String, 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);
        let trap_log = log.clone();
        let stderr_task = tokio::spawn(async move {
            let mut reader = BufReader::new(host_err);
            let mut bytes = Vec::new();
            while let Ok(n) = reader.read_until(b'\n', &mut bytes).await {
                if n == 0 {
                    break;
                }
                let line = String::from_utf8_lossy(&bytes)
                    .trim_end_matches(['\r', '\n'])
                    .to_string();
                bytes.clear();
                if let Some((journal, pid, plugin)) = &log {
                    journal.write(pid, plugin, "stderr", &line);
                } else if err_tx.send(line).await.is_err() {
                    break;
                }
            }
        });
        let rt = self.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 = prepared
                    .map_err(wasmtime::Error::msg)?
                    .instantiate_async(&mut store)
                    .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:?}");
                        if let Some((journal, pid, plugin)) = &trap_log {
                            journal.write(pid, plugin, "core", &full);
                        }
                        Exit::Trap(short_trap(&full))
                    }
                },
            }
        });

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

#[cfg(test)]
mod trap_text_tests {
    use super::*;
    #[test]
    fn short_trap_display_is_utf8_safe() {
        let full = format!("backtrace\nwasm trap: {}\nmore detail", "😀".repeat(1000));
        let status = short_trap(&full);
        assert!(status.len() <= 300);
        assert!(status.starts_with("wasm trap:"));
        assert_eq!(short_trap(""), "trap");
        assert_eq!(short_trap("first line\nsecond"), "first line");
    }
}