//! A utility for writing scripts for use as test executables intended to match the //! subcommands of Bazel build actions so `rustdoc --test`, which builds and tests //! code in a single call, can be run as a test target in a hermetic manner. use std::cmp::Reverse; use std::collections::{BTreeMap, BTreeSet}; use std::env; use std::fs; use std::io::{BufRead, BufReader}; use std::path::{Path, PathBuf}; #[derive(Debug)] struct Options { /// A list of environment variable keys to parse from the build action env. env_keys: BTreeSet, /// A list of substrings to strip from [Options::action_argv]. strip_substrings: Vec, /// The path where the script should be written. output: PathBuf, /// If Bazel generated a params file, we may need to strip roots from it. /// This is the path where we will output our stripped params file. optional_output_params_file: PathBuf, /// The `argv` of the configured rustdoc build action. action_argv: Vec, } /// Parse command line arguments fn parse_args() -> Options { let args: Vec = env::args().collect(); let (writer_args, action_args) = { let split = args .iter() .position(|arg| arg == "--") .expect("Unable to find split identifier `--`"); // Converting each set into a vector makes them easier to parse in // the absence of nightly features let (writer, action) = args.split_at(split); (writer.to_vec(), action.to_vec()) }; // Remove the leading `--` which is expected to be the first // item in `action_args` debug_assert_eq!(action_args[0], "--"); let action_argv = action_args[1..].to_vec(); let output = writer_args .iter() .find(|arg| arg.starts_with("--output=")) .and_then(|arg| arg.splitn(2, '=').last()) .map(PathBuf::from) .expect("Missing `--output` argument"); let optional_output_params_file = writer_args .iter() .find(|arg| arg.starts_with("--optional_test_params=")) .and_then(|arg| arg.splitn(2, '=').last()) .map(PathBuf::from) .expect("Missing `--optional_test_params` argument"); let (strip_substring_args, writer_args): (Vec, Vec) = writer_args .into_iter() .partition(|arg| arg.starts_with("--strip_substring=")); let mut strip_substrings: Vec = strip_substring_args .into_iter() .map(|arg| { arg.splitn(2, '=') .last() .expect("--strip_substring arguments must have assignments using `=`") .to_owned() }) .collect(); // Strip substrings should always be in reverse order of the length of each // string so when filtering we know that the longer strings are checked // first in order to avoid cases where shorter strings might match longer ones. strip_substrings.sort_by_key(|b| Reverse(b.len())); strip_substrings.dedup(); let env_keys = writer_args .into_iter() .filter(|arg| arg.starts_with("--action_env=")) .map(|arg| { arg.splitn(2, '=') .last() .expect("--env arguments must have assignments using `=`") .to_owned() }) .collect(); Options { env_keys, strip_substrings, output, optional_output_params_file, action_argv, } } /// Expand the Bazel Arg file and write it into our manually defined params file fn expand_params_file(mut options: Options) -> Options { let params_extension = if cfg!(target_family = "windows") { ".rustdoc_test.bat-0.params" } else { ".rustdoc_test.sh-0.params" }; // We always need to produce the params file, we might overwrite this later though fs::write(&options.optional_output_params_file, b"unused") .expect("Failed to write params file"); // extract the path for the params file, if it exists let params_path = match options.action_argv.pop() { // Found the params file! Some(arg) if arg.starts_with('@') && arg.ends_with(params_extension) => { let path_str = arg .strip_prefix('@') .expect("Checked that there is an @ prefix"); PathBuf::from(path_str) } // No params file present, exit early Some(arg) => { options.action_argv.push(arg); return options; } None => return options, }; // read the params file let params_file = fs::File::open(params_path).expect("Failed to read the rustdoc params file"); let content: Vec<_> = BufReader::new(params_file) .lines() .map(|line| line.expect("failed to parse param as String")) // Remove any substrings found in the argument .map(|arg| { let mut stripped_arg = arg; options .strip_substrings .iter() .for_each(|substring| stripped_arg = stripped_arg.replace(substring, "")); stripped_arg }) .collect(); // add all arguments fs::write(&options.optional_output_params_file, content.join("\n")) .expect("Failed to write test runner"); // append the path of our new params file let formatted_params_path = format!( "@{}", options .optional_output_params_file .to_str() .expect("invalid UTF-8") ); options.action_argv.push(formatted_params_path); options } /// Write a unix compatible test runner fn write_test_runner_unix( path: &Path, env: &BTreeMap, argv: &[String], strip_substrings: &[String], ) { let mut content = vec![ "#!/usr/bin/env bash".to_owned(), "".to_owned(), // TODO: Instead of creating a symlink to mimic the behavior of // --legacy_external_runfiles, this rule should be able to correcrtly // sanitize the action args to run in a runfiles without this link. "if [[ ! -e 'external' ]]; then ln -s ../ external ; fi".to_owned(), "".to_owned(), "exec env - \\".to_owned(), ]; content.extend(env.iter().map(|(key, val)| format!("{key}='{val}' \\"))); let argv_str = argv .iter() // Remove any substrings found in the argument .map(|arg| { let mut stripped_arg = arg.to_owned(); strip_substrings .iter() .for_each(|substring| stripped_arg = stripped_arg.replace(substring, "")); stripped_arg }) .map(|arg| format!("'{arg}'")) .collect::>() .join(" "); content.extend(vec![argv_str, "".to_owned()]); fs::write(path, content.join("\n")).expect("Failed to write test runner"); } /// Write a windows compatible test runner fn write_test_runner_windows( path: &Path, env: &BTreeMap, argv: &[String], strip_substrings: &[String], ) { let env_str = env .iter() .map(|(key, val)| format!("$env:{key}='{val}'")) .collect::>() .join(" ; "); let argv_str = argv .iter() // Remove any substrings found in the argument .map(|arg| { let mut stripped_arg = arg.to_owned(); strip_substrings .iter() .for_each(|substring| stripped_arg = stripped_arg.replace(substring, "")); stripped_arg }) .map(|arg| format!("'{arg}'")) .collect::>() .join(" "); let content = [ "@ECHO OFF".to_owned(), "".to_owned(), // TODO: Instead of creating a symlink to mimic the behavior of // --legacy_external_runfiles, this rule should be able to correcrtly // sanitize the action args to run in a runfiles without this link. "powershell.exe -c \"if (!(Test-Path .\\external)) { New-Item -Path .\\external -ItemType SymbolicLink -Value ..\\ }\"" .to_owned(), "".to_owned(), format!("powershell.exe -c \"{env_str} ; & {argv_str}\""), "".to_owned(), ]; fs::write(path, content.join("\n")).expect("Failed to write test runner"); } #[cfg(target_family = "unix")] fn set_executable(path: &Path) { use std::os::unix::prelude::PermissionsExt; let mut perm = fs::metadata(path) .expect("Failed to get test runner metadata") .permissions(); perm.set_mode(0o755); fs::set_permissions(path, perm).expect("Failed to set permissions on test runner"); } #[cfg(target_family = "windows")] fn set_executable(_path: &Path) { // Windows determines whether or not a file is executable via the PATHEXT // environment variable. This function is a no-op for this platform. } fn write_test_runner( path: &Path, env: &BTreeMap, argv: &[String], strip_substrings: &[String], ) { if cfg!(target_family = "unix") { write_test_runner_unix(path, env, argv, strip_substrings); } else if cfg!(target_family = "windows") { write_test_runner_windows(path, env, argv, strip_substrings); } set_executable(path); } fn main() { let opt = parse_args(); let opt = expand_params_file(opt); let env: BTreeMap = env::vars() .filter(|(key, _)| opt.env_keys.iter().any(|k| k == key)) .collect(); write_test_runner(&opt.output, &env, &opt.action_argv, &opt.strip_substrings); }