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}