add memory util mem::make_boxed

phip1611 · phip1611 · commit d6f5de824742 · 2022-11-20T19:38:28.000+01:00
diff --git a/uefi/src/lib.rs b/uefi/src/lib.rs
@@ -79,3 +79,6 @@ pub mod global_allocator;
 
 #[cfg(feature = "logger")]
 pub mod logger;
+
+#[cfg(feature = "alloc")]
+pub(crate) mod mem;
diff --git a/uefi/src/mem.rs b/uefi/src/mem.rs
@@ -0,0 +1,130 @@
+//! This is a utility module with helper methods for allocations/memory.
+
+use crate::ResultExt;
+use crate::{Result, Status};
+use ::alloc::{alloc, boxed::Box};
+use core::alloc::Layout;
+use core::fmt::Debug;
+use core::ptr::NonNull;
+use core::slice;
+use uefi::data_types::Align;
+
+/// Helper to return owned versions of certain UEFI data structures on the heap, hence, in a box.
+/// This works as wrapper for low-level UEFI functions that
+/// - consume a mutable reference to a buffer that needs to be filled with some data,
+/// - return a mutable typed reference (which points to the same memory as the input buffer) on
+///   success, and
+/// - return the required buffer size on error.
+pub fn make_boxed<'a, Data: Align + ?Sized + Debug + 'a>(
+    mut fetch_data_fn: impl FnMut(&'a mut [u8]) -> Result<&'a mut Data, Option<usize>>,
+) -> Result<Box<Data>> {
+    let required_size = match fetch_data_fn(&mut [])
+        .expect_err("succeeded unexpectedly")
+        .split()
+    {
+        (s, None) => Err::<usize, crate::result::Error>(s.into()),
+        (_, Some(required_size)) => Ok(required_size),
+    }
+    .unwrap();
+
+    // We add trailing padding because the size of a rust structure must
+    // always be a multiple of alignment.
+    let layout = Layout::from_size_align(required_size, Data::alignment())
+        .unwrap()
+        .pad_to_align();
+
+    // Allocate the buffer.
+    let heap_buf: NonNull<u8> = unsafe {
+        let ptr = alloc::alloc(layout);
+        match NonNull::new(ptr) {
+            None => return Err(Status::OUT_OF_RESOURCES.into()),
+            Some(ptr) => ptr,
+        }
+    };
+
+    // Get the file info using the allocated buffer for storage.
+    let data: Result<&mut Data> = {
+        let buffer = unsafe { slice::from_raw_parts_mut(heap_buf.as_ptr(), layout.size()) };
+        fetch_data_fn(buffer).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.as_ptr(), layout) };
+            return Err(err);
+        }
+    };
+
+    let data = unsafe { Box::from_raw(data) };
+
+    Ok(data)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::ResultExt;
+    use core::mem::{align_of, size_of};
+
+    #[derive(Debug)]
+    #[repr(C)]
+    struct SomeData([u8; 4]);
+
+    impl Align for SomeData {
+        fn alignment() -> usize {
+            align_of::<Self>()
+        }
+    }
+
+    /// Function that behaves like the other UEFI functions. It takes a
+    /// mutable reference to a buffer memory that represents a [`SomeData`]
+    /// instance.
+    fn uefi_function_stub_read(buf: &mut [u8]) -> crate::Result<&mut SomeData, Option<usize>> {
+        if buf.len() < 4 {
+            return Status::BUFFER_TOO_SMALL.into_with(|| panic!(), |_| Some(4));
+        };
+
+        buf[0] = 1;
+        buf[1] = 2;
+        buf[2] = 3;
+        buf[3] = 4;
+
+        let data = unsafe { buf.as_mut_ptr().cast::<SomeData>().as_mut().unwrap() };
+
+        Ok(data)
+    }
+
+    // Some basic checks so that miri reports everything is fine.
+    #[test]
+    fn some_data_type_size_constraints() {
+        assert_eq!(size_of::<SomeData>(), 4);
+        assert_eq!(align_of::<SomeData>(), 1);
+    }
+
+    #[test]
+    fn basic_stub_read() {
+        assert_eq!(
+            uefi_function_stub_read(&mut []).status(),
+            Status::BUFFER_TOO_SMALL
+        );
+        assert_eq!(
+            *uefi_function_stub_read(&mut []).unwrap_err().data(),
+            Some(4)
+        );
+
+        let mut buf: [u8; 4] = [0; 4];
+        let data = uefi_function_stub_read(&mut buf).unwrap();
+
+        assert_eq!(&data.0, &[1, 2, 3, 4])
+    }
+
+    #[test]
+    fn make_boxed_utility() {
+        let fetch_data_fn = |buf| uefi_function_stub_read(buf);
+        let data: Box<SomeData> = make_boxed(fetch_data_fn).unwrap();
+
+        assert_eq!(&data.0, &[1, 2, 3, 4])
+    }
+}
diff --git a/uefi/src/proto/media/file/dir.rs b/uefi/src/proto/media/file/dir.rs
@@ -4,14 +4,7 @@ use crate::Result;
 use core::ffi::c_void;
 
 #[cfg(feature = "alloc")]
