uefi: alloc_pages() and alloc_pool() now return NonNull<[u8]>

uefi-test-runner/src/boot/memory.rs

@@ -3,6 +3,7 @@
 use alloc::vec::Vec;
 use uefi::boot::{self, AllocateType};
 use uefi::mem::memory_map::{MemoryMap, MemoryMapMut, MemoryType};
+use uefi_raw::table::boot::PAGE_SIZE;
 
 pub fn test() {
     info!("Testing memory functions");
@@ -17,32 +18,42 @@ pub fn test() {
 }
 
 fn test_allocate_pages() {
-    let num_pages = 1;
-    let ptr =
+    let num_pages = 3;
+    let mut ptr =
         boot::allocate_pages(AllocateType::AnyPages, MemoryType::LOADER_DATA, num_pages).unwrap();
-    let addr = ptr.as_ptr() as usize;
-    assert_eq!(addr % 4096, 0, "Page pointer is not page-aligned");
+
+    let buffer = unsafe { ptr.as_mut() };
+    assert_eq!(
+        buffer.as_ptr().align_offset(PAGE_SIZE),
+        0,
+        "Page pointer is not page-aligned"
+    );
 
     // Verify the page can be written to.
     {
-        let ptr = ptr.as_ptr();
-        unsafe { ptr.write_volatile(0xff) };
-        unsafe { ptr.add(4095).write_volatile(0xff) };
+        buffer[0] = 0xff;
+        buffer[4095] = 0xff;
+        buffer[5095] = 0xff;
+        assert_eq!(buffer[0], 0xff);
+        assert_eq!(buffer[4095], 0xff);
+        assert_eq!(buffer[5095], 0xff);
     }
 
-    unsafe { boot::free_pages(ptr, num_pages) }.unwrap();
+    unsafe { boot::free_pages(ptr.cast(), num_pages) }.unwrap();
 }
 
 fn test_allocate_pool() {
-    let ptr = boot::allocate_pool(MemoryType::LOADER_DATA, 10).unwrap();
+    let mut ptr = boot::allocate_pool(MemoryType::LOADER_DATA, 10).unwrap();
+    let buffer = unsafe { ptr.as_mut() };
 
     // Verify the allocation can be written to.
     {
-        let ptr = ptr.as_ptr();
-        unsafe { ptr.write_volatile(0xff) };
-        unsafe { ptr.add(9).write_volatile(0xff) };
+        buffer[0] = 0xff;
+        buffer[9] = 0xff;
+        assert_eq!(buffer[0], 0xff);
+        assert_eq!(buffer[9], 0xff);
     }
-    unsafe { boot::free_pool(ptr) }.unwrap();
+    unsafe { boot::free_pool(ptr.cast()) }.unwrap();
 }
 
 // Simple test to ensure our custom allocator works with the `alloc` crate.

uefi-test-runner/src/boot/misc.rs

@@ -222,13 +222,14 @@ fn test_install_configuration_table() {
     let initial_table_count = system::with_config_table(|t| t.len());
 
     // Create the entry data.
-    let config: NonNull<u8> = boot::allocate_pool(MemoryType::RUNTIME_SERVICES_DATA, 1).unwrap();
-    unsafe { config.write(123u8) };
+    let mut config_ptr = boot::allocate_pool(MemoryType::RUNTIME_SERVICES_DATA, 1).unwrap();
+    let buffer = unsafe { config_ptr.as_mut() };
+    buffer[0] = 123;
 
     // Install the table.
     const TABLE_GUID: Guid = guid!("4bec53c4-5fc1-48a1-ab12-df214907d29f");
     unsafe {
-        boot::install_configuration_table(&TABLE_GUID, config.as_ptr().cast()).unwrap();
+        boot::install_configuration_table(&TABLE_GUID, config_ptr.as_ptr().cast()).unwrap();
     }
 
     // Verify the installation.
@@ -244,6 +245,6 @@ fn test_install_configuration_table() {
     // Uninstall the table and free the memory.
     unsafe {
         boot::install_configuration_table(&TABLE_GUID, ptr::null()).unwrap();
-        boot::free_pool(config).unwrap();
+        boot::free_pool(config_ptr.cast()).unwrap();
     }
 }

uefi/CHANGELOG.md

@@ -25,6 +25,12 @@
   `proto::device_path::text` to `proto::device_path`.
 - **Breaking:** `exit_boot_services` now consumes a `Option<MemoryType>` which
   defaults to the recommended value of `MemoryType::LOADER_DATA`.
+- **Breaking**: `allocate_pages` now returns `NonNull<[u8]>` to align it with
+  the Rust allocator API. There is an example in the documentation of that
+  function.
+- **Breaking**: `allocate_pool` now returns `NonNull<[u8]>` to align it with
+  the Rust allocator API. There is an example in the documentation of that
+  function.
 - `boot::memory_map()` will never return `Status::BUFFER_TOO_SMALL` from now on,
   as this is considered a hard internal error where users can't do anything
   about it anyway. It will panic instead.

uefi/src/allocator.rs

@@ -68,10 +68,12 @@ unsafe impl GlobalAlloc for Allocator {
             // only guaranteed to provide eight-byte alignment. Allocate extra
             // space so that we can return an appropriately-aligned pointer
             // within the allocation.
-            let full_alloc_ptr = if let Ok(ptr) = boot::allocate_pool(memory_type, size + align) {
-                ptr.as_ptr()
-            } else {
-                return ptr::null_mut();
+            let full_alloc_ptr = match boot::allocate_pool(memory_type, size + align) {
+                Ok(ptr) => ptr.cast::<u8>().as_ptr(),
+                Err(e) => {
+                    log::error!("Failed to allocate pool: {:?}", e);
+                    return ptr::null_mut();
+                }
             };
 
             // Calculate the offset needed to get an aligned pointer within the
@@ -100,7 +102,8 @@ unsafe impl GlobalAlloc for Allocator {
             // `allocate_pool` always provides eight-byte alignment, so we can
             // use `allocate_pool` directly.
             boot::allocate_pool(memory_type, size)
-                .map(|ptr| ptr.as_ptr())
+                .map(|mut ptr: NonNull<[u8]>| unsafe { ptr.as_mut() })
+                .map(|ptr: &mut [u8]| ptr.as_mut_ptr())
                 .unwrap_or(ptr::null_mut())
         }
     }

uefi/src/boot.rs

@@ -120,42 +120,83 @@ pub unsafe fn raise_tpl(tpl: Tpl) -> TplGuard {
     }
 }
 
-/// Allocates memory pages from the system.
+/// Allocates a consecutive set of memory pages using the UEFI allocator.
+///
+/// The caller is responsible to free the memory using [`free_pages`].
 ///
 /// UEFI OS loaders should allocate memory of the type `LoaderData`.
 ///
+/// # Example
+///```rust,no_run
+/// use uefi::boot::{self, AllocateType};
+/// use uefi_raw::table::boot::MemoryType;
+///
+/// let num_pages = 3;
+/// let mut ptr = boot::allocate_pages(
+///    AllocateType::AnyPages,
+///    MemoryType::LOADER_DATA,
+///    num_pages
+/// ).unwrap();
+///
+/// // ⚠️ Creating the reference is safe, but reading the uninitialized memory
+/// // causes Undefined Behavior (UB)! Please make sure to initialize the memory
+/// // first by:
+/// //   - using `core::ptr::write`,
+/// //   - directly writing to slice indices,
+/// //   - zeroing the memory,
+/// //   - using `.copy_from_slice()`,
+/// //   - or a similar operation.
+/// let buffer: &mut [u8] = unsafe { ptr.as_mut() };
+/// // Now initialize the content of the buffer, cast it, etc.
+/// // Please follow Rust guidelines on safety and UB! ⚠️
+///
+/// // free the allocation
+/// unsafe { boot::free_pages(ptr.cast(), num_pages) }.unwrap();
+/// ```
+///
+/// # Safety
+/// Using this function is safe but reading on initialized memory is not.
+/// Please look into the example code.
+///
 /// # Errors
 ///
 /// * [`Status::OUT_OF_RESOURCES`]: allocation failed.
 /// * [`Status::INVALID_PARAMETER`]: `mem_ty` is [`MemoryType::PERSISTENT_MEMORY`],
 ///   [`MemoryType::UNACCEPTED`], or in the range [`MemoryType::MAX`]`..=0x6fff_ffff`.
 /// * [`Status::NOT_FOUND`]: the requested pages could not be found.
-pub fn allocate_pages(ty: AllocateType, mem_ty: MemoryType, count: usize) -> Result<NonNull<u8>> {
+pub fn allocate_pages(
+    allocation_type: AllocateType,
+    memory_type: MemoryType,
+    count: usize,
+) -> Result<NonNull<[u8]>> {
     let bt = boot_services_raw_panicking();
     let bt = unsafe { bt.as_ref() };
 
-    let (ty, initial_addr) = match ty {
+    let (allocation_type_efi, start_address) = match allocation_type {
         AllocateType::AnyPages => (0, 0),
         AllocateType::MaxAddress(addr) => (1, addr),
         AllocateType::Address(addr) => (2, addr),
     };
 
-    let mut addr1 = initial_addr;
-    unsafe { (bt.allocate_pages)(ty, mem_ty, count, &mut addr1) }.to_result()?;
+    let mut addr1 = start_address;
+    unsafe { (bt.allocate_pages)(allocation_type_efi, memory_type, count, &mut addr1) }
+        .to_result()?;
 
     // The UEFI spec allows `allocate_pages` to return a valid allocation at
     // address zero. Rust does not allow writes through a null pointer (which
     // Rust defines as address zero), so this is not very useful. Only return
     // the allocation if the address is non-null.
     if let Some(ptr) = NonNull::new(addr1 as *mut u8) {
-        return Ok(ptr);
+        let slice = NonNull::slice_from_raw_parts(ptr, count * PAGE_SIZE);
+        return Ok(slice);
     }
 
     // Attempt a second allocation. The first allocation (at address zero) has
     // not yet been freed, so if this allocation succeeds it should be at a
     // non-zero address.
-    let mut addr2 = initial_addr;
-    let r = unsafe { (bt.allocate_pages)(ty, mem_ty, count, &mut addr2) }.to_result();
+    let mut addr2 = start_address;
+    let r = unsafe { (bt.allocate_pages)(allocation_type_efi, memory_type, count, &mut addr2) }
+        .to_result();
 
     // Free the original allocation (ignoring errors).
     let _unused = unsafe { (bt.free_pages)(addr1, count) };
@@ -164,7 +205,8 @@ pub fn allocate_pages(ty: AllocateType, mem_ty: MemoryType, count: usize) -> Res
     // address zero. Otherwise, return a pointer to the second allocation.
     r?;
     if let Some(ptr) = NonNull::new(addr2 as *mut u8) {
-        Ok(ptr)
+        let slice = NonNull::slice_from_raw_parts(ptr, count * PAGE_SIZE);
+        Ok(slice)
     } else {
         Err(Status::OUT_OF_RESOURCES.into())
     }
@@ -189,22 +231,67 @@ pub unsafe fn free_pages(ptr: NonNull<u8>, count: usize) -> Result {
     unsafe { (bt.free_pages)(addr, count) }.to_result()
 }
 
-/// Allocates from a memory pool. The pointer will be 8-byte aligned.
+/// Allocates a consecutive region of bytes using the UEFI allocator. The buffer
+/// will be 8-byte aligned.
+///
+/// The caller is responsible to free the memory using [`free_pool`].
+///
+/// # Arguments
+/// - `memory_type`: The [`MemoryType`] used to persist the allocation in the
+///   UEFI memory map. Typically, UEFI OS loaders should allocate memory of
+///   type [`MemoryType::LOADER_DATA`].
+///- `size`: Amount of bytes to allocate.
+///
+/// # Example
+///```rust,no_run
+/// use uefi::boot::{self, AllocateType};
+/// use uefi_raw::table::boot::MemoryType;
+///
+/// let mut ptr = boot::allocate_pool(
+///    MemoryType::LOADER_DATA,
+///    42
+/// ).unwrap();
+///
+/// // ⚠️ Creating the reference is safe, but reading the uninitialized memory
+/// // causes Undefined Behavior (UB)! Please make sure to initialize the memory
+/// // first by:
+/// //   - using `core::ptr::write`,
+/// //   - directly writing to slice indices,
+/// //   - zeroing the memory,
+/// //   - using `.copy_from_slice()`,
+/// //   - or a similar operation.
+/// let buffer: &mut [u8] = unsafe { ptr.as_mut() };
+/// // Now initialize the content of the buffer, cast it, etc.
+/// // Please follow Rust guidelines on safety and UB! ⚠️
+///
+/// // free the allocation
+/// unsafe { boot::free_pool(ptr.cast()) }.unwrap();
+/// ```
+///
+/// # Safety
+/// Using this function is safe but reading on initialized memory is not.
+/// Please look into the example code.
+///
 ///
 /// # Errors
 ///
 /// * [`Status::OUT_OF_RESOURCES`]: allocation failed.
 /// * [`Status::INVALID_PARAMETER`]: `mem_ty` is [`MemoryType::PERSISTENT_MEMORY`],
 ///   [`MemoryType::UNACCEPTED`], or in the range [`MemoryType::MAX`]`..=0x6fff_ffff`.
-pub fn allocate_pool(mem_ty: MemoryType, size: usize) -> Result<NonNull<u8>> {
+pub fn allocate_pool(memory_type: MemoryType, size: usize) -> Result<NonNull<[u8]>> {
     let bt = boot_services_raw_panicking();
     let bt = unsafe { bt.as_ref() };
 
     let mut buffer = ptr::null_mut();
-    let ptr =
-        unsafe { (bt.allocate_pool)(mem_ty, size, &mut buffer) }.to_result_with_val(|| buffer)?;
+    let ptr = unsafe { (bt.allocate_pool)(memory_type, size, &mut buffer) }
+        .to_result_with_val(|| buffer)?;
 
-    Ok(NonNull::new(ptr).expect("allocate_pool must not return a null pointer if successful"))
+    if let Some(ptr) = NonNull::new(ptr) {
+        let slice = NonNull::slice_from_raw_parts(ptr, size);
+        Ok(slice)
+    } else {
+        Err(Status::OUT_OF_RESOURCES.into())
+    }
 }
 
 /// Frees memory allocated by [`allocate_pool`].

uefi/src/mem/memory_map/impl_.rs

@@ -281,7 +281,7 @@ impl MemoryMapBackingMemory {
     pub(crate) fn new(memory_type: MemoryType) -> crate::Result<Self> {
         let memory_map_meta = boot::memory_map_size();
         let len = Self::safe_allocation_size_hint(memory_map_meta);
-        let ptr = boot::allocate_pool(memory_type, len)?.as_ptr();
+        let ptr = boot::allocate_pool(memory_type, len)?.cast::<u8>().as_ptr();
 
         // Should be fine as UEFI always has  allocations with a guaranteed
         // alignment of 8 bytes.