podbox/

process.rs

use std::ffi::OsString;
use std::io::{self, IoSlice, IoSliceMut};
use std::os::fd::{AsRawFd, FromRawFd, OwnedFd, RawFd};
use std::os::unix::net::UnixStream;
use std::os::unix::process::CommandExt;
use std::process::{Command, ExitStatus, Output};
use std::time::Duration;

use nix::sys::signal::{Signal, kill};
use nix::sys::socket::{ControlMessage, ControlMessageOwned, MsgFlags, recvmsg, sendmsg};
use nix::unistd::Pid;

/// Build a `Vec<OsString>` from a slice of `&str`/`&String` literals.
pub fn args<S: AsRef<str>>(items: &[S]) -> Vec<OsString> {
    items.iter().map(|s| OsString::from(s.as_ref())).collect()
}

/// Replace the current process with the given binary and arguments.
///
/// Uses `CommandExt::exec()` so the shell gets a real TTY.
/// On success this function never returns; on failure it returns an error.
pub fn exec_replace(bin: &str, args: &[OsString]) -> anyhow::Error {
    let mut cmd = Command::new(bin);
    cmd.args(args);
    let err = cmd.exec();
    anyhow::Error::from(err).context(format!("failed to exec {}", bin))
}

/// Run a command, capturing stdout and stderr.
pub fn run_piped(bin: &str, args: &[OsString]) -> anyhow::Result<Output> {
    let output = Command::new(bin)
        .args(args)
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()?
        .wait_with_output()?;
    Ok(output)
}

/// Spawn a command attached to the current terminal.
pub fn spawn_interactive(bin: &str, args: &[OsString]) -> anyhow::Result<ExitStatus> {
    let status = Command::new(bin).args(args).status()?;
    Ok(status)
}

/// Run a command with a timeout, capturing stdout and stderr.
///
/// The child process receives SIGKILL after `timeout` if it has not exited.
pub fn run_piped_timeout(
    bin: &str,
    args: &[OsString],
    timeout: Duration,
) -> anyhow::Result<Output> {
    let child = Command::new(bin)
        .args(args)
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()?;
    wait_child_timeout(child, timeout)
}

/// Run an interactive command with a timeout (SIGTERM, then SIGKILL).
///
/// Unlike `run_piped_timeout`, this sends SIGTERM first with a 5-second
/// grace period before SIGKILL, giving well-behaved processes a chance
/// to clean up.
pub fn spawn_interactive_timeout(
    bin: &str,
    args: &[OsString],
    timeout: Duration,
) -> anyhow::Result<ExitStatus> {
    let mut child = Command::new(bin).args(args).spawn()?;
    let pid = Pid::from_raw(child.id().cast_signed());
    let (tx, rx) = std::sync::mpsc::channel();
    std::thread::spawn(move || {
        if rx.recv_timeout(timeout).is_err() {
            let _ = kill(pid, Signal::SIGTERM);
            std::thread::sleep(Duration::from_secs(5));
            let _ = kill(pid, Signal::SIGKILL);
        }
    });
    let status = child.wait()?;
    let _ = tx.send(());
    Ok(status)
}

/// Wait for a child process to complete, enforcing a timeout via SIGKILL.
pub fn wait_child_timeout(child: std::process::Child, timeout: Duration) -> anyhow::Result<Output> {
    let pid = Pid::from_raw(child.id().cast_signed());
    let (tx, rx) = std::sync::mpsc::channel();
    std::thread::spawn(move || {
        if rx.recv_timeout(timeout).is_err() {
            let _ = kill(pid, Signal::SIGKILL);
        }
    });
    let output = child.wait_with_output()?;
    let _ = tx.send(());
    Ok(output)
}

/// Open a pidfd for a given PID (Linux 5.3+).
///
/// Returns `Err` on old kernels or when the PID does not exist.
pub fn open_pidfd(pid: i32) -> io::Result<OwnedFd> {
    let ret = unsafe { nix::libc::syscall(nix::libc::SYS_pidfd_open, pid, 0) };
    if ret < 0 {
        Err(io::Error::last_os_error())
    } else {
        let fd = i32::try_from(ret).map_err(|_| {
            io::Error::new(
                io::ErrorKind::InvalidInput,
                "pidfd_open returned invalid fd",
            )
        })?;
        // SAFETY: fd is a non-negative fd returned by the kernel.
        Ok(unsafe { OwnedFd::from_raw_fd(fd) })
    }
}

/// Send a raw file descriptor over a connected Unix stream via `SCM_RIGHTS`.
///
/// Sends one dummy byte alongside the descriptor so the receiver can detect EOF.
/// Retries on `EINTR` to prevent spurious session drops.
pub fn send_fd(stream: &UnixStream, fd: RawFd) -> io::Result<()> {
    let raw_fd = stream.as_raw_fd();
    let cmsg = ControlMessage::ScmRights(&[fd]);
    let iov = [IoSlice::new(&[0u8])];
    loop {
        match sendmsg::<()>(raw_fd, &iov, &[cmsg], MsgFlags::empty(), None) {
            Ok(_) => return Ok(()),
            Err(nix::errno::Errno::EINTR) => {}
            Err(e) => return Err(io::Error::from(e)),
        }
    }
}

/// Receive a raw file descriptor from a connected Unix stream via `SCM_RIGHTS`.
///
/// Returns `None` when the sender has closed the connection (EOF).
/// Retries on `EINTR` to prevent spurious session drops.
pub fn recv_fd(stream: &UnixStream) -> io::Result<Option<RawFd>> {
    let raw_fd = stream.as_raw_fd();
    let mut buf = [0u8; 1];
    let mut iov = [IoSliceMut::new(&mut buf)];
    let mut cmsg_buf = vec![0u8; 256];
    let msg = loop {
        match recvmsg::<()>(raw_fd, &mut iov, Some(&mut cmsg_buf), MsgFlags::empty()) {
            Ok(m) => break m,
            Err(nix::errno::Errno::EINTR) => {}
            Err(e) => return Err(io::Error::from(e)),
        }
    };

    if msg.bytes == 0 {
        return Ok(None);
    }

    if let Ok(cmsgs) = msg.cmsgs() {
        for cmsg in cmsgs {
            if let ControlMessageOwned::ScmRights(fds) = cmsg {
                if let Some(&fd) = fds.first() {
                    return Ok(Some(fd));
                }
            }
        }
    }
    Ok(None)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn args_builds_osstring_vec() {
        let v = args(&["foo", "bar", "baz"]);
        assert_eq!(v.len(), 3);
        assert_eq!(v[0], "foo");
        assert_eq!(v[1], "bar");
        assert_eq!(v[2], "baz");
    }

    #[test]
    fn args_accepts_mixed_types() {
        let s = String::from("hello");
        let v = args(&["a", &s, "c"]);
        assert_eq!(v[1], "hello");
    }
}