1use crate::{
4 Component, Error, Operation, QueryPolicy,
5 process::OwnedChild,
6 query::{Stream, reader},
7};
8#[cfg(unix)]
9const TERMINATE: i32 = signal_hook::consts::signal::SIGTERM;
10#[cfg(not(unix))]
11const TERMINATE: i32 = 0;
12
13use std::{
14 io::{self, Write},
15 process::{Command, Output, Stdio},
16 sync::atomic::{AtomicBool, Ordering},
17 sync::mpsc,
18 thread,
19 time::Instant,
20};
21
22pub fn capture(
27 component: Component,
28 command: &mut Command,
29 scope: &Operation,
30 policy: QueryPolicy,
31 on_failure: impl FnOnce(),
32) -> Result<Output, Error> {
33 run(component, command, Some(scope), policy, on_failure)
34}
35
36pub fn run(
38 component: Component,
39 command: &mut Command,
40 scope: Option<&Operation>,
41 policy: QueryPolicy,
42 on_failure: impl FnOnce(),
43) -> Result<Output, Error> {
44 if scope.is_some_and(Operation::is_cancelled) {
45 return Err(Error::Cancelled);
46 }
47 let child = command
48 .stdin(Stdio::null())
49 .stdout(Stdio::piped())
50 .stderr(Stdio::piped())
51 .spawn()
52 .map_err(|source| Error::Process { component, source })?;
53 let mut child = OwnedChild(child);
54 let result = receive(component, &mut child, scope, None, policy);
55 if result.is_err() && !matches!(child.0.try_wait(), Ok(Some(_))) {
56 on_failure();
57 }
58 result
59}
60
61pub fn capture_input(
66 component: Component,
67 command: &mut Command,
68 input: Vec<u8>,
69 policy: QueryPolicy,
70 cancelled: &AtomicBool,
71 closed: &AtomicBool,
72) -> Result<Output, Error> {
73 if cancelled.load(Ordering::Acquire) || closed.load(Ordering::Acquire) {
74 return Err(Error::Cancelled);
75 }
76 let child = command
77 .stdin(Stdio::piped())
78 .stdout(Stdio::piped())
79 .stderr(Stdio::piped())
80 .spawn()
81 .map_err(|source| Error::Process { component, source })?;
82 let mut child = OwnedChild(child);
83 let mut stdin = child.0.stdin.take().ok_or(Error::InvalidResponse)?;
84 let writer = thread::Builder::new()
85 .name("eska-component-input".into())
86 .spawn(move || stdin.write_all(&input))
87 .map_err(|source| Error::Process { component, source })?;
88 let result = receive(
89 component,
90 &mut child,
91 None,
92 Some((cancelled, closed)),
93 policy,
94 );
95 drop(child);
96 let written = writer
97 .join()
98 .map_err(|_| Error::Process {
99 component,
100 source: io::Error::other("input writer failed"),
101 })
102 .and_then(|result| result.map_err(|source| Error::Process { component, source }));
103 let output = result?;
104 if output.status.success() {
105 written?;
106 }
107 Ok(output)
108}
109
110fn receive(
112 component: Component,
113 child: &mut OwnedChild,
114 scope: Option<&Operation>,
115 request: Option<(&AtomicBool, &AtomicBool)>,
116 policy: QueryPolicy,
117) -> Result<Output, Error> {
118 let stdout = child.0.stdout.take().ok_or(Error::InvalidResponse)?;
119 let stderr = child.0.stderr.take().ok_or(Error::InvalidResponse)?;
120 let (sender, receiver) = mpsc::channel();
121 reader(
122 component,
123 stdout,
124 Stream::Stdout,
125 sender.clone(),
126 policy.output_limit,
127 )?;
128 reader(
129 component,
130 stderr,
131 Stream::Stderr,
132 sender,
133 policy.output_limit,
134 )?;
135 let started = Instant::now();
136 let mut stopping = None;
137 let mut status = None;
138 let mut stdout = None;
139 let mut stderr = None;
140 loop {
141 while let Ok((stream, result)) = receiver.try_recv() {
142 match stream {
143 Stream::Stdout => stdout = Some(result?),
144 Stream::Stderr => stderr = Some(result?),
145 }
146 }
147 let cancelled = scope.is_some_and(Operation::is_cancelled)
148 || request.is_some_and(|(cancelled, closed)| {
149 cancelled.load(Ordering::Acquire) || closed.load(Ordering::Acquire)
150 });
151 if stopping.is_none() && (cancelled || started.elapsed() >= policy.timeout) {
152 if policy.cancellation_grace.is_zero() {
153 return Err(reason(cancelled));
154 }
155 signal(
156 child,
157 scope
158 .and_then(Operation::received_signal)
159 .unwrap_or(TERMINATE),
160 )?;
161 stopping = Some((Instant::now(), cancelled));
162 }
163 if status.is_none() {
164 status = child
165 .0
166 .try_wait()
167 .map_err(|source| Error::Process { component, source })?;
168 }
169 if let Some((started, cancelled)) = stopping {
170 if status.is_some() || started.elapsed() >= policy.cancellation_grace {
171 return Err(reason(cancelled));
172 }
173 } else if let Some(status) = status
174 && stdout.is_some()
175 && stderr.is_some()
176 {
177 return Ok(Output {
178 status,
179 stdout: stdout.take().ok_or(Error::InvalidResponse)?,
180 stderr: stderr.take().ok_or(Error::InvalidResponse)?,
181 });
182 }
183 thread::sleep(std::time::Duration::from_millis(10));
184 }
185}
186
187const fn reason(cancelled: bool) -> Error {
189 if cancelled {
190 Error::Cancelled
191 } else {
192 Error::Timeout
193 }
194}
195
196#[cfg(unix)]
198fn signal(child: &OwnedChild, signal: i32) -> Result<(), Error> {
199 use nix::{
200 errno::Errno,
201 sys::signal::{Signal, kill},
202 unistd::Pid,
203 };
204 let pid =
205 i32::try_from(child.0.id()).map_err(|error| Error::Signal(std::io::Error::other(error)))?;
206 let signal = Signal::try_from(signal).map_err(|error| Error::Signal(error.into()))?;
207 match kill(Pid::from_raw(pid), signal) {
208 Ok(()) | Err(Errno::ESRCH) => Ok(()),
209 Err(error) => Err(Error::Signal(error.into())),
210 }
211}
212
213#[cfg(not(unix))]
215fn signal(child: &mut OwnedChild, _: i32) -> Result<(), Error> {
216 child.0.kill().map_err(Error::Signal)
217}
218
219#[cfg(all(test, unix))]
220mod tests {
221 use super::*;
222 use crate::test_support::TestDir;
223 use std::{sync::Arc, time::Duration};
224
225 #[test]
227 fn cancelling_a_blocked_input_reaps_the_child_and_writer() {
228 let directory = TestDir::new();
229 let cancelled = Arc::new(AtomicBool::new(false));
230 let signal = Arc::clone(&cancelled);
231 let trigger = thread::spawn(move || {
232 thread::sleep(Duration::from_millis(50));
233 signal.store(true, Ordering::Release);
234 });
235 let mut command = Command::new("python3");
236 command
237 .args(["-c", "import time; time.sleep(30)"])
238 .current_dir(&directory.0);
239 let start = Instant::now();
240 let result = capture_input(
241 Component::Fmt,
242 &mut command,
243 vec![b'x'; 1024 * 1024],
244 QueryPolicy {
245 timeout: Duration::from_secs(3),
246 output_limit: 1024,
247 cancellation_grace: Duration::ZERO,
248 },
249 &cancelled,
250 &AtomicBool::new(false),
251 );
252 trigger.join().unwrap();
253 assert!(matches!(result, Err(Error::Cancelled)));
254 assert!(start.elapsed() < Duration::from_secs(2));
255 }
256
257 #[test]
259 fn input_and_output_are_complete_with_bounded_streams() {
260 let directory = TestDir::new();
261 let input = vec![b'a'; 256 * 1024];
262 let mut command = Command::new("python3");
263 command.args(["-c","import sys; data=sys.stdin.buffer.read(); sys.stderr.buffer.write(b'checked'); sys.stdout.buffer.write(data)"]).current_dir(&directory.0);
264 let output = capture_input(
265 Component::Fmt,
266 &mut command,
267 input.clone(),
268 QueryPolicy {
269 timeout: Duration::from_secs(3),
270 output_limit: 1024 * 1024,
271 cancellation_grace: Duration::ZERO,
272 },
273 &AtomicBool::new(false),
274 &AtomicBool::new(false),
275 )
276 .unwrap();
277 assert!(output.status.success());
278 assert_eq!(output.stdout, input);
279 assert_eq!(output.stderr, b"checked");
280 }
281}