uefi: Implementing wrappers for EFI Shell env functions

This commit implements wrappers for the following EFI Shell Protocol
functions: set_env() and get_env()

Author: RenTrieu

uefi/src/proto/mod.rs

@@ -51,6 +51,7 @@ pub mod rng;
 #[cfg(feature = "alloc")]
 pub mod scsi;
 pub mod security;
+#[cfg(feature = "alloc")]
 pub mod shell;
 pub mod shell_params;
 pub mod shim;

uefi/src/proto/shell/mod.rs

@@ -6,6 +6,7 @@ use crate::proto::unsafe_protocol;
 use crate::{CStr16, Char16, Error, Result, Status, StatusExt};
 use core::ptr;
 use uefi_raw::protocol::shell::ShellProtocol;
+use alloc::vec::Vec;
 
 /// Shell Protocol
 #[derive(Debug)]
@@ -54,4 +55,82 @@ impl Shell {
         let dir_ptr: *const Char16 = directory.map_or(ptr::null(), |x| x.as_ptr());
         unsafe { (self.0.set_cur_dir)(fs_ptr.cast(), dir_ptr.cast()) }.to_result()
     }
+
+    /// Gets the environment variable or list of environment variables
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - The environment variable name of which to retrieve the
+    ///   value
+    ///   If None, will return all defined shell environment
+    ///   variables
+    ///
+    /// # Returns
+    ///
+    /// * `Some(Vec<env_value>)` - Value of the environment variable
+    /// * `Some(Vec<env_names>)` - Vector of environment variable names
+    /// * `None` - Environment variable doesn't exist
+    #[must_use]
+    pub fn get_env<'a>(&'a self, name: Option<&CStr16>) -> Option<Vec<&'a CStr16>> {
+        let mut env_vec = Vec::new();
+        match name {
+            Some(n) => {
+                let name_ptr: *const Char16 = core::ptr::from_ref::<CStr16>(n).cast();
+                let var_val = unsafe { (self.0.get_env)(name_ptr.cast()) };
+                if var_val.is_null() {
+                    return None;
+                } else {
+                    unsafe { env_vec.push(CStr16::from_ptr(var_val.cast())) };
+                }
+            }
+            None => {
+                let cur_env_ptr = unsafe { (self.0.get_env)(ptr::null()) };
+
+                let mut cur_start = cur_env_ptr;
+                let mut cur_len = 0;
+
+                let mut i = 0;
+                let mut null_count = 0;
+                unsafe {
+                    while null_count <= 1 {
+                        if (*(cur_env_ptr.add(i))) == Char16::from_u16_unchecked(0).into() {
+                            if cur_len > 0 {
+                                env_vec.push(CStr16::from_char16_with_nul_unchecked(
+                                    &(*ptr::slice_from_raw_parts(cur_start.cast(), cur_len + 1)),
+                                ));
+                            }
+                            cur_len = 0;
+                            null_count += 1;
+                        } else {
+                            if null_count > 0 {
+                                cur_start = cur_env_ptr.add(i);
+                            }
+                            null_count = 0;
+                            cur_len += 1;
+                        }
+                        i += 1;
+                    }
+                }
+            }
+        }
+        Some(env_vec)
+    }
+
+    /// Sets the environment variable
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - The environment variable for which to set the value
+    /// * `value` - The new value of the environment variable
+    /// * `volatile` - Indicates whether or not the variable is volatile or
+    ///   not
+    ///
+    /// # Returns
+    ///
+    /// * `Status::SUCCESS` The variable was successfully set
+    pub fn set_env(&self, name: &CStr16, value: &CStr16, volatile: bool) -> Status {
+        let name_ptr: *const Char16 = core::ptr::from_ref::<CStr16>(name).cast();
+        let value_ptr: *const Char16 = core::ptr::from_ref::<CStr16>(value).cast();
+        unsafe { (self.0.set_env)(name_ptr.cast(), value_ptr.cast(), volatile) }
+    }
 }