runtime: Add variable key functions

The interface is slightly different than in `uefi/src/table/runtime.rs`,
where we just have a function to get all keys at once in a `Vec<VariableKey>`.

For the new implementation, two interfaces are
provided:

`get_next_variable_key` is a low-level interface that does not require
the `alloc` feature; it's a simple wrapper around `GetNextVariableName`.

`variable_keys` is a high-level interface that returns an
iterator. (Following the conventions of Rust's std lib, the iterator
type is named `VariableKeys`.) This requires `alloc`.

Author: nicholasbishop

uefi/src/runtime.rs

@@ -9,7 +9,10 @@ use crate::{table, CStr16, Error, Result, Status, StatusExt};
 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, core::mem,
+};
 
 #[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,145 @@ 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<'n, 'v>(
+    name: &'n mut [u16],
+    vendor: &'v mut VariableVendor,
+) -> Result<(), Option<usize>> {
+    let rt = runtime_services_raw_panicking();
+    let rt = unsafe { rt.as_ref() };
+
+    let mut name_size_in_bytes = name.len() * mem::size_of::<u16>();
+
+    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")]
+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;
+                    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.