Skip to main content

strop_core/process/
capture.rs

1//! Bounded capture for short configuration commands. The owning worker drains
2//! both pipes concurrently and never relinquishes cancellation to a running child.
3use super::OwnedProcess;
4use crate::worker::{CancelToken, Failure, FailureKind};
5use std::io::{self, Read};
6use std::process::{Command, ExitStatus, Stdio};
7use std::sync::mpsc::{channel, RecvTimeoutError};
8use std::time::{Duration, Instant};
9
10const LIMIT: u64 = 64 * 1024;
11const DEADLINE: Duration = Duration::from_secs(30);
12const POLL: Duration = Duration::from_millis(20);
13
14pub struct CommandOutput {
15    pub status: ExitStatus,
16    pub stdout: Vec<u8>,
17    pub stderr: Vec<u8>,
18}
19enum Stream {
20    Stdout(io::Result<Vec<u8>>),
21    Stderr(io::Result<Vec<u8>>),
22}
23fn read_pipe(pipe: impl Read) -> io::Result<Vec<u8>> {
24    let mut bytes = Vec::new();
25    pipe.take(LIMIT + 1).read_to_end(&mut bytes)?;
26    if bytes.len() as u64 > LIMIT {
27        return Err(io::Error::other(
28            "configuration command exceeded 64 KiB output limit",
29        ));
30    }
31    Ok(bytes)
32}
33
34/// At most 64 KiB per pipe and 30 seconds. No terminal input/output is inherited.
35pub fn capture(command: &mut Command, token: &CancelToken) -> Result<CommandOutput, Failure> {
36    std::thread::scope(|scope| {
37        command
38            .stdin(Stdio::null())
39            .stdout(Stdio::piped())
40            .stderr(Stdio::piped());
41        // Owned here, INSIDE scope: unwinding kills pipes before scope joins.
42        let mut process = OwnedProcess::spawn(command, token)?;
43        let stdout = process
44            .take_stdout()
45            .ok_or_else(|| Failure::new(FailureKind::Protocol, "missing stdout"))?;
46        let stderr = process
47            .take_stderr()
48            .ok_or_else(|| Failure::new(FailureKind::Protocol, "missing stderr"))?;
49        let (tx, rx) = channel();
50        let out_tx = tx.clone();
51        std::thread::Builder::new()
52            .name("capture-stdout".into())
53            .spawn_scoped(scope, move || {
54                let _ = out_tx.send(Stream::Stdout(read_pipe(stdout)));
55            })
56            .map_err(|error| Failure::new(FailureKind::ThreadStart, error.to_string()))?;
57        std::thread::Builder::new()
58            .name("capture-stderr".into())
59            .spawn_scoped(scope, move || {
60                let _ = tx.send(Stream::Stderr(read_pipe(stderr)));
61            })
62            .map_err(|error| Failure::new(FailureKind::ThreadStart, error.to_string()))?;
63        let deadline = Instant::now() + DEADLINE;
64        let mut stdout = None;
65        let mut stderr = None;
66        loop {
67            if token.is_cancelled() {
68                return Err(Failure::new(
69                    FailureKind::Unavailable,
70                    "configuration command cancelled",
71                ));
72            }
73            if Instant::now() >= deadline {
74                return Err(Failure::new(
75                    FailureKind::Wait,
76                    "configuration command timed out after 30 seconds",
77                ));
78            }
79            let exited = process.has_exited()?;
80            if exited {
81                process.terminate()?; // descendants cannot retain the pipes
82                match (stdout.take(), stderr.take()) {
83                    (Some(stdout), Some(stderr)) => {
84                        return Ok(CommandOutput {
85                            status: process.wait()?,
86                            stdout,
87                            stderr,
88                        });
89                    }
90                    (out, err) => {
91                        stdout = out;
92                        stderr = err;
93                    }
94                }
95            }
96            let event = match rx.recv_timeout(POLL) {
97                Ok(event) => event,
98                Err(RecvTimeoutError::Timeout) => continue,
99                Err(RecvTimeoutError::Disconnected) if !exited => {
100                    std::thread::park_timeout(POLL);
101                    continue;
102                }
103                Err(error) => {
104                    return Err(Failure::new(FailureKind::Disconnected, error.to_string()))
105                }
106            };
107            let (slot, result) = match event {
108                Stream::Stdout(result) => (&mut stdout, result),
109                Stream::Stderr(result) => (&mut stderr, result),
110            };
111            *slot = Some(result.map_err(|error| Failure::new(FailureKind::Io, error.to_string()))?);
112        }
113    })
114}