1use std::borrow::Cow;
24use std::io;
25use std::os::fd::{AsRawFd, OwnedFd, RawFd};
26use std::time::{Duration, Instant};
27
28use rustix::process::{Signal, pidfd_send_signal};
29
30use crate::plan::Plan;
31use crate::workspace::Workspace;
32
33#[must_use]
35#[derive(Debug, Clone)]
36pub struct Output {
37 pub stdout: Vec<u8>,
38 pub stderr: Vec<u8>,
39 pub status: Status,
40 pub duration: Duration,
41 pub exit_code: Option<i32>,
42 pub signal: Option<i32>,
43}
44
45impl Output {
46 #[inline]
47 pub fn success(&self) -> bool {
48 self.status == Status::Exited && self.exit_code == Some(0)
49 }
50
51 #[inline]
52 pub fn stdout_str(&self) -> Cow<'_, str> {
53 String::from_utf8_lossy(&self.stdout)
54 }
55
56 #[inline]
57 pub fn stderr_str(&self) -> Cow<'_, str> {
58 String::from_utf8_lossy(&self.stderr)
59 }
60}
61
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
64pub enum Status {
65 Exited,
66 Signaled,
67 Timeout,
68 OutputLimitExceeded,
69}
70
71#[allow(clippy::cast_possible_truncation)]
74pub fn monitor(pidfd: OwnedFd, workspace: &Workspace, plan: &Plan) -> io::Result<Output> {
75 let start = Instant::now();
76 let deadline = start + plan.timeout;
77
78 let mut stdout_buf = Vec::new();
79 let mut stderr_buf = Vec::new();
80
81 let stdout_fd = workspace.pipes.stdout.read.as_raw_fd();
82 let stderr_fd = workspace.pipes.stderr.read.as_raw_fd();
83 let pidfd_raw = pidfd.as_raw_fd();
84
85 set_nonblocking(stdout_fd)?;
86 set_nonblocking(stderr_fd)?;
87
88 let mut status = Status::Exited;
89 let mut exit_code = None;
90 let mut signal = None;
91 let mut buf = [0u8; 4096];
92
93 loop {
94 let timeout_remaining = deadline.saturating_duration_since(Instant::now());
95 if timeout_remaining.is_zero() {
96 pidfd_send_signal(&pidfd, Signal::KILL).ok();
97 status = Status::Timeout;
98 wait_for_exit(pidfd_raw)?;
99 break;
100 }
101
102 let poll_timeout = timeout_remaining.as_millis().min(100) as i32;
104 let mut fds = [
105 libc::pollfd {
106 fd: stdout_fd,
107 events: libc::POLLIN,
108 revents: 0,
109 },
110 libc::pollfd {
111 fd: stderr_fd,
112 events: libc::POLLIN,
113 revents: 0,
114 },
115 libc::pollfd {
116 fd: pidfd_raw,
117 events: libc::POLLIN,
118 revents: 0,
119 },
120 ];
121
122 let ret = unsafe { libc::poll(fds.as_mut_ptr(), 3, poll_timeout) };
123 if ret < 0 {
124 let err = io::Error::last_os_error();
125 if err.kind() == io::ErrorKind::Interrupted {
126 continue;
127 }
128 return Err(err);
129 }
130
131 if fds[0].revents & libc::POLLIN != 0 {
132 if let Ok(n) = read_nonblocking(stdout_fd, &mut buf) {
133 if n > 0 {
134 if stdout_buf.len() + n > plan.max_output as usize {
135 status = Status::OutputLimitExceeded;
136 pidfd_send_signal(&pidfd, Signal::KILL).ok();
137 wait_for_exit(pidfd_raw)?;
138 break;
139 }
140 stdout_buf.extend_from_slice(&buf[..n]);
141 }
142 }
143 }
144
145 if fds[1].revents & libc::POLLIN != 0 {
146 if let Ok(n) = read_nonblocking(stderr_fd, &mut buf) {
147 if n > 0 {
148 if stderr_buf.len() + n > plan.max_output as usize {
149 status = Status::OutputLimitExceeded;
150 pidfd_send_signal(&pidfd, Signal::KILL).ok();
151 wait_for_exit(pidfd_raw)?;
152 break;
153 }
154 stderr_buf.extend_from_slice(&buf[..n]);
155 }
156 }
157 }
158
159 if fds[2].revents & libc::POLLIN != 0 {
160 let (ec, sig) = wait_for_exit(pidfd_raw)?;
161 exit_code = ec;
162 signal = sig;
163 if sig.is_some() {
164 status = Status::Signaled;
165 }
166 break;
167 }
168
169 if (fds[0].revents & libc::POLLHUP != 0) && (fds[1].revents & libc::POLLHUP != 0) {
170 let (ec, sig) = wait_for_exit(pidfd_raw)?;
171 exit_code = ec;
172 signal = sig;
173 if sig.is_some() {
174 status = Status::Signaled;
175 }
176 break;
177 }
178 }
179
180 drain_remaining(stdout_fd, &mut stdout_buf, &mut buf, plan.max_output);
181 drain_remaining(stderr_fd, &mut stderr_buf, &mut buf, plan.max_output);
182
183 Ok(Output {
184 stdout: stdout_buf,
185 stderr: stderr_buf,
186 status,
187 duration: start.elapsed(),
188 exit_code,
189 signal,
190 })
191}
192
193#[allow(clippy::cast_sign_loss)]
196pub fn write_stdin(workspace: &Workspace, data: &[u8]) -> io::Result<()> {
197 let fd = workspace.pipes.stdin.write.as_raw_fd();
198 let mut written = 0;
199 while written < data.len() {
200 let ret = unsafe {
201 libc::write(
202 fd,
203 data[written..].as_ptr().cast::<libc::c_void>(),
204 data.len() - written,
205 )
206 };
207 if ret < 0 {
208 return Err(io::Error::last_os_error());
209 }
210 written += ret as usize;
211 }
212 Ok(())
213}
214
215#[inline]
216pub(crate) fn set_nonblocking(fd: RawFd) -> io::Result<()> {
217 let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) };
218 if flags < 0 {
219 return Err(io::Error::last_os_error());
220 }
221 let ret = unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) };
222 if ret < 0 {
223 Err(io::Error::last_os_error())
224 } else {
225 Ok(())
226 }
227}
228
229#[allow(clippy::cast_sign_loss)]
231#[inline]
232fn read_nonblocking(fd: RawFd, buf: &mut [u8]) -> io::Result<usize> {
233 let ret = unsafe { libc::read(fd, buf.as_mut_ptr().cast::<libc::c_void>(), buf.len()) };
234 if ret < 0 {
235 Err(io::Error::last_os_error())
236 } else {
237 Ok(ret as usize)
238 }
239}
240
241#[allow(clippy::cast_possible_truncation)]
243fn drain_remaining(fd: RawFd, output: &mut Vec<u8>, buf: &mut [u8], max_output: u64) {
244 let max = max_output as usize;
245 loop {
246 if output.len() >= max {
247 break;
249 }
250 match read_nonblocking(fd, buf) {
251 Ok(0) | Err(_) => break,
252 Ok(n) => {
253 let remaining = max.saturating_sub(output.len());
255 let to_add = n.min(remaining);
256 output.extend_from_slice(&buf[..to_add]);
257 }
258 }
259 }
260}
261
262#[allow(clippy::cast_sign_loss)]
264pub(crate) fn wait_for_exit(pidfd: RawFd) -> io::Result<(Option<i32>, Option<i32>)> {
265 let mut siginfo: libc::siginfo_t = unsafe { std::mem::zeroed() };
266 let ret = unsafe {
267 libc::waitid(
268 libc::P_PIDFD,
269 pidfd as libc::id_t,
270 &mut siginfo,
271 libc::WEXITED,
272 )
273 };
274 if ret < 0 {
275 return Err(io::Error::last_os_error());
276 }
277
278 let code = siginfo.si_code;
279 let status = unsafe { siginfo.si_status() };
280
281 match code {
282 libc::CLD_EXITED => Ok((Some(status), None)),
283 libc::CLD_KILLED | libc::CLD_DUMPED => Ok((None, Some(status))),
284 _ => Ok((None, None)),
285 }
286}
287
288#[cfg(test)]
289mod tests {
290 use super::*;
291
292 #[test]
293 fn output_success() {
294 let output = Output {
295 stdout: vec![],
296 stderr: vec![],
297 status: Status::Exited,
298 duration: Duration::from_millis(100),
299 exit_code: Some(0),
300 signal: None,
301 };
302 assert!(output.success());
303 }
304
305 #[test]
306 fn output_failure() {
307 let output = Output {
308 stdout: vec![],
309 stderr: vec![],
310 status: Status::Exited,
311 duration: Duration::from_millis(100),
312 exit_code: Some(1),
313 signal: None,
314 };
315 assert!(!output.success());
316 }
317}