Add read_entry_boxed_in and get_boxed_info_in that use the allocator_api.

This commit adds new functionality behind the `unsafe` feature that enhances some
existing functions with `allocator_api`-flavored allocators. Existing code will work
with and without the unstable feature.

Author: phip1611

CHANGELOG.md

@@ -1,14 +1,22 @@
 # Changelog
 
 ## uefi - [Unreleased]
+
+### Added
 - Implementations for the trait `EqStrUntilNul` now allow `?Sized` inputs. This means that
   you can write `some_cstr16.eq_str_until_nul("test")` instead of
   `some_cstr16.eq_str_until_nul(&"test")` now.
 - Added `TryFrom<core::ffi::CStr>` implementation for `CStr8`.
 - Added `Directory::read_entry_boxed` which works similar to `File::get_boxed_info`. This allows
-  easier iteration over the entries in a directory.
+  easier iteration over the entries in a directory. (requires the **alloc** feature)
+- Added `Directory::read_entry_boxed_in` and `File::get_boxed_info_in` that use the `allocator_api`
+  feature. (requires the **unstable** and **alloc** features)
 - Added an `core::error::Error` implementation for `Error` to ease
-  integration with error-handling crates.
+  integration with error-handling crates. (requires the **unstable** feature)
+
+### Changed
+
+### Removed
 
 ## uefi-macros - [Unreleased]
 

uefi-test-runner/Cargo.toml

@@ -6,7 +6,8 @@ publish = false
 edition = "2021"
 
 [dependencies]
-uefi = { path = "../uefi", features = ['alloc'] }
+# TODO we should let the uefi-test-runner run with and without unstable.
+uefi = { path = "../uefi", features = ["alloc", "unstable"] }
 uefi-services = { path = "../uefi-services" }
 
 log = { version = "0.4.17", default-features = false }

uefi/Cargo.toml

@@ -19,6 +19,7 @@ logger = []
 # were observed on the VirtualBox UEFI implementation (see uefi-rs#121).
 # In those cases, this feature can be excluded by removing the default features.
 panic-on-logger-errors = []
+# Generic gate to code that uses unstable features of Rust. You usually need a nightly toolchain.
 unstable = []
 
 [dependencies]

uefi/src/lib.rs

@@ -38,6 +38,8 @@
 //!   feature is not enabled, but once a couple more required features
 //!   are stabilized we intend to make the `uefi` crate work on the
 //!   stable channel by default.
+//!   As example, in conjunction with the `alloc`-feature, this gate allows
+//!   the `allocator_api` on certain functions.
 //!
 //! The `global_allocator` and `logger` features require special
 //! handling to perform initialization and tear-down. The
@@ -62,6 +64,7 @@
 #![feature(ptr_metadata)]
 #![cfg_attr(feature = "alloc", feature(vec_into_raw_parts))]
 #![cfg_attr(feature = "unstable", feature(error_in_core))]
+#![cfg_attr(all(feature = "unstable", feature = "alloc"), feature(allocator_api))]
 #![cfg_attr(docsrs, feature(doc_auto_cfg))]
 #![no_std]
 // Enable some additional warnings and lints.
@@ -101,5 +104,6 @@ pub mod global_allocator;
 #[cfg(feature = "logger")]
 pub mod logger;
 
+// As long as this is behind "alloc", we can simplify cfg-feature attributes in this module.
 #[cfg(feature = "alloc")]
 pub(crate) mod mem;

uefi/src/mem.rs

@@ -2,13 +2,19 @@
 
 use crate::ResultExt;
 use crate::{Result, Status};
-use ::alloc::{alloc, boxed::Box};
+use ::alloc::boxed::Box;
 use core::alloc::Layout;
 use core::fmt::Debug;
 use core::slice;
 use uefi::data_types::Align;
 use uefi::Error;
 
+#[cfg(not(feature = "unstable"))]
+use ::alloc::alloc::{alloc, dealloc};
+
+#[cfg(feature = "unstable")]
+use {core::alloc::Allocator, core::ptr::NonNull};
+
 /// Helper to return owned versions of certain UEFI data structures on the heap in a [`Box`]. This
 /// function is intended to wrap low-level UEFI functions of this crate that
 /// - can consume an empty buffer without a panic to get the required buffer size in the errors
@@ -17,12 +23,24 @@ use uefi::Error;
 ///   buffer size is sufficient, and
 /// - return a mutable typed reference that points to the same memory as the input buffer on
 ///   success.
