uefi: allocate_pages() now returns NonNull<[u8]>

This aligns the signature with the Rust allocator API.

Author: phip1611

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,20 +18,28 @@ 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() {

uefi/CHANGELOG.md

@@ -16,6 +16,9 @@
   `&self`.
 - **Breaking:** The `pxe::Mode` struct is now opaque. Use method calls to access
   mode data instead of direct field access.
+- **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.
 - `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/boot.rs

@@ -120,42 +120,70 @@ 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();
+///
+/// let buffer: &mut [u8] = unsafe { ptr.as_mut() };
+/// // now do something with your buffer
+///
+/// // free the allocation
+/// unsafe { boot::free_pages(ptr.cast(), num_pages) }.unwrap();
+/// ```
+///
 /// # 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 +192,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())
     }
@@ -191,6 +220,31 @@ pub unsafe fn free_pages(ptr: NonNull<u8>, count: usize) -> Result {
 
 /// Allocates from a memory pool. The pointer 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();
+///
+/// let buffer: &mut [u8] = unsafe { ptr.as_mut() };
+/// // now do something with your buffer
+///
+/// // free the allocation
+/// unsafe { boot::free_pool(ptr.cast()) }.unwrap();
+/// ```
+///
 /// # Errors
 ///
 /// * [`Status::OUT_OF_RESOURCES`]: allocation failed.