uefi: Modifying Vars to return environment variable values in addition to names

RenTrieu · RenTrieu · commit 6e70d4fcf19a · 2025-08-11T12:45:04.000-07:00
diff --git a/uefi-test-runner/src/proto/shell.rs b/uefi-test-runner/src/proto/shell.rs
@@ -5,12 +5,12 @@ use uefi::proto::shell::Shell;
 use uefi::{boot, cstr16};
 
 /// Test `var()`, `vars()`, and `set_var()`
-pub fn test_env(shell: &ScopedProtocol<Shell>) {
+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(), cstr16!("path"),);
+    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(), cstr16!("nonesting"),);
+    assert_eq!(cur_env_vec.next().unwrap().0, cstr16!("nonesting"));
     let cur_env_vec = shell.vars();
     let default_len = cur_env_vec.count();
 
@@ -27,15 +27,15 @@ pub fn test_env(shell: &ScopedProtocol<Shell>) {
     assert_eq!(cur_env_str, test_val);
 
     let mut found_var = false;
-    for env_var in cur_env_vec {
+    for (env_var, _) in cur_env_vec {
         if env_var == test_var {
             found_var = true;
         }
     }
     assert!(!found_var);
     let cur_env_vec = shell.vars();
     let mut found_var = false;
-    for env_var in cur_env_vec {
+    for (env_var, _) in cur_env_vec {
         if env_var == test_var {
             found_var = true;
         }
@@ -53,7 +53,7 @@ pub fn test_env(shell: &ScopedProtocol<Shell>) {
 
     let cur_env_vec = shell.vars();
     let mut found_var = false;
-    for env_var in cur_env_vec {
+    for (env_var, _) in cur_env_vec {
         if env_var == test_var {
             found_var = true;
         }
@@ -63,98 +63,6 @@ pub fn test_env(shell: &ScopedProtocol<Shell>) {
     assert_eq!(cur_env_vec.count(), default_len);
 }
 
-/// Test `current_dir()` and `set_current_dir()`
-pub fn test_cur_dir(shell: &ScopedProtocol<Shell>) {
-    /* Test setting and getting current file system and current directory */
-    let fs_var = cstr16!("fs0:");
-    let dir_var = cstr16!("/");
-    let status = shell.set_current_dir(Some(fs_var), Some(dir_var));
-    assert!(status.is_ok());
-
-    let cur_fs_str = shell
-        .current_dir(Some(fs_var))
-        .expect("Could not get the current file system mapping");
-    let expected_fs_str = cstr16!("FS0:\\");
-    assert_eq!(cur_fs_str, expected_fs_str);
-
-    // Changing current file system
-    let fs_var = cstr16!("fs1:");
-    let dir_var = cstr16!("/");
-    let status = shell.set_current_dir(Some(fs_var), Some(dir_var));
-    assert!(status.is_ok());
-
-    let cur_fs_str = shell
-        .current_dir(Some(fs_var))
-        .expect("Could not get the current file system mapping");
-    assert_ne!(cur_fs_str, expected_fs_str);
-    let expected_fs_str = cstr16!("FS1:\\");
-    assert_eq!(cur_fs_str, expected_fs_str);
-
-    // Changing current file system and current directory
-    let fs_var = cstr16!("fs0:");
-    let dir_var = cstr16!("efi/");
-    let status = shell.set_current_dir(Some(fs_var), Some(dir_var));
-    assert!(status.is_ok());
-
-    let cur_fs_str = shell
-        .current_dir(Some(fs_var))
-        .expect("Could not get the current file system mapping");
-    assert_ne!(cur_fs_str, expected_fs_str);
-    let expected_fs_str = cstr16!("FS0:\\efi");
-    assert_eq!(cur_fs_str, expected_fs_str);
-
-    /* Test current working directory cases */
-
-    // At this point, the current working file system has not been set
-    // So we expect a NULL output
-    assert!(shell.current_dir(None).is_none());
-
-    // Setting the current working file system and current working directory
-    let dir_var = cstr16!("fs0:/");
-    let status = shell.set_current_dir(None, Some(dir_var));
-    assert!(status.is_ok());
-    let cur_fs_str = shell
-        .current_dir(Some(fs_var))
-        .expect("Could not get the current file system mapping");
-    let expected_fs_str = cstr16!("FS0:");
-    assert_eq!(cur_fs_str, expected_fs_str);
-
-    let cur_fs_str = shell
-        .current_dir(None)
-        .expect("Could not get the current file system mapping");
-    assert_eq!(cur_fs_str, expected_fs_str);
-
-    // Changing current working directory
-    let dir_var = cstr16!("/efi");
-    let status = shell.set_current_dir(None, Some(dir_var));
-    assert!(status.is_ok());
-    let cur_fs_str = shell
-        .current_dir(Some(fs_var))
-        .expect("Could not get the current file system mapping");
-    let expected_fs_str = cstr16!("FS0:\\efi");
-    assert_eq!(cur_fs_str, expected_fs_str);
-    let cur_fs_str = shell
-        .current_dir(None)
-        .expect("Could not get the current file system mapping");
-    assert_eq!(cur_fs_str, expected_fs_str);
-
-    // Changing current directory in a non-current working file system
-    let fs_var = cstr16!("fs0:");
-    let dir_var = cstr16!("efi/tools");
-    let status = shell.set_current_dir(Some(fs_var), Some(dir_var));
-    assert!(status.is_ok());
-    let cur_fs_str = shell
-        .current_dir(None)
-        .expect("Could not get the current file system mapping");
-    assert_ne!(cur_fs_str, expected_fs_str);
-
-    let expected_fs_str = cstr16!("FS0:\\efi\\tools");
-    let cur_fs_str = shell
-        .current_dir(Some(fs_var))
-        .expect("Could not get the current file system mapping");
-    assert_eq!(cur_fs_str, expected_fs_str);
-}
-
 pub fn test() {
     info!("Running shell protocol tests");
 
@@ -163,6 +71,5 @@ pub fn test() {
     let shell =
         boot::open_protocol_exclusive::<Shell>(handle).expect("Failed to open Shell protocol");
 
-    test_env(&shell);
-    test_cur_dir(&shell);
+    test_var(&shell);
 }
diff --git a/uefi/src/proto/shell/mod.rs b/uefi/src/proto/shell/mod.rs
@@ -17,30 +17,45 @@ use crate::{CStr16, Char16, Result, StatusExt};
 #[unsafe_protocol(ShellProtocol::GUID)]
 pub struct Shell(ShellProtocol);
 
+/// Trait for implementing the var function
+pub trait ShellVar {
+    /// 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> {
+pub struct Vars<'a, T: ShellVar> {
     /// Char16 containing names of environment variables
-    inner: *const Char16,
+    names: *const Char16,
+    /// Reference to Shell Protocol
+    protocol: *const T,
     /// Placeholder to attach a lifetime to `Vars`
     placeholder: PhantomData<&'a CStr16>,
 }
 
-impl<'a> Iterator for Vars<'a> {
-    type Item = &'a CStr16;
+impl<'a, T: ShellVar + '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.inner) };
+        let s = unsafe { CStr16::from_ptr(self.names) };
         if s.is_empty() {
             None
         } else {
-            self.inner = unsafe { self.inner.add(s.num_chars() + 1) };
-            Some(s)
+            self.names = unsafe { self.names.add(s.num_chars() + 1) };
+            Some((s, unsafe { self.protocol.as_ref().unwrap().var(s) }))
         }
     }
 }
 
+impl ShellVar for Shell {
+    /// Gets the value of the specified environment variable
+    fn var(&self, name: &CStr16) -> Option<&CStr16> {
+        self.var(name)
+    }
+}
+
 impl Shell {
     /// Gets the value of the specified environment variable
     ///
@@ -67,10 +82,11 @@ impl Shell {
 
     /// Gets an iterator over the names of all environment variables
     #[must_use]
-    pub fn vars(&self) -> Vars<'_> {
+    pub fn vars(&self) -> Vars<'_, Self> {
         let env_ptr = unsafe { (self.0.get_env)(ptr::null()) };
         Vars {
-            inner: env_ptr.cast::<Char16>(),
+            names: env_ptr.cast::<Char16>(),
+            protocol: self,
             placeholder: PhantomData,
         }
     }
@@ -92,61 +108,38 @@ impl Shell {
         let value_ptr: *const Char16 = value.as_ptr();
         unsafe { (self.0.set_env)(name_ptr.cast(), value_ptr.cast(), volatile) }.to_result()
     }
-
-    /// Returns the current directory on the specified device
-    ///
-    /// # Arguments
-    ///
-    /// * `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 current_dir(&self, file_system_mapping: Option<&CStr16>) -> Option<&CStr16> {
-        let mapping_ptr: *const Char16 = file_system_mapping.map_or(ptr::null(), CStr16::as_ptr);
-        let cur_dir = unsafe { (self.0.get_cur_dir)(mapping_ptr.cast()) };
-        if cur_dir.is_null() {
-            None
-        } else {
-            unsafe { Some(CStr16::from_ptr(cur_dir.cast())) }
-        }
-    }
-
-    /// Changes the current directory on the specified device
-    ///
-    /// # Arguments
-    ///
-    /// * `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
-    ///
-    /// # Errors
-    ///
-    /// * `Status::EFI_NOT_FOUND` - The directory does not exist
-    pub fn set_current_dir(
-        &self,
-        file_system: Option<&CStr16>,
-        directory: Option<&CStr16>,
-    ) -> Result {
-        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()));
-        unsafe { (self.0.set_cur_dir)(fs_ptr.cast(), dir_ptr.cast()) }.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> ShellVar 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() {
@@ -155,9 +148,11 @@ mod tests {
         vars_mock.push(0);
         vars_mock.push(0);
         let mut vars = Vars {
-            inner: vars_mock.as_ptr().cast(),
+            names: vars_mock.as_ptr().cast(),
+            protocol: &ShellMock::new(Vec::new(), Vec::new()),
             placeholder: PhantomData,
         };
+
         assert!(vars.next().is_none());
 
         // One environment variable in Vars
@@ -168,10 +163,14 @@ mod tests {
         vars_mock.push(0);
         vars_mock.push(0);
         let vars = Vars {
-            inner: vars_mock.as_ptr().cast(),
+            names: vars_mock.as_ptr().cast(),
+            protocol: &ShellMock::new(Vec::from([cstr16!("foo")]), Vec::from([cstr16!("value")])),
             placeholder: PhantomData,
         };
-        assert_eq!(vars.collect::<Vec<_>>(), Vec::from([cstr16!("foo")]));
+        assert_eq!(
+            vars.collect::<Vec<_>>(),
+            Vec::from([(cstr16!("foo"), Some(cstr16!("value")))])
+        );
 
         // Multiple environment variables in Vars
         let mut vars_mock = Vec::<u16>::new();
@@ -192,12 +191,20 @@ mod tests {
         vars_mock.push(0);
 
         let vars = Vars {
-            inner: vars_mock.as_ptr().cast(),
+            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")]),
+            ),
             placeholder: PhantomData,
         };
         assert_eq!(
             vars.collect::<Vec<_>>(),
-            Vec::from([cstr16!("foo1"), cstr16!("bar"), cstr16!("baz2")])
+            Vec::from([
+                (cstr16!("foo1"), Some(cstr16!("value"))),
+                (cstr16!("bar"), Some(cstr16!("one"))),
+                (cstr16!("baz2"), Some(cstr16!("two")))
+            ])
         );
     }
 }