-pub fn make_boxed<
+///
+/// # Feature `unstable` / `allocator_api`
+/// By default, this function works with Rust's default allocation mechanism. If you activate the
+/// `unstable`-feature, it uses the `allocator_api` instead. In that case, the function takes an
+/// additional parameter describing the specific [`Allocator`]. You can use [`alloc::alloc::Global`]
+/// as default.
+pub(crate) fn make_boxed<
     'a,
+    // The UEFI data structure.
     Data: Align + ?Sized + Debug + 'a,
     F: FnMut(&'a mut [u8]) -> Result<&'a mut Data, Option<usize>>,
+    #[cfg(feature = "unstable")] A: Allocator,
 >(
+    // A function to read the UEFI data structure into a provided buffer.
     mut fetch_data_fn: F,
+    #[cfg(feature = "unstable")]
+    // Allocator of the `allocator_api` feature. You can use `Global` as default.
+    allocator: A,
 ) -> Result<Box<Data>> {
     let required_size = match fetch_data_fn(&mut []).map_err(Error::split) {
         // This is the expected case: the empty buffer passed in is too
@@ -40,13 +58,23 @@ pub fn make_boxed<
         .unwrap()
         .pad_to_align();
 
-    // Allocate the buffer.
-    let heap_buf: *mut u8 = unsafe {
-        let ptr = alloc::alloc(layout);
-        if ptr.is_null() {
-            return Err(Status::OUT_OF_RESOURCES.into());
+    // Allocate the buffer on the heap.
+    let heap_buf: *mut u8 = {
+        #[cfg(not(feature = "unstable"))]
+        {
+            let ptr = unsafe { alloc(layout) };
+            if ptr.is_null() {
+                return Err(Status::OUT_OF_RESOURCES.into());
+            }
+            ptr
         }
-        ptr
+
+        #[cfg(feature = "unstable")]
+        allocator
+            .allocate(layout)
+            .map_err(|_| <Status as Into<Error>>::into(Status::OUT_OF_RESOURCES))?
+            .as_ptr()
+            .cast::<u8>()
     };
 
     // Read the data into the provided buffer.
@@ -59,7 +87,14 @@ pub fn make_boxed<
     let data: &mut Data = match data {
         Ok(data) => data,
         Err(err) => {
-            unsafe { alloc::dealloc(heap_buf, layout) };
+            #[cfg(not(feature = "unstable"))]
+            unsafe {
+                dealloc(heap_buf, layout)
+            };
+            #[cfg(feature = "unstable")]
+            unsafe {
+                allocator.deallocate(NonNull::new(heap_buf).unwrap(), layout)
+            }
             return Err(err);
         }
     };
@@ -73,6 +108,8 @@ pub fn make_boxed<
 mod tests {
     use super::*;
     use crate::ResultExt;
+    #[cfg(feature = "unstable")]
+    use alloc::alloc::Global;
     use core::mem::{align_of, size_of};
 
     /// Some simple dummy type to test [`make_boxed`].
@@ -166,14 +203,27 @@ mod tests {
         assert_eq!(&data.0 .0, &[1, 2, 3, 4]);
     }
 
+    /// This unit tests checks the [`make_boxed`] utility. The test has different code and behavior
+    /// depending on whether the "unstable" feature is active or not.
     #[test]
     fn test_make_boxed_utility() {
         let fetch_data_fn = |buf| uefi_function_stub_read(buf);
+
+        #[cfg(not(feature = "unstable"))]
         let data: Box<SomeData> = make_boxed(fetch_data_fn).unwrap();
+
+        #[cfg(feature = "unstable")]
+        let data: Box<SomeData> = make_boxed(fetch_data_fn, Global).unwrap();
         assert_eq!(&data.0, &[1, 2, 3, 4]);
 
         let fetch_data_fn = |buf| uefi_function_stub_read(buf);
+
+        #[cfg(not(feature = "unstable"))]
         let data: Box<SomeDataAlign16> = make_boxed(fetch_data_fn).unwrap();
+
+        #[cfg(feature = "unstable")]
+        let data: Box<SomeDataAlign16> = make_boxed(fetch_data_fn, Global).unwrap();
+
         assert_eq!(&data.0 .0, &[1, 2, 3, 4]);
     }
 }

uefi/src/proto/media/file/dir.rs

@@ -2,9 +2,10 @@ use super::{File, FileHandle, FileInfo, FromUefi, RegularFile};
 use crate::data_types::Align;
 use crate::Result;
 use core::ffi::c_void;
-
 #[cfg(feature = "alloc")]
 use {crate::mem::make_boxed, alloc::boxed::Box};
+#[cfg(all(feature = "unstable", feature = "alloc"))]
+use {alloc::alloc::Global, core::alloc::Allocator};
 
 /// A `FileHandle` that is also a directory.
 ///
@@ -62,8 +63,7 @@ impl Directory {
         })
     }
 
-    /// Wrapper around [`Self::read_entry`] that returns an owned copy of the data. It has the same
-    /// implications and requirements. On failure, the payload of `Err` is `()´.
+    /// Wrapper around [`Self::read_entry_boxed_in`] that uses the [`Global`] allocator.
     #[cfg(feature = "alloc")]
     pub fn read_entry_boxed(&mut self) -> Result<Option<Box<FileInfo>>> {
         let read_entry_res = self.read_entry(&mut []);
@@ -80,7 +80,42 @@ impl Directory {
                     maybe_info.expect("Should have more entries")
                 })
         };
+
+        #[cfg(not(feature = "unstable"))]
         let file_info = make_boxed::<FileInfo, _>(fetch_data_fn)?;
+
+        #[cfg(feature = "unstable")]
+        let file_info = make_boxed::<FileInfo, _, _>(fetch_data_fn, Global)?;
+
+        Ok(Some(file_info))
+    }
+
+    /// Wrapper around [`Self::read_entry`] that returns an owned copy of the data. It has the same
+    /// implications and requirements. On failure, the payload of `Err` is `()´.
+    ///
+    /// It allows to use a custom allocator via the `allocator_api` feature.
+    #[cfg(all(feature = "unstable", feature = "alloc"))]
+    pub fn read_entry_boxed_in<A: Allocator>(
+        &mut self,
+        allocator: A,
+    ) -> Result<Option<Box<FileInfo>>> {
+        let read_entry_res = self.read_entry(&mut []);
+
+        // If no more entries are available, return early.
+        if let Ok(None) = read_entry_res {
+            return Ok(None);
+        }
+
+        let fetch_data_fn = |buf| {
+            self.read_entry(buf)
+                // this is safe, as above, we checked that there are more entries
+                .map(|maybe_info: Option<&mut FileInfo>| {
+                    maybe_info.expect("Should have more entries")
+                })
+        };
+
+        let file_info = make_boxed::<FileInfo, _, A>(fetch_data_fn, allocator)?;
+
         Ok(Some(file_info))
     }
 

uefi/src/proto/media/file/mod.rs

@@ -16,6 +16,8 @@ use core::ffi::c_void;
 use core::fmt::Debug;
 use core::mem;
 use core::ptr;
+#[cfg(all(feature = "unstable", feature = "alloc"))]
+use {alloc::alloc::Global, core::alloc::Allocator};
 #[cfg(feature = "alloc")]
 use {alloc::boxed::Box, uefi::mem::make_boxed};
 
@@ -161,11 +163,25 @@ pub trait File: Sized {
         (self.imp().flush)(self.imp()).into()
     }
 
+    /// Wrapper around [`Self::get_boxed_info_in`] that uses the [`Global`] allocator.
     #[cfg(feature = "alloc")]
-    /// Read the dynamically allocated info for a file.
     fn get_boxed_info<Info: FileProtocolInfo + ?Sized + Debug>(&mut self) -> Result<Box<Info>> {
         let fetch_data_fn = |buf| self.get_info::<Info>(buf);
+        #[cfg(not(feature = "unstable"))]
         let file_info = make_boxed::<Info, _>(fetch_data_fn)?;
+        #[cfg(feature = "unstable")]
+        let file_info = make_boxed::<Info, _, _>(fetch_data_fn, Global)?;
+        Ok(file_info)
+    }
+
+    /// Read the dynamically allocated info for a file.
+    #[cfg(all(feature = "unstable", feature = "alloc"))]
+    fn get_boxed_info_in<Info: FileProtocolInfo + ?Sized + Debug, A: Allocator>(
+        &mut self,
+        allocator: A,
+    ) -> Result<Box<Info>> {
+        let fetch_data_fn = |buf| self.get_info::<Info>(buf);
+        let file_info = make_boxed::<Info, _, A>(fetch_data_fn, allocator)?;
         Ok(file_info)
     }