uefi: Changing get_envs() to return an Iterator

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

Author: RenTrieu

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

@@ -2,7 +2,7 @@
 
 use uefi::boot::ScopedProtocol;
 use uefi::proto::shell::Shell;
-use uefi::{Error, Status, boot, cstr16, CStr16};
+use uefi::{Error, Status, boot, cstr16};
 
 /// Test `current_dir()` and `set_current_dir()`
 pub fn test_current_dir(shell: &ScopedProtocol<Shell>) {
@@ -100,27 +100,20 @@ pub fn test_current_dir(shell: &ScopedProtocol<Shell>) {
     assert_eq!(cur_fs_str, expected_fs_str);
 }
 
-/// Test ``get_env()``, ``get_envs()``, 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 */
+    let mut cur_env_vec = shell.get_envs();
+    assert_eq!(cur_env_vec.next().unwrap(), cstr16!("path"),);
+    // check pre-defined shell variables; see UEFI Shell spec
+    assert_eq!(cur_env_vec.next().unwrap(), cstr16!("nonesting"),);
     let cur_env_vec = shell.get_envs();
-    assert_eq!(
-        *cur_env_vec.first().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()
-    );
-    let default_len = cur_env_vec.len();
+    let default_len = cur_env_vec.count();
 
     /* 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();
+    let cur_env_vec = shell.get_envs();
+    let test_var = cstr16!("test_var");
+    let test_val = cstr16!("test_val");
     assert!(shell.get_env(test_var).is_none());
     let status = shell.set_env(test_var, test_val, false);
     assert_eq!(status, Status::SUCCESS);
@@ -129,20 +122,41 @@ pub fn test_env(shell: &ScopedProtocol<Shell>) {
         .expect("Could not get environment variable");
     assert_eq!(cur_env_str, test_val);
 
-    assert!(!cur_env_vec.contains(&test_var));
+    let mut found_var = false;
+    for env_var in cur_env_vec {
+        if env_var == test_var {
+            found_var = true;
+        }
+    }
+    assert!(!found_var);
+    let cur_env_vec = shell.get_envs();
+    let mut found_var = false;
+    for env_var in cur_env_vec {
+        if env_var == test_var {
+            found_var = true;
+        }
+    }
+    assert!(found_var);
+
     let cur_env_vec = shell.get_envs();
-    assert!(cur_env_vec.contains(&test_var));
-    assert_eq!(cur_env_vec.len(), default_len + 1);
+    assert_eq!(cur_env_vec.count(), default_len + 1);
 
     /* Test deleting environment variable */
-    let test_val = CStr16::from_str_with_buf("", &mut test_val_buf).unwrap();
+    let test_val = cstr16!("");
     let status = shell.set_env(test_var, test_val, false);
     assert_eq!(status, Status::SUCCESS);
     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);
+    let mut found_var = false;
+    for env_var in cur_env_vec {
+        if env_var == test_var {
+            found_var = true;
+        }
+    }
+    assert!(!found_var);
+    let cur_env_vec = shell.get_envs();
+    assert_eq!(cur_env_vec.count(), default_len);
 }
 
 pub fn test() {

uefi/src/proto/shell/mod.rs

@@ -4,16 +4,45 @@
 
 use crate::proto::unsafe_protocol;
 use crate::{CStr16, Char16, Error, Result, Status, StatusExt};
+use core::marker::PhantomData;
 use core::ptr;
 use uefi_raw::protocol::shell::ShellProtocol;
-use alloc::vec::Vec;
 
 /// Shell Protocol
 #[derive(Debug)]
 #[repr(transparent)]
 #[unsafe_protocol(ShellProtocol::GUID)]
 pub struct Shell(ShellProtocol);
 
+/// Iterator over the names of environmental variables obtained from the Shell protocol.
+#[derive(Debug)]
+pub struct Vars<'a> {
+    /// Char16 containing names of environment variables
+    inner: *const Char16,
+    /// Placeholder to attach a lifetime to `Vars`
+    placeholder: PhantomData<&'a CStr16>,
+}
+
+impl<'a> Iterator for Vars<'a> {
+    type Item = &'a CStr16;
+    // We iterate a list of NUL terminated CStr16s.
+    // The list is terminated with a double NUL.
+    fn next(&mut self) -> Option<Self::Item> {
+        let cur_start = self.inner;
+        let mut cur_len = 0;
+        unsafe {
+            if *(cur_start) == Char16::from_u16_unchecked(0) {
+                return None;
+            }
+            while *(cur_start.add(cur_len)) != Char16::from_u16_unchecked(0) {
+                cur_len += 1;
+            }
+            self.inner = self.inner.add(cur_len + 1);
+            Some(CStr16::from_ptr(cur_start))
+        }
+    }
+}
+
 impl Shell {
     /// Returns the current directory on the specified device.
     ///
@@ -85,36 +114,12 @@ impl Shell {
     ///
     /// * `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(
-                            &(*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;
-            }
+    pub fn get_envs(&self) -> Vars {
+        let env_ptr = unsafe { (self.0.get_env)(ptr::null()) };
+        Vars {
+            inner: env_ptr.cast::<Char16>(),
+            placeholder: PhantomData,
         }
-        env_vec
     }
 
     /// Sets the environment variable
@@ -135,3 +140,64 @@ impl Shell {
         unsafe { (self.0.set_env)(name_ptr.cast(), value_ptr.cast(), volatile) }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use alloc::vec::Vec;
+    use uefi::cstr16;
+
+    /// Testing Vars struct
+    #[test]
+    fn test_vars() {
+        // Empty Vars
+        let mut vars_mock = Vec::<u16>::new();
+        vars_mock.push(0);
+        vars_mock.push(0);
+        let mut vars = Vars {
+            inner: vars_mock.as_ptr().cast(),
+            placeholder: PhantomData,
+        };
+        assert!(vars.next().is_none());
+
+        // One environment variable in Vars
+        let mut vars_mock = Vec::<u16>::new();
+        vars_mock.push(b'f' as u16);
+        vars_mock.push(b'o' as u16);
+        vars_mock.push(b'o' as u16);
+        vars_mock.push(0);
+        vars_mock.push(0);
+        let vars = Vars {
+            inner: vars_mock.as_ptr().cast(),
+            placeholder: PhantomData,
+        };
+        assert_eq!(vars.collect::<Vec<_>>(), Vec::from([cstr16!("foo")]));
+
+        // Multiple environment variables in Vars
+        let mut vars_mock = Vec::<u16>::new();
+        vars_mock.push(b'f' as u16);
+        vars_mock.push(b'o' as u16);
+        vars_mock.push(b'o' as u16);
+        vars_mock.push(b'1' as u16);
+        vars_mock.push(0);
+        vars_mock.push(b'b' as u16);
+        vars_mock.push(b'a' as u16);
+        vars_mock.push(b'r' as u16);
+        vars_mock.push(0);
+        vars_mock.push(b'b' as u16);
+        vars_mock.push(b'a' as u16);
+        vars_mock.push(b'z' as u16);
+        vars_mock.push(b'2' as u16);
+        vars_mock.push(0);
+        vars_mock.push(0);
+
+        let vars = Vars {
+            inner: vars_mock.as_ptr().cast(),
+            placeholder: PhantomData,
+        };
+        assert_eq!(
+            vars.collect::<Vec<_>>(),
+            Vec::from([cstr16!("foo1"), cstr16!("bar"), cstr16!("baz2")])
+        );
+    }
+}