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

This aligns the signature with the Rust allocator API.

Author: phip1611

Diff for: uefi-test-runner/src/boot/memory.rs

@@ -43,15 +43,17 @@ fn test_allocate_pages() {
 }
 
 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.

Diff for: 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();
     }
 }

Diff for: uefi/CHANGELOG.md

@@ -19,6 +19,9 @@
 - **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.

Diff for: 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())
         }
     }

Diff for: uefi/src/boot.rs

@@ -218,7 +218,8 @@ 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`].
 ///
@@ -250,15 +251,20 @@ pub unsafe fn free_pages(ptr: NonNull<u8>, count: usize) -> Result {
 /// * [`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)?;
+        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`].

Diff for: 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.