strop_core/process/
capture.rs1use 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
34pub 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 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()?; 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}