uefi: Splitting get_env() into two separate functions

The UEFI get_env() implementation is used for getting individual environment
variable values and also environment variable names. This is not intuitive
so this commit splits the function into two dedicated ones: one for accessing
values and the other for accessing names.

Co-authored-by: Philipp Schuster <phip1611@gmail.com>

Author: RenTrieu

uefi-test-runner/src/proto/shell.rs

@@ -5,14 +5,12 @@ use uefi::proto::shell::Shell;
 use uefi::{CStr16, boot};
 use uefi_raw::Status;
 
-/// Test ``get_env()`` and ``set_env()``
+/// Test ``get_env()``, ``get_envs()``, and ``set_env()``
 pub fn test_env(shell: &ScopedProtocol<Shell>) {
     let mut test_buf = [0u16; 128];
 
-    /* Test retrieving list of environment variable names (null input) */
-    let cur_env_vec = shell
-        .get_env(None)
-        .expect("Could not get environment variable");
+    /* Test retrieving list of environment variable names */
+    let cur_env_vec = shell.get_envs();
     assert_eq!(
         *cur_env_vec.first().unwrap(),
         CStr16::from_str_with_buf("path", &mut test_buf).unwrap()
@@ -21,27 +19,35 @@ pub fn test_env(shell: &ScopedProtocol<Shell>) {
         *cur_env_vec.get(1).unwrap(),
         CStr16::from_str_with_buf("nonesting", &mut test_buf).unwrap()
     );
+    let default_len = cur_env_vec.len();
 
     /* Test setting and getting a specific environment variable */
     let mut test_env_buf = [0u16; 32];
     let test_var = CStr16::from_str_with_buf("test_var", &mut test_env_buf).unwrap();
     let mut test_val_buf = [0u16; 32];
     let test_val = CStr16::from_str_with_buf("test_val", &mut test_val_buf).unwrap();
-    assert!(shell.get_env(Some(test_var)).is_none());
+    assert!(shell.get_env(test_var).is_none());
     let status = shell.set_env(test_var, test_val, false);
     assert_eq!(status, Status::SUCCESS);
-    let cur_env_str = *shell
-        .get_env(Some(test_var))
-        .expect("Could not get environment variable")
-        .first()
-        .unwrap();
+    let cur_env_str = shell
+        .get_env(test_var)
+        .expect("Could not get environment variable");
     assert_eq!(cur_env_str, test_val);
 
+    assert!(!cur_env_vec.contains(&test_var));
+    let cur_env_vec = shell.get_envs();
+    assert!(cur_env_vec.contains(&test_var));
+    assert_eq!(cur_env_vec.len(), default_len + 1);
+
     /* Test deleting environment variable */
     let test_val = CStr16::from_str_with_buf("", &mut test_val_buf).unwrap();
     let status = shell.set_env(test_var, test_val, false);
     assert_eq!(status, Status::SUCCESS);
-    assert!(shell.get_env(Some(test_var)).is_none());
+    assert!(shell.get_env(test_var).is_none());
+
+    let cur_env_vec = shell.get_envs();
+    assert!(!cur_env_vec.contains(&test_var));
+    assert_eq!(cur_env_vec.len(), default_len);
 }
 
 /// Test ``get_cur_dir()`` and ``set_cur_dir()``

uefi/src/proto/shell/mod.rs

@@ -10,75 +10,76 @@ use uefi_raw::Status;
 
 use core::ptr;
 
-pub use uefi_raw::protocol::shell::ShellProtocol;
+use uefi_raw::protocol::shell::ShellProtocol;
 
 use crate::{CStr16, Char16};
 
 /// Shell Protocol
 #[derive(Debug)]
 #[repr(transparent)]
-#[unsafe_protocol(uefi_raw::protocol::shell::ShellProtocol::GUID)]
-pub struct Shell(uefi_raw::protocol::shell::ShellProtocol);
+#[unsafe_protocol(ShellProtocol::GUID)]
+pub struct Shell(ShellProtocol);
 
 impl Shell {
-    /// Gets the environment variable or list of environment variables
+    /// Gets the value of the specified environment variable
     ///
     /// # Arguments
     ///
     /// * `name` - The environment variable name of which to retrieve the
-    ///   value
-    ///   If None, will return all defined shell environment
-    ///   variables
+    ///   value.
     ///
     /// # 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
+    /// * `Some(<env_value>)` - &CStr16 containing the value of the
+    ///   environment variable
+    /// * `None` - If environment variable does not 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()) };
+    pub fn get_env(&self, name: &CStr16) -> Option<&CStr16> {
+        let name_ptr: *const Char16 = core::ptr::from_ref::<CStr16>(name).cast();
+        let var_val = unsafe { (self.0.get_env)(name_ptr.cast()) };
+        if var_val.is_null() {
+            None
+        } else {
+            unsafe { Some(CStr16::from_ptr(var_val.cast())) }
+        }
+    }
 
-                let mut cur_start = cur_env_ptr;
-                let mut cur_len = 0;
+    /// Gets the list of environment variables
+    ///
+    /// # Returns
+    ///
+    /// * `Vec<env_names>` - Vector of environment variable names
+    #[must_use]
+    pub fn get_envs(&self) -> Vec<&CStr16> {
+        let mut env_vec: Vec<&CStr16> = Vec::new();
+        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;
+        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(
+                            &(*ptr::slice_from_raw_parts(cur_start.cast(), cur_len + 1)),
+                        ).unwrap());
                     }
+                    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)
+        env_vec
     }
 
     /// Sets the environment variable
@@ -87,12 +88,12 @@ impl Shell {
     ///
     /// * `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
+    /// * `volatile` - Indicates whether the variable is volatile or
     ///   not
     ///
     /// # Returns
     ///
-    /// * `Status::SUCCESS` The variable was successfully set
+    /// * `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();
@@ -105,12 +106,13 @@ impl Shell {
     ///
     /// * `file_system_mapping` - The file system mapping for which to get
     ///   the current directory
+    ///
     /// # Returns
     ///
     /// * `Some(cwd)` - CStr16 containing the current working directory
     /// * `None` - Could not retrieve current directory
     #[must_use]
-    pub fn get_cur_dir<'a>(&'a self, file_system_mapping: Option<&CStr16>) -> Option<&'a CStr16> {
+    pub fn get_cur_dir(&self, file_system_mapping: Option<&CStr16>) -> Option<&CStr16> {
         let mapping_ptr: *const Char16 = file_system_mapping.map_or(ptr::null(), |x| (x.as_ptr()));
         let cur_dir = unsafe { (self.0.get_cur_dir)(mapping_ptr.cast()) };
         if cur_dir.is_null() {
@@ -127,13 +129,14 @@ impl Shell {
     /// * `file_system` - Pointer to the file system's mapped name.
     /// * `directory` - Points to the directory on the device specified by
     ///   `file_system`.
+    ///
     /// # Returns
     ///
-    /// * `Status::SUCCESS` The directory was successfully set
+    /// * `Status::SUCCESS` - The directory was successfully set
     ///
     /// # Errors
     ///
-    /// * `Status::EFI_NOT_FOUND` The directory does not exist
+    /// * `Status::EFI_NOT_FOUND` - The directory does not exist
     pub fn set_cur_dir(&self, file_system: Option<&CStr16>, directory: Option<&CStr16>) -> Status {
         let fs_ptr: *const Char16 = file_system.map_or(ptr::null(), |x| (x.as_ptr()));
         let dir_ptr: *const Char16 = directory.map_or(ptr::null(), |x| (x.as_ptr()));