diff --git a/util/process_wrapper/BUILD.bazel b/util/process_wrapper/BUILD.bazel index 1871748071..55074701c0 100644 --- a/util/process_wrapper/BUILD.bazel +++ b/util/process_wrapper/BUILD.bazel @@ -35,6 +35,6 @@ rust_binary( rust_test( name = "process_wrapper_test", - srcs = glob(["*.rs"]), + deps = [":process_wrapper_bin"], crate = "process_wrapper_bin", ) diff --git a/util/process_wrapper/flags.rs b/util/process_wrapper/flags.rs index 113cd4c6e6..50260ce0e5 100644 --- a/util/process_wrapper/flags.rs +++ b/util/process_wrapper/flags.rs @@ -119,7 +119,12 @@ impl<'a> Flags<'a> { all.sort(); let mut help_text = String::new(); - writeln!(&mut help_text, "Help for {}", program_name).unwrap(); + writeln!( + &mut help_text, + "Help for {}: [options] -- [extra arguments]", + program_name + ) + .unwrap(); for line in all { writeln!(&mut help_text, "\t{}", line).unwrap(); } diff --git a/util/process_wrapper/main.rs b/util/process_wrapper/main.rs index 991766d2fa..d49fa5c2b5 100644 --- a/util/process_wrapper/main.rs +++ b/util/process_wrapper/main.rs @@ -13,168 +13,24 @@ // limitations under the License. mod flags; +mod options; mod util; -use std::collections::HashMap; -use std::env; use std::fs::{copy, OpenOptions}; use std::process::{exit, Command, Stdio}; -use flags::{Flags, ParseOutcome}; -use util::*; - -fn environment_block( - environment_file_block: HashMap, - stamp_mappings: &[(String, String)], - subst_mappings: &[(String, String)], -) -> HashMap { - // Taking all environment variables from the current process - // and sending them down to the child process - let mut environment_variables: HashMap = std::env::vars().collect(); - // Have the last values added take precedence over the first. - // This is simpler than needing to track duplicates and explicitly override - // them. - environment_variables.extend(environment_file_block.into_iter()); - for (f, replace_with) in stamp_mappings { - for value in environment_variables.values_mut() { - let from = format!("{{{}}}", f); - let new = value.replace(from.as_str(), replace_with); - *value = new; - } - } - for (f, replace_with) in subst_mappings { - for value in environment_variables.values_mut() { - let from = format!("${{{}}}", f); - let new = value.replace(from.as_str(), replace_with); - *value = new; - } - } - environment_variables -} - -fn prepare_args(mut args: Vec, subst_mappings: &[(String, String)]) -> Vec { - for (f, replace_with) in subst_mappings { - for arg in args.iter_mut() { - let from = format!("${{{}}}", f); - let new = arg.replace(from.as_str(), replace_with); - *arg = new; - } - } - args -} +use crate::options::options; fn main() { - // Process argument list until -- is encountered. - // Everything after is sent to the child process. - let mut subst_mapping_raw = None; - let mut volatile_status_file_raw = None; - let mut env_file_raw = None; - let mut arg_file_raw = None; - let mut touch_file_raw = None; - let mut copy_output_raw = None; - let mut stdout_file_raw = None; - let mut stderr_file_raw = None; - let mut flags = Flags::new(); - flags.define_repeated_flag("--subst", "", &mut subst_mapping_raw); - flags.define_flag("--volatile-status-file", "", &mut volatile_status_file_raw); - flags.define_repeated_flag("--env-file", "", &mut env_file_raw); - flags.define_repeated_flag("--arg-file", "", &mut arg_file_raw); - flags.define_flag("--touch-file", "", &mut touch_file_raw); - flags.define_repeated_flag("--copy-output", "", &mut copy_output_raw); - flags.define_flag( - "--stdout-file", - "Redirect subprocess stdout in this file.", - &mut stdout_file_raw, - ); - flags.define_flag( - "--stderr-file", - "Redirect subprocess stderr in this file.", - &mut stderr_file_raw, - ); - let mut child_args = match flags - .parse(env::args().collect()) - .expect("flag parse error") - { - ParseOutcome::Help(help) => { - eprintln!("{}", help); - return; - } - ParseOutcome::Parsed(p) => p, + let opts = match options() { + Err(err) => panic!("process wrapper error: {}", err), + Ok(v) => v, }; - let subst_mappings: Vec<(String, String)> = subst_mapping_raw - .unwrap_or_default() - .into_iter() - .map(|arg| { - let (key, val) = arg.split_once('=').unwrap_or_else(|| { - panic!( - "process wrapper error: empty key for substitution '{}'", - arg - ) - }); - let v = if val == "${pwd}" { - std::env::current_dir() - .unwrap() - .to_str() - .unwrap() - .to_string() - } else { - val.to_owned() - }; - (key.to_owned(), v) - }) - .collect(); - let stamp_mappings = - volatile_status_file_raw.map_or_else(Vec::new, |s| read_stamp_status_to_array(s).unwrap()); - let environment_file_block: HashMap = env_file_raw - .unwrap_or_default() - .into_iter() - .flat_map(|path| -> Vec<(String, String)> { - let lines = read_file_to_array(path).unwrap(); - lines - .into_iter() - .map(|l| { - let splits = l - .split_once('=') - .expect("process wrapper error: environment file invalid"); - (splits.0.to_owned(), splits.1.to_owned()) - }) - .collect() - }) - .collect(); - let mut file_arguments: Vec = arg_file_raw - .unwrap_or_default() - .into_iter() - .flat_map(|path| read_file_to_array(path).unwrap()) - .collect(); - // Process --copy-output - let copy_output = copy_output_raw.map(|co| { - if co.len() != 2 { - panic!( - "process wrapper error: \"--copy-output\" needs exactly 2 parameters, {} provided", - co.len() - ) - } - let copy_source = &co[0]; - let copy_dest = &co[1]; - if copy_source == copy_dest { - panic!( - "process wrapper error: \"--copy-output\" source ({}) and dest ({}) need to be different.", - copy_source, copy_dest - ) - } - (copy_source.to_owned(), copy_dest.to_owned()) - }); - child_args.append(&mut file_arguments); - let child_args = prepare_args(child_args, &subst_mappings); - - let (exec_path, args) = child_args.split_first().expect("process wrapper error: at least one argument after -- is required (the child process path) but none found."); - let vars = environment_block(environment_file_block, &stamp_mappings, &subst_mappings); - - let status = Command::new(exec_path) - .args(args) + let status = Command::new(opts.executable) + .args(opts.child_arguments) .env_clear() - .envs(vars) - .stdout(if let Some(stdout_file) = stdout_file_raw { + .envs(opts.child_environment) + .stdout(if let Some(stdout_file) = opts.stdout_file { OpenOptions::new() .create(true) .truncate(true) @@ -185,7 +41,7 @@ fn main() { } else { Stdio::inherit() }) - .stderr(if let Some(stderr_file) = stderr_file_raw { + .stderr(if let Some(stderr_file) = opts.stderr_file { OpenOptions::new() .create(true) .truncate(true) @@ -200,14 +56,14 @@ fn main() { .expect("process wrapper error: failed to spawn child process"); if status.success() { - if let Some(tf) = touch_file_raw { + if let Some(tf) = opts.touch_file { OpenOptions::new() .create(true) .write(true) .open(tf) .expect("process wrapper error: failed to create touch file"); } - if let Some((copy_source, copy_dest)) = copy_output { + if let Some((copy_source, copy_dest)) = opts.copy_output { copy(©_source, ©_dest).unwrap_or_else(|_| { panic!( "process wrapper error: failed to copy {} into {}", diff --git a/util/process_wrapper/options.rs b/util/process_wrapper/options.rs new file mode 100644 index 0000000000..25a80aaff4 --- /dev/null +++ b/util/process_wrapper/options.rs @@ -0,0 +1,229 @@ +use std::collections::HashMap; +use std::env; +use std::fmt; +use std::process::exit; + +use crate::flags::{FlagParseError, Flags, ParseOutcome}; +use crate::util::*; + +#[derive(Debug)] +pub(crate) enum OptionError { + FlagError(FlagParseError), + Generic(String), +} + +impl fmt::Display for OptionError { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + Self::FlagError(e) => write!(f, "error parsing flags: {}", e), + Self::Generic(s) => write!(f, "{}", s), + } + } +} + +#[derive(Debug)] +pub(crate) struct Options { + // Contains the path to the child executable + pub(crate) executable: String, + // Contains arguments for the child process fetched from files. + pub(crate) child_arguments: Vec, + // Contains environment variables for the child process fetched from files. + pub(crate) child_environment: HashMap, + // If set, create the specified file after the child process successfully + // terminated its execution. + pub(crate) touch_file: Option, + // If set to (source, dest) copies the source file to dest. + pub(crate) copy_output: Option<(String, String)>, + // If set, redirects the child process stdout to this file. + pub(crate) stdout_file: Option, + // If set, redirects the child process stderr to this file. + pub(crate) stderr_file: Option, +} + +pub(crate) fn options() -> Result { + // Process argument list until -- is encountered. + // Everything after is sent to the child process. + let mut subst_mapping_raw = None; + let mut volatile_status_file_raw = None; + let mut env_file_raw = None; + let mut arg_file_raw = None; + let mut touch_file = None; + let mut copy_output_raw = None; + let mut stdout_file = None; + let mut stderr_file = None; + let mut flags = Flags::new(); + flags.define_repeated_flag("--subst", "", &mut subst_mapping_raw); + flags.define_flag("--volatile-status-file", "", &mut volatile_status_file_raw); + flags.define_repeated_flag( + "--env-file", + "File(s) containing environment variables to pass to the child process.", + &mut env_file_raw, + ); + flags.define_repeated_flag( + "--arg-file", + "File(s) containing command line arguments to pass to the child process.", + &mut arg_file_raw, + ); + flags.define_flag( + "--touch-file", + "Create this file after the child process runs successfully.", + &mut touch_file, + ); + flags.define_repeated_flag("--copy-output", "", &mut copy_output_raw); + flags.define_flag( + "--stdout-file", + "Redirect subprocess stdout in this file.", + &mut stdout_file, + ); + flags.define_flag( + "--stderr-file", + "Redirect subprocess stderr in this file.", + &mut stderr_file, + ); + + let mut child_args = match flags + .parse(env::args().collect()) + .map_err(OptionError::FlagError)? + { + ParseOutcome::Help(help) => { + eprintln!("{}", help); + exit(0); + } + ParseOutcome::Parsed(p) => p, + }; + let current_dir = std::env::current_dir() + .unwrap() + .to_str() + .unwrap() + .to_owned(); + let subst_mappings = subst_mapping_raw + .unwrap_or_default() + .into_iter() + .map(|arg| { + let (key, val) = arg.split_once('=').ok_or_else(|| { + OptionError::Generic(format!( + "process wrapper error: empty key for substitution '{}'", + arg + )) + })?; + let v = if val == "${pwd}" { + current_dir.as_str() + } else { + val + } + .to_owned(); + Ok((key.to_owned(), v)) + }) + .collect::, OptionError>>()?; + let stamp_mappings = + volatile_status_file_raw.map_or_else(Vec::new, |s| read_stamp_status_to_array(s).unwrap()); + + let environment_file_block = env_from_files(env_file_raw.unwrap_or_default())?; + let mut file_arguments = args_from_file(arg_file_raw.unwrap_or_default())?; + // Process --copy-output + let copy_output = copy_output_raw + .map(|co| { + if co.len() != 2 { + return Err(OptionError::Generic(format!( + "\"--copy-output\" needs exactly 2 parameters, {} provided", + co.len() + ))); + } + let copy_source = &co[0]; + let copy_dest = &co[1]; + if copy_source == copy_dest { + return Err(OptionError::Generic(format!( + "\"--copy-output\" source ({}) and dest ({}) need to be different.", + copy_source, copy_dest + ))); + } + Ok((copy_source.to_owned(), copy_dest.to_owned())) + }) + .transpose()?; + + // Prepare the environment variables, unifying those read from files with the ones + // of the current process. + let vars = environment_block(environment_file_block, &stamp_mappings, &subst_mappings); + // Append all the arguments fetched from files to those provided via command line. + child_args.append(&mut file_arguments); + let child_args = prepare_args(child_args, &subst_mappings); + // Split the executable path from the rest of the arguments. + let (exec_path, args) = child_args.split_first().ok_or_else(|| { + OptionError::Generic( + "at least one argument after -- is required (the child process path)".to_owned(), + ) + })?; + + Ok(Options { + executable: exec_path.to_owned(), + child_arguments: args.to_vec(), + child_environment: vars, + touch_file, + copy_output, + stdout_file, + stderr_file, + }) +} + +fn args_from_file(paths: Vec) -> Result, OptionError> { + let mut args = vec![]; + for path in paths.into_iter() { + let mut lines = read_file_to_array(path).map_err(OptionError::Generic)?; + args.append(&mut lines); + } + Ok(args) +} + +fn env_from_files(paths: Vec) -> Result, OptionError> { + let mut env_vars = HashMap::new(); + for path in paths.into_iter() { + let lines = read_file_to_array(path).map_err(OptionError::Generic)?; + for line in lines.into_iter() { + let (k, v) = line.split_once('=').ok_or_else(|| { + OptionError::Generic("process wrapper error: environment file invalid".to_owned()) + })?; + env_vars.insert(k.to_owned(), v.to_owned()); + } + } + Ok(env_vars) +} + +fn prepare_args(mut args: Vec, subst_mappings: &[(String, String)]) -> Vec { + for (f, replace_with) in subst_mappings { + for arg in args.iter_mut() { + let from = format!("${{{}}}", f); + let new = arg.replace(from.as_str(), replace_with); + *arg = new; + } + } + args +} + +fn environment_block( + environment_file_block: HashMap, + stamp_mappings: &[(String, String)], + subst_mappings: &[(String, String)], +) -> HashMap { + // Taking all environment variables from the current process + // and sending them down to the child process + let mut environment_variables: HashMap = std::env::vars().collect(); + // Have the last values added take precedence over the first. + // This is simpler than needing to track duplicates and explicitly override + // them. + environment_variables.extend(environment_file_block.into_iter()); + for (f, replace_with) in stamp_mappings { + for value in environment_variables.values_mut() { + let from = format!("{{{}}}", f); + let new = value.replace(from.as_str(), replace_with); + *value = new; + } + } + for (f, replace_with) in subst_mappings { + for value in environment_variables.values_mut() { + let from = format!("${{{}}}", f); + let new = value.replace(from.as_str(), replace_with); + *value = new; + } + } + environment_variables +} diff --git a/util/process_wrapper/util.rs b/util/process_wrapper/util.rs index 0b5d37ccce..f7adfdc5c8 100644 --- a/util/process_wrapper/util.rs +++ b/util/process_wrapper/util.rs @@ -15,13 +15,22 @@ use std::fs::File; use std::io::{BufRead, BufReader, Read}; +pub(crate) fn read_file_to_array(path: String) -> Result, String> { + let file = File::open(path).map_err(|e| e.to_string())?; + read_to_array(file) +} + +pub(crate) fn read_stamp_status_to_array(path: String) -> Result, String> { + let file = File::open(path).map_err(|e| e.to_string())?; + stamp_status_to_array(file) +} + fn read_to_array(reader: impl Read) -> Result, String> { let reader = BufReader::new(reader); let mut ret = vec![]; let mut escaped_line = String::new(); for l in reader.lines() { let line = l.map_err(|e| e.to_string())?; - // Skip empty lines if any if line.is_empty() { continue; } @@ -30,7 +39,7 @@ fn read_to_array(reader: impl Read) -> Result, String> { // escape let end_backslash_count = line.chars().rev().take_while(|&c| c == '\\').count(); // a 0 or even number of backslashes do not lead to a new line escape - let escape = !(end_backslash_count == 0 || end_backslash_count % 2 == 0); + let escape = end_backslash_count % 2 == 1; // remove backslashes and add back two for every one let l = line.trim_end_matches('\\'); escaped_line.push_str(l); @@ -46,11 +55,6 @@ fn read_to_array(reader: impl Read) -> Result, String> { Ok(ret) } -pub(crate) fn read_file_to_array(path: String) -> Result, String> { - let file = File::open(path).map_err(|e| e.to_string())?; - read_to_array(file) -} - fn stamp_status_to_array(reader: impl Read) -> Result, String> { let escaped_lines = read_to_array(reader)?; escaped_lines @@ -64,11 +68,6 @@ fn stamp_status_to_array(reader: impl Read) -> Result, Str .collect() } -pub(crate) fn read_stamp_status_to_array(path: String) -> Result, String> { - let file = File::open(path).map_err(|e| e.to_string())?; - stamp_status_to_array(file) -} - #[cfg(test)] mod test { use super::*;