Skip to content

Commit

Permalink
fix: passing environment variables to containers via temmporary --env…
Browse files Browse the repository at this point in the history
…-file
  • Loading branch information
evilsocket committed Nov 1, 2024
1 parent ce5ea1b commit caf2bb8
Show file tree
Hide file tree
Showing 3 changed files with 115 additions and 4 deletions.
78 changes: 78 additions & 0 deletions src/book/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -82,8 +82,34 @@ impl Container {
.to_string(),
app_in_path: true,
args: vec!["run".to_string(), "--rm".to_string()],
env: BTreeMap::new(),
temp_env_file: None,
};

// handle environment variables if present
if !cmdline.env.is_empty() {
let mut env_contents = String::new();
for (key, value) in &cmdline.env {
env_contents.push_str(&format!("{}={}\n", key, value));
}

// create temp file
let temp_file = tempfile::NamedTempFile::new()
.map_err(|e| anyhow::anyhow!("failed to create temp env file: {}", e))?;

// write env vars
std::fs::write(temp_file.path(), env_contents)
.map_err(|e| anyhow::anyhow!("failed to write env file: {}", e))?;

// add env-file arg
dockerized
.args
.push(format!("--env-file={}", temp_file.path().display()));

// keep temp file alive until docker run completes
dockerized.temp_env_file = Some(temp_file);
}

