ignition_core/deploy/transport/
process.rs1use std::ffi::{OsStr, OsString};
28use std::process::Stdio;
29use std::time::Duration;
30
31use tokio::io::{AsyncReadExt, AsyncWriteExt};
32
33pub const STDOUT_CAP_BYTES: usize = 32 * 1024 * 1024;
36pub const STDERR_CAP_BYTES: usize = 64 * 1024;
38pub const OVERFLOW_CODE: i32 = -2;
40pub const SPAWN_FAILED_CODE: i32 = 127;
42
43#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct ProcessOutput {
46 pub stdout: Vec<u8>,
48 pub stderr: Vec<u8>,
50 pub code: i32,
53 pub timed_out: bool,
55}
56
57#[async_trait::async_trait]
59pub trait ProcessRunner: Send + Sync {
60 async fn run(
63 &self,
64 program: &OsStr,
65 args: &[OsString],
66 stdin: Option<Vec<u8>>,
67 timeout: Duration,
68 ) -> ProcessOutput;
69}
70
71pub struct TokioProcessRunner;
73
74#[async_trait::async_trait]
75impl ProcessRunner for TokioProcessRunner {
76 async fn run(
77 &self,
78 program: &OsStr,
79 args: &[OsString],
80 stdin: Option<Vec<u8>>,
81 timeout: Duration,
82 ) -> ProcessOutput {
83 run_capped(program, args, stdin, timeout, STDOUT_CAP_BYTES).await
84 }
85}
86
87pub fn utf8_locale() -> &'static str {
90 if cfg!(target_os = "macos") {
91 "en_US.UTF-8"
92 } else {
93 "C.UTF-8"
94 }
95}
96
97struct Overflow;
99
100async fn read_capped<R>(reader: Option<R>, cap: usize) -> Result<Vec<u8>, Overflow>
102where
103 R: tokio::io::AsyncRead + Unpin,
104{
105 let Some(mut reader) = reader else {
106 return Ok(Vec::new());
107 };
108 let mut out = Vec::new();
109 let mut chunk = [0u8; 8192];
110 loop {
111 match reader.read(&mut chunk).await {
112 Ok(0) | Err(_) => return Ok(out),
113 Ok(n) => {
114 if out.len() + n > cap {
115 return Err(Overflow);
116 }
117 out.extend_from_slice(&chunk[..n]);
118 }
119 }
120 }
121}
122
123async fn read_tail<R>(reader: Option<R>, cap: usize) -> Vec<u8>
125where
126 R: tokio::io::AsyncRead + Unpin,
127{
128 let Some(mut reader) = reader else {
129 return Vec::new();
130 };
131 let mut out: Vec<u8> = Vec::new();
132 let mut chunk = [0u8; 8192];
133 loop {
134 match reader.read(&mut chunk).await {
135 Ok(0) | Err(_) => return out,
136 Ok(n) => {
137 out.extend_from_slice(&chunk[..n]);
138 if out.len() > cap {
139 let excess = out.len() - cap;
140 out.drain(..excess);
141 }
142 }
143 }
144 }
145}
146
147fn failed(code: i32, timed_out: bool, note: String) -> ProcessOutput {
148 ProcessOutput {
149 stdout: Vec::new(),
150 stderr: note.into_bytes(),
151 code,
152 timed_out,
153 }
154}
155
156pub(crate) async fn run_capped(
157 program: &OsStr,
158 args: &[OsString],
159 stdin: Option<Vec<u8>>,
160 timeout: Duration,
161 stdout_cap: usize,
162) -> ProcessOutput {
163 let mut command = tokio::process::Command::new(program);
164 command
165 .args(args)
166 .env("LC_ALL", utf8_locale())
167 .stdin(Stdio::piped())
168 .stdout(Stdio::piped())
169 .stderr(Stdio::piped())
170 .kill_on_drop(true);
171 let mut child = match command.spawn() {
172 Ok(child) => child,
173 Err(err) => {
174 return failed(
175 SPAWN_FAILED_CODE,
176 false,
177 format!(
178 "failed to spawn {}: {:?}",
179 program.to_string_lossy(),
180 err.kind()
181 ),
182 );
183 }
184 };
185 let stdin_handle = child.stdin.take();
186 let stdout = child.stdout.take();
187 let stderr = child.stderr.take();
188
189 let io = async {
190 let write = async {
191 if let Some(mut handle) = stdin_handle {
192 if let Some(bytes) = &stdin {
193 let _ = handle.write_all(bytes).await;
196 }
197 let _ = handle.shutdown().await;
198 }
200 Ok::<(), Overflow>(())
201 };
202 let err = async { Ok::<Vec<u8>, Overflow>(read_tail(stderr, STDERR_CAP_BYTES).await) };
203 let ((), out, err) = tokio::try_join!(write, read_capped(stdout, stdout_cap), err)?;
204 let status = child.wait().await;
205 Ok::<_, Overflow>((out, err, status))
206 };
207
208 match tokio::time::timeout(timeout, io).await {
209 Ok(Ok((stdout, stderr, status))) => ProcessOutput {
210 stdout,
211 stderr,
212 code: status.ok().and_then(|s| s.code()).unwrap_or(-1),
213 timed_out: false,
214 },
215 Ok(Err(Overflow)) => {
216 let _ = child.kill().await;
218 failed(
219 OVERFLOW_CODE,
220 false,
221 format!("stdout exceeded {stdout_cap} bytes; the child was killed"),
222 )
223 }
224 Err(_elapsed) => {
225 let _ = child.kill().await;
226 ProcessOutput {
227 stdout: Vec::new(),
228 stderr: Vec::new(),
229 code: -1,
230 timed_out: true,
231 }
232 }
233 }
234}
235
236#[cfg(test)]
237mod tests {
238 use super::*;
239
240 fn os(s: &str) -> OsString {
241 OsString::from(s)
242 }
243
244 #[cfg(unix)]
247 #[tokio::test]
248 async fn children_run_in_a_utf8_locale() {
249 let out = run_capped(
250 &os("sh"),
251 &[os("-c"), os("printf %s \"$LC_ALL\"")],
252 None,
253 Duration::from_secs(5),
254 1024,
255 )
256 .await;
257 assert_eq!(out.code, 0);
258 assert_eq!(String::from_utf8_lossy(&out.stdout), utf8_locale());
259 assert!(utf8_locale().ends_with("UTF-8"));
260 }
261
262 #[cfg(unix)]
263 #[tokio::test]
264 async fn timeout_kills_child() {
265 let dir = tempfile::tempdir().expect("tempdir");
266 let pidfile = dir.path().join("pid");
267 let script = format!("echo $$ > '{}'; exec sleep 30", pidfile.display());
268 let started = std::time::Instant::now();
269 let out = TokioProcessRunner
270 .run(
271 OsStr::new("sh"),
272 &[os("-c"), os(&script)],
273 None,
274 Duration::from_millis(300),
275 )
276 .await;
277 assert!(out.timed_out, "{out:?}");
278 assert!(started.elapsed() < Duration::from_secs(5));
279 let pid = std::fs::read_to_string(&pidfile)
280 .expect("pid file")
281 .trim()
282 .to_string();
283 let alive = std::process::Command::new("kill")
284 .args(["-0", &pid])
285 .stderr(std::process::Stdio::null())
286 .status()
287 .expect("kill")
288 .success();
289 assert!(!alive, "child {pid} survived the timeout");
290 }
291
292 #[cfg(windows)]
293 #[tokio::test]
294 async fn timeout_kills_child() {
295 let started = std::time::Instant::now();
296 let out = TokioProcessRunner
297 .run(
298 OsStr::new("ping"),
299 &[os("-n"), os("30"), os("127.0.0.1")],
300 None,
301 Duration::from_millis(300),
302 )
303 .await;
304 assert!(out.timed_out, "{out:?}");
305 assert!(started.elapsed() < Duration::from_secs(5));
306 }
307
308 #[cfg(unix)]
309 #[tokio::test]
310 async fn stdin_is_closed() {
311 let empty = TokioProcessRunner
312 .run(OsStr::new("cat"), &[], None, Duration::from_secs(10))
313 .await;
314 assert_eq!((empty.code, empty.timed_out), (0, false), "{empty:?}");
315 assert!(empty.stdout.is_empty());
316 let fed = TokioProcessRunner
317 .run(
318 OsStr::new("cat"),
319 &[],
320 Some(b"abc".to_vec()),
321 Duration::from_secs(10),
322 )
323 .await;
324 assert_eq!(fed.stdout, b"abc");
325 assert_eq!(fed.code, 0);
326 }
327
328 #[tokio::test]
329 async fn spawn_failure_is_127() {
330 let out = TokioProcessRunner
331 .run(
332 OsStr::new("ign-no-such-program-22-05"),
333 &[],
334 None,
335 Duration::from_secs(5),
336 )
337 .await;
338 assert_eq!(out.code, 127);
339 assert!(!out.timed_out);
340 let reason = String::from_utf8_lossy(&out.stderr).into_owned();
341 assert!(reason.contains("failed to spawn"), "{reason}");
342 }
343
344 #[cfg(unix)]
345 #[tokio::test]
346 async fn stdout_over_the_cap_kills_and_fails() {
347 let out = run_capped(
348 OsStr::new("sh"),
349 &[os("-c"), os("yes | head -c 1000000; exec sleep 30")],
350 None,
351 Duration::from_secs(20),
352 1024,
353 )
354 .await;
355 assert_eq!(out.code, OVERFLOW_CODE, "{out:?}");
356 assert!(!out.timed_out);
357 assert!(out.stdout.is_empty());
358 }
359
360 #[cfg(unix)]
361 #[tokio::test]
362 async fn stderr_keeps_only_its_tail() {
363 let out = TokioProcessRunner
364 .run(
365 OsStr::new("sh"),
366 &[
367 os("-c"),
368 os("head -c 200000 /dev/zero | tr '\\0' x >&2; echo END >&2"),
369 ],
370 None,
371 Duration::from_secs(20),
372 )
373 .await;
374 assert_eq!(out.code, 0);
375 assert!(out.stderr.len() <= STDERR_CAP_BYTES);
376 assert!(out.stderr.ends_with(b"END\n"));
377 }
378
379 #[cfg(unix)]
380 #[tokio::test]
381 async fn nonzero_exit_code_is_reported() {
382 let out = TokioProcessRunner
383 .run(
384 OsStr::new("sh"),
385 &[os("-c"), os("echo oops >&2; exit 3")],
386 None,
387 Duration::from_secs(10),
388 )
389 .await;
390 assert_eq!(out.code, 3);
391 assert_eq!(out.stderr, b"oops\n");
392 }
393}