uefi: Implementing var(), vars(), and set_var() wrapper functions

RenTrieu · phip1611 · RenTrieu · commit 75f22a06ee74 · 2025-10-24T12:13:14.000-07:00
The var() and vars() functions wrap the UEFI get_env() implementation to
retrieve a single and multiple environment variables respectively.
The set_var() function wraps the UEFI set_env() variable to set the value
of an environment variable.

Co-authored-by: Philipp Schuster &lt;phip1611@gmail.com&gt;
diff --git a/uefi-test-runner/src/proto/shell.rs b/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,43 +100,46 @@ pub fn test_current_dir(shell: &ScopedProtocol<Shell>) {
     assert_eq!(cur_fs_str, expected_fs_str);
 }
 
-/// Test ``get_env()`` 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");
-    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()
-    );
+/// Test `var()`, `vars()`, and `set_var()`
+pub fn test_var(shell: &ScopedProtocol<Shell>) {
+    /* Test retrieving list of environment variable names */
+    let mut cur_env_vec = shell.vars();
+    assert_eq!(cur_env_vec.next().unwrap().0, cstr16!("path"));
+    // check pre-defined shell variables; see UEFI Shell spec
+    assert_eq!(cur_env_vec.next().unwrap().0, cstr16!("nonesting"));
+    let cur_env_vec = shell.vars();
+    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();
-    assert!(shell.get_env(Some(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 test_var = cstr16!("test_var");
+    let test_val = cstr16!("test_val");
+
+    let found_var = shell.vars().any(|(env_var, _)| env_var == test_var);
+    assert!(!found_var);
+    assert!(shell.var(test_var).is_none());
+
+    let status = shell.set_var(test_var, test_val, false);
+    assert!(status.is_ok());
+    let cur_env_str = shell
+        .var(test_var)
+        .expect("Could not get environment variable");
     assert_eq!(cur_env_str, test_val);
 
+    let found_var = shell.vars().any(|(env_var, _)| env_var == test_var);
+    assert!(found_var);
+    let cur_env_vec = shell.vars();
+    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 status = shell.set_env(test_var, test_val, false);
-    assert_eq!(status, Status::SUCCESS);
-    assert!(shell.get_env(Some(test_var)).is_none());
+    let test_val = cstr16!("");
+    let status = shell.set_var(test_var, test_val, false);
+    assert!(status.is_ok());
+    assert!(shell.var(test_var).is_none());
+
+    let found_var = shell.vars().any(|(env_var, _)| env_var == test_var);
+    assert!(!found_var);
+    let cur_env_vec = shell.vars();
+    assert_eq!(cur_env_vec.count(), default_len);
 }
 
 pub fn test() {
@@ -148,5 +151,5 @@ pub fn test() {
         boot::open_protocol_exclusive::<Shell>(handle).expect("Failed to open Shell protocol");
 
     test_current_dir(&shell);
-    test_env(&shell);
+    test_var(&shell);
 }
diff --git a/uefi/src/proto/shell/mod.rs b/uefi/src/proto/shell/mod.rs
@@ -4,16 +4,56 @@
 
 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);
 
+/// Trait for implementing the var function
+pub trait ShellVarProvider {
+    /// Gets the value of the specified environment variable
+    fn var(&self, name: &CStr16) -> Option<&CStr16>;
+}
+
+/// Iterator over the names of environmental variables obtained from the Shell protocol.
+#[derive(Debug)]
+pub struct Vars<'a, T: ShellVarProvider> {
+    /// Char16 containing names of environment variables
+    names: *const Char16,
+    /// Reference to Shell Protocol
+    protocol: *const T,
+    /// Marker to attach a lifetime to `Vars`
+    _marker: PhantomData<&'a CStr16>,
+}
+
+impl<'a, T: ShellVarProvider + 'a> Iterator for Vars<'a, T> {
+    type Item = (&'a CStr16, Option<&'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 s = unsafe { CStr16::from_ptr(self.names) };
+        if s.is_empty() {
+            None
+        } else {
+            self.names = unsafe { self.names.add(s.num_chars() + 1) };
+            Some((s, unsafe { self.protocol.as_ref().unwrap().var(s) }))
+        }
+    }
+}
+
+impl ShellVarProvider for Shell {
+    /// Gets the value of the specified environment variable
+    fn var(&self, name: &CStr16) -> Option<&CStr16> {
+        self.var(name)
+    }
+}
+
 impl Shell {
     /// Returns the current directory on the specified device.
     ///
@@ -56,64 +96,42 @@ impl Shell {
         unsafe { (self.0.set_cur_dir)(fs_ptr.cast(), dir_ptr.cast()) }.to_result()
     }
 
-    /// 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()) };
-
-                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;
-                    }
-                }
-            }
+    pub fn var(&self, name: &CStr16) -> Option<&CStr16> {
+        let name_ptr: *const Char16 = name.as_ptr();
+        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())) }
+        }
+    }
+
+    /// Gets an iterator over the names of all environment variables
+    ///
+    /// # Returns
+    ///
+    /// * `Vars` - Iterator over the names of the environment variables
+    #[must_use]
+    pub fn vars(&self) -> Vars<'_, Self> {
+        let env_ptr = unsafe { (self.0.get_env)(ptr::null()) };
+        Vars {
+            names: env_ptr.cast::<Char16>(),
+            protocol: self,
+            _marker: PhantomData,
         }