// add volumes if any
if let Some(volumes) = &self.volumes {
for volume in volumes {
Expand Down Expand Up @@ -397,6 +423,8 @@ mod tests {
app: "original_app".to_string(),
app_in_path: true,
args: vec!["arg1".to_string(), "arg2".to_string()],
env: BTreeMap::new(),
temp_env_file: None,
};

let wrapped_cmdline = container.wrap(original_cmdline).unwrap();
Expand Down Expand Up @@ -515,4 +543,54 @@ functions:
assert!(result.get_function("function1").is_ok());
assert!(result.get_function("function2").is_err());
}

#[test]
fn test_wrap_with_env() {
let env: BTreeMap<String, String> = {
let mut env = BTreeMap::new();
env.insert("TEST_VAR".to_string(), "test_value".to_string());
env
};

let command_line =
CommandLine::from_vec_with_env(&vec!["echo".to_string(), "test".to_string()], env)
.unwrap();

let container = Container {
source: ContainerSource::Image("test_image".to_string()),
args: None,
volumes: None,
force: false,
preserve_app: true,
platform: None,
};

let wrapped = container.wrap(command_line).unwrap();

// Find the env-file argument
let env_file_arg = wrapped
.args
.iter()
.find(|arg| arg.starts_with("--env-file="))
.expect("--env-file argument not found")
.clone();

// Extract the file path
let env_file_path = env_file_arg
.strip_prefix("--env-file=")
.expect("Failed to strip --env-file= prefix");

let env_file = std::path::Path::new(env_file_path);
assert!(env_file.exists());

// Read the env file contents
let env_file_contents = std::fs::read_to_string(env_file).expect("Failed to read env file");

// Verify it contains the expected environment variable
assert!(env_file_contents.contains("TEST_VAR=test_value"));

// Clean up the env file
drop(wrapped);
assert!(!env_file.exists(), "env file was not deleted");
}
}
12 changes: 9 additions & 3 deletions src/book/runtime.rs
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,7 @@ impl<'a> FunctionRef<'a> {
) -> anyhow::Result<CommandLine> {
// determine the command line to execute
let command_line = self.function.execution.get_command_line()?;
let mut env = BTreeMap::new();

// interpolate the arguments
let command_line = {
Expand All @@ -199,7 +200,7 @@ impl<'a> FunctionRef<'a> {
if var_name.starts_with("env.") || var_name.starts_with("ENV.") {
let env_var_name = var_name.replace("env.", "").replace("ENV.", "");
let env_var = std::env::var(&env_var_name);
if let Ok(value) = env_var {
let env_var_value = if let Ok(value) = env_var {
value
} else if var_default.is_some() {
var_default.unwrap().to_string()
Expand All @@ -208,7 +209,12 @@ impl<'a> FunctionRef<'a> {
"environment variable {} not set",
env_var_name
));
}
};

// add the environment variable to the command line for later use
env.insert(env_var_name, env_var_value.to_owned());

env_var_value
} else if let Some(value) = arguments.get(var_name) {
// if the value is empty and there's a default value, use the default value
if value.is_empty() && var_default.is_some() {
Expand All @@ -231,7 +237,7 @@ impl<'a> FunctionRef<'a> {
};

// final parsing
CommandLine::from_vec(&command_line)
CommandLine::from_vec_with_env(&command_line, env)
}
}

Expand Down
29 changes: 28 additions & 1 deletion src/runtime/cmd.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,15 @@
use std::fmt;
use std::{collections::BTreeMap, fmt};

#[derive(Debug)]
pub struct CommandLine {
pub sudo: bool,
pub app: String,
pub app_in_path: bool,
pub args: Vec<String>,
pub env: BTreeMap<String, String>,

// used to keep a valid reference to this while the command is running
pub temp_env_file: Option<tempfile::NamedTempFile>,
}

impl CommandLine {
Expand Down Expand Up @@ -47,9 +51,20 @@ impl CommandLine {
app,
args,
app_in_path,
env: BTreeMap::new(),
temp_env_file: None,
})
}

pub fn from_vec_with_env(
vec: &Vec<String>,
env: BTreeMap<String, String>,
) -> anyhow::Result<Self> {
let mut cmd = Self::from_vec(vec)?;
cmd.env = env;
Ok(cmd)
}

pub async fn execute(&self) -> anyhow::Result<String> {
let output = tokio::process::Command::new(&self.app)
.args(&self.args)
Expand Down Expand Up @@ -111,6 +126,8 @@ mod tests {
app: "ls".to_string(),
args: vec!["-l".to_string(), "-a".to_string()],
app_in_path: true,
env: BTreeMap::new(),
temp_env_file: None,
};
assert_eq!(format!("{}", cmd), "ls -l -a");

Expand All @@ -119,6 +136,8 @@ mod tests {
app: "apt".to_string(),
args: vec!["install".to_string(), "package".to_string()],
app_in_path: true,
env: BTreeMap::new(),
temp_env_file: None,
};
assert_eq!(format!("{}", cmd_with_sudo), "sudo apt install package");
}
Expand All @@ -130,6 +149,8 @@ mod tests {
app: "echo".to_string(),
args: vec!["-n".to_string(), "Hello, World!".to_string()],
app_in_path: true,
env: BTreeMap::new(),
temp_env_file: None,
};
let result = cmd.execute().await.unwrap();
assert_eq!(result, "Hello, World!");
Expand All @@ -142,6 +163,8 @@ mod tests {
app: "ls".to_string(),
args: vec!["nonexistent_file".to_string()],
app_in_path: true,
env: BTreeMap::new(),
temp_env_file: None,
};
let result = cmd.execute().await.unwrap();
assert!(result.contains("EXIT CODE:"));
Expand All @@ -158,6 +181,8 @@ mod tests {
"echo 'Hello' && echo 'Error' >&2".to_string(),
],
app_in_path: true,
env: BTreeMap::new(),
temp_env_file: None,
};
let result = cmd.execute().await.unwrap();
assert!(result.contains("Hello"));
Expand All @@ -171,6 +196,8 @@ mod tests {
app: "".to_string(),
args: vec!["arg1".to_string(), "arg2".to_string()],
app_in_path: true,
env: BTreeMap::new(),
temp_env_file: None,
};
let result = cmd.execute().await;
assert!(result.is_err());
Expand Down

0 comments on commit caf2bb8

Please sign in to comment.