Skip to main content

podbox/
process.rs

1use std::ffi::OsString;
2use std::io::{self, BufRead, IoSlice, IoSliceMut, Write};
3use std::os::fd::{AsRawFd, FromRawFd, OwnedFd, RawFd};
4use std::os::unix::net::UnixStream;
5use std::os::unix::process::CommandExt;
6use std::process::{Command, ExitStatus, Output};
7use std::time::Duration;
8
9use nix::sys::signal::{Signal, kill};
10use nix::sys::socket::{ControlMessage, ControlMessageOwned, MsgFlags, recvmsg, sendmsg};
11use nix::unistd::Pid;
12
13/// Build a `Vec<OsString>` from a slice of `&str`/`&String` literals.
14pub fn args<S: AsRef<str>>(items: &[S]) -> Vec<OsString> {
15    items.iter().map(|s| OsString::from(s.as_ref())).collect()
16}
17
18/// Replace the current process with the given binary and arguments.
19///
20/// Uses `CommandExt::exec()` so the shell gets a real TTY.
21/// On success this function never returns; on failure it returns an error.
22pub fn exec_replace(bin: &str, args: &[OsString]) -> anyhow::Error {
23    let mut cmd = Command::new(bin);
24    cmd.args(args);
25    let err = cmd.exec();
26    anyhow::Error::from(err).context(format!("failed to exec {bin}"))
27}
28
29/// Run a command, capturing stdout and stderr.
30pub fn run_piped(bin: &str, args: &[OsString]) -> anyhow::Result<Output> {
31    let output = Command::new(bin)
32        .args(args)
33        .stdout(std::process::Stdio::piped())
34        .stderr(std::process::Stdio::piped())
35        .spawn()?
36        .wait_with_output()?;
37    Ok(output)
38}
39
40/// Spawn a command attached to the current terminal.
41pub fn spawn_interactive(bin: &str, args: &[OsString]) -> anyhow::Result<ExitStatus> {
42    let status = Command::new(bin).args(args).status()?;
43    Ok(status)
44}
45
46/// Run a command while teeing stdout+stderr into a log file.
47///
48/// Every child line is appended to `log` as produced. When `mirror` is true
49/// (verbose mode) lines are also echoed to our stdout so long operations
50/// stream live; otherwise output is captured silently and the caller shows a
51/// tail from the log on failure.
52pub fn run_with_log(
53    bin: &str,
54    args: &[OsString],
55    log: &mut std::fs::File,
56    mirror: bool,
57) -> anyhow::Result<ExitStatus> {
58    let mut child = Command::new(bin)
59        .args(args)
60        .stdout(std::process::Stdio::piped())
61        .stderr(std::process::Stdio::piped())
62        .spawn()
63        .map_err(|e| anyhow::Error::new(e).context(format!("failed to execute {bin}")))?;
64
65    let mut out = child.stdout.take().expect("stdout piped");
66    let mut err = child.stderr.take().expect("stderr piped");
67    let mut log_out = log.try_clone()?;
68    let mut log_err = log.try_clone()?;
69
70    let t_out = std::thread::spawn(move || tee_stream(&mut out, &mut log_out, false, mirror));
71    let t_err = std::thread::spawn(move || tee_stream(&mut err, &mut log_err, true, mirror));
72
73    let status = child.wait()?;
74    let _ = t_out.join();
75    let _ = t_err.join();
76    let _ = log.flush();
77    Ok(status)
78}
79
80/// Copy one child pipe into the log until EOF; optionally mirror to stderr.
81fn tee_stream<R: io::Read, W: io::Write>(src: &mut R, dst: &mut W, to_stderr: bool, mirror: bool) {
82    let mut reader = io::BufReader::new(src);
83    let mut line = String::new();
84    loop {
85        line.clear();
86        match reader.read_line(&mut line) {
87            Ok(0) | Err(_) => break,
88            Ok(_) => {
89                if dst.write_all(line.as_bytes()).is_err() {
90                    break;
91                }
92                if mirror {
93                    let _ = if to_stderr {
94                        io::stderr().write_all(line.as_bytes())
95                    } else {
96                        io::stdout().write_all(line.as_bytes())
97                    };
98                }
99            }
100        }
101    }
102}
103
104/// Run a command with a timeout, capturing stdout and stderr.
105///
106/// The child process receives SIGKILL after `timeout` if it has not exited.
107pub fn run_piped_timeout(
108    bin: &str,
109    args: &[OsString],
110    timeout: Duration,
111) -> anyhow::Result<Output> {
112    let child = Command::new(bin)
113        .args(args)
114        .stdout(std::process::Stdio::piped())
115        .stderr(std::process::Stdio::piped())
116        .spawn()?;
117    wait_child_timeout(child, timeout)
118}
119
120/// Run an interactive command with a timeout (SIGTERM, then SIGKILL).
121///
122/// Unlike `run_piped_timeout`, this sends SIGTERM first with a 5-second
123/// grace period before SIGKILL, giving well-behaved processes a chance
124/// to clean up.
125pub fn spawn_interactive_timeout(
126    bin: &str,
127    args: &[OsString],
128    timeout: Duration,
129) -> anyhow::Result<ExitStatus> {
130    let mut child = Command::new(bin).args(args).spawn()?;
131    let pid = Pid::from_raw(child.id().cast_signed());
132    let (tx, rx) = std::sync::mpsc::channel();
133    std::thread::spawn(move || {
134        if rx.recv_timeout(timeout).is_err() {
135            let _ = kill(pid, Signal::SIGTERM);
136            std::thread::sleep(Duration::from_secs(5));
137            let _ = kill(pid, Signal::SIGKILL);
138        }
139    });
140    let status = child.wait()?;
141    let _ = tx.send(());
142    Ok(status)
143}
144
145/// Wait for a child process to complete, enforcing a timeout via SIGKILL.
146pub fn wait_child_timeout(child: std::process::Child, timeout: Duration) -> anyhow::Result<Output> {
147    let pid = Pid::from_raw(child.id().cast_signed());
148    let (tx, rx) = std::sync::mpsc::channel();
149    std::thread::spawn(move || {
150        if rx.recv_timeout(timeout).is_err() {
151            let _ = kill(pid, Signal::SIGKILL);
152        }
153    });
154    let output = child.wait_with_output()?;
155    let _ = tx.send(());
156    Ok(output)
157}
158
159/// Open a pidfd for a given PID (Linux 5.3+).
160///
161/// Returns `Err` on old kernels or when the PID does not exist.
162pub fn open_pidfd(pid: i32) -> io::Result<OwnedFd> {
163    let pid = rustix::process::Pid::from_raw(pid)
164        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "invalid PID"))?;
165    rustix::process::pidfd_open(pid, rustix::process::PidfdFlags::empty()).map_err(io::Error::from)
166}
167
168/// Send a raw file descriptor over a connected Unix stream via `SCM_RIGHTS`.
169///
170/// Sends one dummy byte alongside the descriptor so the receiver can detect EOF.
171/// Retries on `EINTR` to prevent spurious session drops.
172pub fn send_fd(stream: &UnixStream, fd: RawFd) -> io::Result<()> {
173    let raw_fd = stream.as_raw_fd();
174    let cmsg = ControlMessage::ScmRights(&[fd]);
175    let iov = [IoSlice::new(&[0u8])];
176    loop {
177        match sendmsg::<()>(raw_fd, &iov, &[cmsg], MsgFlags::empty(), None) {
178            Ok(_) => return Ok(()),
179            Err(nix::errno::Errno::EINTR) => {}
180            Err(e) => return Err(io::Error::from(e)),
181        }
182    }
183}
184
185/// Adopt a raw descriptor received via `SCM_RIGHTS` into an owned handle.
186///
187/// The kernel duplicates ancillary-data descriptors into the receiving
188/// process on delivery, transferring their single reference — whoever
189/// calls this must adopt them (here) or close them (leak). This is the
190/// single audit point for that ownership transfer in this crate.
191pub fn adopt_scm_fd(raw: RawFd) -> OwnedFd {
192    // SAFETY: `raw` arrived via SCM_RIGHTS over a Unix socket; the kernel
193    // duplicated it into this process on delivery and we now hold exactly
194    // one reference to a valid open descriptor. No safe wrapper exists for
195    // adopting externally-sourced fds.
196    #[allow(unsafe_code)]
197    unsafe {
198        OwnedFd::from_raw_fd(raw)
199    }
200}
201
202/// Receive a raw file descriptor from a connected Unix stream via `SCM_RIGHTS`.
203///
204/// Returns `None` when the sender has closed the connection (EOF).
205/// Retries on `EINTR` to prevent spurious session drops.
206pub fn recv_fd(stream: &UnixStream) -> io::Result<Option<RawFd>> {
207    let raw_fd = stream.as_raw_fd();
208    let mut buf = [0u8; 1];
209    let mut iov = [IoSliceMut::new(&mut buf)];
210    let mut cmsg_buf = vec![0u8; 256];
211    let msg = loop {
212        match recvmsg::<()>(raw_fd, &mut iov, Some(&mut cmsg_buf), MsgFlags::empty()) {
213            Ok(m) => break m,
214            Err(nix::errno::Errno::EINTR) => {}
215            Err(e) => return Err(io::Error::from(e)),
216        }
217    };
218
219    if msg.bytes == 0 {
220        return Ok(None);
221    }
222
223    if let Ok(cmsgs) = msg.cmsgs() {
224        for cmsg in cmsgs {
225            if let ControlMessageOwned::ScmRights(fds) = cmsg {
226                if let Some(&fd) = fds.first() {
227                    return Ok(Some(fd));
228                }
229            }
230        }
231    }
232    Ok(None)
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238
239    #[test]
240    fn args_builds_osstring_vec() {
241        let v = args(&["foo", "bar", "baz"]);
242        assert_eq!(v.len(), 3);
243        assert_eq!(v[0], "foo");
244        assert_eq!(v[1], "bar");
245        assert_eq!(v[2], "baz");
246    }
247
248    #[test]
249    fn args_accepts_mixed_types() {
250        let s = String::from("hello");
251        let v = args(&["a", &s, "c"]);
252        assert_eq!(v[1], "hello");
253    }
254}