-        Some(env_vec)
     }
 
     /// Sets the environment variable
@@ -122,15 +140,114 @@ 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
-    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) }
+    /// * `Status::SUCCESS` - The variable was successfully set
+    pub fn set_var(&self, name: &CStr16, value: &CStr16, volatile: bool) -> Result {
+        let name_ptr: *const Char16 = name.as_ptr();
+        let value_ptr: *const Char16 = value.as_ptr();
+        unsafe { (self.0.set_env)(name_ptr.cast(), value_ptr.cast(), volatile) }.to_result()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use alloc::collections::BTreeMap;
+    use alloc::vec::Vec;
+    use uefi::cstr16;
+
+    struct ShellMock<'a> {
+        inner: BTreeMap<&'a CStr16, &'a CStr16>,
+    }
+
+    impl<'a> ShellMock<'a> {
+        fn new(names: Vec<&'a CStr16>, values: Vec<&'a CStr16>) -> ShellMock<'a> {
+            let mut inner_map = BTreeMap::new();
+            for (name, val) in names.iter().zip(values.iter()) {
+                inner_map.insert(*name, *val);
+            }
+            ShellMock { inner: inner_map }
+        }
+    }
+    impl<'a> ShellVarProvider for ShellMock<'a> {
+        fn var(&self, name: &CStr16) -> Option<&CStr16> {
+            if let Some(val) = self.inner.get(name) {
+                Some(*val)
+            } else {
+                None
+            }
+        }
+    }
+
+    /// 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 {
+            names: vars_mock.as_ptr().cast(),
+            protocol: &ShellMock::new(Vec::new(), Vec::new()),
+            _marker: 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 {
+            names: vars_mock.as_ptr().cast(),
+            protocol: &ShellMock::new(Vec::from([cstr16!("foo")]), Vec::from([cstr16!("value")])),
+            _marker: PhantomData,
+        };
+        assert_eq!(
+            vars.collect::<Vec<_>>(),
+            Vec::from([(cstr16!("foo"), Some(cstr16!("value")))])
+        );
+
+        // 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 {
+            names: vars_mock.as_ptr().cast(),
+            protocol: &ShellMock::new(
+                Vec::from([cstr16!("foo1"), cstr16!("bar"), cstr16!("baz2")]),
+                Vec::from([cstr16!("value"), cstr16!("one"), cstr16!("two")]),
+            ),
+            _marker: PhantomData,
+        };
+        assert_eq!(
+            vars.collect::<Vec<_>>(),
+            Vec::from([
+                (cstr16!("foo1"), Some(cstr16!("value"))),
+                (cstr16!("bar"), Some(cstr16!("one"))),
+                (cstr16!("baz2"), Some(cstr16!("two")))
+            ])
+        );
     }
 }
