kcode-rust-bins 3.0.1

Author, locally publish, and run small Rust binaries
Documentation
use std::io::{Read, Write};
use std::path::Path;
use std::process::{Child, Command, ExitStatus, Stdio};
use std::thread;
use std::time::{Duration, Instant};

use crate::model::{Error, Result, RunRequest, RunResult};
use crate::protocol::{read_output, write_message};
use crate::sandbox::ensure_runtime_image;

const DIAGNOSTIC_LIMIT: usize = 1024 * 1024;

pub(crate) fn run(
    executable_path: &Path,
    version: String,
    request: RunRequest,
    timeout: Duration,
) -> Result<RunResult> {
    if timeout.is_zero() {
        return Err(Error::new(
            "invalid_timeout",
            "binary timeout must be greater than zero",
        ));
    }
    let RunRequest { text, objects } = request;

    let runtime_image = ensure_runtime_image()?;
    let mut command = runtime_command(executable_path, &runtime_image);
    command
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());
    let mut child = command
        .spawn()
        .map_err(|error| Error::new("execution", format!("could not start Podman: {error}")))?;
    let mut stdin = child
        .stdin
        .take()
        .ok_or_else(|| Error::new("execution", "Podman stdin was not piped"))?;
    let mut stdout = child
        .stdout
        .take()
        .ok_or_else(|| Error::new("execution", "Podman stdout was not piped"))?;
    let mut stderr = child
        .stderr
        .take()
        .ok_or_else(|| Error::new("execution", "Podman stderr was not piped"))?;
    let writer = thread::spawn(move || {
        let result = write_message(&mut stdin, &text, &objects).and_then(|()| {
            stdin
                .flush()
                .map_err(|error| Error::new("protocol", format!("flush stdin: {error}")))
        });
        drop(stdin);
        result
    });
    let reader = thread::spawn(move || {
        let result = read_output(&mut stdout);
        if result.is_err() {
            let _ = std::io::copy(&mut stdout, &mut std::io::sink());
        }
        result
    });
    let diagnostics_reader = thread::spawn(move || drain_diagnostics(&mut stderr));

    let wait = wait_for_child(&mut child, timeout);
    let diagnostics = join_thread(diagnostics_reader, "stderr reader")?;
    let output = join_thread(reader, "stdout reader")?;
    let input = join_thread(writer, "stdin writer")?;
    let status = wait?;

    if !status.success() {
        return Err(Error::new(
            "execution",
            format!(
                "tool exited with status {}; diagnostics: {}",
                status
                    .code()
                    .map_or_else(|| "signal".to_owned(), |code| code.to_string()),
                diagnostics.trim()
            ),
        ));
    }
    input?;
    let output = output?;

    Ok(RunResult {
        version,
        text: output.text,
        objects: output.objects,
        diagnostics,
    })
}

fn runtime_command(executable: &Path, image: &str) -> Command {
    let mut volume = executable.as_os_str().to_os_string();
    volume.push(":/tool:ro,Z");
    let mut command = Command::new("podman");
    command
        .arg("run")
        .arg("--rm")
        .arg("--interactive")
        .arg("--read-only")
        .arg("--cap-drop=all")
        .arg("--security-opt=no-new-privileges")
        .arg("--userns=keep-id")
        .arg("--pids-limit=256")
        .arg("--memory=2g")
        .arg("--cpus=4")
        .arg("--tmpfs=/tmp:rw,nosuid,nodev,size=536870912")
        .arg("--pull=never")
        .arg("--volume")
        .arg(volume)
        .arg("--entrypoint=/tool")
        .arg(image);
    command
}

fn wait_for_child(child: &mut Child, timeout: Duration) -> Result<ExitStatus> {
    let deadline = Instant::now() + timeout;
    loop {
        match child.try_wait() {
            Ok(Some(status)) => return Ok(status),
            Ok(None) if Instant::now() < deadline => thread::sleep(Duration::from_millis(25)),
            Ok(None) => {
                let _ = child.kill();
                let _ = child.wait();
                return Err(Error::new(
                    "execution",
                    format!("tool exceeded {} second timeout", timeout.as_secs()),
                ));
            }
            Err(error) => {
                let _ = child.kill();
                let _ = child.wait();
                return Err(Error::new(
                    "execution",
                    format!("could not wait for tool: {error}"),
                ));
            }
        }
    }
}

fn drain_diagnostics(reader: &mut impl Read) -> String {
    let mut retained = Vec::new();
    let mut buffer = [0_u8; 8192];
    let mut truncated = false;
    loop {
        match reader.read(&mut buffer) {
            Ok(0) => break,
            Ok(read) => {
                let available = DIAGNOSTIC_LIMIT.saturating_sub(retained.len());
                let keep = available.min(read);
                retained.extend_from_slice(&buffer[..keep]);
                truncated |= keep < read;
            }
            Err(error) => {
                retained.extend_from_slice(format!("\n[stderr read failed: {error}]").as_bytes());
                break;
            }
        }
    }
    let mut diagnostics = String::from_utf8_lossy(&retained).into_owned();
    if truncated {
        diagnostics.push_str("\n[diagnostics truncated]");
    }
    diagnostics
}

fn join_thread<T>(handle: thread::JoinHandle<T>, label: &str) -> Result<T> {
    handle
        .join()
        .map_err(|_| Error::new("execution", format!("{label} panicked")))
}

#[cfg(test)]
mod tests {
    use std::path::Path;

    use super::{DIAGNOSTIC_LIMIT, drain_diagnostics, runtime_command};

    #[test]
    fn runtime_is_open_networked_and_isolated() {
        let command = runtime_command(Path::new("/tmp/tool"), "runtime-image");
        let arguments = command
            .get_args()
            .map(|argument| argument.to_string_lossy().into_owned())
            .collect::<Vec<_>>();
        assert!(arguments.contains(&"--read-only".to_owned()));
        assert!(arguments.contains(&"--cap-drop=all".to_owned()));
        assert!(arguments.contains(&"--security-opt=no-new-privileges".to_owned()));
        assert!(arguments.contains(&"--memory=2g".to_owned()));
        assert!(
            !arguments
                .iter()
                .any(|argument| argument == "--network=none")
        );
    }

    #[test]
    fn diagnostics_are_bounded_but_drained() {
        let bytes = vec![b'x'; DIAGNOSTIC_LIMIT + 100];
        let diagnostics = drain_diagnostics(&mut bytes.as_slice());
        assert!(diagnostics.len() < DIAGNOSTIC_LIMIT + 100);
        assert!(diagnostics.ends_with("[diagnostics truncated]"));
    }
}