planarian_backend_linux/
cmd.rs

1//! Running a command with a timeout and bounded output.
2
3use std::io::Read;
4use std::process::{Command, Stdio};
5use std::time::{Duration, Instant};
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct Output {
9    pub status: Option<i32>,
10    pub stdout: String,
11    pub stderr: String,
12    pub timed_out: bool,
13    /// More output than `max_output` came and was dropped: what is here is
14    /// incomplete, and never a basis for a plan (rule 21).
15    pub truncated: bool,
16}
17
18impl Output {
19    pub fn ok(&self) -> bool {
20        self.status == Some(0) && !self.timed_out
21    }
22
23    /// The last few lines of stderr (or stdout), for an error message.
24    pub fn tail(&self) -> String {
25        let src = if self.stderr.trim().is_empty() { &self.stdout } else { &self.stderr };
26        let lines: Vec<&str> = src.lines().collect();
27        lines[lines.len().saturating_sub(4)..].join(" | ")
28    }
29}
30
31pub const MAX_OUTPUT: usize = 1 << 20;
32
33/// Run `argv` with a timeout; kill it (and, through its process group, its
34/// children) when the timeout expires.
35pub fn run(argv: &[&str], timeout: Duration) -> Output {
36    run_limited(argv, timeout, MAX_OUTPUT)
37}
38
39pub fn run_limited(argv: &[&str], timeout: Duration, max_output: usize) -> Output {
40    let fail =
41        |e: String| Output { status: None, stdout: String::new(), stderr: e, timed_out: false, truncated: false };
42    let Some((prog, args)) = argv.split_first() else { return fail("empty command".into()) };
43    let mut cmd = Command::new(prog);
44    cmd.args(args).stdin(Stdio::null()).stdout(Stdio::piped()).stderr(Stdio::piped());
45    std::os::unix::process::CommandExt::process_group(&mut cmd, 0);
46    let mut child = match cmd.spawn() {
47        Ok(c) => c,
48        Err(e) => return fail(format!("{prog}: {e}")),
49    };
50    let mut out = child.stdout.take().expect("piped");
51    let mut err = child.stderr.take().expect("piped");
52    let out_t = std::thread::spawn(move || {
53        let mut buf = Vec::new();
54        let _ = (&mut out).take(max_output as u64).read_to_end(&mut buf);
55        let dropped = std::io::copy(&mut out, &mut std::io::sink()).unwrap_or(0);
56        (buf, dropped > 0)
57    });
58    let err_t = std::thread::spawn(move || {
59        let mut buf = Vec::new();
60        let _ = (&mut err).take(max_output as u64).read_to_end(&mut buf);
61        let dropped = std::io::copy(&mut err, &mut std::io::sink()).unwrap_or(0);
62        (buf, dropped > 0)
63    });
64    let deadline = Instant::now() + timeout;
65    let mut timed_out = false;
66    let status = loop {
67        match child.try_wait() {
68            Ok(Some(s)) => break s.code(),
69            Ok(None) if Instant::now() >= deadline => {
70                timed_out = true;
71                if let Some(pid) = rustix::process::Pid::from_raw(child.id() as i32) {
72                    let _ = rustix::process::kill_process_group(pid, rustix::process::Signal::KILL);
73                }
74                let _ = child.kill();
75                break child.wait().ok().and_then(|s| s.code());
76            }
77            Ok(None) => std::thread::sleep(Duration::from_millis(20)),
78            Err(e) => return fail(e.to_string()),
79        }
80    };
81    let (out, out_cut) = out_t.join().unwrap_or_default();
82    let (err, err_cut) = err_t.join().unwrap_or_default();
83    let stdout = String::from_utf8_lossy(&out).into_owned();
84    let stderr = String::from_utf8_lossy(&err).into_owned();
85    Output { status, stdout, stderr, timed_out, truncated: out_cut || err_cut }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn runs_and_times_out() {
94        let o = run(&["sh", "-c", "echo hi; echo err >&2; exit 3"], Duration::from_secs(5));
95        assert_eq!(o.status, Some(3));
96        assert_eq!(o.stdout.trim(), "hi");
97        assert_eq!(o.tail(), "err");
98        let o = run(&["sh", "-c", "sleep 5 & sleep 5"], Duration::from_millis(200));
99        assert!(o.timed_out && !o.ok());
100        let o = run_limited(&["sh", "-c", "yes | head -c 100000"], Duration::from_secs(5), 10);
101        assert_eq!(o.stdout.len(), 10);
102        assert!(!run(&["/nonexistent"], Duration::from_secs(1)).ok());
103    }
104}
planarian_backend_linux/cmd.rs

planarian_backend_linux/
cmd.rs
