std: uefi: Add basic Env variables

- Implement environment variable functions
- Using EFI Shell protocol.

Signed-off-by: Ayush Singh <ayush@beagleboard.org>

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 4ced7065c82..abc8e69a285 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 4eb7698b43a..8160dcf35a1 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(()) }
+    }
+}