Skip to content

Commit

Permalink
process wrapper: move argument parsing in a separate mod.
Browse files Browse the repository at this point in the history
  • Loading branch information
@gigaroby
gigaroby committed Mar 7, 2022
1 parent e46a2ce commit 6fe93df
Show file tree
Hide file tree
Showing 5 changed files with 259 additions and 170 deletions.
2 changes: 1 addition & 1 deletion util/process_wrapper/BUILD.bazel
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,6 @@ rust_binary(

rust_test(
name = "process_wrapper_test",
srcs = glob(["*.rs"]),
deps = [":process_wrapper_bin"],
crate = "process_wrapper_bin",
)
7 changes: 6 additions & 1 deletion util/process_wrapper/flags.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
Expand Down
168 changes: 12 additions & 156 deletions util/process_wrapper/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, String>,
stamp_mappings: &[(String, String)],
subst_mappings: &[(String, String)],
) -> HashMap<String, String> {
// Taking all environment variables from the current process
// and sending them down to the child process
let mut environment_variables: HashMap<String, String> = 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<String>, subst_mappings: &[(String, String)]) -> Vec<String> {
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<String, String> = 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<String> = 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)
Expand All @@ -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)
Expand All @@ -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(&copy_source, &copy_dest).unwrap_or_else(|_| {
panic!(
"process wrapper error: failed to copy {} into {}",
Expand Down
Loading

0 comments on commit 6fe93df

Please sign in to comment.