From 81af1a3bb39e24504fdf378afc06d1c5a565836f Mon Sep 17 00:00:00 2001
From: Nicholas Bishop <nicholasbishop@google.com>
Date: Fri, 4 Oct 2024 18:19:51 -0400
Subject: [PATCH] uefi: Improve the VariableKey type and iterator

Prior to this commit, the `name` field of `VariableKey` was not public, so it
was not possible to construct `VariableKey` outside of this crate. This is a
problem for unit tests that want to mock `runtime::variable_keys`; there's no
way for the unit test to construct the iterator elements.

Fix by making `name` public. Also change the `name` type from a `Vec<u16>` to a
`CString16`; this makes the type easier to work with, since in all cases
variable names should be UCS-2. The `VariableKeys` iterator now yields an error
for variables with non-UCS-2 names (but such errors do not stop iteration; you
can simply continue on to the next variable key).

Also deprecate the `VariableKey::name()` method, since it just returns the same
thing as the `name` field now.
---
 uefi-test-runner/src/runtime/vars.rs |  2 +-
 uefi/CHANGELOG.md                    |  6 ++++++
 uefi/src/runtime.rs                  | 31 ++++++++++++----------------
 3 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/uefi-test-runner/src/runtime/vars.rs b/uefi-test-runner/src/runtime/vars.rs
index 930141c02..d46a8f978 100644
--- a/uefi-test-runner/src/runtime/vars.rs
+++ b/uefi-test-runner/src/runtime/vars.rs
@@ -48,7 +48,7 @@ fn test_variables() {
     let find_by_key = || {
         runtime::variable_keys().any(|k| {
             let k = k.as_ref().unwrap();
-            k.name().unwrap() == NAME && &k.vendor == VENDOR
+            k.name == NAME && &k.vendor == VENDOR
         })
     };
     assert!(find_by_key());
diff --git a/uefi/CHANGELOG.md b/uefi/CHANGELOG.md
index 9dafd9b3a..beb67633a 100644
--- a/uefi/CHANGELOG.md
+++ b/uefi/CHANGELOG.md
@@ -17,6 +17,12 @@ details of the deprecated items that were removed in this release.
   deprecated conversion from `uefi::table::boot::ScopedProtocol` has been
   removed.
 - Fixed `boot::open_protocol` to properly handle a null interface pointer.
+- `VariableKey` now has a public `name` field. This `name` field always contains
+  a valid string, so the `VariableKey::name()` method has been deprecated. Since
+  all fields of `VariableKey` are now public, the type can be constructed by
+  users.
+- The `VariableKeys` iterator will now yield an error item if a variable name is
+  not UCS-2.
 
 
 # uefi - 0.32.0 (2024-09-09)
diff --git a/uefi/src/runtime.rs b/uefi/src/runtime.rs
index 1d7d4545b..b8dfd9878 100644
--- a/uefi/src/runtime.rs
+++ b/uefi/src/runtime.rs
@@ -15,7 +15,8 @@ use uefi_raw::table::boot::MemoryDescriptor;
 
 #[cfg(feature = "alloc")]
 use {
-    crate::mem::make_boxed, crate::Guid, alloc::borrow::ToOwned, alloc::boxed::Box, alloc::vec::Vec,
+    crate::mem::make_boxed, crate::CString16, crate::Guid, alloc::borrow::ToOwned,
+    alloc::boxed::Box, alloc::vec::Vec,
 };
 
 #[cfg(all(feature = "unstable", feature = "alloc"))]
@@ -303,15 +304,13 @@ impl Iterator for VariableKeys {
 
         match result {
             Ok(()) => {
-                // Copy the name buffer, truncated after the first null
-                // character (if one is present).
-                let name = if let Some(nul_pos) = self.name.iter().position(|c| *c == 0) {
-                    self.name[..=nul_pos].to_owned()
-                } else {
-                    self.name.clone()
+                // Convert the name to a `CStr16`, yielding an error if invalid.
+                let Ok(name) = CStr16::from_u16_until_nul(&self.name) else {
+                    return Some(Err(Status::UNSUPPORTED.into()));
                 };
+
                 Some(Ok(VariableKey {
-                    name,
+                    name: name.to_owned(),
                     vendor: self.vendor,
                 }))
             }
@@ -868,30 +867,26 @@ impl TryFrom<&[u8]> for Time {
 #[cfg(feature = "alloc")]
 #[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
 pub struct VariableKey {
-    pub(crate) name: Vec<u16>,
     /// Unique identifier for the vendor.
     pub vendor: VariableVendor,
+
+    /// Name of the variable, unique with the vendor namespace.
+    pub name: CString16,
 }
 
 #[cfg(feature = "alloc")]
 impl VariableKey {
     /// Name of the variable.
+    #[deprecated = "Use the VariableKey.name field instead"]
     pub fn name(&self) -> core::result::Result<&CStr16, crate::data_types::FromSliceWithNulError> {
-        CStr16::from_u16_with_nul(&self.name)
+        Ok(&self.name)
     }
 }
 
 #[cfg(feature = "alloc")]
 impl Display for VariableKey {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
-        write!(f, "VariableKey {{ name: ")?;
-
-        match self.name() {
-            Ok(name) => write!(f, "\"{name}\"")?,
-            Err(err) => write!(f, "Err({err:?})")?,
-        }
-
-        write!(f, ", vendor: ")?;
+        write!(f, "VariableKey {{ name: \"{}\", vendor: ", self.name)?;
 
         if self.vendor == VariableVendor::GLOBAL_VARIABLE {
             write!(f, "GLOBAL_VARIABLE")?;