From 588bfb4d50980f71434dc34b524d44545dd2d721 Mon Sep 17 00:00:00 2001
From: Ayush Singh <ayush@beagleboard.org>
Date: Tue, 15 Oct 2024 15:18:57 +0530
Subject: [PATCH 1/2] std: uefi: Add basic Env variables

- Implement environment variable functions
- Using EFI Shell protocol.

Signed-off-by: Ayush Singh <ayush@beagleboard.org>
---
 library/std/src/sys/pal/uefi/helpers.rs |  23 ++++-
 library/std/src/sys/pal/uefi/os.rs      | 124 +++++++++++++++++++++---
 2 files changed, 132 insertions(+), 15 deletions(-)

diff --git a/library/std/src/sys/pal/uefi/helpers.rs b/library/std/src/sys/pal/uefi/helpers.rs
index 4ced7065c826d..abc8e69a285f3 100644
--- a/library/std/src/sys/pal/uefi/helpers.rs
+++ b/library/std/src/sys/pal/uefi/helpers.rs
@@ -10,7 +10,7 @@
 //! - More information about protocols can be found [here](https://edk2-docs.gitbook.io/edk-ii-uefi-driver-writer-s-guide/3_foundation/36_protocols_and_handles)
 
 use r_efi::efi::{self, Guid};
-use r_efi::protocols::{device_path, device_path_to_text};
+use r_efi::protocols::{device_path, device_path_to_text, shell};
 
 use crate::ffi::{OsStr, OsString};
 use crate::io::{self, const_io_error};
@@ -424,3 +424,24 @@ pub(crate) fn os_string_to_raw(s: &OsStr) -> Option<Box<[r_efi::efi::Char16]>> {
     let temp = s.encode_wide().chain(Some(0)).collect::<Box<[r_efi::efi::Char16]>>();
     if temp[..temp.len() - 1].contains(&0) { None } else { Some(temp) }
 }
+
+pub(crate) fn open_shell() -> Option<NonNull<shell::Protocol>> {
+    static LAST_VALID_HANDLE: AtomicPtr<crate::ffi::c_void> =
+        AtomicPtr::new(crate::ptr::null_mut());
+
+    if let Some(handle) = NonNull::new(LAST_VALID_HANDLE.load(Ordering::Acquire)) {
+        if let Ok(protocol) = open_protocol::<shell::Protocol>(handle, shell::PROTOCOL_GUID) {
+            return Some(protocol);
+        }
+    }
+
+    let handles = locate_handles(shell::PROTOCOL_GUID).ok()?;
+    for handle in handles {
+        if let Ok(protocol) = open_protocol::<shell::Protocol>(handle, shell::PROTOCOL_GUID) {
+            LAST_VALID_HANDLE.store(handle.as_ptr(), Ordering::Release);
+            return Some(protocol);
+        }
+    }
+
+    None
+}
diff --git a/library/std/src/sys/pal/uefi/os.rs b/library/std/src/sys/pal/uefi/os.rs
index 4eb7698b43aa9..8160dcf35a1cc 100644
--- a/library/std/src/sys/pal/uefi/os.rs
+++ b/library/std/src/sys/pal/uefi/os.rs
@@ -192,44 +192,58 @@ pub fn current_exe() -> io::Result<PathBuf> {
     helpers::device_path_to_text(protocol).map(PathBuf::from)
 }
 
-pub struct Env(!);
+pub struct EnvStrDebug<'a> {
+    iter: &'a [(OsString, OsString)],
+}
+
+impl fmt::Debug for EnvStrDebug<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut list = f.debug_list();
+        for (a, b) in self.iter {
+            list.entry(&(a.to_str().unwrap(), b.to_str().unwrap()));
+        }
+        list.finish()
+    }
+}
+
+pub struct Env(crate::vec::IntoIter<(OsString, OsString)>);
 
 impl Env {
     // FIXME(https://github.com/rust-lang/rust/issues/114583): Remove this when <OsStr as Debug>::fmt matches <str as Debug>::fmt.
     pub fn str_debug(&self) -> impl fmt::Debug + '_ {
-        let Self(inner) = self;
-        match *inner {}
+        EnvStrDebug { iter: self.0.as_slice() }
     }
 }
 
 impl Iterator for Env {
     type Item = (OsString, OsString);
+
     fn next(&mut self) -> Option<(OsString, OsString)> {
-        self.0
+        self.0.next()
     }
 }
 
 impl fmt::Debug for Env {
-    fn fmt(&self, _: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let Self(inner) = self;
-        match *inner {}
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
     }
 }
 
 pub fn env() -> Env {
-    panic!("not supported on this platform")
+    let env = uefi_env::get_all().expect("not supported on this platform");
+    Env(env.into_iter())
 }
 
