Merge pull request #559 from phip1611/dir-info-boxed

phip1611 · web-flow · commit 3a709611e8ea · 2022-11-24T09:43:48.000+01:00
`Directory::read_entry_boxed` plus common abstraction `make_boxed`
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,8 @@
   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.
 
 ## uefi-macros - [Unreleased]
 
diff --git a/uefi-test-runner/src/proto/media/known_disk.rs b/uefi-test-runner/src/proto/media/known_disk.rs
@@ -1,4 +1,5 @@
 use alloc::string::ToString;
+use core::cell::RefCell;
 use core::ptr::NonNull;
 use uefi::prelude::*;
 use uefi::proto::media::block::BlockIO;
@@ -21,20 +22,61 @@ fn test_existing_dir(directory: &mut Directory) {
 
     assert!(dir.is_directory().unwrap());
 
-    let mut dir = dir.into_directory().expect("not a directory");
-
-    // Collect and validate the directory entries.
-    let mut entry_names = vec![];
-    let mut buf = vec![0; 200];
-    loop {
-        let entry = dir.read_entry(&mut buf).expect("failed to read directory");
-        if let Some(entry) = entry {
-            entry_names.push(entry.file_name().to_string());
-        } else {
-            break;
+    let dir = dir.into_directory().expect("Should be a directory");
+
+    let dir = RefCell::new(dir);
+
+    // Backing memory to read the file info data into.
+    let mut stack_buf = [0; 200];
+
+    // The file names that the test read from the directory.
+    let entry_names = RefCell::new(vec![]);
+
+    // Expected file names in the directory.
+    const EXPECTED: &[&str] = &[".", "..", "test_input.txt"];
+
+    // Reads the whole directory with provided backing memory.
+    let mut test_read_dir_stack_mem = || {
+        let mut dir = dir.borrow_mut();
+        let mut entry_names = entry_names.borrow_mut();
+        loop {
+            let entry = dir
+                .read_entry(&mut stack_buf)
+                .expect("failed to read directory");
+            if let Some(entry) = entry {
+                entry_names.push(entry.file_name().to_string());
+            } else {
+                break;
+            }
         }
+        assert_eq!(&*entry_names, EXPECTED);
+    };
+
+    // Reads the whole directory but returns owned memory on the heap.
+    let test_read_dir_heap_mem = || {
+        let mut dir = dir.borrow_mut();
+        let mut entry_names = entry_names.borrow_mut();
+        loop {
+            let entry = dir.read_entry_boxed().expect("failed to read directory");
+            if let Some(entry) = entry {
+                entry_names.push(entry.file_name().to_string());
+            } else {
+                break;
+            }
+        }
+        assert_eq!(&*entry_names, EXPECTED);
+    };
+
+    // Tests all read dir test functions three times.
+    for _ in 0..3 {
+        entry_names.borrow_mut().clear();
+        dir.borrow_mut().reset_entry_readout().unwrap();
+        test_read_dir_stack_mem();
+
+        entry_names.borrow_mut().clear();
+        dir.borrow_mut().reset_entry_readout().unwrap();
+        test_read_dir_heap_mem();
     }
-    assert_eq!(entry_names, [".", "..", "test_input.txt"]);
 }
 
 /// Test that deleting a file opened in read-only mode fails with a
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,179 @@
+//! 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::slice;
+use uefi::data_types::Align;
+use uefi::Error;
+
+/// 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
+///   payload,
+/// - consume a mutable reference to a buffer that will be filled with some data if the provided
+///   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<
+    '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>> {
+    let required_size = match fetch_data_fn(&mut []).map_err(Error::split) {
+        // This is the expected case: the empty buffer passed in is too
+        // small, so we get the required size.
+        Err((Status::BUFFER_TOO_SMALL, Some(required_size))) => Ok(required_size),
+        // Propagate any other error.
+        Err((status, _)) => Err(Error::from(status)),
+        // Success is unexpected, return an error.
+        Ok(_) => Err(Error::from(Status::UNSUPPORTED)),
+    }?;
+
+    // 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: *mut u8 = unsafe {
+        let ptr = alloc::alloc(layout);
+        if ptr.is_null() {
+            return Err(Status::OUT_OF_RESOURCES.into());
+        }
+        ptr
+    };
+
+    // 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()
+    };
+
+    // 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) };
+            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};
+
+    /// Some simple dummy type to test [`make_boxed`].
+    #[derive(Debug)]
+    #[repr(C)]
+    struct SomeData([u8; 4]);
+
+    impl Align for SomeData {
+        fn alignment() -> usize {
+            align_of::<Self>()
+        }
+    }
+
+    /// Type wrapper that ensures an alignment of 16 for the underlying data.
+    #[derive(Debug)]
+    #[repr(C, align(16))]
+    struct Align16<T>(T);
+
+    /// Version of [`SomeData`] that has an alignment of 16.
+    type SomeDataAlign16 = Align16<SomeData>;
+
+    impl Align for SomeDataAlign16 {
+        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<Data: Align>(buf: &mut [u8]) -> Result<&mut Data, Option<usize>> {
+        let required_size = size_of::<Data>();
+
+        if buf.len() < required_size {
+            // We can use an zero-length buffer to find the required size.
+            return Status::BUFFER_TOO_SMALL.into_with(|| panic!(), |_| Some(required_size));
+        };
+
+        // assert alignment
+        assert_eq!(
+            buf.as_ptr() as usize % Data::alignment(),
+            0,
+            "The buffer must be correctly aligned!"
+        );
+
+        buf[0] = 1;
+        buf[1] = 2;
+        buf[2] = 3;
+        buf[3] = 4;
+
+        let data = unsafe { &mut *buf.as_mut_ptr().cast::<Data>() };
+
+        Ok(data)
+    }
+
+    // Some basic sanity checks so that we can catch problems early that miri would detect
+    // otherwise.
+    #[test]
+    fn test_some_data_type_size_constraints() {
+        assert_eq!(size_of::<SomeData>(), 4);
+        assert_eq!(SomeData::alignment(), 1);
+        assert_eq!(
+            size_of::<SomeDataAlign16>(),
+            16,
+            "The size must be 16 instead of 4, as in Rust the runtime size is a multiple of the alignment."
+        );
+        assert_eq!(SomeDataAlign16::alignment(), 16);
+    }
+
+    // Tests `uefi_function_stub_read` which is the foundation for the `test_make_boxed_utility`
+    // test.
+    #[test]
+    fn test_basic_stub_read() {
+        assert_eq!(
+            uefi_function_stub_read::<SomeData>(&mut []).status(),
+            Status::BUFFER_TOO_SMALL
+        );
+        assert_eq!(
+            *uefi_function_stub_read::<SomeData>(&mut [])
+                .unwrap_err()
+                .data(),
+            Some(4)
+        );
+
+        let mut buf: [u8; 4] = [0; 4];
+        let data: &mut SomeData = uefi_function_stub_read(&mut buf).unwrap();
+        assert_eq!(&data.0, &[1, 2, 3, 4]);
+
+        let mut buf: Align16<[u8; 16]> = Align16([0; 16]);
+        let data: &mut SomeDataAlign16 = uefi_function_stub_read(&mut buf.0).unwrap();
+        assert_eq!(&data.0 .0, &[1, 2, 3, 4]);
+    }
+
+    #[test]
+    fn test_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]);
+
+        let fetch_data_fn = |buf| uefi_function_stub_read(buf);
+        let data: Box<SomeDataAlign16> = make_boxed(fetch_data_fn).unwrap();
+        assert_eq!(&data.0 .0, &[1, 2, 3, 4]);
+    }
+}
diff --git a/uefi/src/proto/media/file/dir.rs b/uefi/src/proto/media/file/dir.rs
@@ -3,6 +3,9 @@ use crate::data_types::Align;
 use crate::Result;
 use core::ffi::c_void;
 
+#[cfg(feature = "alloc")]
+use {crate::mem::make_boxed, alloc::boxed::Box};
+
 /// A `FileHandle` that is also a directory.
 ///
 /// Use `File::into_type` or `Directory::new` to create a `Directory`. In
@@ -20,7 +23,7 @@ impl Directory {
         Self(RegularFile::new(handle))
     }
 
-    /// Read the next directory entry
+    /// Read the next directory entry.
     ///
     /// Try to read the next directory entry into `buffer`. If the buffer is too small, report the
     /// required buffer size as part of the error. If there are no more directory entries, return
@@ -56,6 +59,28 @@ 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 `()´.
+    #[cfg(feature = "alloc")]
+    pub fn read_entry_boxed(&mut self) -> 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, _>(fetch_data_fn)?;
+        Ok(Some(file_info))
+    }
+
     /// Start over the process of enumerating directory entries
     pub fn reset_entry_readout(&mut self) -> Result {
         self.0.set_position(0)
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};
@@ -166,53 +162,11 @@ pub trait File: Sized {
     }
 
     #[cfg(feature = "alloc")]
-    /// Get the dynamically allocated info for a file
+    /// Read 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.
diff --git a/xtask/src/cargo.rs b/xtask/src/cargo.rs
