Merge pull request #1252 from rust-osdev/bishop-runtime-mod-varkeys-2

Add variable key functions to `uefi::runtime`

Author: phip1611

uefi-test-runner/src/runtime/vars.rs

@@ -51,6 +51,15 @@ fn test_variables(rt: &RuntimeServices) {
         info!("First variable: {}", key);
     }
 
+    // Test that the `runtime::variable_keys` iterator gives exactly the same
+    // list as the `RuntimeServices::variable_keys` function.
+    assert_eq!(
+        runtime::variable_keys()
+            .map(|k| k.unwrap())
+            .collect::<alloc::vec::Vec<_>>(),
+        variable_keys
+    );
+
     info!("Testing delete_variable()");
     rt.delete_variable(NAME, VENDOR)
         .expect("failed to delete variable");
@@ -86,6 +95,15 @@ fn test_variables_freestanding() {
     assert_eq!(&*data, VALUE);
     assert_eq!(attrs, ATTRS);
 
+    // Test that the variable is present in the `variable_keys` iterator.
+    let find_by_key = || {
+        runtime::variable_keys().any(|k| {
+            let k = k.as_ref().unwrap();
+            k.name().unwrap() == NAME && &k.vendor == VENDOR
+        })
+    };
+    assert!(find_by_key());
+
     // Delete the variable and verify it can no longer be read.
     runtime::delete_variable(NAME, VENDOR).expect("failed to delete variable");
     assert_eq!(
@@ -94,6 +112,8 @@ fn test_variables_freestanding() {
             .status(),
         Status::NOT_FOUND
     );
+    // Variable is no longer present in the `variable_keys` iterator.
+    assert!(!find_by_key());
 }
 
 fn test_variable_info(rt: &RuntimeServices) {

uefi/CHANGELOG.md

@@ -8,6 +8,7 @@
 - Add standard derives for `ConfigTableEntry`.
 - `PcrEvent`/`PcrEventInputs` impl `Align`, `Eq`, and `PartialEq`.
 - Added `PcrEvent::new_in_box` and `PcrEventInputs::new_in_box`.
+- `VariableKey` impls `Clone`, `Eq`, `PartialEq`, `Ord`, `PartialOrd`, and `Hash`.
 
 ## Changed
 - **Breaking:** `uefi::helpers::init` no longer takes an argument.

uefi/src/runtime.rs

@@ -6,10 +6,13 @@
 //! of the UEFI specification for details.
 
 use crate::{table, CStr16, Error, Result, Status, StatusExt};
+use core::mem;
 use core::ptr::{self, NonNull};
 
 #[cfg(feature = "alloc")]
-use {crate::mem::make_boxed, alloc::boxed::Box};
+use {
+    crate::mem::make_boxed, crate::Guid, alloc::borrow::ToOwned, alloc::boxed::Box, alloc::vec::Vec,
+};
 
 #[cfg(all(feature = "unstable", feature = "alloc"))]
 use alloc::alloc::Global;
@@ -18,6 +21,9 @@ pub use crate::table::runtime::{Daylight, Time, TimeCapabilities, TimeError, Tim
 pub use uefi_raw::capsule::{CapsuleBlockDescriptor, CapsuleFlags, CapsuleHeader};
 pub use uefi_raw::table::runtime::{ResetType, VariableAttributes, VariableVendor};
 
+#[cfg(feature = "alloc")]
+pub use crate::table::runtime::VariableKey;
+
 fn runtime_services_raw_panicking() -> NonNull<uefi_raw::table::runtime::RuntimeServices> {
     let st = table::system_table_raw_panicking();
     // SAFETY: valid per requirements of `set_system_table`.
@@ -141,6 +147,146 @@ pub fn get_variable_boxed(
     }
 }
 
+/// Gets each variable key (name and vendor) one at a time.
+///
+/// This is used to iterate over variable keys. See [`variable_keys`] for a more
+/// convenient interface that requires the `alloc` feature.
+///
+/// To get the first variable key, `name` must be initialized to start with a
+/// null character. The `vendor` value is arbitrary. On success, the first
+/// variable's name and vendor will be written out to `name` and `vendor`. Keep
+/// calling `get_next_variable_key` with the same `name` and `vendor` references
+/// to get the remaining variable keys.
+///
+/// All variable names should be valid strings, but this may not be enforced by
+/// firmware. To convert to a string, truncate at the first null and call
+/// [`CStr16::from_u16_with_nul`].
+///
+/// # Errors
+///
+/// * [`Status::NOT_FOUND`]: indicates end of iteration, the last variable keys
+///   was retrieved by the previous call to `get_next_variable_key`.
+/// * [`Status::BUFFER_TOO_SMALL`]: `name` is not large enough. The required
+///   size (in `u16` characters, not bytes) will be returned in the error data.
+/// * [`Status::INVALID_PARAMETER`]: `name` does not contain a null character, or
+///   the `name` and `vendor` are not an existing variable.
+/// * [`Status::DEVICE_ERROR`]: variable could not be read due to a hardware error.
+/// * [`Status::UNSUPPORTED`]: this platform does not support variable storage
+///   after exiting boot services.
+pub fn get_next_variable_key(
+    name: &mut [u16],
+    vendor: &mut VariableVendor,
+) -> Result<(), Option<usize>> {
+    let rt = runtime_services_raw_panicking();
+    let rt = unsafe { rt.as_ref() };
+
+    let mut name_size_in_bytes = mem::size_of_val(name);
+
+    let status = unsafe {
+        (rt.get_next_variable_name)(&mut name_size_in_bytes, name.as_mut_ptr(), &mut vendor.0)
+    };
+    match status {
+        Status::SUCCESS => Ok(()),
+        Status::BUFFER_TOO_SMALL => Err(Error::new(
+            status,
+            Some(name_size_in_bytes / mem::size_of::<u16>()),
+        )),
+        _ => Err(Error::new(status, None)),
+    }
+}
+
+/// Get an iterator over all UEFI variables.
+///
+/// See [`VariableKeys`] for details.
+#[cfg(feature = "alloc")]
+#[must_use]
+pub fn variable_keys() -> VariableKeys {
+    VariableKeys::new()
+}
+
+/// Iterator over all UEFI variables.
+///
+/// Each iteration yields a `Result<`[`VariableKey`]`>`. Error values:
+///
+/// * [`Status::DEVICE_ERROR`]: variable could not be read due to a hardware error.
+/// * [`Status::UNSUPPORTED`]: this platform does not support variable storage
+///   after exiting boot services.
+#[cfg(feature = "alloc")]
+#[derive(Debug)]
+pub struct VariableKeys {
+    name: Vec<u16>,
+    vendor: VariableVendor,
+    is_done: bool,
+}
+
+#[cfg(feature = "alloc")]
+impl VariableKeys {
+    fn new() -> Self {
+        // Create a the name buffer with a reasonable default capacity, and
+        // initialize it to an empty null-terminated string.
+        let mut name = Vec::with_capacity(32);
+        name.push(0);
+
+        Self {
+            // Give the name buffer a reasonable default capacity.
+            name,
+            // The initial vendor GUID is arbitrary.
+            vendor: VariableVendor(Guid::default()),
+            is_done: false,
+        }
+    }
+}
+
+#[cfg(feature = "alloc")]
+impl Iterator for VariableKeys {
+    type Item = Result<VariableKey>;
+
+    fn next(&mut self) -> Option<Result<VariableKey>> {
+        if self.is_done {
+            return None;
+        }
+
+        let mut result = get_next_variable_key(&mut self.name, &mut self.vendor);
+
+        // If the name buffer was too small, resize it to be big enough and call
+        // `get_next_variable_key` again.
+        if let Err(err) = &result {
+            if let Some(required_size) = err.data() {
+                self.name.resize(*required_size, 0u16);
+                result = get_next_variable_key(&mut self.name, &mut self.vendor);
+            }
+        }
+
+        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()
+                };
+                Some(Ok(VariableKey {
+                    name,
+                    vendor: self.vendor,
+                }))
+            }
+            Err(err) => {
+                if err.status() == Status::NOT_FOUND {
+                    // This status indicates the end of the list. The final variable
+                    // has already been yielded at this point, so return `None`.
+                    self.is_done = true;
+                    None
+                } else {
+                    // Return the error and end iteration.
+                    self.is_done = true;
+                    Some(Err(err.to_err_without_payload()))
+                }
+            }
+        }
+    }
+}
+
 /// Sets the value of a variable. This can be used to create a new variable,
 /// update an existing variable, or (when the size of `data` is zero)
 /// delete a variable.

uefi/src/table/runtime.rs

@@ -647,9 +647,9 @@ impl TryFrom<&[u8]> for Time {
 
 /// Unique key for a variable.
 #[cfg(feature = "alloc")]
-#[derive(Debug)]
+#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
 pub struct VariableKey {
-    name: Vec<u16>,
+    pub(crate) name: Vec<u16>,
     /// Unique identifier for the vendor.
     pub vendor: VariableVendor,
 }