-pub fn getenv(_: &OsStr) -> Option<OsString> {
-    None
+pub fn getenv(key: &OsStr) -> Option<OsString> {
+    uefi_env::get(key)
 }
 
-pub unsafe fn setenv(_: &OsStr, _: &OsStr) -> io::Result<()> {
-    Err(io::const_io_error!(io::ErrorKind::Unsupported, "cannot set env vars on this platform"))
+pub unsafe fn setenv(key: &OsStr, val: &OsStr) -> io::Result<()> {
+    uefi_env::set(key, val)
 }
 
-pub unsafe fn unsetenv(_: &OsStr) -> io::Result<()> {
-    Err(io::const_io_error!(io::ErrorKind::Unsupported, "cannot unset env vars on this platform"))
+pub unsafe fn unsetenv(key: &OsStr) -> io::Result<()> {
+    uefi_env::unset(key)
 }
 
 pub fn temp_dir() -> PathBuf {
@@ -294,3 +308,85 @@ mod uefi_shell {
         None
     }
 }
+
+mod uefi_env {
+    use crate::ffi::{OsStr, OsString};
+    use crate::io;
+    use crate::os::uefi::ffi::OsStringExt;
+    use crate::ptr::NonNull;
+    use crate::sys::{helpers, unsupported_err};
+
+    pub(crate) fn get(key: &OsStr) -> Option<OsString> {
+        let shell = helpers::open_shell()?;
+        let mut key_ptr = helpers::os_string_to_raw(key)?;
+        unsafe { get_raw(shell, key_ptr.as_mut_ptr()) }
+    }
+
+    pub(crate) fn set(key: &OsStr, val: &OsStr) -> io::Result<()> {
+        let mut key_ptr = helpers::os_string_to_raw(key)
+            .ok_or(io::const_io_error!(io::ErrorKind::InvalidInput, "Invalid Key"))?;
+        let mut val_ptr = helpers::os_string_to_raw(val)
+            .ok_or(io::const_io_error!(io::ErrorKind::InvalidInput, "Invalid Value"))?;
+        unsafe { set_raw(key_ptr.as_mut_ptr(), val_ptr.as_mut_ptr()) }
+    }
+
+    pub(crate) fn unset(key: &OsStr) -> io::Result<()> {
+        let mut key_ptr = helpers::os_string_to_raw(key)
+            .ok_or(io::const_io_error!(io::ErrorKind::InvalidInput, "Invalid Key"))?;
+        unsafe { set_raw(key_ptr.as_mut_ptr(), crate::ptr::null_mut()) }
+    }
+
+    pub(crate) fn get_all() -> io::Result<Vec<(OsString, OsString)>> {
+        let shell = helpers::open_shell().ok_or(unsupported_err())?;
+
+        let mut vars = Vec::new();
+        let val = unsafe { ((*shell.as_ptr()).get_env)(crate::ptr::null_mut()) };
+
+        if val.is_null() {
+            return Ok(vars);
+        }
+
+        let mut start = 0;
+
+        // UEFI Shell returns all keys seperated by NULL.
+        // End of string is denoted by two NULLs
+        for i in 0.. {
+            if unsafe { *val.add(i) } == 0 {
+                // Two NULL signal end of string
+                if i == start {
+                    break;
+                }
+
+                let key = OsString::from_wide(unsafe {
+                    crate::slice::from_raw_parts(val.add(start), i - start)
+                });
+                // SAFETY: val.add(start) is always NULL terminated
+                let val = unsafe { get_raw(shell, val.add(start)) }
+                    .ok_or(io::const_io_error!(io::ErrorKind::InvalidInput, "Invalid Value"))?;
+
+                vars.push((key, val));
+                start = i + 1;
+            }
+        }
+
+        Ok(vars)
+    }
+
+    unsafe fn get_raw(
+        shell: NonNull<r_efi::efi::protocols::shell::Protocol>,
+        key_ptr: *mut r_efi::efi::Char16,
+    ) -> Option<OsString> {
+        let val = unsafe { ((*shell.as_ptr()).get_env)(key_ptr) };
+        helpers::os_string_from_raw(val)
+    }
+
+    unsafe fn set_raw(
+        key_ptr: *mut r_efi::efi::Char16,
+        val_ptr: *mut r_efi::efi::Char16,
+    ) -> io::Result<()> {
+        let shell = helpers::open_shell().ok_or(unsupported_err())?;
+        let r =
+            unsafe { ((*shell.as_ptr()).set_env)(key_ptr, val_ptr, r_efi::efi::Boolean::FALSE) };
+        if r.is_error() { Err(io::Error::from_raw_os_error(r.as_usize())) } else { Ok(()) }
+    }
+}

From 753536aba86310d85e71bb49062d883938602417 Mon Sep 17 00:00:00 2001
From: Ayush Singh <ayush@beagleboard.org>
Date: Wed, 16 Oct 2024 09:10:54 +0530
Subject: [PATCH 2/2] std: uefi: Use common function for UEFI shell

- Since in almost all cases, there will only be 1 UEFI shell, share the
  shell handle between all functions that require it.

Signed-off-by: Ayush Singh <ayush@beagleboard.org>
---
 library/std/src/sys/pal/uefi/os.rs | 38 ++----------------------------
 1 file changed, 2 insertions(+), 36 deletions(-)

diff --git a/library/std/src/sys/pal/uefi/os.rs b/library/std/src/sys/pal/uefi/os.rs
index 8160dcf35a1cc..27395f7c3c0b3 100644
--- a/library/std/src/sys/pal/uefi/os.rs
+++ b/library/std/src/sys/pal/uefi/os.rs
@@ -125,7 +125,7 @@ pub fn error_string(errno: RawOsError) -> String {
 }
 
 pub fn getcwd() -> io::Result<PathBuf> {
-    match uefi_shell::open_shell() {
+    match helpers::open_shell() {
         Some(shell) => {
             // SAFETY: path_ptr is managed by UEFI shell and should not be deallocated
             let path_ptr = unsafe { ((*shell.as_ptr()).get_cur_dir)(crate::ptr::null_mut()) };
@@ -144,7 +144,7 @@ pub fn getcwd() -> io::Result<PathBuf> {
 }
 
 pub fn chdir(p: &path::Path) -> io::Result<()> {
-    let shell = uefi_shell::open_shell().ok_or(unsupported_err())?;
+    let shell = helpers::open_shell().ok_or(unsupported_err())?;
 
     let mut p = helpers::os_string_to_raw(p.as_os_str())
         .ok_or(io::const_io_error!(io::ErrorKind::InvalidData, "Invalid path"))?;
@@ -275,40 +275,6 @@ pub fn getpid() -> u32 {
     panic!("no pids on this platform")
 }
 
-mod uefi_shell {
-    use r_efi::protocols::shell;
-
-    use super::super::helpers;
-    use crate::ptr::NonNull;
-    use crate::sync::atomic::{AtomicPtr, Ordering};
-
-    pub fn open_shell() -> Option<NonNull<shell::Protocol>> {
-        static LAST_VALID_HANDLE: AtomicPtr<crate::ffi::c_void> =
-            AtomicPtr::new(crate::ptr::null_mut());
-
-        if let Some(handle) = NonNull::new(LAST_VALID_HANDLE.load(Ordering::Acquire)) {
-            if let Ok(protocol) = helpers::open_protocol::<shell::Protocol>(
-                handle,
-                r_efi::protocols::shell::PROTOCOL_GUID,
-            ) {
-                return Some(protocol);
-            }
-        }
-
-        let handles = helpers::locate_handles(shell::PROTOCOL_GUID).ok()?;
-        for handle in handles {
-            if let Ok(protocol) =
-                helpers::open_protocol::<shell::Protocol>(handle, shell::PROTOCOL_GUID)
-            {
-                LAST_VALID_HANDLE.store(handle.as_ptr(), Ordering::Release);
-                return Some(protocol);
-            }
-        }
-
-        None
-    }
-}
-
 mod uefi_env {
     use crate::ffi::{OsStr, OsString};
     use crate::io;

<!DOCTYPE html PUBLIC '-//W3C//DTD XHTML 1.0 Transitional//EN' 'http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd'>
<html xmlns='http://www.w3.org/1999/xhtml'>
<head>
<title>pFad - Phonifier reborn</title>
<meta http-equiv='Content-Type' content='text/html; charset=utf-8' />
</head>
<body>
<h1>Pfad - The Proxy pFad of &#169; 2024 Garber Painting. All rights reserved.</h1>


<!-- Disclaimer -->
<p>Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.</p>
<br>
<p>Alternative Proxies:</p><p><a href="http://clevelandohioweatherforecast.com/php-proxy/index.php?q=https://patch-diff.githubusercontent.com/raw/rust-lang/rust/pull/127462.patch" target="_blank">Alternative Proxy</a></p><p><a href="http://clevelandohioweatherforecast.com/pFad/index.php?u=https://patch-diff.githubusercontent.com/raw/rust-lang/rust/pull/127462.patch" target="_blank">pFad Proxy</a></p><p><a href="http://clevelandohioweatherforecast.com/pFad/v3index.php?u=https://patch-diff.githubusercontent.com/raw/rust-lang/rust/pull/127462.patch" target="_blank">pFad v3 Proxy</a></p><p><a href="http://clevelandohioweatherforecast.com/pFad/v4index.php?u=https://patch-diff.githubusercontent.com/raw/rust-lang/rust/pull/127462.patch" target="_blank">pFad v4 Proxy</a></p></body>
</html>