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
13pub fn args<S: AsRef<str>>(items: &[S]) -> Vec<OsString> {
15 items.iter().map(|s| OsString::from(s.as_ref())).collect()
16}
17
18pub 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
29pub 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
40pub fn spawn_interactive(bin: &str, args: &[OsString]) -> anyhow::Result<ExitStatus> {
42 let status = Command::new(bin).args(args).status()?;
43 Ok(status)
44}
45
46pub 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
80fn 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
104pub 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
120pub 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
145pub 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
159pub 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
168pub 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
185pub fn adopt_scm_fd(raw: RawFd) -> OwnedFd {
192 #[allow(unsafe_code)]
197 unsafe {
198 OwnedFd::from_raw_fd(raw)
199 }
200}
201
202pub 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}