Added read_entry_boxed_in and get_boxed_info_in that use the allocator_api

Author: phip1611

CHANGELOG.md

@@ -8,7 +8,9 @@
 - Added `Directory::read_entry_boxed` which works similar to `File::get_boxed_info`. This allows
   easier iteration over the entries in a directory.
 - Added an `core::error::Error` implementation for `Error` to ease
-  integration with error-handling crates.
+  integration with error-handling crates. (requires the **unstable** feature)
+- Added `Directory::read_entry_boxed_in` and `File::get_boxed_info_in` that use the `allocator_api`
+  feature. (requires the **unstable** feature)
 
 ## uefi-macros - [Unreleased]
 

uefi/src/lib.rs

@@ -62,6 +62,8 @@
 #![feature(ptr_metadata)]
 #![cfg_attr(feature = "alloc", feature(vec_into_raw_parts))]
 #![cfg_attr(feature = "unstable", feature(error_in_core))]
+#![cfg_attr(feature = "unstable", feature(allocator_api))]
+#![cfg_attr(feature = "unstable", feature(slice_ptr_get))]
 #![cfg_attr(docsrs, feature(doc_auto_cfg))]
 #![no_std]
 // Enable some additional warnings and lints.

uefi/src/mem.rs

@@ -2,13 +2,25 @@
 
 use crate::ResultExt;
 use crate::{Result, Status};
-use ::alloc::{alloc, boxed::Box};
-use core::alloc::Layout;
+use ::alloc::{alloc::Global, boxed::Box};
+use core::alloc::{Allocator, Layout};
 use core::fmt::Debug;
-use core::slice;
+use core::ptr::NonNull;
 use uefi::data_types::Align;
 use uefi::Error;
 
+/// Wrapper around [`make_boxed_in`] that uses [`Global`]/the system allocator.
+#[allow(unused)]
+pub fn make_boxed<
+    'a,
+    Data: Align + ?Sized + Debug + 'a,
+    F: FnMut(&'a mut [u8]) -> Result<&'a mut Data, Option<usize>>,
+>(
+    mut fetch_data_fn: F,
+) -> Result<Box<Data>> {
+    make_boxed_in::<'a, Global, _, _>(Global, fetch_data_fn)
+}
+
 /// 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,11 +29,15 @@ 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<
+///
+/// It takes a custom allocator via the `allocator_api`.
+pub fn make_boxed_in<
     'a,
+    A: Allocator,
     Data: Align + ?Sized + Debug + 'a,
     F: FnMut(&'a mut [u8]) -> Result<&'a mut Data, Option<usize>>,
 >(
+    allocator: A,
     mut fetch_data_fn: F,
 ) -> Result<Box<Data>> {
     let required_size = match fetch_data_fn(&mut []).map_err(Error::split) {
@@ -41,25 +57,21 @@ pub fn make_boxed<
         .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());
-        }
-        ptr
-    };
+    let mut heap_buf: NonNull<[u8]> = allocator
+        .allocate(layout)
+        .map_err(|_| <Status as Into<Error>>::into(Status::OUT_OF_RESOURCES))?;
 
     // Read the data into the provided buffer.
     let data: Result<&mut Data> = {
-        let buffer = unsafe { slice::from_raw_parts_mut(heap_buf, required_size) };
-        fetch_data_fn(buffer).discard_errdata()
+        let slice = unsafe { heap_buf.as_mut() };
+        fetch_data_fn(slice).discard_errdata()
     };
 
     // If an error occurred, deallocate the memory before returning.
     let data: &mut Data = match data {
         Ok(data) => data,
         Err(err) => {
-            unsafe { alloc::dealloc(heap_buf, layout) };
+            unsafe { allocator.deallocate(heap_buf.as_non_null_ptr(), layout) };
             return Err(err);
         }
     };

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

@@ -4,7 +4,7 @@ use crate::Result;
 use core::ffi::c_void;
 
 #[cfg(feature = "alloc")]
-use {crate::mem::make_boxed, alloc::boxed::Box};
+use {crate::mem::make_boxed_in, alloc::alloc::Global, alloc::boxed::Box, core::alloc::Allocator};
 
 /// A `FileHandle` that is also a directory.
 ///
@@ -61,10 +61,21 @@ impl Directory {
         })
     }
 
+    /// 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>>> {
+        self.read_entry_boxed_in(Global)
+    }
+
     /// 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(feature = "alloc")]
-    pub fn read_entry_boxed(&mut self) -> Result<Option<Box<FileInfo>>> {
+    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.
@@ -79,7 +90,7 @@ impl Directory {
                     maybe_info.expect("Should have more entries")
                 })
         };
-        let file_info = make_boxed::<FileInfo, _>(fetch_data_fn)?;
+        let file_info = make_boxed_in::<A, FileInfo, _>(allocator, fetch_data_fn)?;
         Ok(Some(file_info))
     }
 

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

@@ -17,7 +17,7 @@ use core::fmt::Debug;
 use core::mem;
 use core::ptr;
 #[cfg(feature = "alloc")]
-use {alloc::boxed::Box, uefi::mem::make_boxed};
+use {alloc::alloc::Global, alloc::boxed::Box, core::alloc::Allocator, uefi::mem::make_boxed_in};
 
 pub use self::info::{FileInfo, FileProtocolInfo, FileSystemInfo, FileSystemVolumeLabel, FromUefi};
 pub use self::{dir::Directory, regular::RegularFile};
@@ -162,10 +162,19 @@ pub trait File: Sized {
     }
 
     #[cfg(feature = "alloc")]
-    /// Read the dynamically allocated info for a file.
+    /// Wrapper around [`Self::get_boxed_info_in`] that uses the [`Global`] allocator.
     fn get_boxed_info<Info: FileProtocolInfo + ?Sized + Debug>(&mut self) -> Result<Box<Info>> {
+        self.get_boxed_info_in(Global)
+    }
+
+    #[cfg(feature = "alloc")]
+    /// Read the dynamically allocated info for a file.
+    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, _>(fetch_data_fn)?;
+        let file_info = make_boxed_in::<A, Info, _>(allocator, fetch_data_fn)?;
         Ok(file_info)
     }