| use crate::process_error::ProcessError; |
| use crate::read2; |
| |
| use anyhow::{bail, Context, Result}; |
| use jobserver::Client; |
| use shell_escape::escape; |
| use tempfile::NamedTempFile; |
| |
| use std::collections::BTreeMap; |
| use std::env; |
| use std::ffi::{OsStr, OsString}; |
| use std::fmt; |
| use std::io::{self, Write}; |
| use std::iter::once; |
| use std::path::Path; |
| use std::process::{Command, ExitStatus, Output, Stdio}; |
| |
| /// A builder object for an external process, similar to [`std::process::Command`]. |
| #[derive(Clone, Debug)] |
| pub struct ProcessBuilder { |
| /// The program to execute. |
| program: OsString, |
| /// A list of arguments to pass to the program. |
| args: Vec<OsString>, |
| /// Any environment variables that should be set for the program. |
| env: BTreeMap<String, Option<OsString>>, |
| /// The directory to run the program from. |
| cwd: Option<OsString>, |
| /// A list of wrappers that wrap the original program when calling |
| /// [`ProcessBuilder::wrapped`]. The last one is the outermost one. |
| wrappers: Vec<OsString>, |
| /// The `make` jobserver. See the [jobserver crate] for |
| /// more information. |
| /// |
| /// [jobserver crate]: https://docs.rs/jobserver/ |
| jobserver: Option<Client>, |
| /// `true` to include environment variable in display. |
| display_env_vars: bool, |
| /// `true` to retry with an argfile if hitting "command line too big" error. |
| /// See [`ProcessBuilder::retry_with_argfile`] for more information. |
| retry_with_argfile: bool, |
| /// Data to write to stdin. |
| stdin: Option<Vec<u8>>, |
| } |
| |
| impl fmt::Display for ProcessBuilder { |
| fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
| write!(f, "`")?; |
| |
| if self.display_env_vars { |
| for (key, val) in self.env.iter() { |
| if let Some(val) = val { |
| let val = escape(val.to_string_lossy()); |
| if cfg!(windows) { |
| write!(f, "set {}={}&& ", key, val)?; |
| } else { |
| write!(f, "{}={} ", key, val)?; |
| } |
| } |
| } |
| } |
| |
| write!(f, "{}", self.get_program().to_string_lossy())?; |
| |
| for arg in self.get_args() { |
| write!(f, " {}", escape(arg.to_string_lossy()))?; |
| } |
| |
| write!(f, "`") |
| } |
| } |
| |
| impl ProcessBuilder { |
| /// Creates a new [`ProcessBuilder`] with the given executable path. |
| pub fn new<T: AsRef<OsStr>>(cmd: T) -> ProcessBuilder { |
| ProcessBuilder { |
| program: cmd.as_ref().to_os_string(), |
| args: Vec::new(), |
| cwd: None, |
| env: BTreeMap::new(), |
| wrappers: Vec::new(), |
| jobserver: None, |
| display_env_vars: false, |
| retry_with_argfile: false, |
| stdin: None, |
| } |
| } |
| |
| /// (chainable) Sets the executable for the process. |
| pub fn program<T: AsRef<OsStr>>(&mut self, program: T) -> &mut ProcessBuilder { |
| self.program = program.as_ref().to_os_string(); |
| self |
| } |
| |
| /// (chainable) Adds `arg` to the args list. |
| pub fn arg<T: AsRef<OsStr>>(&mut self, arg: T) -> &mut ProcessBuilder { |
| self.args.push(arg.as_ref().to_os_string()); |
| self |
| } |
| |
| /// (chainable) Adds multiple `args` to the args list. |
| pub fn args<T: AsRef<OsStr>>(&mut self, args: &[T]) -> &mut ProcessBuilder { |
| self.args |
| .extend(args.iter().map(|t| t.as_ref().to_os_string())); |
| self |
| } |
| |
| /// (chainable) Replaces the args list with the given `args`. |
| pub fn args_replace<T: AsRef<OsStr>>(&mut self, args: &[T]) -> &mut ProcessBuilder { |
| if let Some(program) = self.wrappers.pop() { |
| // User intend to replace all args, so we |
| // - use the outermost wrapper as the main program, and |
| // - cleanup other inner wrappers. |
| self.program = program; |
| self.wrappers = Vec::new(); |
| } |
| self.args = args.iter().map(|t| t.as_ref().to_os_string()).collect(); |
| self |
| } |
| |
| /// (chainable) Sets the current working directory of the process. |
| pub fn cwd<T: AsRef<OsStr>>(&mut self, path: T) -> &mut ProcessBuilder { |
| self.cwd = Some(path.as_ref().to_os_string()); |
| self |
| } |
| |
| /// (chainable) Sets an environment variable for the process. |
| pub fn env<T: AsRef<OsStr>>(&mut self, key: &str, val: T) -> &mut ProcessBuilder { |
| self.env |
| .insert(key.to_string(), Some(val.as_ref().to_os_string())); |
| self |
| } |
| |
| /// (chainable) Unsets an environment variable for the process. |
| pub fn env_remove(&mut self, key: &str) -> &mut ProcessBuilder { |
| self.env.insert(key.to_string(), None); |
| self |
| } |
| |
| /// Gets the executable name. |
| pub fn get_program(&self) -> &OsString { |
| self.wrappers.last().unwrap_or(&self.program) |
| } |
| |
| /// Gets the program arguments. |
| pub fn get_args(&self) -> impl Iterator<Item = &OsString> { |
| self.wrappers |
| .iter() |
| .rev() |
| .chain(once(&self.program)) |
| .chain(self.args.iter()) |
| .skip(1) // Skip the main `program |
| } |
| |
| /// Gets the current working directory for the process. |
| pub fn get_cwd(&self) -> Option<&Path> { |
| self.cwd.as_ref().map(Path::new) |
| } |
| |
| /// Gets an environment variable as the process will see it (will inherit from environment |
| /// unless explicitally unset). |
| pub fn get_env(&self, var: &str) -> Option<OsString> { |
| self.env |
| .get(var) |
| .cloned() |
| .or_else(|| Some(env::var_os(var))) |
| .and_then(|s| s) |
| } |
| |
| /// Gets all environment variables explicitly set or unset for the process (not inherited |
| /// vars). |
| pub fn get_envs(&self) -> &BTreeMap<String, Option<OsString>> { |
| &self.env |
| } |
| |
| /// Sets the `make` jobserver. See the [jobserver crate][jobserver_docs] for |
| /// more information. |
| /// |
| /// [jobserver_docs]: https://docs.rs/jobserver/0.1.6/jobserver/ |
| pub fn inherit_jobserver(&mut self, jobserver: &Client) -> &mut Self { |
| self.jobserver = Some(jobserver.clone()); |
| self |
| } |
| |
| /// Enables environment variable display. |
| pub fn display_env_vars(&mut self) -> &mut Self { |
| self.display_env_vars = true; |
| self |
| } |
| |
| /// Enables retrying with an argfile if hitting "command line too big" error |
| /// |
| /// This is primarily for the `@path` arg of rustc and rustdoc, which treat |
| /// each line as an command-line argument, so `LF` and `CRLF` bytes are not |
| /// valid as an argument for argfile at this moment. |
| /// For example, `RUSTDOCFLAGS="--crate-version foo\nbar" cargo doc` is |
| /// valid when invoking from command-line but not from argfile. |
| /// |
| /// To sum up, the limitations of the argfile are: |
| /// |
| /// - Must be valid UTF-8 encoded. |
| /// - Must not contain any newlines in each argument. |
| /// |
| /// Ref: |
| /// |
| /// - <https://doc.rust-lang.org/rustdoc/command-line-arguments.html#path-load-command-line-flags-from-a-path> |
| /// - <https://doc.rust-lang.org/rustc/command-line-arguments.html#path-load-command-line-flags-from-a-path> |
| pub fn retry_with_argfile(&mut self, enabled: bool) -> &mut Self { |
| self.retry_with_argfile = enabled; |
| self |
| } |
| |
| /// Sets a value that will be written to stdin of the process on launch. |
| pub fn stdin<T: Into<Vec<u8>>>(&mut self, stdin: T) -> &mut Self { |
| self.stdin = Some(stdin.into()); |
| self |
| } |
| |
| fn should_retry_with_argfile(&self, err: &io::Error) -> bool { |
| self.retry_with_argfile && imp::command_line_too_big(err) |
| } |
| |
| /// Like [`Command::status`] but with a better error message. |
| pub fn status(&self) -> Result<ExitStatus> { |
| self._status() |
| .with_context(|| ProcessError::could_not_execute(self)) |
| } |
| |
| fn _status(&self) -> io::Result<ExitStatus> { |
| if !debug_force_argfile(self.retry_with_argfile) { |
| let mut cmd = self.build_command(); |
| match cmd.spawn() { |
| Err(ref e) if self.should_retry_with_argfile(e) => {} |
| Err(e) => return Err(e), |
| Ok(mut child) => return child.wait(), |
| } |
| } |
| let (mut cmd, argfile) = self.build_command_with_argfile()?; |
| let status = cmd.spawn()?.wait(); |
| close_tempfile_and_log_error(argfile); |
| status |
| } |
| |
| /// Runs the process, waiting for completion, and mapping non-success exit codes to an error. |
| pub fn exec(&self) -> Result<()> { |
| let exit = self.status()?; |
| if exit.success() { |
| Ok(()) |
| } else { |
| Err(ProcessError::new( |
| &format!("process didn't exit successfully: {}", self), |
| Some(exit), |
| None, |
| ) |
| .into()) |
| } |
| } |
| |
| /// Replaces the current process with the target process. |
| /// |
| /// On Unix, this executes the process using the Unix syscall `execvp`, which will block |
| /// this process, and will only return if there is an error. |
| /// |
| /// On Windows this isn't technically possible. Instead we emulate it to the best of our |
| /// ability. One aspect we fix here is that we specify a handler for the Ctrl-C handler. |
| /// In doing so (and by effectively ignoring it) we should emulate proxying Ctrl-C |
| /// handling to the application at hand, which will either terminate or handle it itself. |
| /// According to Microsoft's documentation at |
| /// <https://docs.microsoft.com/en-us/windows/console/ctrl-c-and-ctrl-break-signals>. |
| /// the Ctrl-C signal is sent to all processes attached to a terminal, which should |
| /// include our child process. If the child terminates then we'll reap them in Cargo |
| /// pretty quickly, and if the child handles the signal then we won't terminate |
| /// (and we shouldn't!) until the process itself later exits. |
| pub fn exec_replace(&self) -> Result<()> { |
| imp::exec_replace(self) |
| } |
| |
| /// Like [`Command::output`] but with a better error message. |
| pub fn output(&self) -> Result<Output> { |
| self._output() |
| .with_context(|| ProcessError::could_not_execute(self)) |
| } |
| |
| fn _output(&self) -> io::Result<Output> { |
| if !debug_force_argfile(self.retry_with_argfile) { |
| let mut cmd = self.build_command(); |
| match piped(&mut cmd, self.stdin.is_some()).spawn() { |
| Err(ref e) if self.should_retry_with_argfile(e) => {} |
| Err(e) => return Err(e), |
| Ok(mut child) => { |
| if let Some(stdin) = &self.stdin { |
| child.stdin.take().unwrap().write_all(stdin)?; |
| } |
| return child.wait_with_output(); |
| } |
| } |
| } |
| let (mut cmd, argfile) = self.build_command_with_argfile()?; |
| let mut child = piped(&mut cmd, self.stdin.is_some()).spawn()?; |
| if let Some(stdin) = &self.stdin { |
| child.stdin.take().unwrap().write_all(stdin)?; |
| } |
| let output = child.wait_with_output(); |
| close_tempfile_and_log_error(argfile); |
| output |
| } |
| |
| /// Executes the process, returning the stdio output, or an error if non-zero exit status. |
| pub fn exec_with_output(&self) -> Result<Output> { |
| let output = self.output()?; |
| if output.status.success() { |
| Ok(output) |
| } else { |
| Err(ProcessError::new( |
| &format!("process didn't exit successfully: {}", self), |
| Some(output.status), |
| Some(&output), |
| ) |
| .into()) |
| } |
| } |
| |
| /// Executes a command, passing each line of stdout and stderr to the supplied callbacks, which |
| /// can mutate the string data. |
| /// |
| /// If any invocations of these function return an error, it will be propagated. |
| /// |
| /// If `capture_output` is true, then all the output will also be buffered |
| /// and stored in the returned `Output` object. If it is false, no caching |
| /// is done, and the callbacks are solely responsible for handling the |
| /// output. |
| pub fn exec_with_streaming( |
| &self, |
| on_stdout_line: &mut dyn FnMut(&str) -> Result<()>, |
| on_stderr_line: &mut dyn FnMut(&str) -> Result<()>, |
| capture_output: bool, |
| ) -> Result<Output> { |
| let mut stdout = Vec::new(); |
| let mut stderr = Vec::new(); |
| |
| let mut callback_error = None; |
| let mut stdout_pos = 0; |
| let mut stderr_pos = 0; |
| |
| let spawn = |mut cmd| { |
| if !debug_force_argfile(self.retry_with_argfile) { |
| match piped(&mut cmd, false).spawn() { |
| Err(ref e) if self.should_retry_with_argfile(e) => {} |
| Err(e) => return Err(e), |
| Ok(child) => return Ok((child, None)), |
| } |
| } |
| let (mut cmd, argfile) = self.build_command_with_argfile()?; |
| Ok((piped(&mut cmd, false).spawn()?, Some(argfile))) |
| }; |
| |
| let status = (|| { |
| let cmd = self.build_command(); |
| let (mut child, argfile) = spawn(cmd)?; |
| let out = child.stdout.take().unwrap(); |
| let err = child.stderr.take().unwrap(); |
| read2(out, err, &mut |is_out, data, eof| { |
| let pos = if is_out { |
| &mut stdout_pos |
| } else { |
| &mut stderr_pos |
| }; |
| let idx = if eof { |
| data.len() |
| } else { |
| match data[*pos..].iter().rposition(|b| *b == b'\n') { |
| Some(i) => *pos + i + 1, |
| None => { |
| *pos = data.len(); |
| return; |
| } |
| } |
| }; |
| |
| let new_lines = &data[..idx]; |
| |
| for line in String::from_utf8_lossy(new_lines).lines() { |
| if callback_error.is_some() { |
| break; |
| } |
| let callback_result = if is_out { |
| on_stdout_line(line) |
| } else { |
| on_stderr_line(line) |
| }; |
| if let Err(e) = callback_result { |
| callback_error = Some(e); |
| break; |
| } |
| } |
| |
| if capture_output { |
| let dst = if is_out { &mut stdout } else { &mut stderr }; |
| dst.extend(new_lines); |
| } |
| |
| data.drain(..idx); |
| *pos = 0; |
| })?; |
| let status = child.wait(); |
| if let Some(argfile) = argfile { |
| close_tempfile_and_log_error(argfile); |
| } |
| status |
| })() |
| .with_context(|| ProcessError::could_not_execute(self))?; |
| let output = Output { |
| status, |
| stdout, |
| stderr, |
| }; |
| |
| { |
| let to_print = if capture_output { Some(&output) } else { None }; |
| if let Some(e) = callback_error { |
| let cx = ProcessError::new( |
| &format!("failed to parse process output: {}", self), |
| Some(output.status), |
| to_print, |
| ); |
| bail!(anyhow::Error::new(cx).context(e)); |
| } else if !output.status.success() { |
| bail!(ProcessError::new( |
| &format!("process didn't exit successfully: {}", self), |
| Some(output.status), |
| to_print, |
| )); |
| } |
| } |
| |
| Ok(output) |
| } |
| |
| /// Builds the command with an `@<path>` argfile that contains all the |
| /// arguments. This is primarily served for rustc/rustdoc command family. |
| fn build_command_with_argfile(&self) -> io::Result<(Command, NamedTempFile)> { |
| use std::io::Write as _; |
| |
| let mut tmp = tempfile::Builder::new() |
| .prefix("cargo-argfile.") |
| .tempfile()?; |
| |
| let mut arg = OsString::from("@"); |
| arg.push(tmp.path()); |
| let mut cmd = self.build_command_without_args(); |
| cmd.arg(arg); |
| tracing::debug!("created argfile at {} for {self}", tmp.path().display()); |
| |
| let cap = self.get_args().map(|arg| arg.len() + 1).sum::<usize>(); |
| let mut buf = Vec::with_capacity(cap); |
| for arg in &self.args { |
| let arg = arg.to_str().ok_or_else(|| { |
| io::Error::new( |
| io::ErrorKind::Other, |
| format!( |
| "argument for argfile contains invalid UTF-8 characters: `{}`", |
| arg.to_string_lossy() |
| ), |
| ) |
| })?; |
| if arg.contains('\n') { |
| return Err(io::Error::new( |
| io::ErrorKind::Other, |
| format!("argument for argfile contains newlines: `{arg}`"), |
| )); |
| } |
| writeln!(buf, "{arg}")?; |
| } |
| tmp.write_all(&mut buf)?; |
| Ok((cmd, tmp)) |
| } |
| |
| /// Builds a command from `ProcessBuilder` for everything but not `args`. |
| fn build_command_without_args(&self) -> Command { |
| let mut command = { |
| let mut iter = self.wrappers.iter().rev().chain(once(&self.program)); |
| let mut cmd = Command::new(iter.next().expect("at least one `program` exists")); |
| cmd.args(iter); |
| cmd |
| }; |
| if let Some(cwd) = self.get_cwd() { |
| command.current_dir(cwd); |
| } |
| for (k, v) in &self.env { |
| match *v { |
| Some(ref v) => { |
| command.env(k, v); |
| } |
| None => { |
| command.env_remove(k); |
| } |
| } |
| } |
| if let Some(ref c) = self.jobserver { |
| c.configure(&mut command); |
| } |
| command |
| } |
| |
| /// Converts `ProcessBuilder` into a `std::process::Command`, and handles |
| /// the jobserver, if present. |
| /// |
| /// Note that this method doesn't take argfile fallback into account. The |
| /// caller should handle it by themselves. |
| pub fn build_command(&self) -> Command { |
| let mut command = self.build_command_without_args(); |
| for arg in &self.args { |
| command.arg(arg); |
| } |
| command |
| } |
| |
| /// Wraps an existing command with the provided wrapper, if it is present and valid. |
| /// |
| /// # Examples |
| /// |
| /// ```rust |
| /// use cargo_util::ProcessBuilder; |
| /// // Running this would execute `rustc` |
| /// let cmd = ProcessBuilder::new("rustc"); |
| /// |
| /// // Running this will execute `sccache rustc` |
| /// let cmd = cmd.wrapped(Some("sccache")); |
| /// ``` |
| pub fn wrapped(mut self, wrapper: Option<impl AsRef<OsStr>>) -> Self { |
| if let Some(wrapper) = wrapper.as_ref() { |
| let wrapper = wrapper.as_ref(); |
| if !wrapper.is_empty() { |
| self.wrappers.push(wrapper.to_os_string()); |
| } |
| } |
| self |
| } |
| } |
| |
| /// Forces the command to use `@path` argfile. |
| /// |
| /// You should set `__CARGO_TEST_FORCE_ARGFILE` to enable this. |
| fn debug_force_argfile(retry_enabled: bool) -> bool { |
| cfg!(debug_assertions) && env::var("__CARGO_TEST_FORCE_ARGFILE").is_ok() && retry_enabled |
| } |
| |
| /// Creates new pipes for stderr, stdout, and optionally stdin. |
| fn piped(cmd: &mut Command, pipe_stdin: bool) -> &mut Command { |
| cmd.stdout(Stdio::piped()) |
| .stderr(Stdio::piped()) |
| .stdin(if pipe_stdin { |
| Stdio::piped() |
| } else { |
| Stdio::null() |
| }) |
| } |
| |
| fn close_tempfile_and_log_error(file: NamedTempFile) { |
| file.close().unwrap_or_else(|e| { |
| tracing::warn!("failed to close temporary file: {e}"); |
| }); |
| } |
| |
| #[cfg(unix)] |
| mod imp { |
| use super::{close_tempfile_and_log_error, debug_force_argfile, ProcessBuilder, ProcessError}; |
| use anyhow::Result; |
| use std::io; |
| use std::os::unix::process::CommandExt; |
| |
| pub fn exec_replace(process_builder: &ProcessBuilder) -> Result<()> { |
| let mut error; |
| let mut file = None; |
| if debug_force_argfile(process_builder.retry_with_argfile) { |
| let (mut command, argfile) = process_builder.build_command_with_argfile()?; |
| file = Some(argfile); |
| error = command.exec() |
| } else { |
| let mut command = process_builder.build_command(); |
| error = command.exec(); |
| if process_builder.should_retry_with_argfile(&error) { |
| let (mut command, argfile) = process_builder.build_command_with_argfile()?; |
| file = Some(argfile); |
| error = command.exec() |
| } |
| } |
| if let Some(file) = file { |
| close_tempfile_and_log_error(file); |
| } |
| |
| Err(anyhow::Error::from(error).context(ProcessError::new( |
| &format!("could not execute process {}", process_builder), |
| None, |
| None, |
| ))) |
| } |
| |
| pub fn command_line_too_big(err: &io::Error) -> bool { |
| err.raw_os_error() == Some(libc::E2BIG) |
| } |
| } |
| |
| #[cfg(windows)] |
| mod imp { |
| use super::{ProcessBuilder, ProcessError}; |
| use anyhow::Result; |
| use std::io; |
| use windows_sys::Win32::Foundation::{BOOL, FALSE, TRUE}; |
| use windows_sys::Win32::System::Console::SetConsoleCtrlHandler; |
| |
| unsafe extern "system" fn ctrlc_handler(_: u32) -> BOOL { |
| // Do nothing; let the child process handle it. |
| TRUE |
| } |
| |
| pub fn exec_replace(process_builder: &ProcessBuilder) -> Result<()> { |
| unsafe { |
| if SetConsoleCtrlHandler(Some(ctrlc_handler), TRUE) == FALSE { |
| return Err(ProcessError::new("Could not set Ctrl-C handler.", None, None).into()); |
| } |
| } |
| |
| // Just execute the process as normal. |
| process_builder.exec() |
| } |
| |
| pub fn command_line_too_big(err: &io::Error) -> bool { |
| use windows_sys::Win32::Foundation::ERROR_FILENAME_EXCED_RANGE; |
| err.raw_os_error() == Some(ERROR_FILENAME_EXCED_RANGE as i32) |
| } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use super::ProcessBuilder; |
| use std::fs; |
| |
| #[test] |
| fn argfile_build_succeeds() { |
| let mut cmd = ProcessBuilder::new("echo"); |
| cmd.args(["foo", "bar"].as_slice()); |
| let (cmd, argfile) = cmd.build_command_with_argfile().unwrap(); |
| |
| assert_eq!(cmd.get_program(), "echo"); |
| let cmd_args: Vec<_> = cmd.get_args().map(|s| s.to_str().unwrap()).collect(); |
| assert_eq!(cmd_args.len(), 1); |
| assert!(cmd_args[0].starts_with("@")); |
| assert!(cmd_args[0].contains("cargo-argfile.")); |
| |
| let buf = fs::read_to_string(argfile.path()).unwrap(); |
| assert_eq!(buf, "foo\nbar\n"); |
| } |
| |
| #[test] |
| fn argfile_build_fails_if_arg_contains_newline() { |
| let mut cmd = ProcessBuilder::new("echo"); |
| cmd.arg("foo\n"); |
| let err = cmd.build_command_with_argfile().unwrap_err(); |
| assert_eq!( |
| err.to_string(), |
| "argument for argfile contains newlines: `foo\n`" |
| ); |
| } |
| |
| #[test] |
| fn argfile_build_fails_if_arg_contains_invalid_utf8() { |
| let mut cmd = ProcessBuilder::new("echo"); |
| |
| #[cfg(windows)] |
| let invalid_arg = { |
| use std::os::windows::prelude::*; |
| std::ffi::OsString::from_wide(&[0x0066, 0x006f, 0xD800, 0x006f]) |
| }; |
| |
| #[cfg(unix)] |
| let invalid_arg = { |
| use std::os::unix::ffi::OsStrExt; |
| std::ffi::OsStr::from_bytes(&[0x66, 0x6f, 0x80, 0x6f]).to_os_string() |
| }; |
| |
| cmd.arg(invalid_arg); |
| let err = cmd.build_command_with_argfile().unwrap_err(); |
| assert_eq!( |
| err.to_string(), |
| "argument for argfile contains invalid UTF-8 characters: `fo�o`" |
| ); |
| } |
| } |