snix_build/sandbox/
event.rs

1use std::os::unix::process::ExitStatusExt;
2
3use futures::StreamExt;
4use futures::stream::BoxStream;
5use tokio_util::io::ReaderStream;
6use tracing::warn;
7
8/// A common set of events emitted by sandboxed builds.
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub enum SandboxEvent {
11    Stdout(Vec<u8>),
12    Stderr(Vec<u8>),
13    ExitCode(i32),
14}
15
16/// Stream stdout, stderr, and exit code from a spawned child process.
17///
18/// Keeps `guard` alive for the duration of the stream (e.g. FUSE mount, inputs guard).
19pub fn stream_process<G: 'static + Send>(
20    mut child: tokio::process::Child,
21    guard: G,
22) -> std::io::Result<BoxStream<'static, SandboxEvent>> {
23    let stdout = child
24        .stdout
25        .take()
26        .ok_or_else(|| std::io::Error::other("stdout pipe was not captured"))?;
27    let stderr = child
28        .stderr
29        .take()
30        .ok_or_else(|| std::io::Error::other("stderr pipe was not captured"))?;
31
32    let stream = async_stream::stream! {
33        let _guard = guard;
34        let stdout_stream = ReaderStream::new(stdout)
35            .map(|res| res.map(|b| SandboxEvent::Stdout(b.to_vec())));
36        let stderr_stream = ReaderStream::new(stderr)
37            .map(|res| res.map(|b| SandboxEvent::Stderr(b.to_vec())));
38
39        let output_stream = futures::stream::select(stdout_stream, stderr_stream);
40        tokio::pin!(output_stream);
41        while let Some(res) = output_stream.next().await {
42            match res {
43                Ok(event) => yield event,
44                Err(e) => {
45                    warn!("error reading sandbox output: {e}");
46                    break;
47                }
48            }
49        }
50
51        let exit_code = match child.wait().await {
52            Ok(status) => status
53                .code()
54                .unwrap_or_else(|| 128 + status.signal().unwrap_or(0)),
55            Err(_) => 1,
56        };
57
58        yield SandboxEvent::ExitCode(exit_code);
59    }
60    .boxed();
61
62    Ok(stream)
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68    use std::process::Stdio;
69    use tokio::process::Command;
70
71    #[tokio::test]
72    async fn test_stream_process() {
73        let mut cmd = Command::new("/bin/sh");
74        cmd.args(["-c", "echo hello; echo err >&2; exit 42"]);
75        cmd.stdout(Stdio::piped())
76            .stderr(Stdio::piped())
77            .kill_on_drop(true);
78
79        let child = cmd.spawn().expect("failed to spawn /bin/sh");
80
81        let (drop_tx, drop_rx) = tokio::sync::oneshot::channel();
82        struct DropGuard(Option<tokio::sync::oneshot::Sender<()>>);
83        impl Drop for DropGuard {
84            fn drop(&mut self) {
85                if let Some(tx) = self.0.take() {
86                    let _ = tx.send(());
87                }
88            }
89        }
90
91        let mut stream = stream_process(child, DropGuard(Some(drop_tx)))
92            .expect("failed to start stream_process");
93        let mut stdout = Vec::new();
94        let mut stderr = Vec::new();
95        let mut exit_code = None;
96
97        while let Some(event) = stream.next().await {
98            match event {
99                SandboxEvent::Stdout(bytes) => stdout.extend(bytes),
100                SandboxEvent::Stderr(bytes) => stderr.extend(bytes),
101                SandboxEvent::ExitCode(code) => exit_code = Some(code),
102            }
103        }
104
105        assert_eq!(stdout, b"hello\n");
106        assert_eq!(stderr, b"err\n");
107        assert_eq!(exit_code, Some(42));
108        assert_eq!(
109            drop_rx.await,
110            Ok(()),
111            "guard was dropped after stream ended"
112        );
113    }
114}
snix_build/sandbox/event.rs

snix_build/sandbox/
event.rs
