Rollup merge of #136418 - Ayush1325:command-env, r=jhpratt

uefi: process: Add support for command environment variables

Set environment variables before launching the process and restore the prior variables after the program exists.

This is the same implementation as the one used by UEFI Shell Execute [0].

[0]: https://github.com/tianocore/edk2/blob/2d2642f4832ebc45cb7d5ba9430b933d953b94f2/ShellPkg/Application/Shell/ShellProtocol.c#L1700

1 files changed, 62 insertions, 3 deletions

diff --git a/library/std/src/sys/pal/uefi/process.rs b/library/std/src/sys/pal/uefi/process.rs
index 3077a72eac6..0757f1cb490 100644
--- a/library/std/src/sys/pal/uefi/process.rs
+++ b/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,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<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)
+}