From 1ed804ddb9c9c018aeb386669f81e42a21721d90 Mon Sep 17 00:00:00 2001
From: Ayush Singh <ayush@beagleboard.org>
Date: Sun, 19 Jan 2025 22:23:35 +0530
Subject: [PATCH] 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

Signed-off-by: Ayush Singh <ayush@beagleboard.org>
---
 library/std/src/sys/pal/uefi/process.rs | 65 +++++++++++++++++++++++--
 1 file 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 3077a72eac661..0757f1cb490d4 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)
+}