uefi: process: Add support for command environment variables

library/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<OsString>,
     stdout: Option<Stdio>,
     stderr: Option<Stdio>,
+    env: CommandEnv,
 }
 
 // passed back to std::process with the pipes connected to the child, if any
@@ -40,15 +42,21 @@ 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) {
         self.args.push(arg.to_os_string());
     }
 
     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<BTreeMap<EnvKey, (Option<OsString>, Option<OsString>)>> {
+    if env.is_unchanged() {
+        return None;
+    }
+
+    let mut result = BTreeMap::<EnvKey, (Option<OsString>, Option<OsString>)>::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<OsString> = 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)
+}