sausage/spikes/s2-shell/src/process.rs

//! Starts a uutils component in its own store and exchanges bounded pipe chunks.
use super::State;
use sausage_vfs::{
    Host, Mounts, Process,
    adapter::{self, Adapter, View},
};
use serde::Deserialize;
use std::{
    path::PathBuf,
    sync::{
        Arc,
        atomic::{AtomicI32, Ordering},
    },
};
use tokio::net::UnixStream;
use wasmtime::{
    Engine, Store, StoreContextMut,
    component::{Component, Linker, ResourceTable},
};
use wasmtime_wasi::{
    WasiCtxBuilder,
    cli::{AsyncStdinStream, AsyncStdoutStream},
};

pub struct Runtime {
    pub engine: Engine,
    pub utility: Component,
    pub world: PathBuf,
}
pub struct Child {
    input: Option<UnixStream>,
    output: Option<UnixStream>,
    error: Option<UnixStream>,
    code: Arc<AtomicI32>,
}
#[derive(Deserialize)]
struct Request {
    argv: Vec<String>,
    env: Vec<(String, String)>,
    cwd: String,
}
fn pair() -> wasmtime::Result<(UnixStream, UnixStream)> {
    let (a, b) = std::os::unix::net::UnixStream::pair()?;
    a.set_nonblocking(true)?;
    b.set_nonblocking(true)?;
    Ok((UnixStream::from_std(a)?, UnixStream::from_std(b)?))
}
fn drain(pipe: &mut Option<UnixStream>) -> wasmtime::Result<Vec<u8>> {
    let Some(stream) = pipe else {
        return Ok(Vec::new());
    };
    let mut bytes = vec![0; 8192];
    match stream.try_read(&mut bytes) {
        Ok(0) => {
            *pipe = None;
            bytes.clear();
        }
        Ok(n) => bytes.truncate(n),
        Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => bytes.clear(),
        Err(e) => return Err(e.into()),
    }
    Ok(bytes)
}
pub fn add(linker: &mut Linker<State>) -> wasmtime::Result<()> {
    let mut host = linker.instance("s2:process/host")?;
    host.func_wrap(
        "spawn",
        |mut store: StoreContextMut<'_, State>, (json,): (String,)| {
            let request: Request = serde_json::from_str(&json)?;
            let runtime = store.data().runtime.clone();
            let name = std::path::Path::new(&request.argv[0])
                .file_name()
                .ok_or_else(|| wasmtime::format_err!("missing command name"))?;
            wasmtime::ensure!(
                runtime.world.join("bin").join(name).is_file(),
                "command is absent from /bin"
            );
            let (input, child_input) = pair()?;
            let (output, child_output) = pair()?;
            let (error, child_error) = pair()?;
            let code = Arc::new(AtomicI32::new(-1));
            let result_code = code.clone();
            tokio::spawn(async move {
                let result = run(runtime, request, child_input, child_output, child_error).await;
                let status = match result {
                    Ok(status) => status,
                    Err(error) => {
                        eprintln!("child component: {error:#}");
                        125
                    }
                };
                result_code.store(status, Ordering::Release);
            });
            let state = store.data_mut();
            let pid = state.children.len() as u32;
            state.children.push(Child {
                input: Some(input),
                output: Some(output),
                error: Some(error),
                code,
            });
            Ok((pid,))
        },
    )?;
    host.func_wrap(
        "exchange",
        |mut store: StoreContextMut<'_, State>,
         (pid, bytes, eof, closed): (u32, Vec<u8>, bool, bool)| {
            let child = &mut store.data_mut().children[pid as usize];
            let consumed = match child.input.as_mut() {
                Some(input) if !bytes.is_empty() => match input.try_write(&bytes) {
                    Ok(n) => n,
                    Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => 0,
                    Err(e) if e.kind() == std::io::ErrorKind::BrokenPipe => {
                        child.input = None;
                        bytes.len()
                    }
                    Err(e) => return Err(e.into()),
                },
                _ => bytes.len(),
            };
            if eof && consumed == bytes.len() {
                child.input = None;
            }
            if closed {
                child.output = None;
            }
            let out = drain(&mut child.output)?;
            let err = drain(&mut child.error)?;
            let code = child.code.load(Ordering::Acquire);
            let done = code >= 0 && child.output.is_none() && child.error.is_none();
            Ok(((consumed as u32, done, code.max(0) as u32, out, err),))
        },
    )?;
    Ok(())
}
async fn run(
    runtime: Arc<Runtime>,
    request: Request,
    input: UnixStream,
    output: UnixStream,
    error: UnixStream,
) -> wasmtime::Result<i32> {
    let mut mounts = Mounts::default();
    mounts.mount(
        "/",
        Arc::new(Host::new(&runtime.world, Process::default())?),
        false,
    );
    let mut linker = Linker::new(&runtime.engine);
    wasmtime_wasi::p2::add_to_linker_async(&mut linker)?;
    adapter::add_to_linker(&mut linker, |s: &mut State| View {
        adapter: &mut s.fs,
        table: &mut s.table,
    })?;
    let mut wasi = WasiCtxBuilder::new();
    let (reader, _) = input.into_split();
    let (_, writer) = output.into_split();
    let (_, err_writer) = error.into_split();
    wasi.args(&request.argv)
        .envs(&request.env)
        .env("S2_CWD", &request.cwd)
        .initial_cwd(&request.cwd)
        .stdin(AsyncStdinStream::new(reader))
        .stdout(AsyncStdoutStream::new(8192, writer))
        .stderr(AsyncStdoutStream::new(8192, err_writer));
    let mut store = Store::new(
        &runtime.engine,
        State {
            wasi: wasi.build(),
            table: ResourceTable::new(),
            fs: Adapter::new(mounts),
            runtime: runtime.clone(),
            children: Vec::new(),
        },
    );
    let command = wasmtime_wasi::p2::bindings::Command::instantiate_async(
        &mut store,
        &runtime.utility,
        &linker,
    )
    .await?;
    match command.wasi_cli_run().call_run(&mut store).await {
        Ok(Ok(())) => Ok(0),
        Ok(Err(())) => Ok(1),
        Err(error) => match error.downcast_ref::<wasmtime_wasi::I32Exit>() {
            Some(exit) => Ok(exit.0),
            None => Err(error),
        },
    }
}