Skip to main content

compiletest/build_helper/src/
arg_file_command.rs

1//! This module is explictly not `mod`ed as it's shared across multiple crates
2//! like bootstrap and compiletest via `#[path]` moduled declarations.
3//! It's important to keep this file isolated from the rest of build_helper so it can be compiled
4//! without build_helper.
5
6// Roughly match the `std::process::Command` API
7#![allow(dead_code, unreachable_pub)]
8
9use std::ffi::{OsStr, OsString};
10use std::io::Write;
11use std::path::Path;
12use std::process::{Command, CommandEnvs};
13
14use tempfile::NamedTempFile;
15
16/// A wrapper around [`Command`] that adds support for arg files.
17/// This is useful as we have some commands that can get very long and at times
18/// hit the OS limit (usually Windows)
19///
20/// This implementation is based off the `ProcessBuilder` implementation in Cargo
21/// but simplified.
22///
23/// NOTE: In most scenarios we want to avoid arg files as it makes debugging more complicated
24///       so we try to avoid it if the command is not close to the OS limit.
25#[derive(Debug)]
26pub struct ArgFileCommand {
27    command: Command,
28    args: Vec<OsString>,
29}
30
31impl ArgFileCommand {
32    pub fn new<S: AsRef<OsStr>>(program: S) -> Self {
33        let command = Command::new(program);
34        Self { command, args: Vec::new() }
35    }
36    pub fn arg<S: AsRef<OsStr>>(&mut self, arg: S) -> &mut Self {
37        self.args.push(arg.as_ref().to_os_string());
38        self
39    }
40
41    pub fn args<I, S>(&mut self, args: I) -> &mut Self
42    where
43        I: IntoIterator<Item = S>,
44        S: AsRef<OsStr>,
45    {
46        self.args.extend(args.into_iter().map(|s| s.as_ref().to_os_string()));
47        self
48    }
49
50    pub fn env<K, V>(&mut self, key: K, val: V) -> &mut Self
51    where
52        K: AsRef<OsStr>,
53        V: AsRef<OsStr>,
54    {
55        self.command.env(key, val);
56        self
57    }
58
59    pub fn get_envs(&self) -> CommandEnvs<'_> {
60        self.command.get_envs()
61    }
62
63    pub fn env_remove<K: AsRef<OsStr>>(&mut self, key: K) -> &mut Self {
64        self.command.env_remove(key);
65        self
66    }
67
68    pub fn current_dir<P: AsRef<Path>>(&mut self, dir: P) -> &mut Self {
69        self.command.current_dir(dir);
70        self
71    }
72
73    pub fn stdin(&mut self, stdin: std::process::Stdio) -> &mut Self {
74        self.command.stdin(stdin);
75        self
76    }
77
78    pub fn build(mut self) -> std::io::Result<(Command, Option<NamedTempFile>)> {
79        // On Windows there is a hard limit of ~32KB, so we cut off at 30KB to
80        // give some buffer just incase.
81        #[cfg(windows)]
82        let threshold: usize = 30 * 1024;
83        // On unix the limit is defined by ARG_MAX. If its not explicitly set we set it to 1MB
84        // which is fairly large but lower than the ~2MB that it defaults to on most systems.
85        #[cfg(unix)]
86        let threshold: usize =
87            std::env::var("ARG_MAX").ok().and_then(|v| v.parse().ok()).unwrap_or(1024 * 1024);
88
89        let total_arg_len: usize = self.args.iter().map(|a| a.len() + 1).sum();
90        if total_arg_len <= threshold {
91            self.command.args(self.args);
92            return Ok((self.command, None));
93        }
94
95        let mut tmp = tempfile::Builder::new().prefix("bootstrap-argfile.").tempfile()?;
96
97        let mut arg = OsString::from("@");
98        arg.push(tmp.path());
99        self.command.arg(arg);
100
101        let mut buf = Vec::with_capacity(total_arg_len);
102        for arg in &self.args {
103            let arg = arg.to_str().ok_or_else(|| {
104                std::io::Error::other(format!(
105                    "argument for argfile contains invalid UTF-8 characters: `{}`",
106                    arg.to_string_lossy()
107                ))
108            })?;
109            if arg.contains('\n') {
110                return Err(std::io::Error::other(format!(
111                    "argument for argfile contains newlines: `{arg}`"
112                )));
113            }
114            writeln!(buf, "{arg}")?;
115        }
116        tmp.write_all(&buf)?;
117        tmp.flush()?;
118
119        Ok((self.command, Some(tmp)))
120    }
121}