Added initial implementation of get_env()

Author: RenTrieu

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

@@ -1,39 +1,42 @@
 use uefi::boot;
 use uefi::proto::shell::{Shell, ShellFileHandle};
 use uefi::CStr16;
+use uefi::data_types::Char16;
 
 pub fn test() {
     info!("Running shell protocol tests");
 
     let handle = boot::get_handle_for_protocol::<Shell>().expect("No Shell handles");
-    // let mut fs = uefi::fs::FileSystem::new(sfs
-    // let (fs_handle, mut sfs) = find_test_disk();
 
     let mut shell =
         boot::open_protocol_exclusive::<Shell>(handle).expect("Failed to open Shell protocol");
 
     // create some files
-    let mut test_buf = [0u16; 12];
-    let test_str = CStr16::from_str_with_buf("test", &mut test_buf).unwrap();
+    // let mut test_buf = [0u16; 12];
+    // let test_str = CStr16::from_str_with_buf("test", &mut test_buf).unwrap();
 
-    let cur_env_str = shell.get_env(None).expect("Could not get environment variable");
-    info!("cur_env_str size: {}", cur_env_str.num_chars());
-    info!("cur_env_str: {}", cur_env_str);
+    /* Test retrieving list of environment variable names (null input) */
+    let cur_env_vec = shell.get_env(None).expect("Could not get environment variable").vec().unwrap();
+    let mut test_buf = [0u16; 64];
+    assert_eq!(*cur_env_vec.get(0).unwrap(), CStr16::from_str_with_buf("path", &mut test_buf).unwrap());
+    assert_eq!(*cur_env_vec.get(1).unwrap(), CStr16::from_str_with_buf("nonesting", &mut test_buf).unwrap());
 
-    /* Testing setting and getting a specific environment variable */
+    // Debug statements TODO: Remove
+    info!("cur_env_vec size: {}", cur_env_vec.len());
+    for (i, env_var) in cur_env_vec.iter().enumerate() {
+        info!("i: {}, env_var: {}", i, env_var);
+    }
+
+    /* 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());
     shell.set_env(test_var, test_val, false);
-    let cur_env_str = shell.get_env(Some(test_var)).expect("Could not get environment variable");
+    let cur_env_str = shell.get_env(Some(test_var)).expect("Could not get environment variable").val().unwrap();
     assert_eq!(cur_env_str, test_val);
 
-    // for (i, c) in cur_env_str.iter().enumerate() {
-    //     info!("cur_env_str: i: {}, c: {}", i, c);
-    // }
-
     // let mut cur_fs_buf = [0u16; 32];
     // let cur_fs_str = CStr16::from_str_with_buf("", &mut cur_fs_buf).unwrap();
     // info!("cur_fs_str size 1: {}", cur_fs_str.num_chars());

uefi/src/proto/shell/mod.rs

@@ -1,5 +1,9 @@
 //! EFI Shell Protocol v2.2
 
+#![cfg(feature = "alloc")]
+
+use alloc::vec::Vec;
+
 use core::{ffi::c_void, marker::PhantomData, mem::MaybeUninit, ptr};
 
 use uefi_macros::unsafe_protocol;
@@ -96,6 +100,36 @@ impl core::fmt::Debug for Shell {
     }
 }
 
+/// Enum describing output options for get_env()
+///
+/// `EnvOutput::Val` - Value for a given variable
+/// `EnvOutput::VarVec` - Vec of current environment variables
+#[derive(Debug)]
+pub enum EnvOutput<'a> {
+    /// Value for a given variable
+    Val(&'a CStr16),
+    /// Vec of current environment variable names
+    Vec(Vec<&'a CStr16>),
+}
+
+impl<'a> EnvOutput<'a> {
+    /// Extracts the env var value from EnvOutput
+    pub fn val(self) -> Option<&'a CStr16> {
+        match self {
+            EnvOutput::Val(v) => Some(v),
+            _ => None,
+        }
+    }
+
+    /// Extracts the vector of variable names from EnvOutput
+    pub fn vec(self) -> Option<Vec<&'a CStr16>> {
+        match self {
+            EnvOutput::Vec(v) => Some(v),
+            _ => None,
+        }
+    }
+}
+
 impl Shell {
     /// TODO
     pub fn execute(
@@ -131,13 +165,49 @@ impl Shell {
     ///
     /// * `Some(env_value)` - Value of the environment variable
     /// * `None` - Environment variable doesn't exist
-    pub fn get_env<'a>(&'a self, name: Option<&CStr16>) -> Option<&'a CStr16> {
-        let name_ptr: *const Char16 = name.map_or(core::ptr::null(), |x| (x as *const CStr16).cast());
-        let var_val = (self.get_env)(name_ptr);
-        if var_val.is_null() {
-            None
-        } else {
-            unsafe { Some(CStr16::from_ptr(var_val)) }
+    pub fn get_env<'a>(&'a self, name: Option<&CStr16>) -> Option<EnvOutput<'a>> {
+        match name {
+            Some(n) => {
+                let name_ptr: *const Char16 = (n as *const CStr16).cast();
+                let var_val = (self.get_env)(name_ptr);
+                if var_val.is_null() {
+                    None
+                } else {
+                    unsafe { Some(EnvOutput::Val(CStr16::from_ptr(var_val))) }
+                }
+            }
+            None => {
+                let mut env_vec = Vec::new();
+                let cur_env_ptr = (self.get_env)(core::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) {
+                            if cur_len > 0 {
+                                // TODO: Optimize this to directly convert a
+                                // Char16 slice to CStr16
+                                env_vec.push(CStr16::from_ptr(cur_start));
+                            }
+                            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(EnvOutput::Vec(env_vec))
+            }
         }
     }