-use {
-    crate::{ResultExt, Status},
-    ::alloc::boxed::Box,
-    alloc::alloc,
-    core::alloc::Layout,
-    core::ptr::NonNull,
-    core::slice,
-};
+use {crate::mem::make_boxed, alloc::boxed::Box};
 
 /// A `FileHandle` that is also a directory.
 ///
@@ -77,51 +70,15 @@ impl Directory {
             return Ok(None);
         }
 
-        let required_size = match read_entry_res
-            .expect_err("zero sized read unexpectedly succeeded")
-            .split()
-        {
-            // Early return if something has failed.
-            (s, None) => return Err(s.into()),
-            (_, Some(required_size)) => required_size,
+        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")
+                })
         };
-
-        // We add trailing padding because the size of a rust structure must
-        // always be a multiple of alignment.
-        let layout = Layout::from_size_align(required_size, FileInfo::alignment())
-            .unwrap()
-            .pad_to_align();
-
-        // Allocate the buffer.
-        let heap_buf: NonNull<u8> = unsafe {
-            let ptr = alloc::alloc(layout);
-            match NonNull::new(ptr) {
-                None => return Err(Status::OUT_OF_RESOURCES.into()),
-                Some(ptr) => ptr,
-            }
-        };
-
-        // Get the file info using the allocated buffer for storage.
-        let info = {
-            let buffer = unsafe { slice::from_raw_parts_mut(heap_buf.as_ptr(), layout.size()) };
-            self.read_entry(buffer).discard_errdata()
-        };
-
-        // If an error occurred, deallocate the memory before returning.
-        let info = match info {
-            Ok(info) => info,
-            Err(err) => {
-                unsafe { alloc::dealloc(heap_buf.as_ptr(), layout) };
-                return Err(err);
-            }
-        };
-
-        // Wrap the file info in a box so that it will be deallocated on
-        // drop. This is valid because the memory was allocated with the
-        // global allocator.
-        let info = info.map(|info| unsafe { Box::from_raw(info) });
-
-        Ok(info)
+        let file_info = make_boxed::<FileInfo>(fetch_data_fn)?;
+        Ok(Some(file_info))
     }
 
     /// Start over the process of enumerating directory entries
diff --git a/uefi/src/proto/media/file/mod.rs b/uefi/src/proto/media/file/mod.rs
@@ -17,11 +17,7 @@ use core::fmt::Debug;
 use core::mem;
 use core::ptr;
 #[cfg(feature = "alloc")]
-use {
-    crate::ResultExt,
-    ::alloc::{alloc, alloc::Layout, boxed::Box},
-    core::slice,
-};
+use {alloc::boxed::Box, uefi::mem::make_boxed};
 
 pub use self::info::{FileInfo, FileProtocolInfo, FileSystemInfo, FileSystemVolumeLabel, FromUefi};
 pub use self::{dir::Directory, regular::RegularFile};
@@ -168,51 +164,9 @@ pub trait File: Sized {
     #[cfg(feature = "alloc")]
     /// Get the dynamically allocated info for a file
     fn get_boxed_info<Info: FileProtocolInfo + ?Sized + Debug>(&mut self) -> Result<Box<Info>> {
-        // Initially try get_info with an empty array, this should always fail
-        // as all Info types at least need room for a null-terminator.
-        let size = match self
-            .get_info::<Info>(&mut [])
-            .expect_err("zero sized get_info unexpectedly succeeded")
-            .split()
-        {
-            (s, None) => return Err(s.into()),
-            (_, Some(size)) => size,
-        };
-
-        // We add trailing padding because the size of a rust structure must
-        // always be a multiple of alignment.
-        let layout = Layout::from_size_align(size, Info::alignment())
-            .unwrap()
-            .pad_to_align();
-
-        // Allocate the buffer.
-        let data: *mut u8 = unsafe {
-            let data = alloc::alloc(layout);
-            if data.is_null() {
-                return Err(Status::OUT_OF_RESOURCES.into());
-            }
-            data
-        };
-
-        // Get the file info using the allocated buffer for storage.
-        let info = {
-            let buffer = unsafe { slice::from_raw_parts_mut(data, layout.size()) };
-            self.get_info::<Info>(buffer).discard_errdata()
-        };
-
-        // If an error occurred, deallocate the memory before returning.
-        let info = match info {
-            Ok(info) => info,
-            Err(err) => {
-                unsafe { alloc::dealloc(data, layout) };
-                return Err(err);
-            }
-        };
-
-        // Wrap the file info in a box so that it will be deallocated on
-        // drop. This is valid because the memory was allocated with the
-        // global allocator.
-        unsafe { Ok(Box::from_raw(info)) }
+        let fetch_data_fn = |buf| self.get_info::<Info>(buf);
+        let file_info = make_boxed::<Info>(fetch_data_fn)?;
+        Ok(file_info)
     }
 
     /// Returns if the underlying file is a regular file.
