1 //! A utility for writing scripts for use as test executables intended to match the
2 //! subcommands of Bazel build actions so `rustdoc --test`, which builds and tests
3 //! code in a single call, can be run as a test target in a hermetic manner.
4
5 use std::cmp::Reverse;
6 use std::collections::{BTreeMap, BTreeSet};
7 use std::env;
8 use std::fs;
9 use std::io::{BufRead, BufReader};
10 use std::path::{Path, PathBuf};
11
12 #[derive(Debug)]
13 struct Options {
14 /// A list of environment variable keys to parse from the build action env.
15 env_keys: BTreeSet<String>,
16
17 /// A list of substrings to strip from [Options::action_argv].
18 strip_substrings: Vec<String>,
19
20 /// The path where the script should be written.
21 output: PathBuf,
22
23 /// If Bazel generated a params file, we may need to strip roots from it.
24 /// This is the path where we will output our stripped params file.
25 optional_output_params_file: PathBuf,
26
27 /// The `argv` of the configured rustdoc build action.
28 action_argv: Vec<String>,
29 }
30
31 /// Parse command line arguments
parse_args() -> Options32 fn parse_args() -> Options {
33 let args: Vec<String> = env::args().collect();
34 let (writer_args, action_args) = {
35 let split = args
36 .iter()
37 .position(|arg| arg == "--")
38 .expect("Unable to find split identifier `--`");
39
40 // Converting each set into a vector makes them easier to parse in
41 // the absence of nightly features
42 let (writer, action) = args.split_at(split);
43 (writer.to_vec(), action.to_vec())
44 };
45
46 // Remove the leading `--` which is expected to be the first
47 // item in `action_args`
48 debug_assert_eq!(action_args[0], "--");
49 let action_argv = action_args[1..].to_vec();
50
51 let output = writer_args
52 .iter()
53 .find(|arg| arg.starts_with("--output="))
54 .and_then(|arg| arg.splitn(2, '=').last())
55 .map(PathBuf::from)
56 .expect("Missing `--output` argument");
57
58 let optional_output_params_file = writer_args
59 .iter()
60 .find(|arg| arg.starts_with("--optional_test_params="))
61 .and_then(|arg| arg.splitn(2, '=').last())
62 .map(PathBuf::from)
63 .expect("Missing `--optional_test_params` argument");
64
65 let (strip_substring_args, writer_args): (Vec<String>, Vec<String>) = writer_args
66 .into_iter()
67 .partition(|arg| arg.starts_with("--strip_substring="));
68
69 let mut strip_substrings: Vec<String> = strip_substring_args
70 .into_iter()
71 .map(|arg| {
72 arg.splitn(2, '=')
73 .last()
74 .expect("--strip_substring arguments must have assignments using `=`")
75 .to_owned()
76 })
77 .collect();
78
79 // Strip substrings should always be in reverse order of the length of each
80 // string so when filtering we know that the longer strings are checked
81 // first in order to avoid cases where shorter strings might match longer ones.
82 strip_substrings.sort_by_key(|b| Reverse(b.len()));
83 strip_substrings.dedup();
84
85 let env_keys = writer_args
86 .into_iter()
87 .filter(|arg| arg.starts_with("--action_env="))
88 .map(|arg| {
89 arg.splitn(2, '=')
90 .last()
91 .expect("--env arguments must have assignments using `=`")
92 .to_owned()
93 })
94 .collect();
95
96 Options {
97 env_keys,
98 strip_substrings,
99 output,
100 optional_output_params_file,
101 action_argv,
102 }
103 }
104
105 /// Expand the Bazel Arg file and write it into our manually defined params file
expand_params_file(mut options: Options) -> Options106 fn expand_params_file(mut options: Options) -> Options {
107 let params_extension = if cfg!(target_family = "windows") {
108 ".rustdoc_test.bat-0.params"
109 } else {
110 ".rustdoc_test.sh-0.params"
111 };
112
113 // We always need to produce the params file, we might overwrite this later though
114 fs::write(&options.optional_output_params_file, b"unused")
115 .expect("Failed to write params file");
116
117 // extract the path for the params file, if it exists
118 let params_path = match options.action_argv.pop() {
119 // Found the params file!
120 Some(arg) if arg.starts_with('@') && arg.ends_with(params_extension) => {
121 let path_str = arg
122 .strip_prefix('@')
123 .expect("Checked that there is an @ prefix");
124 PathBuf::from(path_str)
125 }
126 // No params file present, exit early
127 Some(arg) => {
128 options.action_argv.push(arg);
129 return options;
130 }
131 None => return options,
132 };
133
134 // read the params file
135 let params_file = fs::File::open(params_path).expect("Failed to read the rustdoc params file");
136 let content: Vec<_> = BufReader::new(params_file)
137 .lines()
138 .map(|line| line.expect("failed to parse param as String"))
139 // Remove any substrings found in the argument
140 .map(|arg| {
141 let mut stripped_arg = arg;
142 options
143 .strip_substrings
144 .iter()
145 .for_each(|substring| stripped_arg = stripped_arg.replace(substring, ""));
146 stripped_arg
147 })
148 .collect();
149
150 // add all arguments
151 fs::write(&options.optional_output_params_file, content.join("\n"))
152 .expect("Failed to write test runner");
153
154 // append the path of our new params file
155 let formatted_params_path = format!(
156 "@{}",
157 options
158 .optional_output_params_file
159 .to_str()
160 .expect("invalid UTF-8")
161 );
162 options.action_argv.push(formatted_params_path);
163
164 options
165 }
166
167 /// Write a unix compatible test runner
write_test_runner_unix( path: &Path, env: &BTreeMap<String, String>, argv: &[String], strip_substrings: &[String], )168 fn write_test_runner_unix(
169 path: &Path,
170 env: &BTreeMap<String, String>,
171 argv: &[String],
172 strip_substrings: &[String],
173 ) {
174 let mut content = vec![
175 "#!/usr/bin/env bash".to_owned(),
176 "".to_owned(),
177 // TODO: Instead of creating a symlink to mimic the behavior of
178 // --legacy_external_runfiles, this rule should be able to correcrtly
179 // sanitize the action args to run in a runfiles without this link.
180 "if [[ ! -e 'external' ]]; then ln -s ../ external ; fi".to_owned(),
181 "".to_owned(),
182 "exec env - \\".to_owned(),
183 ];
184
185 content.extend(env.iter().map(|(key, val)| format!("{key}='{val}' \\")));
186
187 let argv_str = argv
188 .iter()
189 // Remove any substrings found in the argument
190 .map(|arg| {
191 let mut stripped_arg = arg.to_owned();
192 strip_substrings
193 .iter()
194 .for_each(|substring| stripped_arg = stripped_arg.replace(substring, ""));
195 stripped_arg
196 })
197 .map(|arg| format!("'{arg}'"))
198 .collect::<Vec<String>>()
199 .join(" ");
200
201 content.extend(vec![argv_str, "".to_owned()]);
202
203 fs::write(path, content.join("\n")).expect("Failed to write test runner");
204 }
205
206 /// Write a windows compatible test runner
207 fn write_test_runner_windows(
208 path: &Path,
209 env: &BTreeMap<String, String>,
210 argv: &[String],
211 strip_substrings: &[String],
212 ) {
213 let env_str = env
214 .iter()
215 .map(|(key, val)| format!("$env:{key}='{val}'"))
216 .collect::<Vec<String>>()
217 .join(" ; ");
218
219 let argv_str = argv
220 .iter()
221 // Remove any substrings found in the argument
222 .map(|arg| {
223 let mut stripped_arg = arg.to_owned();
224 strip_substrings
225 .iter()
226 .for_each(|substring| stripped_arg = stripped_arg.replace(substring, ""));
227 stripped_arg
228 })
229 .map(|arg| format!("'{arg}'"))
230 .collect::<Vec<String>>()
231 .join(" ");
232
233 let content = [
234 "@ECHO OFF".to_owned(),
235 "".to_owned(),
236 // TODO: Instead of creating a symlink to mimic the behavior of
237 // --legacy_external_runfiles, this rule should be able to correcrtly
238 // sanitize the action args to run in a runfiles without this link.
239 "powershell.exe -c \"if (!(Test-Path .\\external)) { New-Item -Path .\\external -ItemType SymbolicLink -Value ..\\ }\""
240 .to_owned(),
241 "".to_owned(),
242 format!("powershell.exe -c \"{env_str} ; & {argv_str}\""),
243 "".to_owned(),
244 ];
245
246 fs::write(path, content.join("\n")).expect("Failed to write test runner");
247 }
248
249 #[cfg(target_family = "unix")]
set_executable(path: &Path)250 fn set_executable(path: &Path) {
251 use std::os::unix::prelude::PermissionsExt;
252
253 let mut perm = fs::metadata(path)
254 .expect("Failed to get test runner metadata")
255 .permissions();
256
257 perm.set_mode(0o755);
258 fs::set_permissions(path, perm).expect("Failed to set permissions on test runner");
259 }
260
261 #[cfg(target_family = "windows")]
set_executable(_path: &Path)262 fn set_executable(_path: &Path) {
263 // Windows determines whether or not a file is executable via the PATHEXT
264 // environment variable. This function is a no-op for this platform.
265 }
266
write_test_runner( path: &Path, env: &BTreeMap<String, String>, argv: &[String], strip_substrings: &[String], )267 fn write_test_runner(
268 path: &Path,
269 env: &BTreeMap<String, String>,
270 argv: &[String],
271 strip_substrings: &[String],
272 ) {
273 if cfg!(target_family = "unix") {
274 write_test_runner_unix(path, env, argv, strip_substrings);
275 } else if cfg!(target_family = "windows") {
276 write_test_runner_windows(path, env, argv, strip_substrings);
277 }
278
279 set_executable(path);
280 }
281
main()282 fn main() {
283 let opt = parse_args();
284 let opt = expand_params_file(opt);
285
286 let env: BTreeMap<String, String> = env::vars()
287 .filter(|(key, _)| opt.env_keys.iter().any(|k| k == key))
288 .collect();
289
290 write_test_runner(&opt.output, &env, &opt.action_argv, &opt.strip_substrings);
291 }
292