From a01bde2a4b206f98b7421ff0678b779b0f694ef0 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sun, 28 Jan 2024 01:08:01 -0500 Subject: [PATCH] env var policy --- Cargo.lock | 34 ++++- crates/lib/subprocess_common/Cargo.toml | 3 + .../command_adapter/command_adapter_args.rs | 0 .../src/command_adapter/mod.rs | 2 - .../src/command_runner/command_args.rs | 7 + .../command_runner.rs} | 55 +++---- .../src/command_runner/command_runner_args.rs | 28 ++++ .../src/command_runner/env_var_policy.rs | 138 ++++++++++++++++++ .../src/command_runner/mod.rs | 4 + crates/lib/subprocess_common/src/lib.rs | 2 +- 10 files changed, 235 insertions(+), 38 deletions(-) delete mode 100644 crates/lib/subprocess_common/src/command_adapter/command_adapter_args.rs delete mode 100644 crates/lib/subprocess_common/src/command_adapter/mod.rs create mode 100644 crates/lib/subprocess_common/src/command_runner/command_args.rs rename crates/lib/subprocess_common/src/{command_adapter/command_adapter.rs => command_runner/command_runner.rs} (79%) create mode 100644 crates/lib/subprocess_common/src/command_runner/command_runner_args.rs create mode 100644 crates/lib/subprocess_common/src/command_runner/env_var_policy.rs create mode 100644 crates/lib/subprocess_common/src/command_runner/mod.rs diff --git a/Cargo.lock b/Cargo.lock index f012171a25..daaded6471 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3632,7 +3632,7 @@ dependencies = [ "serde", "serde_derive", "serde_json", - "serial_test", + "serial_test 2.0.0", "sqlx", "tokens", "tokio", @@ -3905,7 +3905,7 @@ dependencies = [ "serde", "serde_derive", "serde_json", - "serial_test", + "serial_test 2.0.0", "sqlx", "tokens", "tokio", @@ -5526,7 +5526,21 @@ dependencies = [ "lazy_static", "log", "parking_lot 0.12.1", - "serial_test_derive", + "serial_test_derive 2.0.0", +] + +[[package]] +name = "serial_test" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "953ad9342b3aaca7cb43c45c097dd008d4907070394bd0751a0aa8817e5a018d" +dependencies = [ + "dashmap", + "futures 0.3.29", + "lazy_static", + "log", + "parking_lot 0.12.1", + "serial_test_derive 3.0.0", ] [[package]] @@ -5540,6 +5554,17 @@ dependencies = [ "syn 2.0.39", ] +[[package]] +name = "serial_test_derive" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b93fb4adc70021ac1b47f7d45e8cc4169baaa7ea58483bc5b721d19a26202212" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.39", +] + [[package]] name = "server_environment" version = "0.0.1" @@ -6127,7 +6152,7 @@ name = "storyteller_root" version = "0.0.1" dependencies = [ "my-workspace-hack", - "serial_test", + "serial_test 2.0.0", "tempdir", ] @@ -6214,6 +6239,7 @@ dependencies = [ "filesys", "log", "my-workspace-hack", + "serial_test 3.0.0", "subprocess", ] diff --git a/crates/lib/subprocess_common/Cargo.toml b/crates/lib/subprocess_common/Cargo.toml index ae80fbbf75..0d2f32df83 100644 --- a/crates/lib/subprocess_common/Cargo.toml +++ b/crates/lib/subprocess_common/Cargo.toml @@ -29,3 +29,6 @@ subprocess = "0.2.9" #serde_derive = "1.0.125" #serde_json = "1.0.64" #toml = "0.5.8" + +[dev-dependencies] +serial_test = "3.0.0" diff --git a/crates/lib/subprocess_common/src/command_adapter/command_adapter_args.rs b/crates/lib/subprocess_common/src/command_adapter/command_adapter_args.rs deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/crates/lib/subprocess_common/src/command_adapter/mod.rs b/crates/lib/subprocess_common/src/command_adapter/mod.rs deleted file mode 100644 index 7a73486070..0000000000 --- a/crates/lib/subprocess_common/src/command_adapter/mod.rs +++ /dev/null @@ -1,2 +0,0 @@ -pub mod command_adapter; -mod command_adapter_args; \ No newline at end of file diff --git a/crates/lib/subprocess_common/src/command_runner/command_args.rs b/crates/lib/subprocess_common/src/command_runner/command_args.rs new file mode 100644 index 0000000000..a96a4a1508 --- /dev/null +++ b/crates/lib/subprocess_common/src/command_runner/command_args.rs @@ -0,0 +1,7 @@ + +/// Trait for arguments that can be turned into a command line string. +pub trait CommandArgs { + /// Convert to the call signature. + /// NB: We're not returning a HashMap since positions may need to be respected, or some args may have no values. + fn to_command_string(&self) -> String; +} diff --git a/crates/lib/subprocess_common/src/command_adapter/command_adapter.rs b/crates/lib/subprocess_common/src/command_runner/command_runner.rs similarity index 79% rename from crates/lib/subprocess_common/src/command_adapter/command_adapter.rs rename to crates/lib/subprocess_common/src/command_runner/command_runner.rs index 0c1de09858..15328a9dbb 100644 --- a/crates/lib/subprocess_common/src/command_adapter/command_adapter.rs +++ b/crates/lib/subprocess_common/src/command_runner/command_runner.rs @@ -1,7 +1,7 @@ use std::env; use std::ffi::OsString; use std::fs::File; -use std::path::{Path, PathBuf}; +use std::path::PathBuf; use std::time::Duration; use log::info; @@ -11,11 +11,12 @@ use errors::AnyhowResult; use filesys::path_to_string::path_to_string; use crate::command_exit_status::CommandExitStatus; +use crate::command_runner::command_runner_args::{FileOrCreate, RunAsSubprocessArgs}; use crate::docker_options::DockerOptions; use crate::executable_or_command::ExecutableOrShellCommand; #[derive(Clone)] -pub struct CommandAdapter { +pub struct CommandRunner { /// The script, executable, or command to run. pub executable_or_command: ExecutableOrShellCommand, @@ -35,13 +36,8 @@ pub struct CommandAdapter { pub maybe_execution_timeout: Option, } -pub struct RunAsSubprocessArgs<'a> { - /// If set, write stderr to this file. - /// The file must not already exist and must not already be open. - pub maybe_stderr_output_file: Option<&'a Path>, -} -impl CommandAdapter { +impl CommandRunner { /// Run the command with the `subprocess` crate utilities pub fn run_with_subprocess(&self, args: RunAsSubprocessArgs<'_>) -> AnyhowResult { @@ -74,7 +70,9 @@ impl CommandAdapter { } } - // TODO: Command arguments + let command_arguments = args.args.to_command_string(); + command.push_str(&command_arguments); + command.push_str(" "); if let Some(docker_options) = self.maybe_docker_options.as_ref() { command = docker_options.to_command_string(&command); @@ -109,11 +107,24 @@ impl CommandAdapter { config.env = Some(env_vars); } - if let Some(stderr_output_file) = args.maybe_stderr_output_file { - info!("stderr will be written to file: {:?}", stderr_output_file); + match args.maybe_stderr_output_file { + Some(FileOrCreate::NewFileWithName(stderr_output_file)) => { + info!("stderr will be written to file: {:?}", stderr_output_file); - let stderr_file = File::create(stderr_output_file)?; - config.stderr = Redirection::File(stderr_file); + let stderr_file = File::create(stderr_output_file)?; + config.stderr = Redirection::File(stderr_file); + }, + _ => {}, + } + + match args.maybe_stdout_output_file { + Some(FileOrCreate::NewFileWithName(stdout_output_file)) => { + info!("stdout will be written to file: {:?}", stdout_output_file); + + let stdout_file = File::create(stdout_output_file)?; + config.stdout = Redirection::File(stdout_file); + }, + _ => {}, } let mut popen_handle = Popen::create(&command_parts, config)?; @@ -132,24 +143,6 @@ impl CommandAdapter { } }; - //let mut popen = Popen::create( - // &["echo", "hello"], - // PopenConfig { - // stdout: Redirection::Pipe, - // ..Default::default() - // } - //)?; - - //let stdout = popen.stdout.take().unwrap(); - //let mut stdout_reader = BufReader::new(stdout); - - //let mut output = String::new(); - //stdout_reader.read_to_string(&mut output)?; - - //let status = popen.wait()?; - //println!("status: {:?}", status); - //println!("output: {:?}", output); - match maybe_exit_status { None => { // NB: If the program didn't successfully terminate, kill it. diff --git a/crates/lib/subprocess_common/src/command_runner/command_runner_args.rs b/crates/lib/subprocess_common/src/command_runner/command_runner_args.rs new file mode 100644 index 0000000000..97c1f954bf --- /dev/null +++ b/crates/lib/subprocess_common/src/command_runner/command_runner_args.rs @@ -0,0 +1,28 @@ +use std::path::Path; + +use crate::command_runner::command_args::CommandArgs; + +#[derive(Clone, Debug)] +pub enum FileOrCreate<'a> { + /// Open the file at the path. Overwrite if it already exists. + NewFileWithName(&'a Path), + + // Future option(1): new file + // /// Create a new file (and presumably return the name) + // NewFile, + + // Future option(2): append-only file + // Future option(3): already open file handle +} + + +pub struct RunAsSubprocessArgs<'a> { + /// The args for the process we're going to call + pub args: Box, + + /// If set, write stderr to this file. + pub maybe_stderr_output_file: Option>, + + /// If set, write stdout to this file. + pub maybe_stdout_output_file: Option>, +} diff --git a/crates/lib/subprocess_common/src/command_runner/env_var_policy.rs b/crates/lib/subprocess_common/src/command_runner/env_var_policy.rs new file mode 100644 index 0000000000..bae94f1281 --- /dev/null +++ b/crates/lib/subprocess_common/src/command_runner/env_var_policy.rs @@ -0,0 +1,138 @@ +use std::collections::HashSet; +use std::env; +use std::ffi::OsString; + +#[derive(Clone, Debug)] +pub enum EnvVarPolicy { + /// Copy no env vars to the subprocess + CopyNone, + /// Copy all env vars to the subprocess + CopyAll, + /// Copy only the following env vars to the subprocess + CopyOnly(HashSet), + /// Copy all except the following env vars to the subprocess + CopyExcept(HashSet), +} + +impl EnvVarPolicy { + /// Build the list of env vars to send to the subprocess, if any, based on the policy. + pub fn get_env_vars_for_subprocess(&self) -> Option> { + match self { + EnvVarPolicy::CopyNone => { + None + } + EnvVarPolicy::CopyAll => { + Some(env::vars().into_iter() + .map(|(key, value)| (OsString::from(key), OsString::from(value))) + .collect::>()) + } + EnvVarPolicy::CopyOnly(names) => { + Some(env::vars().into_iter() + .filter(|(key, _value)| names.contains(key)) + .map(|(key, value)| (OsString::from(key), OsString::from(value))) + .collect::>()) + + } + EnvVarPolicy::CopyExcept(names) => { + Some(env::vars().into_iter() + .filter(|(key, _value)| !names.contains(key)) + .map(|(key, value)| (OsString::from(key), OsString::from(value))) + .collect::>()) + } + } + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use crate::command_runner::env_var_policy::EnvVarPolicy; + + // NB: Careful with `set_var` as it escapes into *all* other tests! To make matters + // more complicated, tests run in parallel and in any order. We're using a special + // test variable so we never infect other tests. + fn set_env_vars() { + std::env::set_var("STORYTELLER_TEST_VAR_FOO", "foo"); + std::env::set_var("STORYTELLER_TEST_VAR_BAR", "bar"); + std::env::set_var("STORYTELLER_TEST_VAR_BAZ", "baz"); + std::env::set_var("STORYTELLER_TEST_VAR_BIN", "bin"); + } + + // NB: Careful with `set_var` as it escapes into *all* other tests! To make matters + // more complicated, tests run in parallel and in any order. We're using a special + // test variable so we never infect other tests. + fn unset_env_vars() { + std::env::remove_var("STORYTELLER_TEST_VAR_FOO"); + std::env::remove_var("STORYTELLER_TEST_VAR_BAR"); + std::env::remove_var("STORYTELLER_TEST_VAR_BAZ"); + std::env::remove_var("STORYTELLER_TEST_VAR_BIN"); + } + + #[test] + #[serial_test::serial] + fn copy_none() { + let policy = EnvVarPolicy::CopyNone; + assert_eq!(policy.get_env_vars_for_subprocess(), None); + } + + #[test] + #[serial_test::serial] + fn copy_all() { + let policy = EnvVarPolicy::CopyAll; + + set_env_vars(); + + let args = policy.get_env_vars_for_subprocess().expect("should contain args"); + + // NB: All variables contained + assert!(args.iter().find(|(key, _value)| key == "STORYTELLER_TEST_VAR_FOO").is_some()); + assert!(args.iter().find(|(key, _value)| key == "STORYTELLER_TEST_VAR_BAR").is_some()); + assert!(args.iter().find(|(key, _value)| key == "STORYTELLER_TEST_VAR_BAZ").is_some()); + assert!(args.iter().find(|(key, _value)| key == "STORYTELLER_TEST_VAR_BIN").is_some()); + + unset_env_vars(); + } + + #[test] + #[serial_test::serial] + fn copy_except() { + let policy = EnvVarPolicy::CopyExcept(HashSet::from([ + "STORYTELLER_TEST_VAR_FOO".to_string(), + "STORYTELLER_TEST_VAR_BIN".to_string(), + ])); + + set_env_vars(); + + let args = policy.get_env_vars_for_subprocess().expect("should contain args"); + + // NB: Only two contained + assert!(args.iter().find(|(key, _value)| key == "STORYTELLER_TEST_VAR_FOO").is_none()); + assert!(args.iter().find(|(key, _value)| key == "STORYTELLER_TEST_VAR_BAR").is_some()); + assert!(args.iter().find(|(key, _value)| key == "STORYTELLER_TEST_VAR_BAZ").is_some()); + assert!(args.iter().find(|(key, _value)| key == "STORYTELLER_TEST_VAR_BIN").is_none()); + + unset_env_vars(); + } + + #[test] + #[serial_test::serial] + fn copy_only() { + let policy = EnvVarPolicy::CopyOnly(HashSet::from([ + "STORYTELLER_TEST_VAR_FOO".to_string(), + "STORYTELLER_TEST_VAR_BIN".to_string(), + ])); + + set_env_vars(); + + let args = policy.get_env_vars_for_subprocess().expect("should contain args"); + + // NB: Only two contained + assert!(args.iter().find(|(key, _value)| key == "STORYTELLER_TEST_VAR_FOO").is_some()); + assert!(args.iter().find(|(key, _value)| key == "STORYTELLER_TEST_VAR_BAR").is_none()); + assert!(args.iter().find(|(key, _value)| key == "STORYTELLER_TEST_VAR_BAZ").is_none()); + assert!(args.iter().find(|(key, _value)| key == "STORYTELLER_TEST_VAR_BIN").is_some()); + + unset_env_vars(); + } +} diff --git a/crates/lib/subprocess_common/src/command_runner/mod.rs b/crates/lib/subprocess_common/src/command_runner/mod.rs new file mode 100644 index 0000000000..3b70c2f5fb --- /dev/null +++ b/crates/lib/subprocess_common/src/command_runner/mod.rs @@ -0,0 +1,4 @@ +pub mod command_runner; +mod command_runner_args; +mod command_args; +mod env_var_policy; \ No newline at end of file diff --git a/crates/lib/subprocess_common/src/lib.rs b/crates/lib/subprocess_common/src/lib.rs index e4d252240c..41554a3892 100644 --- a/crates/lib/subprocess_common/src/lib.rs +++ b/crates/lib/subprocess_common/src/lib.rs @@ -13,7 +13,7 @@ #![allow(dead_code)] #![allow(non_snake_case)] -pub mod command_adapter; +pub mod command_runner; pub mod command_exit_status; pub mod docker_options; pub mod executable_or_command;