1use std::collections::BTreeSet;
39use std::io::{self, Read, Write};
40use std::path::PathBuf;
41use std::process::{Child, Command, ExitStatus, Stdio};
42use std::time::{Duration, Instant};
43
44#[derive(Debug, Clone, PartialEq, Eq)]
46pub enum EnvPolicy {
47 InheritExcept { deny: Vec<String> },
56 ClearExcept { allow: Vec<String> },
62}
63
64impl EnvPolicy {
65 pub fn minimal_allow() -> Vec<String> {
69 let base: &[&str] = if cfg!(windows) {
70 &["PATH", "PATHEXT", "SYSTEMROOT", "SYSTEMDRIVE", "COMSPEC", "TEMP", "TMP", "USERPROFILE"]
71 } else {
72 &["PATH", "HOME", "TMPDIR", "LANG", "LC_ALL", "TZ"]
73 };
74 base.iter().map(|s| s.to_string()).collect()
75 }
76}
77
78impl Default for EnvPolicy {
79 fn default() -> Self {
80 EnvPolicy::InheritExcept { deny: secret_env_vars() }
81 }
82}
83
84static SECRET_ENV: std::sync::Mutex<Option<BTreeSet<String>>> = std::sync::Mutex::new(None);
92
93pub fn deny_env_var(name: &str) {
99 if name.trim().is_empty() {
100 return;
101 }
102 let mut guard = SECRET_ENV.lock().unwrap_or_else(|e| e.into_inner());
103 guard.get_or_insert_with(BTreeSet::new).insert(name.to_string());
104}
105
106pub fn secret_env_vars() -> Vec<String> {
108 let guard = SECRET_ENV.lock().unwrap_or_else(|e| e.into_inner());
109 guard.as_ref().map(|s| s.iter().cloned().collect()).unwrap_or_default()
110}
111
112#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
114pub enum StderrMode {
115 #[default]
118 Pipe,
119 Inherit,
122}
123
124#[derive(Debug, Clone)]
126pub struct SpawnPolicy {
127 pub timeout: Option<Duration>,
131 pub max_output_bytes: usize,
135 pub env: EnvPolicy,
136 pub stderr: StderrMode,
137 pub current_dir: Option<PathBuf>,
139}
140
141pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(300);
145
146pub const DEFAULT_MAX_OUTPUT: usize = 64 * 1024 * 1024;
150
151impl Default for SpawnPolicy {
152 fn default() -> Self {
153 SpawnPolicy {
154 timeout: Some(DEFAULT_TIMEOUT),
155 max_output_bytes: DEFAULT_MAX_OUTPUT,
156 env: EnvPolicy::default(),
157 stderr: StderrMode::default(),
158 current_dir: None,
159 }
160 }
161}
162
163impl SpawnPolicy {
164 pub fn deny_vars<I: IntoIterator<Item = String>>(mut self, vars: I) -> Self {
168 let mut denied: BTreeSet<String> = match self.env {
174 EnvPolicy::InheritExcept { deny } => deny.into_iter().collect(),
175 EnvPolicy::ClearExcept { allow } => {
176 self.env = EnvPolicy::ClearExcept { allow };
177 return self;
178 }
179 };
180 denied.extend(vars);
181 self.env = EnvPolicy::InheritExcept { deny: denied.into_iter().collect() };
182 self
183 }
184
185 pub fn timeout(mut self, timeout: Option<Duration>) -> Self {
186 self.timeout = timeout;
187 self
188 }
189
190 pub fn stderr(mut self, mode: StderrMode) -> Self {
191 self.stderr = mode;
192 self
193 }
194}
195
196#[derive(Debug)]
198pub struct SpawnOutput {
199 pub status: ExitStatus,
200 pub stdout: Vec<u8>,
201 pub stderr: Vec<u8>,
202 pub timed_out: bool,
205 pub stdout_truncated: bool,
206 pub stderr_truncated: bool,
207}
208
209impl SpawnOutput {
210 pub fn stderr_text(&self) -> String {
212 String::from_utf8_lossy(&self.stderr).trim().to_string()
213 }
214
215 pub fn failure(&self, what: &str) -> Option<String> {
218 if self.timed_out {
219 return Some(format!("{what} timed out and was killed"));
220 }
221 if !self.status.success() {
222 let err = self.stderr_text();
223 return Some(if err.is_empty() {
224 format!("{what} exited with {}", self.status)
225 } else {
226 format!("{what} exited with {}: {err}", self.status)
227 });
228 }
229 None
230 }
231}
232
233fn spawn_past_busy_exec(cmd: &mut Command) -> io::Result<Child> {
241 let mut attempt = 0u32;
242 loop {
243 match cmd.spawn() {
244 Err(e) if e.kind() == io::ErrorKind::ExecutableFileBusy && attempt < 20 => {
245 attempt += 1;
246 std::thread::sleep(Duration::from_millis(5 * u64::from(attempt)));
247 }
248 other => return other,
249 }
250 }
251}
252
253pub fn run(
263 mut cmd: Command,
264 stdin: Option<&[u8]>,
265 extra_env: &[(&str, &str)],
266 policy: &SpawnPolicy,
267) -> io::Result<SpawnOutput> {
268 match &policy.env {
269 EnvPolicy::InheritExcept { deny } => {
270 for var in deny {
271 cmd.env_remove(var);
272 }
273 }
274 EnvPolicy::ClearExcept { allow } => {
275 cmd.env_clear();
276 for var in allow {
277 if let Ok(val) = std::env::var(var) {
278 cmd.env(var, val);
279 }
280 }
281 }
282 }
283 for (k, v) in extra_env {
284 cmd.env(k, v);
285 }
286 if let Some(dir) = &policy.current_dir {
287 cmd.current_dir(dir);
288 }
289
290 cmd.stdin(if stdin.is_some() { Stdio::piped() } else { Stdio::null() })
291 .stdout(Stdio::piped())
292 .stderr(match policy.stderr {
293 StderrMode::Pipe => Stdio::piped(),
294 StderrMode::Inherit => Stdio::inherit(),
295 });
296
297 let mut child = spawn_past_busy_exec(&mut cmd)?;
298
299 let stdin_thread = child.stdin.take().map(|mut pipe| {
303 let payload = stdin.unwrap_or_default().to_vec();
304 std::thread::spawn(move || {
305 let _ = pipe.write_all(&payload);
306 })
308 });
309
310 let cap = policy.max_output_bytes;
311 let out_thread = child.stdout.take().map(|pipe| std::thread::spawn(move || drain(pipe, cap)));
312 let err_thread = child.stderr.take().map(|pipe| std::thread::spawn(move || drain(pipe, cap)));
313
314 let (status, timed_out) = wait_bounded(&mut child, policy.timeout)?;
315
316 if let Some(t) = stdin_thread {
317 let _ = t.join();
318 }
319 let (stdout, stdout_truncated) = out_thread.and_then(|t| t.join().ok()).unwrap_or((Vec::new(), false));
320 let (stderr, stderr_truncated) = err_thread.and_then(|t| t.join().ok()).unwrap_or((Vec::new(), false));
321
322 Ok(SpawnOutput { status, stdout, stderr, timed_out, stdout_truncated, stderr_truncated })
323}
324
325fn drain<R: Read>(mut src: R, cap: usize) -> (Vec<u8>, bool) {
329 let mut kept = Vec::new();
330 let mut buf = [0u8; 16 * 1024];
331 let mut truncated = false;
332 loop {
333 match src.read(&mut buf) {
334 Ok(0) => break,
335 Ok(n) => {
336 if kept.len() < cap {
337 let room = cap - kept.len();
338 let take = room.min(n);
339 kept.extend_from_slice(&buf[..take]);
340 if take < n {
341 truncated = true;
342 }
343 } else {
344 truncated = true;
345 }
346 }
347 Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
348 Err(_) => break,
349 }
350 }
351 (kept, truncated)
352}
353
354fn wait_bounded(child: &mut Child, timeout: Option<Duration>) -> io::Result<(ExitStatus, bool)> {
360 let Some(limit) = timeout else {
361 return Ok((child.wait()?, false));
362 };
363 let deadline = Instant::now() + limit;
364 let mut nap = Duration::from_millis(1);
365 loop {
366 if let Some(status) = child.try_wait()? {
367 return Ok((status, false));
368 }
369 if Instant::now() >= deadline {
370 let _ = child.kill();
371 let status = child.wait()?;
373 return Ok((status, true));
374 }
375 std::thread::sleep(nap);
376 nap = (nap * 2).min(Duration::from_millis(50));
377 }
378}
379
380#[cfg(test)]
381mod tests {
382 use super::*;
383
384 fn sh(script: &str) -> Command {
385 let mut c = Command::new("/bin/sh");
386 c.arg("-c").arg(script);
387 c
388 }
389
390 #[test]
391 #[cfg_attr(windows, ignore = "uses /bin/sh")]
392 fn captures_stdout_and_exit_status() {
393 let out = run(sh("printf hello"), None, &[], &SpawnPolicy::default()).unwrap();
394 assert!(out.status.success());
395 assert_eq!(out.stdout, b"hello");
396 assert!(!out.timed_out);
397 assert!(!out.stdout_truncated);
398 }
399
400 #[test]
401 #[cfg_attr(windows, ignore = "uses /bin/sh")]
402 fn stdin_reaches_the_child() {
403 let out = run(sh("cat"), Some(b"payload"), &[], &SpawnPolicy::default()).unwrap();
404 assert_eq!(out.stdout, b"payload");
405 }
406
407 #[test]
408 #[cfg_attr(windows, ignore = "uses /bin/sh")]
409 fn timeout_kills_a_hung_child() {
410 let policy = SpawnPolicy::default().timeout(Some(Duration::from_millis(150)));
411 let out = run(sh("sleep 30"), None, &[], &policy).unwrap();
412 assert!(out.timed_out, "expected the child to be killed");
413 assert!(!out.status.success());
414 assert!(out.failure("tool").unwrap().contains("timed out"));
415 }
416
417 #[test]
418 #[cfg_attr(windows, ignore = "uses /bin/sh")]
419 fn output_cap_truncates_without_hanging() {
420 let policy = SpawnPolicy { max_output_bytes: 1024, ..SpawnPolicy::default() };
422 let out = run(sh("head -c 200000 /dev/zero"), None, &[], &policy).unwrap();
423 assert_eq!(out.stdout.len(), 1024);
424 assert!(out.stdout_truncated);
425 assert!(!out.timed_out, "draining past the cap must not stall the child");
426 }
427
428 #[test]
429 #[cfg_attr(windows, ignore = "uses /bin/sh")]
430 fn large_stdin_does_not_deadlock() {
431 let big = vec![b'x'; 4 * 1024 * 1024];
434 let policy = SpawnPolicy::default().timeout(Some(Duration::from_secs(20)));
435 let out = run(sh("cat"), Some(&big), &[], &policy).unwrap();
436 assert!(!out.timed_out, "write-then-read deadlock");
437 assert_eq!(out.stdout.len(), big.len());
438 }
439
440 #[test]
441 #[cfg_attr(windows, ignore = "uses /bin/sh")]
442 fn denied_vars_do_not_reach_the_child() {
443 std::env::set_var("AREEV_TEST_SECRET", "hunter2");
444 let policy = SpawnPolicy::default().deny_vars(["AREEV_TEST_SECRET".to_string()]);
445 let out = run(sh("printf %s \"${AREEV_TEST_SECRET:-absent}\""), None, &[], &policy).unwrap();
446 std::env::remove_var("AREEV_TEST_SECRET");
447 assert_eq!(String::from_utf8_lossy(&out.stdout), "absent");
448 }
449
450 #[test]
451 #[cfg_attr(windows, ignore = "uses /bin/sh")]
452 fn inherited_vars_still_reach_the_child() {
453 std::env::set_var("AREEV_TEST_KEEP", "kept");
456 std::env::set_var("AREEV_TEST_DROP", "dropped");
457 let policy = SpawnPolicy::default().deny_vars(["AREEV_TEST_DROP".to_string()]);
458 let out = run(sh("printf %s \"${AREEV_TEST_KEEP:-absent}\""), None, &[], &policy).unwrap();
459 std::env::remove_var("AREEV_TEST_KEEP");
460 std::env::remove_var("AREEV_TEST_DROP");
461 assert_eq!(String::from_utf8_lossy(&out.stdout), "kept");
462 }
463
464 #[test]
465 #[cfg_attr(windows, ignore = "uses /bin/sh")]
466 fn clear_except_drops_everything_unlisted_but_keeps_extras() {
467 std::env::set_var("AREEV_TEST_AMBIENT", "ambient");
468 let policy = SpawnPolicy {
469 env: EnvPolicy::ClearExcept { allow: EnvPolicy::minimal_allow() },
470 ..SpawnPolicy::default()
471 };
472 let out = run(
473 sh("printf %s \"${AREEV_TEST_AMBIENT:-absent}/${AREEV_EXTRA:-none}\""),
474 None,
475 &[("AREEV_EXTRA", "set")],
476 &policy,
477 )
478 .unwrap();
479 std::env::remove_var("AREEV_TEST_AMBIENT");
480 assert_eq!(String::from_utf8_lossy(&out.stdout), "absent/set");
481 }
482
483 #[test]
484 #[cfg_attr(windows, ignore = "uses /bin/sh")]
485 fn nonzero_exit_reports_stderr() {
486 let out = run(sh("echo boom >&2; exit 3"), None, &[], &SpawnPolicy::default()).unwrap();
487 let msg = out.failure("embed command").unwrap();
488 assert!(msg.contains("embed command"), "{msg}");
489 assert!(msg.contains("boom"), "{msg}");
490 }
491
492 #[test]
493 #[cfg_attr(windows, ignore = "uses /bin/sh")]
494 fn registered_secrets_are_scrubbed_without_the_seam_asking() {
495 std::env::set_var("AREEV_TEST_REGISTERED", "hunter2");
498 deny_env_var("AREEV_TEST_REGISTERED");
499 let out = run(
500 sh("printf %s \"${AREEV_TEST_REGISTERED:-absent}\""),
501 None,
502 &[],
503 &SpawnPolicy::default(),
504 )
505 .unwrap();
506 std::env::remove_var("AREEV_TEST_REGISTERED");
507 assert_eq!(String::from_utf8_lossy(&out.stdout), "absent");
508 }
509
510 #[test]
511 fn deny_env_var_ignores_blanks() {
512 deny_env_var(" ");
513 deny_env_var("");
514 assert!(!secret_env_vars().iter().any(|v| v.trim().is_empty()));
515 }
516
517 #[test]
520 #[cfg(target_os = "linux")]
521 fn an_executable_still_open_for_writing_is_retried_not_failed() {
522 use std::os::unix::fs::PermissionsExt;
523 let dir = std::env::temp_dir().join(format!("areev-etxtbsy-{}", std::process::id()));
524 std::fs::create_dir_all(&dir).unwrap();
525 let path = dir.join("busy.sh");
526 let mut held = std::fs::File::create(&path).unwrap();
527 held.write_all(b"#!/bin/sh\nprintf ok\n").unwrap();
528 std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o700)).unwrap();
529
530 let bare = Command::new(&path).spawn().map(|mut c| c.wait());
534 assert_eq!(bare.err().map(|e| e.kind()), Some(io::ErrorKind::ExecutableFileBusy));
535
536 let release = std::thread::spawn(move || {
537 std::thread::sleep(Duration::from_millis(40));
538 drop(held);
539 });
540 let out = run(Command::new(&path), None, &[], &SpawnPolicy::default()).unwrap();
541 release.join().unwrap();
542 std::fs::remove_dir_all(&dir).ok();
543 assert_eq!(out.stdout, b"ok");
544 }
545
546 #[test]
547 fn spawn_failure_is_an_error_not_a_panic() {
548 let cmd = Command::new("areev-no-such-binary-eaf1");
549 assert!(run(cmd, None, &[], &SpawnPolicy::default()).is_err());
550 }
551}