1use crate::process_error::ProcessError;
2use crate::read2;
3
4use anyhow::{Context, Result, bail};
5use jobserver::Client;
6use shell_escape::escape;
7use tempfile::NamedTempFile;
8
9use std::collections::BTreeMap;
10use std::env;
11use std::ffi::{OsStr, OsString};
12use std::fmt;
13use std::io::{self, Write};
14use std::iter::once;
15use std::path::Path;
16use std::process::{Command, ExitStatus, Output, Stdio};
17
18#[derive(Clone, Debug)]
20pub struct ProcessBuilder {
21 program: OsString,
23 arg0: Option<OsString>,
25 args: Vec<OsString>,
27 env: BTreeMap<String, Option<OsString>>,
29 cwd: Option<OsString>,
31 wrappers: Vec<OsString>,
34 jobserver: Option<Client>,
39 display_env_vars: bool,
41 retry_with_argfile: bool,
44 stdin: Option<Vec<u8>>,
46}
47
48impl fmt::Display for ProcessBuilder {
49 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50 write!(f, "`")?;
51
52 if self.display_env_vars {
53 for (key, val) in self.env.iter() {
54 if let Some(val) = val {
55 let val = escape(val.to_string_lossy());
56 if cfg!(windows) {
57 write!(f, "set {}={}&& ", key, val)?;
58 } else {
59 write!(f, "{}={} ", key, val)?;
60 }
61 }
62 }
63 }
64
65 write!(f, "{}", self.get_program().to_string_lossy())?;
66
67 for arg in self.get_args() {
68 write!(f, " {}", escape(arg.to_string_lossy()))?;
69 }
70
71 write!(f, "`")
72 }
73}
74
75impl ProcessBuilder {
76 pub fn new<T: AsRef<OsStr>>(cmd: T) -> ProcessBuilder {
78 ProcessBuilder {
79 program: cmd.as_ref().to_os_string(),
80 arg0: None,
81 args: Vec::new(),
82 cwd: None,
83 env: BTreeMap::new(),
84 wrappers: Vec::new(),
85 jobserver: None,
86 display_env_vars: false,
87 retry_with_argfile: false,
88 stdin: None,
89 }
90 }
91
92 pub fn program<T: AsRef<OsStr>>(&mut self, program: T) -> &mut ProcessBuilder {
94 self.program = program.as_ref().to_os_string();
95 self
96 }
97
98 pub fn arg0<T: AsRef<OsStr>>(&mut self, arg: T) -> &mut ProcessBuilder {
100 self.arg0 = Some(arg.as_ref().to_os_string());
101 self
102 }
103
104 pub fn arg<T: AsRef<OsStr>>(&mut self, arg: T) -> &mut ProcessBuilder {
106 self.args.push(arg.as_ref().to_os_string());
107 self
108 }
109
110 pub fn args<T: AsRef<OsStr>>(&mut self, args: &[T]) -> &mut ProcessBuilder {
112 self.args
113 .extend(args.iter().map(|t| t.as_ref().to_os_string()));
114 self
115 }
116
117 pub fn args_replace<T: AsRef<OsStr>>(&mut self, args: &[T]) -> &mut ProcessBuilder {
119 if let Some(program) = self.wrappers.pop() {
120 self.program = program;
124 self.wrappers = Vec::new();
125 }
126 self.args = args.iter().map(|t| t.as_ref().to_os_string()).collect();
127 self
128 }
129
130 pub fn cwd<T: AsRef<OsStr>>(&mut self, path: T) -> &mut ProcessBuilder {
132 self.cwd = Some(path.as_ref().to_os_string());
133 self
134 }
135
136 pub fn env<T: AsRef<OsStr>>(&mut self, key: &str, val: T) -> &mut ProcessBuilder {
138 self.env
139 .insert(key.to_string(), Some(val.as_ref().to_os_string()));
140 self
141 }
142
143 pub fn env_remove(&mut self, key: &str) -> &mut ProcessBuilder {
145 self.env.insert(key.to_string(), None);
146 self
147 }
148
149 pub fn get_program(&self) -> &OsString {
151 self.wrappers.last().unwrap_or(&self.program)
152 }
153
154 pub fn get_arg0(&self) -> Option<&OsStr> {
156 self.arg0.as_deref()
157 }
158
159 pub fn get_args(&self) -> impl Iterator<Item = &OsString> {
161 self.wrappers
162 .iter()
163 .rev()
164 .chain(once(&self.program))
165 .chain(self.args.iter())
166 .skip(1) }
168
169 pub fn get_cwd(&self) -> Option<&Path> {
171 self.cwd.as_ref().map(Path::new)
172 }
173
174 pub fn get_env(&self, var: &str) -> Option<OsString> {
177 self.env
178 .get(var)
179 .cloned()
180 .or_else(|| Some(env::var_os(var)))
181 .and_then(|s| s)
182 }
183
184 pub fn get_envs(&self) -> &BTreeMap<String, Option<OsString>> {
187 &self.env
188 }
189
190 pub fn inherit_jobserver(&mut self, jobserver: &Client) -> &mut Self {
195 self.jobserver = Some(jobserver.clone());
196 self
197 }
198
199 pub fn display_env_vars(&mut self) -> &mut Self {
201 self.display_env_vars = true;
202 self
203 }
204
205 pub fn retry_with_argfile(&mut self, enabled: bool) -> &mut Self {
223 self.retry_with_argfile = enabled;
224 self
225 }
226
227 pub fn stdin<T: Into<Vec<u8>>>(&mut self, stdin: T) -> &mut Self {
229 self.stdin = Some(stdin.into());
230 self
231 }
232
233 fn should_retry_with_argfile(&self, err: &io::Error) -> bool {
234 self.retry_with_argfile && imp::command_line_too_big(err)
235 }
236
237 pub fn status(&self) -> Result<ExitStatus> {
239 self._status()
240 .with_context(|| ProcessError::could_not_execute(self))
241 }
242
243 fn _status(&self) -> io::Result<ExitStatus> {
244 if !debug_force_argfile(self.retry_with_argfile) {
245 let mut cmd = self.build_command();
246 match cmd.spawn() {
247 Err(ref e) if self.should_retry_with_argfile(e) => {}
248 Err(e) => return Err(e),
249 Ok(mut child) => return child.wait(),
250 }
251 }
252 let (mut cmd, argfile) = self.build_command_with_argfile()?;
253 let status = cmd.spawn()?.wait();
254 close_tempfile_and_log_error(argfile);
255 status
256 }
257
258 pub fn exec(&self) -> Result<()> {
260 let exit = self.status()?;
261 if exit.success() {
262 Ok(())
263 } else {
264 Err(ProcessError::new(
265 &format!("process didn't exit successfully: {}", self),
266 Some(exit),
267 None,
268 )
269 .into())
270 }
271 }
272
273 pub fn exec_replace(&self) -> Result<()> {
289 imp::exec_replace(self)
290 }
291
292 pub fn output(&self) -> Result<Output> {
294 self._output()
295 .with_context(|| ProcessError::could_not_execute(self))
296 }
297
298 fn _output(&self) -> io::Result<Output> {
299 if !debug_force_argfile(self.retry_with_argfile) {
300 let mut cmd = self.build_command();
301 match piped(&mut cmd, self.stdin.is_some()).spawn() {
302 Err(ref e) if self.should_retry_with_argfile(e) => {}
303 Err(e) => return Err(e),
304 Ok(mut child) => {
305 if let Some(stdin) = &self.stdin {
306 child.stdin.take().unwrap().write_all(stdin)?;
307 }
308 return child.wait_with_output();
309 }
310 }
311 }
312 let (mut cmd, argfile) = self.build_command_with_argfile()?;
313 let mut child = piped(&mut cmd, self.stdin.is_some()).spawn()?;
314 if let Some(stdin) = &self.stdin {
315 child.stdin.take().unwrap().write_all(stdin)?;
316 }
317 let output = child.wait_with_output();
318 close_tempfile_and_log_error(argfile);
319 output
320 }
321
322 pub fn exec_with_output(&self) -> Result<Output> {
324 let output = self.output()?;
325 if output.status.success() {
326 Ok(output)
327 } else {
328 Err(ProcessError::new(
329 &format!("process didn't exit successfully: {}", self),
330 Some(output.status),
331 Some(&output),
332 )
333 .into())
334 }
335 }
336
337 pub fn exec_with_streaming(
347 &self,
348 on_stdout_line: &mut dyn FnMut(&str) -> Result<()>,
349 on_stderr_line: &mut dyn FnMut(&str) -> Result<()>,
350 capture_output: bool,
351 ) -> Result<Output> {
352 let mut stdout = Vec::new();
353 let mut stderr = Vec::new();
354
355 let mut callback_error = None;
356 let mut stdout_pos = 0;
357 let mut stderr_pos = 0;
358
359 let spawn = |mut cmd| {
360 if !debug_force_argfile(self.retry_with_argfile) {
361 match piped(&mut cmd, false).spawn() {
362 Err(ref e) if self.should_retry_with_argfile(e) => {}
363 Err(e) => return Err(e),
364 Ok(child) => return Ok((child, None)),
365 }
366 }
367 let (mut cmd, argfile) = self.build_command_with_argfile()?;
368 Ok((piped(&mut cmd, false).spawn()?, Some(argfile)))
369 };
370
371 let status = (|| {
372 let cmd = self.build_command();
373 let (mut child, argfile) = spawn(cmd)?;
374 let out = child.stdout.take().unwrap();
375 let err = child.stderr.take().unwrap();
376 read2(out, err, &mut |is_out, data, eof| {
377 let pos = if is_out {
378 &mut stdout_pos
379 } else {
380 &mut stderr_pos
381 };
382 let idx = if eof {
383 data.len()
384 } else {
385 match data[*pos..].iter().rposition(|b| *b == b'\n') {
386 Some(i) => *pos + i + 1,
387 None => {
388 *pos = data.len();
389 return;
390 }
391 }
392 };
393
394 let new_lines = &data[..idx];
395
396 for line in String::from_utf8_lossy(new_lines).lines() {
397 if callback_error.is_some() {
398 break;
399 }
400 let callback_result = if is_out {
401 on_stdout_line(line)
402 } else {
403 on_stderr_line(line)
404 };
405 if let Err(e) = callback_result {
406 callback_error = Some(e);
407 break;
408 }
409 }
410
411 if capture_output {
412 let dst = if is_out { &mut stdout } else { &mut stderr };
413 dst.extend(new_lines);
414 }
415
416 data.drain(..idx);
417 *pos = 0;
418 })?;
419 let status = child.wait();
420 if let Some(argfile) = argfile {
421 close_tempfile_and_log_error(argfile);
422 }
423 status
424 })()
425 .with_context(|| ProcessError::could_not_execute(self))?;
426 let output = Output {
427 status,
428 stdout,
429 stderr,
430 };
431
432 {
433 let to_print = if capture_output { Some(&output) } else { None };
434 if let Some(e) = callback_error {
435 let cx = ProcessError::new(
436 &format!("failed to parse process output: {}", self),
437 Some(output.status),
438 to_print,
439 );
440 bail!(anyhow::Error::new(cx).context(e));
441 } else if !output.status.success() {
442 bail!(ProcessError::new(
443 &format!("process didn't exit successfully: {}", self),
444 Some(output.status),
445 to_print,
446 ));
447 }
448 }
449
450 Ok(output)
451 }
452
453 fn build_command_with_argfile(&self) -> io::Result<(Command, NamedTempFile)> {
456 use std::io::Write as _;
457
458 let mut tmp = tempfile::Builder::new()
459 .prefix("cargo-argfile.")
460 .tempfile()?;
461
462 let mut arg = OsString::from("@");
463 arg.push(tmp.path());
464 let mut cmd = self.build_command_without_args();
465 cmd.arg(arg);
466 tracing::debug!("created argfile at {} for {self}", tmp.path().display());
467
468 let cap = self.get_args().map(|arg| arg.len() + 1).sum::<usize>();
469 let mut buf = Vec::with_capacity(cap);
470 for arg in &self.args {
471 let arg = arg.to_str().ok_or_else(|| {
472 io::Error::new(
473 io::ErrorKind::Other,
474 format!(
475 "argument for argfile contains invalid UTF-8 characters: `{}`",
476 arg.to_string_lossy()
477 ),
478 )
479 })?;
480 if arg.contains('\n') {
481 return Err(io::Error::new(
482 io::ErrorKind::Other,
483 format!("argument for argfile contains newlines: `{arg}`"),
484 ));
485 }
486 writeln!(buf, "{arg}")?;
487 }
488 tmp.write_all(&mut buf)?;
489 Ok((cmd, tmp))
490 }
491
492 fn build_command_without_args(&self) -> Command {
494 let mut command = {
495 let mut iter = self.wrappers.iter().rev().chain(once(&self.program));
496 let mut cmd = Command::new(iter.next().expect("at least one `program` exists"));
497 cmd.args(iter);
498 cmd
499 };
500 #[cfg(unix)]
501 if let Some(arg0) = self.get_arg0() {
502 use std::os::unix::process::CommandExt as _;
503 command.arg0(arg0);
504 }
505 if let Some(cwd) = self.get_cwd() {
506 command.current_dir(cwd);
507 }
508 for (k, v) in &self.env {
509 match *v {
510 Some(ref v) => {
511 command.env(k, v);
512 }
513 None => {
514 command.env_remove(k);
515 }
516 }
517 }
518 if let Some(ref c) = self.jobserver {
519 c.configure(&mut command);
520 }
521 command
522 }
523
524 pub fn build_command(&self) -> Command {
530 let mut command = self.build_command_without_args();
531 for arg in &self.args {
532 command.arg(arg);
533 }
534 command
535 }
536
537 pub fn wrapped(mut self, wrapper: Option<impl AsRef<OsStr>>) -> Self {
550 if let Some(wrapper) = wrapper.as_ref() {
551 let wrapper = wrapper.as_ref();
552 if !wrapper.is_empty() {
553 self.wrappers.push(wrapper.to_os_string());
554 }
555 }
556 self
557 }
558}
559
560fn debug_force_argfile(retry_enabled: bool) -> bool {
564 cfg!(debug_assertions) && env::var("__CARGO_TEST_FORCE_ARGFILE").is_ok() && retry_enabled
565}
566
567fn piped(cmd: &mut Command, pipe_stdin: bool) -> &mut Command {
569 cmd.stdout(Stdio::piped())
570 .stderr(Stdio::piped())
571 .stdin(if pipe_stdin {
572 Stdio::piped()
573 } else {
574 Stdio::null()
575 })
576}
577
578fn close_tempfile_and_log_error(file: NamedTempFile) {
579 file.close().unwrap_or_else(|e| {
580 tracing::warn!("failed to close temporary file: {e}");
581 });
582}
583
584#[cfg(unix)]
585mod imp {
586 use super::{ProcessBuilder, ProcessError, close_tempfile_and_log_error, debug_force_argfile};
587 use anyhow::Result;
588 use std::io;
589 use std::os::unix::process::CommandExt;
590
591 pub fn exec_replace(process_builder: &ProcessBuilder) -> Result<()> {
592 let mut error;
593 let mut file = None;
594 if debug_force_argfile(process_builder.retry_with_argfile) {
595 let (mut command, argfile) = process_builder.build_command_with_argfile()?;
596 file = Some(argfile);
597 error = command.exec()
598 } else {
599 let mut command = process_builder.build_command();
600 error = command.exec();
601 if process_builder.should_retry_with_argfile(&error) {
602 let (mut command, argfile) = process_builder.build_command_with_argfile()?;
603 file = Some(argfile);
604 error = command.exec()
605 }
606 }
607 if let Some(file) = file {
608 close_tempfile_and_log_error(file);
609 }
610
611 Err(anyhow::Error::from(error).context(ProcessError::new(
612 &format!("could not execute process {}", process_builder),
613 None,
614 None,
615 )))
616 }
617
618 pub fn command_line_too_big(err: &io::Error) -> bool {
619 err.raw_os_error() == Some(libc::E2BIG)
620 }
621}
622
623#[cfg(windows)]
624mod imp {
625 use super::{ProcessBuilder, ProcessError};
626 use anyhow::Result;
627 use std::io;
628 use windows_sys::Win32::Foundation::{FALSE, TRUE};
629 use windows_sys::Win32::System::Console::SetConsoleCtrlHandler;
630 use windows_sys::core::BOOL;
631
632 unsafe extern "system" fn ctrlc_handler(_: u32) -> BOOL {
633 TRUE
635 }
636
637 pub fn exec_replace(process_builder: &ProcessBuilder) -> Result<()> {
638 unsafe {
639 if SetConsoleCtrlHandler(Some(ctrlc_handler), TRUE) == FALSE {
640 return Err(ProcessError::new("Could not set Ctrl-C handler.", None, None).into());
641 }
642 }
643
644 process_builder.exec()
646 }
647
648 pub fn command_line_too_big(err: &io::Error) -> bool {
649 use windows_sys::Win32::Foundation::ERROR_FILENAME_EXCED_RANGE;
650 err.raw_os_error() == Some(ERROR_FILENAME_EXCED_RANGE as i32)
651 }
652}
653
654#[cfg(test)]
655mod tests {
656 use super::ProcessBuilder;
657 use std::fs;
658
659 #[test]
660 fn argfile_build_succeeds() {
661 let mut cmd = ProcessBuilder::new("echo");
662 cmd.args(["foo", "bar"].as_slice());
663 let (cmd, argfile) = cmd.build_command_with_argfile().unwrap();
664
665 assert_eq!(cmd.get_program(), "echo");
666 let cmd_args: Vec<_> = cmd.get_args().map(|s| s.to_str().unwrap()).collect();
667 assert_eq!(cmd_args.len(), 1);
668 assert!(cmd_args[0].starts_with("@"));
669 assert!(cmd_args[0].contains("cargo-argfile."));
670
671 let buf = fs::read_to_string(argfile.path()).unwrap();
672 assert_eq!(buf, "foo\nbar\n");
673 }
674
675 #[test]
676 fn argfile_build_fails_if_arg_contains_newline() {
677 let mut cmd = ProcessBuilder::new("echo");
678 cmd.arg("foo\n");
679 let err = cmd.build_command_with_argfile().unwrap_err();
680 assert_eq!(
681 err.to_string(),
682 "argument for argfile contains newlines: `foo\n`"
683 );
684 }
685
686 #[test]
687 fn argfile_build_fails_if_arg_contains_invalid_utf8() {
688 let mut cmd = ProcessBuilder::new("echo");
689
690 #[cfg(windows)]
691 let invalid_arg = {
692 use std::os::windows::prelude::*;
693 std::ffi::OsString::from_wide(&[0x0066, 0x006f, 0xD800, 0x006f])
694 };
695
696 #[cfg(unix)]
697 let invalid_arg = {
698 use std::os::unix::ffi::OsStrExt;
699 std::ffi::OsStr::from_bytes(&[0x66, 0x6f, 0x80, 0x6f]).to_os_string()
700 };
701
702 cmd.arg(invalid_arg);
703 let err = cmd.build_command_with_argfile().unwrap_err();
704 assert_eq!(
705 err.to_string(),
706 "argument for argfile contains invalid UTF-8 characters: `fo�o`"
707 );
708 }
709}