1use crate::{CapturedOutput, NativeOutputCapture, OwnedCommand};
11use std::{
12 ffi::OsStr,
13 io,
14 process::{ChildStdin, ExitStatus, Stdio},
15 sync::{
16 atomic::{AtomicBool, Ordering},
17 mpsc, Arc,
18 },
19 thread::JoinHandle,
20 time::{Duration, Instant},
21};
22
23const DRAIN_WAIT: Duration = Duration::from_secs(2);
26
27const STDIN_WAIT: Duration = Duration::from_millis(500);
31
32const STDIN_CANCEL_WAIT: Duration = Duration::from_secs(2);
36
37const CANCEL_RETRY: Duration = Duration::from_millis(25);
39
40#[cfg(windows)]
43const WRITE_CHUNK: usize = 64 * 1024;
44
45#[cfg(unix)]
47const WRITE_POLL: Duration = Duration::from_millis(50);
48
49#[derive(Clone, Copy, Debug, Eq, PartialEq)]
51pub struct RunOptions {
52 pub timeout: Duration,
54 pub stdout_cap: usize,
56 pub stderr_cap: usize,
58 pub no_window: bool,
60}
61
62impl RunOptions {
63 pub fn new(timeout: Duration) -> Self {
65 Self {
66 timeout,
67 stdout_cap: 1024 * 1024,
68 stderr_cap: 1024 * 1024,
69 no_window: true,
70 }
71 }
72
73 pub fn stdout_cap(mut self, bytes: usize) -> Self {
74 self.stdout_cap = bytes;
75 self
76 }
77
78 pub fn stderr_cap(mut self, bytes: usize) -> Self {
79 self.stderr_cap = bytes;
80 self
81 }
82
83 pub fn no_window(mut self, hidden: bool) -> Self {
84 self.no_window = hidden;
85 self
86 }
87}
88
89#[derive(Clone, Debug, Eq, PartialEq)]
91pub struct RunOutput {
92 pub status: ExitStatus,
95 pub timed_out: bool,
97 pub stdout: CapturedOutput,
98 pub stderr: CapturedOutput,
99}
100
101pub fn run_with_stdin<I, S>(
124 program: impl AsRef<OsStr>,
125 args: I,
126 input: &[u8],
127 options: &RunOptions,
128) -> io::Result<RunOutput>
129where
130 I: IntoIterator<Item = S>,
131 S: AsRef<OsStr>,
132{
133 run_inner(program, args, Some(input.to_vec()), options)
134}
135
136pub fn run<I, S>(program: impl AsRef<OsStr>, args: I, options: &RunOptions) -> io::Result<RunOutput>
140where
141 I: IntoIterator<Item = S>,
142 S: AsRef<OsStr>,
143{
144 run_inner(program, args, None, options)
145}
146
147fn run_inner<I, S>(
148 program: impl AsRef<OsStr>,
149 args: I,
150 input: Option<Vec<u8>>,
151 options: &RunOptions,
152) -> io::Result<RunOutput>
153where
154 I: IntoIterator<Item = S>,
155 S: AsRef<OsStr>,
156{
157 let mut command = OwnedCommand::new(program);
158 command
159 .command_mut()
160 .args(args)
161 .stdin(if input.is_some() {
162 Stdio::piped()
163 } else {
164 Stdio::null()
165 })
166 .stdout(Stdio::piped())
167 .stderr(Stdio::piped());
168 if options.no_window {
169 command.windows_hide();
170 } else {
171 command.windows_visible();
172 }
173 let mut child = command.spawn()?;
174 let writer = match (input, child.take_stdin()) {
178 (Some(bytes), Some(stdin)) => Some(StdinWriter::spawn(stdin, bytes)?),
179 _ => None,
180 };
181 let mut stdout = child
182 .take_stdout()
183 .map(|pipe| NativeOutputCapture::stdout(pipe, options.stdout_cap))
184 .transpose()?;
185 let mut stderr = child
186 .take_stderr()
187 .map(|pipe| NativeOutputCapture::stderr(pipe, options.stderr_cap))
188 .transpose()?;
189 let (status, timed_out) = match child.wait_then_kill(options.timeout)? {
192 crate::WaitOutcome::Exited(status) => (status, false),
193 crate::WaitOutcome::Terminated(status) => (status, true),
194 };
195 let stdout = stdout
196 .as_mut()
197 .map_or_else(CapturedOutput::default, |capture| {
198 capture.finish_within(DRAIN_WAIT)
199 });
200 let stderr = stderr
201 .as_mut()
202 .map_or_else(CapturedOutput::default, |capture| {
203 capture.finish_within(DRAIN_WAIT)
204 });
205 if let Some(writer) = writer {
210 writer.finish()?;
211 }
212 Ok(RunOutput {
213 status,
214 timed_out,
215 stdout,
216 stderr,
217 })
218}
219
220struct WriterState {
222 cancel: AtomicBool,
223 #[cfg(windows)]
227 handle: std::sync::Mutex<Option<usize>>,
228 #[cfg(unix)]
233 pipe: std::sync::Mutex<Option<ChildStdin>>,
234}
235
236impl WriterState {
237 fn new() -> Self {
238 Self {
239 cancel: AtomicBool::new(false),
240 #[cfg(windows)]
241 handle: std::sync::Mutex::new(None),
242 #[cfg(unix)]
243 pipe: std::sync::Mutex::new(None),
244 }
245 }
246
247 #[cfg(unix)]
248 fn pipe(&self) -> std::sync::MutexGuard<'_, Option<ChildStdin>> {
249 self.pipe
250 .lock()
251 .unwrap_or_else(std::sync::PoisonError::into_inner)
252 }
253
254 #[cfg(unix)]
256 fn close_pipe(&self) {
257 drop(self.pipe().take());
258 }
259
260 fn cancelled(&self) -> bool {
261 self.cancel.load(Ordering::Acquire)
262 }
263
264 fn cancel(&self) {
267 self.cancel.store(true, Ordering::Release);
268 #[cfg(windows)]
269 {
270 let held = self
271 .handle
272 .lock()
273 .unwrap_or_else(std::sync::PoisonError::into_inner);
274 if let Some(raw) = *held {
275 unsafe { crate::native_capture::platform::cancel_all_io(raw as _) };
279 }
280 }
281 }
282}
283
284#[cfg(test)]
287static LIVE_WRITERS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
288
289#[cfg(test)]
290struct LiveWriter;
291
292#[cfg(test)]
293impl LiveWriter {
294 fn start() -> Self {
295 LIVE_WRITERS.fetch_add(1, Ordering::SeqCst);
296 Self
297 }
298}
299
300#[cfg(test)]
301impl Drop for LiveWriter {
302 fn drop(&mut self) {
303 LIVE_WRITERS.fetch_sub(1, Ordering::SeqCst);
304 }
305}
306
307struct StdinWriter {
312 state: Arc<WriterState>,
313 done: mpsc::Receiver<io::Result<()>>,
314 handle: Option<JoinHandle<()>>,
315}
316
317impl StdinWriter {
318 fn spawn(stdin: ChildStdin, bytes: Vec<u8>) -> io::Result<Self> {
319 let state = Arc::new(WriterState::new());
320 let (sender, done) = mpsc::channel();
321 let shared = Arc::clone(&state);
322 #[cfg(unix)]
323 {
324 *state.pipe() = Some(stdin);
325 }
326 #[cfg(test)]
327 let live = LiveWriter::start();
328 let handle = std::thread::Builder::new()
329 .name("rightkit-run-stdin".into())
330 .spawn(move || {
331 #[cfg(test)]
332 let _live = live;
333 #[cfg(unix)]
336 let result = {
337 let result = write_input(&bytes, &shared);
338 shared.close_pipe();
339 result
340 };
341 #[cfg(windows)]
342 let result = {
343 let mut stdin = stdin;
344 let result = write_input(&mut stdin, &bytes, &shared);
345 drop(stdin);
346 result
347 };
348 let _ = sender.send(result);
349 })?;
350 Ok(Self {
351 state,
352 done,
353 handle: Some(handle),
354 })
355 }
356
357 fn finish(mut self) -> io::Result<()> {
363 let result = match self.done.recv_timeout(STDIN_WAIT) {
364 Err(mpsc::RecvTimeoutError::Timeout) => self.cancel_and_wait(),
365 other => other,
366 };
367 self.settle(result)
368 }
369
370 fn cancel_and_wait(&self) -> Result<io::Result<()>, mpsc::RecvTimeoutError> {
372 let deadline = Instant::now() + STDIN_CANCEL_WAIT;
373 loop {
374 self.state.cancel();
375 match self.done.recv_timeout(CANCEL_RETRY) {
376 Err(mpsc::RecvTimeoutError::Timeout) if Instant::now() < deadline => {}
377 other => return other,
378 }
379 }
380 }
381
382 fn settle(&mut self, result: Result<io::Result<()>, mpsc::RecvTimeoutError>) -> io::Result<()> {
383 let handle = self.handle.take();
384 match result {
385 Ok(result) => {
386 if let Some(handle) = handle {
387 let _ = handle.join();
388 }
389 match result {
390 Err(error) if error.kind() != io::ErrorKind::BrokenPipe => Err(error),
391 _ => Ok(()),
392 }
393 }
394 Err(mpsc::RecvTimeoutError::Timeout) => {
395 #[cfg(unix)]
402 self.state.close_pipe();
403 Err(io::Error::new(
404 io::ErrorKind::TimedOut,
405 "the stdin writer did not stop within its cancellation bound",
406 ))
407 }
408 Err(mpsc::RecvTimeoutError::Disconnected) => {
409 if let Some(handle) = handle {
410 let _ = handle.join();
411 }
412 Err(io::Error::other("the stdin writer thread panicked"))
413 }
414 }
415 }
416}
417
418impl Drop for StdinWriter {
419 fn drop(&mut self) {
420 if self.handle.is_some() {
425 let result = self.cancel_and_wait();
426 let _ = self.settle(result);
427 }
428 }
429}
430
431#[cfg(unix)]
436fn write_input(bytes: &[u8], state: &WriterState) -> io::Result<()> {
437 use std::os::fd::AsRawFd;
438 let fd = {
439 let guard = state.pipe();
440 let Some(pipe) = guard.as_ref() else {
441 return Ok(());
442 };
443 let fd = pipe.as_raw_fd();
444 unsafe {
447 let flags = libc::fcntl(fd, libc::F_GETFL);
448 if flags < 0 || libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) < 0 {
449 return Err(io::Error::last_os_error());
450 }
451 }
452 fd
453 };
454 let mut written = 0;
455 while written < bytes.len() {
456 if state.cancelled() {
457 return Ok(());
458 }
459 let rest = &bytes[written..];
460 let chunk = rest.len().min(1 << 20);
461 let outcome = {
464 let guard = state.pipe();
465 let Some(pipe) = guard.as_ref() else {
466 return Ok(());
468 };
469 let count = unsafe { libc::write(pipe.as_raw_fd(), rest.as_ptr().cast(), chunk) };
472 if count >= 0 {
473 Ok(count as usize)
474 } else {
475 Err(io::Error::last_os_error())
476 }
477 };
478 match outcome {
479 Ok(count) => written += count,
480 Err(error) => match error.kind() {
481 io::ErrorKind::Interrupted => {}
482 io::ErrorKind::WouldBlock => {
483 let mut entry = libc::pollfd {
484 fd,
485 events: libc::POLLOUT,
486 revents: 0,
487 };
488 unsafe { libc::poll(&mut entry, 1, WRITE_POLL.as_millis() as libc::c_int) };
491 }
492 _ => return Err(error),
493 },
494 }
495 }
496 Ok(())
497}
498
499#[cfg(windows)]
503fn write_input(stdin: &mut ChildStdin, bytes: &[u8], state: &WriterState) -> io::Result<()> {
504 use std::{io::Write, os::windows::io::AsRawHandle};
505 let publish = |value: Option<usize>| {
506 *state
507 .handle
508 .lock()
509 .unwrap_or_else(std::sync::PoisonError::into_inner) = value;
510 };
511 publish(Some(stdin.as_raw_handle() as usize));
512 let mut result = Ok(());
513 for chunk in bytes.chunks(WRITE_CHUNK) {
514 if state.cancelled() {
515 break;
516 }
517 if let Err(error) = stdin.write_all(chunk) {
518 result = Err(error);
519 break;
520 }
521 }
522 publish(None);
524 match result {
525 Err(_) if state.cancelled() => Ok(()),
526 other => other,
527 }
528}
529
530#[cfg(test)]
531mod tests {
532 use super::*;
533 use std::sync::atomic::Ordering::SeqCst;
534
535 static WRITER_TESTS: std::sync::Mutex<()> = std::sync::Mutex::new(());
538
539 fn serial() -> std::sync::MutexGuard<'static, ()> {
540 WRITER_TESTS
541 .lock()
542 .unwrap_or_else(std::sync::PoisonError::into_inner)
543 }
544
545 struct KillOnDrop(std::process::Child);
547
548 impl Drop for KillOnDrop {
549 fn drop(&mut self) {
550 let _ = self.0.kill();
551 let _ = self.0.wait();
552 }
553 }
554
555 #[test]
556 fn defaults_hide_the_window_and_cap_both_streams() {
557 let options = RunOptions::new(Duration::from_secs(1));
558 assert!(options.no_window);
559 assert_eq!(options.stdout_cap, 1024 * 1024);
560 assert_eq!(options.stderr_cap, 1024 * 1024);
561 let options = options.stdout_cap(4).stderr_cap(8).no_window(false);
562 assert_eq!((options.stdout_cap, options.stderr_cap), (4, 8));
563 assert!(!options.no_window);
564 }
565
566 #[cfg(unix)]
567 #[test]
568 fn a_quick_program_returns_its_output_and_status() {
569 let options = RunOptions::new(Duration::from_secs(10));
570 let output = run("sh", ["-c", "printf hi; exit 3"], &options).unwrap();
571 assert!(!output.timed_out);
572 assert_eq!(output.status.code(), Some(3));
573 assert_eq!(output.stdout.bytes, b"hi".to_vec());
574 }
575
576 #[cfg(unix)]
577 #[test]
578 fn a_program_running_within_the_deadline_is_not_signalled() {
579 let options = RunOptions::new(Duration::from_secs(10));
580 let output = run("sh", ["-c", "sleep 1; printf done"], &options).unwrap();
581 assert!(!output.timed_out);
582 assert!(output.status.success());
583 assert_eq!(output.stdout.bytes, b"done".to_vec());
584 }
585
586 #[cfg(unix)]
587 #[test]
588 fn a_slow_program_is_terminated_at_the_deadline() {
589 let options = RunOptions::new(Duration::from_millis(200));
590 let output = run("sh", ["-c", "sleep 30"], &options).unwrap();
591 assert!(output.timed_out);
592 }
593
594 #[cfg(unix)]
595 #[test]
596 fn stdin_input_reaches_the_child_and_early_exit_is_not_an_error() {
597 let _serial = serial();
598 let options = RunOptions::new(Duration::from_secs(10));
599 let output = run_with_stdin("sh", ["-c", "cat"], b"hello", &options).unwrap();
600 assert!(output.status.success());
601 assert_eq!(output.stdout.bytes, b"hello".to_vec());
602 let big = vec![b'x'; 1 << 20];
605 let output = run_with_stdin("sh", ["-c", "exit 7"], &big, &options).unwrap();
606 assert_eq!(output.status.code(), Some(7));
607 assert_eq!(LIVE_WRITERS.load(SeqCst), 0);
608 }
609
610 #[test]
614 fn a_reader_that_never_reads_cannot_pin_the_writer() {
615 let _serial = serial();
616 #[cfg(unix)]
617 let mut command = {
618 let mut command = std::process::Command::new("sleep");
619 command.arg("30");
620 command
621 };
622 #[cfg(windows)]
623 let mut command = {
624 let mut command = std::process::Command::new("ping");
625 command.args(["-n", "30", "127.0.0.1"]);
626 command
627 };
628 let mut child = KillOnDrop(
629 command
630 .stdin(Stdio::piped())
631 .stdout(Stdio::null())
632 .stderr(Stdio::null())
633 .spawn()
634 .unwrap(),
635 );
636 let stdin = child.0.stdin.take().unwrap();
637 let writer = StdinWriter::spawn(stdin, vec![b'x'; 2 << 20]).unwrap();
638 std::thread::sleep(Duration::from_millis(200));
640 assert_eq!(LIVE_WRITERS.load(SeqCst), 1, "the writer should be blocked");
641 let started = std::time::Instant::now();
642 writer.finish().unwrap();
643 let elapsed = started.elapsed();
644 assert_eq!(
645 LIVE_WRITERS.load(SeqCst),
646 0,
647 "the writer thread must be gone"
648 );
649 assert!(
650 elapsed < STDIN_WAIT + STDIN_CANCEL_WAIT + Duration::from_secs(3),
651 "finish took {elapsed:?}"
652 );
653 assert!(
654 child.0.try_wait().unwrap().is_none(),
655 "the reader must have stayed alive throughout"
656 );
657 }
658
659 #[cfg(unix)]
664 #[test]
665 fn a_writer_that_never_reports_times_out_and_the_parent_closes_the_write_end() {
666 let _serial = serial();
667 let mut child = KillOnDrop(
668 std::process::Command::new("cat")
669 .stdin(Stdio::piped())
670 .stdout(Stdio::null())
671 .stderr(Stdio::null())
672 .spawn()
673 .unwrap(),
674 );
675 let state = Arc::new(WriterState::new());
676 *state.pipe() = Some(child.0.stdin.take().unwrap());
677 let (_sender, done) = mpsc::channel::<io::Result<()>>();
679 let (release, parked) = mpsc::channel::<()>();
680 let stuck = std::thread::spawn(move || {
681 let _ = parked.recv();
682 });
683 let writer = StdinWriter {
684 state: Arc::clone(&state),
685 done,
686 handle: Some(stuck),
687 };
688 let started = Instant::now();
689 let error = writer.finish().unwrap_err();
690 let elapsed = started.elapsed();
691 assert_eq!(error.kind(), io::ErrorKind::TimedOut);
692 assert!(
693 elapsed >= STDIN_WAIT + STDIN_CANCEL_WAIT
694 && elapsed < STDIN_WAIT + STDIN_CANCEL_WAIT + Duration::from_secs(3),
695 "finish took {elapsed:?}"
696 );
697 assert!(state.cancelled(), "the cancel flag must stay set");
698 assert!(
699 state.pipe().is_none(),
700 "the parent must have closed the write end"
701 );
702 let deadline = Instant::now() + Duration::from_secs(5);
704 let status = loop {
705 if let Some(status) = child.0.try_wait().unwrap() {
706 break status;
707 }
708 assert!(Instant::now() < deadline, "the reader never saw EOF");
709 std::thread::sleep(Duration::from_millis(20));
710 };
711 assert!(status.success());
712 let _ = release.send(());
713 }
714
715 #[cfg(unix)]
718 struct Grandchild(std::path::PathBuf);
719
720 #[cfg(unix)]
721 impl Grandchild {
722 fn new(label: &str) -> Self {
723 let nanos = std::time::SystemTime::now()
724 .duration_since(std::time::UNIX_EPOCH)
725 .map_or(0, |elapsed| elapsed.as_nanos());
726 Self(std::env::temp_dir().join(format!(
727 "rk-run-stdin-{label}-{}-{nanos}.pid",
728 std::process::id()
729 )))
730 }
731
732 fn pid(&self) -> Option<i32> {
733 let pid: i32 = std::fs::read_to_string(&self.0).ok()?.trim().parse().ok()?;
734 (pid > 1).then_some(pid)
735 }
736
737 fn count_path(&self) -> std::path::PathBuf {
739 let mut path = self.0.clone().into_os_string();
740 path.push(".count");
741 path.into()
742 }
743
744 fn alive(&self) -> bool {
745 self.pid()
747 .is_some_and(|pid| unsafe { libc::kill(pid, 0) } == 0)
748 }
749 }
750
751 #[cfg(unix)]
752 impl Drop for Grandchild {
753 fn drop(&mut self) {
754 if let Some(pid) = self.pid() {
755 unsafe { libc::kill(pid, libc::SIGKILL) };
757 }
758 let _ = std::fs::remove_file(&self.0);
759 let _ = std::fs::remove_file(self.count_path());
760 }
761 }
762
763 #[cfg(unix)]
769 fn holder_script(body: &str, tail: &str) -> String {
770 format!(
771 "if command -v setsid >/dev/null 2>&1; then \
772 setsid sh -c 'echo $$ > \"$1\"; {body}' sh \"$1\" <&0 >/dev/null 2>&1 & \
773 elif command -v perl >/dev/null 2>&1; then \
774 perl -MPOSIX -e 'POSIX::setsid(); exec @ARGV' sh -c 'echo $$ > \"$1\"; {body}' sh \"$1\" <&0 >/dev/null 2>&1 & \
775 else exit 99; fi; \
776 i=0; while [ ! -s \"$1\" ] && [ $i -lt 100 ]; do sleep 0.05; i=$((i+1)); done; {tail}"
777 )
778 }
779
780 #[cfg(unix)]
781 #[test]
782 fn a_session_escaped_stdin_holder_does_not_hold_run_past_its_exit() {
783 let _serial = serial();
784 let holder = Grandchild::new("exit");
785 let script = holder_script("sleep 2; wc -c > \"$1.count\"; exec sleep 30", "exit 0");
789 let big = vec![b'x'; 2 << 20];
790 let options = RunOptions::new(Duration::from_secs(10));
791 let started = std::time::Instant::now();
792 let output = run_with_stdin(
793 "sh",
794 ["-c", &script, "sh", holder.0.to_str().unwrap()],
795 &big,
796 &options,
797 )
798 .unwrap();
799 let elapsed = started.elapsed();
800 if output.status.code() == Some(99) {
801 eprintln!("skipped: no setsid or perl to leave the session");
802 return;
803 }
804 assert!(
805 holder.alive(),
806 "the fixture's grandchild must still hold stdin"
807 );
808 assert!(!output.timed_out);
809 assert!(output.status.success(), "{:?}", output.status);
810 assert!(elapsed < Duration::from_secs(6), "run took {elapsed:?}");
812 assert_eq!(
813 LIVE_WRITERS.load(SeqCst),
814 0,
815 "the writer thread must be gone"
816 );
817 let deadline = std::time::Instant::now() + Duration::from_secs(10);
818 let received = loop {
819 if let Some(count) = std::fs::read_to_string(holder.count_path())
820 .ok()
821 .and_then(|text| text.trim().parse::<usize>().ok())
822 {
823 break count;
824 }
825 assert!(
826 std::time::Instant::now() < deadline,
827 "the grandchild never reached EOF: the write end is still open"
828 );
829 std::thread::sleep(Duration::from_millis(50));
830 };
831 assert!(
832 received < big.len() / 2,
833 "{received} bytes reached the grandchild after run returned"
834 );
835 }
836
837 #[cfg(unix)]
838 #[test]
839 fn a_session_escaped_stdin_holder_does_not_hold_run_past_its_timeout() {
840 let _serial = serial();
841 let holder = Grandchild::new("timeout");
842 let script = holder_script("exec sleep 30", "sleep 30");
844 let big = vec![b'x'; 2 << 20];
845 let timeout = Duration::from_secs(1);
846 let options = RunOptions::new(timeout);
847 let started = std::time::Instant::now();
848 let output = run_with_stdin(
849 "sh",
850 ["-c", &script, "sh", holder.0.to_str().unwrap()],
851 &big,
852 &options,
853 )
854 .unwrap();
855 let elapsed = started.elapsed();
856 if !holder.alive() && !output.timed_out {
857 eprintln!("skipped: no setsid or perl to leave the session");
858 return;
859 }
860 assert!(
861 holder.alive(),
862 "the fixture's grandchild must still hold stdin"
863 );
864 assert!(output.timed_out);
865 assert!(
866 elapsed < timeout + Duration::from_secs(6),
867 "run took {elapsed:?}"
868 );
869 assert_eq!(
870 LIVE_WRITERS.load(SeqCst),
871 0,
872 "the writer thread must be gone"
873 );
874 }
875
876 #[cfg(unix)]
877 #[test]
878 fn output_past_the_cap_is_truncated_not_blocking() {
879 let options = RunOptions::new(Duration::from_secs(10)).stdout_cap(4);
880 let output = run("sh", ["-c", "head -c 100000 /dev/zero"], &options).unwrap();
881 assert_eq!(output.stdout.bytes.len(), 4);
882 assert!(output.stdout.truncated);
883 }
884}