vamp/src/shell.rs

#[cfg(unix)]
use std::{
    ffi::OsString,
    io::{self, Read},
    os::{
        fd::OwnedFd,
        unix::{net::UnixStream, process::CommandExt},
    },
    process::{Child, Command, Stdio},
    sync::mpsc::{self, RecvTimeoutError},
    thread,
    time::Duration,
};
use std::{
    path::PathBuf,
    sync::mpsc::{Receiver, Sender, TryRecvError},
    thread::JoinHandle,
};

#[cfg(unix)]
const OUTPUT_LIMIT: usize = 4 * 1024 * 1024;
#[cfg(unix)]
const TRUNCATION_NOTE: &[u8] = b"\n[vamp: output truncated after 4 MiB]\n";

#[derive(Debug)]
pub(crate) struct Output {
    pub(crate) code: Option<i32>,
    pub(crate) stdout: Vec<u8>,
    pub(crate) stderr: Vec<u8>,
}

pub(crate) struct Worker {
    cancel: Sender<()>,
    results: Receiver<Result<Output, String>>,
    thread: Option<JoinHandle<()>>,
    finished: bool,
}

impl Worker {
    pub(crate) fn start(command: String, root: PathBuf) -> Result<Self, String> {
        #[cfg(unix)]
        {
            let shell = std::env::var_os("SHELL")
                .filter(|shell| !shell.is_empty())
                .unwrap_or_else(|| "/bin/sh".into());
            Self::start_with_shell(command, root, shell)
        }
        #[cfg(not(unix))]
        {
            let _ = (command, root);
            Err("Shell commands require Unix process-group cancellation on this platform".into())
        }
    }

    #[cfg(unix)]
    fn start_with_shell(command: String, root: PathBuf, shell: OsString) -> Result<Self, String> {
        let (cancel, cancelled) = mpsc::channel();
        let (send, results) = mpsc::channel();
        let thread = thread::Builder::new()
            .name("vamp-shell".into())
            .spawn(move || {
                let _ = send.send(run(command, root, shell, cancelled));
            })
            .map_err(|error| format!("Cannot start shell worker: {error}"))?;
        Ok(Self {
            cancel,
            results,
            thread: Some(thread),
            finished: false,
        })
    }

    pub(crate) fn poll(&mut self) -> Option<Result<Output, String>> {
        if self.finished {
            return None;
        }
        let result = match self.results.try_recv() {
            Ok(output) => output,
            Err(TryRecvError::Empty) => return None,
            Err(TryRecvError::Disconnected) => Err("Shell worker stopped unexpectedly".into()),
        };
        self.finished = true;
        Some(result)
    }
}

impl Drop for Worker {
    fn drop(&mut self) {
        let _ = self.cancel.send(());
        if let Some(thread) = self.thread.take() {
            let _ = thread.join();
        }
    }
}

#[cfg(unix)]
struct Process(Child);

#[cfg(unix)]
impl Drop for Process {
    fn drop(&mut self) {
        unsafe extern "C" {
            fn kill(pid: i32, signal: i32) -> i32;
        }
        // The shell and its descendants share a private process group; killing
        // only the leader would leave children holding the output descriptors.
        unsafe { kill(-(self.0.id() as i32), 9) };
        let _ = self.0.kill();
        let _ = self.0.wait();
    }
}

#[cfg(unix)]
struct Capture {
    reader: UnixStream,
    bytes: Vec<u8>,
    truncated: bool,
    eof: bool,
}

#[cfg(unix)]
impl Capture {
    fn new() -> Result<(Self, Stdio), String> {
        let (reader, writer) = UnixStream::pair()
            .map_err(|error| format!("Cannot create shell output stream: {error}"))?;
        reader
            .set_nonblocking(true)
            .map_err(|error| format!("Cannot configure shell output stream: {error}"))?;
        Ok((
            Self {
                reader,
                bytes: Vec::new(),
                truncated: false,
                eof: false,
            },
            Stdio::from(OwnedFd::from(writer)),
        ))
    }

    fn read(&mut self) -> io::Result<bool> {
        if self.eof {
            return Ok(false);
        }
        let mut chunk = [0; 16 * 1024];
        match self.reader.read(&mut chunk) {
            Ok(0) => {
                self.eof = true;
                Ok(false)
            }
            Ok(count) => {
                let keep = count.min(OUTPUT_LIMIT.saturating_sub(self.bytes.len()));
                self.bytes.extend_from_slice(&chunk[..keep]);
                self.truncated |= keep != count;
                Ok(true)
            }
            Err(error) if error.kind() == io::ErrorKind::WouldBlock => Ok(false),
            Err(error) if error.kind() == io::ErrorKind::Interrupted => Ok(true),
            Err(error) => Err(error),
        }
    }

    fn finish(mut self) -> Vec<u8> {
        if self.truncated {
            self.bytes.extend_from_slice(TRUNCATION_NOTE);
        }
        self.bytes
    }
}

#[cfg(unix)]
fn run(
    command: String,
    root: PathBuf,
    shell: OsString,
    cancelled: Receiver<()>,
) -> Result<Output, String> {
    if cancelled.try_recv().is_ok() {
        return Err("Shell command cancelled".into());
    }
    let (mut stdout, stdout_writer) = Capture::new()?;
    let (mut stderr, stderr_writer) = Capture::new()?;
    let child = Command::new(shell)
        .arg("-c")
        .arg(command)
        .current_dir(root)
        .stdin(Stdio::null())
        .stdout(stdout_writer)
        .stderr(stderr_writer)
        .process_group(0)
        .spawn()
        .map_err(|error| format!("Cannot start shell command: {error}"))?;
    let mut process = Some(Process(child));
    let mut status = None;
    loop {
        if cancelled.try_recv().is_ok() {
            return Err("Shell command cancelled".into());
        }
        // Read each stream once per turn so a busy writer cannot starve the
        // other stream or cancellation. Overflow is still drained, not retained.
        let out = stdout
            .read()
            .map_err(|error| format!("Cannot read shell stdout: {error}"))?;
        let err = stderr
            .read()
            .map_err(|error| format!("Cannot read shell stderr: {error}"))?;
        if let Some(child) = process.as_mut() {
            status = child
                .0
                .try_wait()
                .map_err(|error| format!("Cannot wait for shell command: {error}"))?;
            if status.is_some() {
                // Background jobs must not outlive the command. Killing the
                // remaining group also closes inherited output descriptors.
                drop(process.take());
            }
        }
        if let Some(status) = status
            && stdout.eof
            && stderr.eof
        {
            return Ok(Output {
                code: status.code(),
                stdout: stdout.finish(),
                stderr: stderr.finish(),
            });
        }
        if !out && !err {
            match cancelled.recv_timeout(Duration::from_millis(10)) {
                Ok(()) | Err(RecvTimeoutError::Disconnected) => {
                    return Err("Shell command cancelled".into());
                }
                Err(RecvTimeoutError::Timeout) => {}
            }
        }
    }
}

#[cfg(all(test, unix))]
#[path = "../tests/unit/shell.rs"]
mod tests;