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>(
82    src: &mut R,
83    dst: &mut W,
84    to_stderr: bool,
85    mirror: bool,
86) {
87    let mut reader = io::BufReader::new(src);
88    let mut line = String::new();
89    loop {
90        line.clear();
91        match reader.read_line(&mut line) {
92            Ok(0) | Err(_) => break,
93            Ok(_) => {
94                if dst.write_all(line.as_bytes()).is_err() {
95                    break;
96                }
97                if mirror {
98                    let _ = if to_stderr {
99                        io::stderr().write_all(line.as_bytes())
100                    } else {
101                        io::stdout().write_all(line.as_bytes())
102                    };
103                }
104            }
105        }
106    }
107}
108
109/// Run a command with a timeout, capturing stdout and stderr.
110///
111/// The child process receives SIGKILL after `timeout` if it has not exited.
112pub fn run_piped_timeout(
113    bin: &str,
114    args: &[OsString],
115    timeout: Duration,
116) -> anyhow::Result<Output> {
117    let child = Command::new(bin)
118        .args(args)
119        .stdout(std::process::Stdio::piped())
120        .stderr(std::process::Stdio::piped())
121        .spawn()?;
122    wait_child_timeout(child, timeout)
123}
124
125/// Run an interactive command with a timeout (SIGTERM, then SIGKILL).
126///
127/// Unlike `run_piped_timeout`, this sends SIGTERM first with a 5-second
128/// grace period before SIGKILL, giving well-behaved processes a chance
129/// to clean up.
130pub fn spawn_interactive_timeout(
131    bin: &str,
132    args: &[OsString],
133    timeout: Duration,
134) -> anyhow::Result<ExitStatus> {
135    let mut child = Command::new(bin).args(args).spawn()?;
136    let pid = Pid::from_raw(child.id().cast_signed());
137    let (tx, rx) = std::sync::mpsc::channel();
138    std::thread::spawn(move || {
139        if rx.recv_timeout(timeout).is_err() {
140            let _ = kill(pid, Signal::SIGTERM);
141            std::thread::sleep(Duration::from_secs(5));
142            let _ = kill(pid, Signal::SIGKILL);
143        }
144    });
145    let status = child.wait()?;
146    let _ = tx.send(());
147    Ok(status)
148}
149
150/// Wait for a child process to complete, enforcing a timeout via SIGKILL.
151pub fn wait_child_timeout(child: std::process::Child, timeout: Duration) -> anyhow::Result<Output> {
152    let pid = Pid::from_raw(child.id().cast_signed());
153    let (tx, rx) = std::sync::mpsc::channel();
154    std::thread::spawn(move || {
155        if rx.recv_timeout(timeout).is_err() {
156            let _ = kill(pid, Signal::SIGKILL);
157        }
158    });
159    let output = child.wait_with_output()?;
160    let _ = tx.send(());
161    Ok(output)
162}
163
164/// Open a pidfd for a given PID (Linux 5.3+).
165///
166/// Returns `Err` on old kernels or when the PID does not exist.
167pub fn open_pidfd(pid: i32) -> io::Result<OwnedFd> {
168    let pid = rustix::process::Pid::from_raw(pid)
169        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "invalid PID"))?;
170    rustix::process::pidfd_open(pid, rustix::process::PidfdFlags::empty()).map_err(io::Error::from)
171}
172
173/// Send a raw file descriptor over a connected Unix stream via `SCM_RIGHTS`.
174///
175/// Sends one dummy byte alongside the descriptor so the receiver can detect EOF.
176/// Retries on `EINTR` to prevent spurious session drops.
177pub fn send_fd(stream: &UnixStream, fd: RawFd) -> io::Result<()> {
178    let raw_fd = stream.as_raw_fd();
179    let cmsg = ControlMessage::ScmRights(&[fd]);
180    let iov = [IoSlice::new(&[0u8])];
181    loop {
182        match sendmsg::<()>(raw_fd, &iov, &[cmsg], MsgFlags::empty(), None) {
183            Ok(_) => return Ok(()),
184            Err(nix::errno::Errno::EINTR) => {}
185            Err(e) => return Err(io::Error::from(e)),
186        }
187    }
188}
189
190/// Adopt a raw descriptor received via `SCM_RIGHTS` into an owned handle.
191///
192/// The kernel duplicates ancillary-data descriptors into the receiving
193/// process on delivery, transferring their single reference — whoever
194/// calls this must adopt them (here) or close them (leak). This is the
195/// single audit point for that ownership transfer in this crate.
196pub fn adopt_scm_fd(raw: RawFd) -> OwnedFd {
197    // SAFETY: `raw` arrived via SCM_RIGHTS over a Unix socket; the kernel
198    // duplicated it into this process on delivery and we now hold exactly
199    // one reference to a valid open descriptor. No safe wrapper exists for
200    // adopting externally-sourced fds.
201    #[allow(unsafe_code)]
202    unsafe {
203        OwnedFd::from_raw_fd(raw)
204    }
205}
206
207/// Receive a raw file descriptor from a connected Unix stream via `SCM_RIGHTS`.
208///
209/// Returns `None` when the sender has closed the connection (EOF).
210/// Retries on `EINTR` to prevent spurious session drops.
211pub fn recv_fd(stream: &UnixStream) -> io::Result<Option<RawFd>> {
212    let raw_fd = stream.as_raw_fd();
213    let mut buf = [0u8; 1];
214    let mut iov = [IoSliceMut::new(&mut buf)];
215    let mut cmsg_buf = vec![0u8; 256];
216    let msg = loop {
217        match recvmsg::<()>(raw_fd, &mut iov, Some(&mut cmsg_buf), MsgFlags::empty()) {
218            Ok(m) => break m,
219            Err(nix::errno::Errno::EINTR) => {}
220            Err(e) => return Err(io::Error::from(e)),
221        }
222    };
223
224    if msg.bytes == 0 {
225        return Ok(None);
226    }
227
228    if let Ok(cmsgs) = msg.cmsgs() {
229        for cmsg in cmsgs {
230            if let ControlMessageOwned::ScmRights(fds) = cmsg {
231                if let Some(&fd) = fds.first() {
232                    return Ok(Some(fd));
233                }
234            }
235        }
236    }
237    Ok(None)
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243
244    #[test]
245    fn args_builds_osstring_vec() {
246        let v = args(&["foo", "bar", "baz"]);
247        assert_eq!(v.len(), 3);
248        assert_eq!(v[0], "foo");
249        assert_eq!(v[1], "bar");
250        assert_eq!(v[2], "baz");
251    }
252
253    #[test]
254    fn args_accepts_mixed_types() {
255        let s = String::from("hello");
256        let v = args(&["a", &s, "c"]);
257        assert_eq!(v[1], "hello");
258    }
259}
podbox/process.rs

podbox/
process.rs
