diff --git a/std/src/sys/pal/uefi/process.rs b/std/src/sys/pal/uefi/process.rs index 3077a72eac661..0757f1cb490d4 100644 --- a/std/src/sys/pal/uefi/process.rs +++ b/std/src/sys/pal/uefi/process.rs @@ -1,6 +1,7 @@ use r_efi::protocols::simple_text_output; use super::helpers; +use crate::collections::BTreeMap; pub use crate::ffi::OsString as EnvKey; use crate::ffi::{OsStr, OsString}; use crate::num::{NonZero, NonZeroI32}; @@ -21,6 +22,7 @@ pub struct Command { args: Vec, stdout: Option, stderr: Option, + env: CommandEnv, } // passed back to std::process with the pipes connected to the child, if any @@ -40,7 +42,13 @@ pub enum Stdio { impl Command { pub fn new(program: &OsStr) -> Command { - Command { prog: program.to_os_string(), args: Vec::new(), stdout: None, stderr: None } + Command { + prog: program.to_os_string(), + args: Vec::new(), + stdout: None, + stderr: None, + env: Default::default(), + } } pub fn arg(&mut self, arg: &OsStr) { @@ -48,7 +56,7 @@ impl Command { } pub fn env_mut(&mut self) -> &mut CommandEnv { - panic!("unsupported") + &mut self.env } pub fn cwd(&mut self, _dir: &OsStr) { @@ -76,7 +84,7 @@ impl Command { } pub fn get_envs(&self) -> CommandEnvs<'_> { - panic!("unsupported") + self.env.iter() } pub fn get_current_dir(&self) -> Option<&Path> { @@ -140,8 +148,30 @@ impl Command { cmd.stderr_inherit() }; + let env = env_changes(&self.env); + + // Set any new vars + if let Some(e) = &env { + for (k, (_, v)) in e { + match v { + Some(v) => crate::env::set_var(k, v), + None => crate::env::remove_var(k), + } + } + } + let stat = cmd.start_image()?; + // Rollback any env changes + if let Some(e) = env { + for (k, (v, _)) in e { + match v { + Some(v) => crate::env::set_var(k, v), + None => crate::env::remove_var(k), + } + } + } + let stdout = cmd.stdout()?; let stderr = cmd.stderr()?; @@ -725,3 +755,32 @@ mod uefi_command_internal { res.into_boxed_slice() } } + +/// Create a map of environment variable changes. Allows efficient setting and rolling back of +/// enviroment variable changes. +/// +/// Entry: (Old Value, New Value) +fn env_changes(env: &CommandEnv) -> Option, Option)>> { + if env.is_unchanged() { + return None; + } + + let mut result = BTreeMap::, Option)>::new(); + + // Check if we want to clear all prior variables + if env.does_clear() { + for (k, v) in crate::env::vars_os() { + result.insert(k.into(), (Some(v), None)); + } + } + + for (k, v) in env.iter() { + let v: Option = v.map(Into::into); + result + .entry(k.into()) + .and_modify(|cur| *cur = (cur.0.clone(), v.clone())) + .or_insert((crate::env::var_os(k), v)); + } + + Some(result) +}