cmd_wrapper/
std_command.rs

1use std::{
2    ffi::OsStr,
3    io::{self, Error, Read, Result, Write},
4    path::Path,
5    process::{Child, Command, ExitStatus, Stdio},
6    sync::mpsc::{self, Receiver},
7    thread::{self, JoinHandle},
8};
9
10use derivative::Derivative;
11use tracing::{error, instrument, warn};
12
13#[derive(Debug)]
14pub struct StdCommand {
15    command: Command,
16}
17
18impl StdCommand {
19    pub fn new<S: Into<String>>(program: S) -> Self {
20        Self {
21            command: Command::new(program.into()),
22        }
23    }
24
25    pub fn arg<S: AsRef<OsStr>>(&mut self, arg: S) -> &mut Self {
26        self.command.arg(arg);
27        self
28    }
29
30    pub fn args<I, S>(&mut self, args: I) -> &mut Self
31    where
32        I: IntoIterator<Item = S>,
33        S: AsRef<OsStr>,
34    {
35        self.command.args(args);
36        self
37    }
38
39    pub fn env<K, V>(&mut self, key: K, val: V) -> &mut Self
40    where
41        K: AsRef<OsStr>,
42        V: AsRef<OsStr>,
43    {
44        self.command.env(key, val);
45        self
46    }
47
48    pub fn envs<I, K, V>(&mut self, vars: I) -> &mut Self
49    where
50        I: IntoIterator<Item = (K, V)>,
51        K: AsRef<OsStr>,
52        V: AsRef<OsStr>,
53    {
54        self.command.envs(vars);
55        self
56    }
57
58    pub fn current_dir<P: AsRef<Path>>(&mut self, dir: P) -> &mut Self {
59        self.command.current_dir(dir);
60        self
61    }
62
63    pub fn stdin<T: Into<Stdio>>(&mut self, stdin: T) -> &mut Self {
64        self.command.stdin(stdin);
65        self
66    }
67
68    pub fn stdout<T: Into<Stdio>>(&mut self, stdout: T) -> &mut Self {
69        self.command.stdout(stdout);
70        self
71    }
72
73    pub fn stderr<T: Into<Stdio>>(&mut self, stderr: T) -> &mut Self {
74        self.command.stderr(stderr);
75        self
76    }
77
78    #[instrument(level = "trace")]
79    pub fn into_inner(self) -> Command {
80        self.command
81    }
82
83    #[instrument(level = "trace")]
84    pub fn child(mut self) -> Result<Child> {
85        self.command.spawn()
86    }
87
88    #[instrument(level = "trace")]
89    pub fn spawn(&mut self) -> Result<StdCommandProcess> {
90        let mut child = self.command.spawn()?;
91
92        let stdin = child
93            .stdin
94            .take()
95            .map(|s| Box::new(s) as Box<dyn Write>)
96            .unwrap_or_else(|| Box::new(SyncSink));
97        let stdout = child
98            .stdout
99            .take()
100            .map(|s| Box::new(s) as Box<dyn Read>)
101            .unwrap_or_else(|| Box::new(io::empty()));
102        let stderr = child
103            .stderr
104            .take()
105            .map(|s| Box::new(s) as Box<dyn Read>)
106            .unwrap_or_else(|| Box::new(io::empty()));
107
108        let pid = child.id();
109
110        let (end_tx, end_rx) = mpsc::channel::<i32>();
111        let end_handler = thread::spawn(move || {
112            let exit_code = match child.wait() {
113                Ok(status) if status.code() == Some(0) => 0,
114                Ok(_) => 1,
115                Err(err) => {
116                    error!("waiting for child process failed: {err}");
117                    -1
118                }
119            };
120
121            if let Err(err) = end_tx.send(exit_code) {
122                warn!("receiver dropped before sending exit code: {err}");
123            }
124        });
125
126        Ok(StdCommandProcess {
127            stdin,
128            stdout,
129            stderr,
130            pid,
131            end_rx,
132            end_handler,
133        })
134    }
135
136    #[instrument(level = "trace")]
137    pub fn output(&mut self) -> Result<String> {
138        let output = self.command.output()?;
139
140        if !output.status.success() {
141            let stderr = String::from_utf8_lossy(&output.stderr);
142            return Err(Error::other(stderr));
143        }
144
145        Ok(String::from_utf8_lossy(&output.stdout).to_string())
146    }
147
148    #[instrument(level = "trace")]
149    pub fn status(&mut self) -> Result<ExitStatus> {
150        self.command.status()
151    }
152
153    #[instrument(level = "trace")]
154    pub fn execute_and_print(&mut self) -> Result<()> {
155        let output = self.command.output()?;
156        if !output.stdout.is_empty() {
157            io::stdout().write_all(&output.stdout)?;
158        }
159        if !output.stderr.is_empty() {
160            io::stderr().write_all(&output.stderr)?;
161        }
162        Ok(())
163    }
164
165    pub fn read_full_stderr_if_any(stderr: &mut impl Read) -> Result<()> {
166        let mut peek_buf = vec![0u8; 1024];
167        let n = stderr.read(&mut peek_buf)?;
168        if n == 0 {
169            return Ok(());
170        }
171
172        let mut full = peek_buf[..n].to_vec();
173        let mut rest = Vec::new();
174        stderr.read_to_end(&mut rest)?;
175        full.extend(rest);
176
177        let msg = String::from_utf8_lossy(&full).trim().to_string();
178        if !msg.is_empty() {
179            return Err(Error::other(msg));
180        }
181
182        Ok(())
183    }
184}
185
186#[derive(Derivative)]
187#[derivative(Debug)]
188pub struct StdCommandProcess {
189    #[derivative(Debug = "ignore")]
190    pub stdin: Box<dyn Write>,
191    #[derivative(Debug = "ignore")]
192    pub stdout: Box<dyn Read>,
193    #[derivative(Debug = "ignore")]
194    pub stderr: Box<dyn Read>,
195    pub pid: u32,
196    pub end_rx: Receiver<i32>,
197    pub end_handler: JoinHandle<()>,
198}
199
200struct SyncSink;
201
202impl Write for SyncSink {
203    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
204        Ok(buf.len())
205    }
206    fn flush(&mut self) -> io::Result<()> {
207        Ok(())
208    }
209}