cargo/util/
process_builder.rs

1use crate::util::{process_error, read2, CargoResult, CargoResultExt};
2use anyhow::bail;
3use jobserver::Client;
4use shell_escape::escape;
5use std::collections::BTreeMap;
6use std::env;
7use std::ffi::{OsStr, OsString};
8use std::fmt;
9use std::iter::once;
10use std::path::Path;
11use std::process::{Command, Output, Stdio};
12
13/// A builder object for an external process, similar to `std::process::Command`.
14#[derive(Clone, Debug)]
15pub struct ProcessBuilder {
16    /// The program to execute.
17    program: OsString,
18    /// A list of arguments to pass to the program.
19    args: Vec<OsString>,
20    /// Any environment variables that should be set for the program.
21    env: BTreeMap<String, Option<OsString>>,
22    /// The directory to run the program from.
23    cwd: Option<OsString>,
24    /// The `make` jobserver. See the [jobserver crate][jobserver_docs] for
25    /// more information.
26    ///
27    /// [jobserver_docs]: https://docs.rs/jobserver/0.1.6/jobserver/
28    jobserver: Option<Client>,
29    /// `true` to include environment variable in display.
30    display_env_vars: bool,
31}
32
33impl fmt::Display for ProcessBuilder {
34    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
35        write!(f, "`")?;
36
37        if self.display_env_vars {
38            for (key, val) in self.env.iter() {
39                if let Some(val) = val {
40                    let val = escape(val.to_string_lossy());
41                    if cfg!(windows) {
42                        write!(f, "set {}={}&& ", key, val)?;
43                    } else {
44                        write!(f, "{}={} ", key, val)?;
45                    }
46                }
47            }
48        }
49
50        write!(f, "{}", self.program.to_string_lossy())?;
51
52        for arg in &self.args {
53            write!(f, " {}", escape(arg.to_string_lossy()))?;
54        }
55
56        write!(f, "`")
57    }
58}
59
60impl ProcessBuilder {
61    /// (chainable) Sets the executable for the process.
62    pub fn program<T: AsRef<OsStr>>(&mut self, program: T) -> &mut ProcessBuilder {
63        self.program = program.as_ref().to_os_string();
64        self
65    }
66
67    /// (chainable) Adds `arg` to the args list.
68    pub fn arg<T: AsRef<OsStr>>(&mut self, arg: T) -> &mut ProcessBuilder {
69        self.args.push(arg.as_ref().to_os_string());
70        self
71    }
72
73    /// (chainable) Adds multiple `args` to the args list.
74    pub fn args<T: AsRef<OsStr>>(&mut self, args: &[T]) -> &mut ProcessBuilder {
75        self.args
76            .extend(args.iter().map(|t| t.as_ref().to_os_string()));
77        self
78    }
79
80    /// (chainable) Replaces the args list with the given `args`.
81    pub fn args_replace<T: AsRef<OsStr>>(&mut self, args: &[T]) -> &mut ProcessBuilder {
82        self.args = args.iter().map(|t| t.as_ref().to_os_string()).collect();
83        self
84    }
85
86    /// (chainable) Sets the current working directory of the process.
87    pub fn cwd<T: AsRef<OsStr>>(&mut self, path: T) -> &mut ProcessBuilder {
88        self.cwd = Some(path.as_ref().to_os_string());
89        self
90    }
91
92    /// (chainable) Sets an environment variable for the process.
93    pub fn env<T: AsRef<OsStr>>(&mut self, key: &str, val: T) -> &mut ProcessBuilder {
94        self.env
95            .insert(key.to_string(), Some(val.as_ref().to_os_string()));
96        self
97    }
98
99    /// (chainable) Unsets an environment variable for the process.
100    pub fn env_remove(&mut self, key: &str) -> &mut ProcessBuilder {
101        self.env.insert(key.to_string(), None);
102        self
103    }
104
105    /// Gets the executable name.
106    pub fn get_program(&self) -> &OsString {
107        &self.program
108    }
109
110    /// Gets the program arguments.
111    pub fn get_args(&self) -> &[OsString] {
112        &self.args
113    }
114
115    /// Gets the current working directory for the process.
116    pub fn get_cwd(&self) -> Option<&Path> {
117        self.cwd.as_ref().map(Path::new)
118    }
119
120    /// Gets an environment variable as the process will see it (will inherit from environment
121    /// unless explicitally unset).
122    pub fn get_env(&self, var: &str) -> Option<OsString> {
123        self.env
124            .get(var)
125            .cloned()
126            .or_else(|| Some(env::var_os(var)))
127            .and_then(|s| s)
128    }
129
130    /// Gets all environment variables explicitly set or unset for the process (not inherited
131    /// vars).
132    pub fn get_envs(&self) -> &BTreeMap<String, Option<OsString>> {
133        &self.env
134    }
135
136    /// Sets the `make` jobserver. See the [jobserver crate][jobserver_docs] for
137    /// more information.
138    ///
139    /// [jobserver_docs]: https://docs.rs/jobserver/0.1.6/jobserver/
140    pub fn inherit_jobserver(&mut self, jobserver: &Client) -> &mut Self {
141        self.jobserver = Some(jobserver.clone());
142        self
143    }
144
145    /// Enables environment variable display.
146    pub fn display_env_vars(&mut self) -> &mut Self {
147        self.display_env_vars = true;
148        self
149    }
150
151    /// Runs the process, waiting for completion, and mapping non-success exit codes to an error.
152    pub fn exec(&self) -> CargoResult<()> {
153        let mut command = self.build_command();
154        let exit = command.status().chain_err(|| {
155            process_error(&format!("could not execute process {}", self), None, None)
156        })?;
157
158        if exit.success() {
159            Ok(())
160        } else {
161            Err(process_error(
162                &format!("process didn't exit successfully: {}", self),
163                Some(exit),
164                None,
165            )
166            .into())
167        }
168    }
169
170    /// Replaces the current process with the target process.
171    ///
172    /// On Unix, this executes the process using the Unix syscall `execvp`, which will block
173    /// this process, and will only return if there is an error.
174    ///
175    /// On Windows this isn't technically possible. Instead we emulate it to the best of our
176    /// ability. One aspect we fix here is that we specify a handler for the Ctrl-C handler.
177    /// In doing so (and by effectively ignoring it) we should emulate proxying Ctrl-C
178    /// handling to the application at hand, which will either terminate or handle it itself.
179    /// According to Microsoft's documentation at
180    /// <https://docs.microsoft.com/en-us/windows/console/ctrl-c-and-ctrl-break-signals>.
181    /// the Ctrl-C signal is sent to all processes attached to a terminal, which should
182    /// include our child process. If the child terminates then we'll reap them in Cargo
183    /// pretty quickly, and if the child handles the signal then we won't terminate
184    /// (and we shouldn't!) until the process itself later exits.
185    pub fn exec_replace(&self) -> CargoResult<()> {
186        imp::exec_replace(self)
187    }
188
189    /// Executes the process, returning the stdio output, or an error if non-zero exit status.
190    pub fn exec_with_output(&self) -> CargoResult<Output> {
191        let mut command = self.build_command();
192
193        let output = command.output().chain_err(|| {
194            process_error(&format!("could not execute process {}", self), None, None)
195        })?;
196
197        if output.status.success() {
198            Ok(output)
199        } else {
200            Err(process_error(
201                &format!("process didn't exit successfully: {}", self),
202                Some(output.status),
203                Some(&output),
204            )
205            .into())
206        }
207    }
208
209    /// Executes a command, passing each line of stdout and stderr to the supplied callbacks, which
210    /// can mutate the string data.
211    ///
212    /// If any invocations of these function return an error, it will be propagated.
213    ///
214    /// If `capture_output` is true, then all the output will also be buffered
215    /// and stored in the returned `Output` object. If it is false, no caching
216    /// is done, and the callbacks are solely responsible for handling the
217    /// output.
218    pub fn exec_with_streaming(
219        &self,
220        on_stdout_line: &mut dyn FnMut(&str) -> CargoResult<()>,
221        on_stderr_line: &mut dyn FnMut(&str) -> CargoResult<()>,
222        capture_output: bool,
223    ) -> CargoResult<Output> {
224        let mut stdout = Vec::new();
225        let mut stderr = Vec::new();
226
227        let mut cmd = self.build_command();
228        cmd.stdout(Stdio::piped())
229            .stderr(Stdio::piped())
230            .stdin(Stdio::null());
231
232        let mut callback_error = None;
233        let status = (|| {
234            let mut child = cmd.spawn()?;
235            let out = child.stdout.take().unwrap();
236            let err = child.stderr.take().unwrap();
237            read2(out, err, &mut |is_out, data, eof| {
238                let idx = if eof {
239                    data.len()
240                } else {
241                    match data.iter().rposition(|b| *b == b'\n') {
242                        Some(i) => i + 1,
243                        None => return,
244                    }
245                };
246                {
247                    // scope for new_lines
248                    let new_lines = if capture_output {
249                        let dst = if is_out { &mut stdout } else { &mut stderr };
250                        let start = dst.len();
251                        let data = data.drain(..idx);
252                        dst.extend(data);
253                        &dst[start..]
254                    } else {
255                        &data[..idx]
256                    };
257                    for line in String::from_utf8_lossy(new_lines).lines() {
258                        if callback_error.is_some() {
259                            break;
260                        }
261                        let callback_result = if is_out {
262                            on_stdout_line(line)
263                        } else {
264                            on_stderr_line(line)
265                        };
266                        if let Err(e) = callback_result {
267                            callback_error = Some(e);
268                        }
269                    }
270                }
271                if !capture_output {
272                    data.drain(..idx);
273                }
274            })?;
275            child.wait()
276        })()
277        .chain_err(|| process_error(&format!("could not execute process {}", self), None, None))?;
278        let output = Output {
279            stdout,
280            stderr,
281            status,
282        };
283
284        {
285            let to_print = if capture_output { Some(&output) } else { None };
286            if let Some(e) = callback_error {
287                let cx = process_error(
288                    &format!("failed to parse process output: {}", self),
289                    Some(output.status),
290                    to_print,
291                );
292                bail!(anyhow::Error::new(cx).context(e));
293            } else if !output.status.success() {
294                bail!(process_error(
295                    &format!("process didn't exit successfully: {}", self),
296                    Some(output.status),
297                    to_print,
298                ));
299            }
300        }
301
302        Ok(output)
303    }
304
305    /// Converts `ProcessBuilder` into a `std::process::Command`, and handles the jobserver, if
306    /// present.
307    pub fn build_command(&self) -> Command {
308        let mut command = Command::new(&self.program);
309        if let Some(cwd) = self.get_cwd() {
310            command.current_dir(cwd);
311        }
312        for arg in &self.args {
313            command.arg(arg);
314        }
315        for (k, v) in &self.env {
316            match *v {
317                Some(ref v) => {
318                    command.env(k, v);
319                }
320                None => {
321                    command.env_remove(k);
322                }
323            }
324        }
325        if let Some(ref c) = self.jobserver {
326            c.configure(&mut command);
327        }
328        command
329    }
330
331    /// Wraps an existing command with the provided wrapper, if it is present and valid.
332    ///
333    /// # Examples
334    ///
335    /// ```rust
336    /// use cargo::util::{ProcessBuilder, process};
337    /// // Running this would execute `rustc`
338    /// let cmd: ProcessBuilder = process("rustc");
339    ///
340    /// // Running this will execute `sccache rustc`
341    /// let cmd = cmd.wrapped(Some("sccache"));
342    /// ```
343    pub fn wrapped(mut self, wrapper: Option<impl AsRef<OsStr>>) -> Self {
344        let wrapper = if let Some(wrapper) = wrapper.as_ref() {
345            wrapper.as_ref()
346        } else {
347            return self;
348        };
349
350        if wrapper.is_empty() {
351            return self;
352        }
353
354        let args = once(self.program).chain(self.args.into_iter()).collect();
355
356        self.program = wrapper.to_os_string();
357        self.args = args;
358
359        self
360    }
361}
362
363/// A helper function to create a `ProcessBuilder`.
364pub fn process<T: AsRef<OsStr>>(cmd: T) -> ProcessBuilder {
365    ProcessBuilder {
366        program: cmd.as_ref().to_os_string(),
367        args: Vec::new(),
368        cwd: None,
369        env: BTreeMap::new(),
370        jobserver: None,
371        display_env_vars: false,
372    }
373}
374
375#[cfg(unix)]
376mod imp {
377    use crate::util::{process_error, ProcessBuilder};
378    use crate::CargoResult;
379    use std::os::unix::process::CommandExt;
380
381    pub fn exec_replace(process_builder: &ProcessBuilder) -> CargoResult<()> {
382        let mut command = process_builder.build_command();
383        let error = command.exec();
384        Err(anyhow::Error::from(error)
385            .context(process_error(
386                &format!("could not execute process {}", process_builder),
387                None,
388                None,
389            ))
390            .into())
391    }
392}
393
394#[cfg(windows)]
395mod imp {
396    use crate::util::{process_error, ProcessBuilder};
397    use crate::CargoResult;
398    use winapi::shared::minwindef::{BOOL, DWORD, FALSE, TRUE};
399    use winapi::um::consoleapi::SetConsoleCtrlHandler;
400
401    unsafe extern "system" fn ctrlc_handler(_: DWORD) -> BOOL {
402        // Do nothing; let the child process handle it.
403        TRUE
404    }
405
406    pub fn exec_replace(process_builder: &ProcessBuilder) -> CargoResult<()> {
407        unsafe {
408            if SetConsoleCtrlHandler(Some(ctrlc_handler), TRUE) == FALSE {
409                return Err(process_error("Could not set Ctrl-C handler.", None, None).into());
410            }
411        }
412
413        // Just execute the process as normal.
414        process_builder.exec()
415    }
416}