From 2f4aece6c2f29da16c82cfc2040085a1aba07dd8 Mon Sep 17 00:00:00 2001
From: Philipp Schuster <phip1611@gmail.com>
Date: Sat, 25 May 2024 15:59:31 +0200
Subject: [PATCH] uefi: BootServices::allocate_pool now returns NonZero<u8>
 instead of *mut u8

This is safe as each null pointer (=failed allocation) is already returned as
non-successful status code by UEFI.
---
 uefi-test-runner/src/boot/misc.rs | 6 ++++--
 uefi/CHANGELOG.md                 | 2 ++
 uefi/src/allocator.rs             | 3 ++-
 uefi/src/table/boot.rs            | 8 ++++++--
 uefi/src/table/system.rs          | 2 +-
 5 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/uefi-test-runner/src/boot/misc.rs b/uefi-test-runner/src/boot/misc.rs
index 0399ff2ba..8b7422be4 100644
--- a/uefi-test-runner/src/boot/misc.rs
+++ b/uefi-test-runner/src/boot/misc.rs
@@ -121,7 +121,8 @@ fn test_install_protocol_interface(bt: &BootServices) {
             mem::size_of::<TestProtocol>(),
         )
         .unwrap()
-        .cast();
+        .cast()
+        .as_ptr();
     unsafe { alloc.write(TestProtocol { data: 123 }) };
 
     let _ = unsafe {
@@ -187,7 +188,8 @@ fn test_install_configuration_table(st: &SystemTable<Boot>) {
     let config = st
         .boot_services()
         .allocate_pool(MemoryType::ACPI_RECLAIM, 1)
-        .expect("Failed to allocate config table");
+        .expect("Failed to allocate config table")
+        .as_ptr();
     unsafe { config.write(42) };
 
     let count = st.config_table().len();
diff --git a/uefi/CHANGELOG.md b/uefi/CHANGELOG.md
index d7c79efdc..0f01003eb 100644
--- a/uefi/CHANGELOG.md
+++ b/uefi/CHANGELOG.md
@@ -15,6 +15,8 @@
 ## Changed
 - `SystemTable::exit_boot_services` is now `unsafe`. See that method's
   documentation for details of obligations for callers.
+- `BootServices::allocate_pool` now returns `NonZero<u8>` instead of
+  `*mut u8`.
 
 ## Removed
 - Removed the `panic-on-logger-errors` feature of the `uefi` crate. Logger
diff --git a/uefi/src/allocator.rs b/uefi/src/allocator.rs
index 3be75c164..d1ef3dc93 100644
--- a/uefi/src/allocator.rs
+++ b/uefi/src/allocator.rs
@@ -86,7 +86,7 @@ unsafe impl GlobalAlloc for Allocator {
             // within the allocation.
             let full_alloc_ptr =
                 if let Ok(ptr) = boot_services.allocate_pool(memory_type, size + align) {
-                    ptr
+                    ptr.as_ptr()
                 } else {
                     return ptr::null_mut();
                 };
@@ -116,6 +116,7 @@ unsafe impl GlobalAlloc for Allocator {
             // use `allocate_pool` directly.
             boot_services
                 .allocate_pool(memory_type, size)
+                .map(|ptr| ptr.as_ptr())
                 .unwrap_or(ptr::null_mut())
         }
     }
diff --git a/uefi/src/table/boot.rs b/uefi/src/table/boot.rs
index 6d9b74289..1a5e44876 100644
--- a/uefi/src/table/boot.rs
+++ b/uefi/src/table/boot.rs
@@ -271,9 +271,13 @@ impl BootServices {
     ///
     /// * [`uefi::Status::OUT_OF_RESOURCES`]
     /// * [`uefi::Status::INVALID_PARAMETER`]
-    pub fn allocate_pool(&self, mem_ty: MemoryType, size: usize) -> Result<*mut u8> {
+    pub fn allocate_pool(&self, mem_ty: MemoryType, size: usize) -> Result<NonNull<u8>> {
         let mut buffer = ptr::null_mut();
-        unsafe { (self.0.allocate_pool)(mem_ty, size, &mut buffer) }.to_result_with_val(|| buffer)
+        let ptr = unsafe { (self.0.allocate_pool)(mem_ty, size, &mut buffer) }
+            .to_result_with_val(|| buffer)?;
+
+        Ok(NonNull::new(ptr)
+            .expect("UEFI should return error if an allocation failed but never a null pointer"))
     }
 
     /// Frees memory allocated from a pool.
diff --git a/uefi/src/table/system.rs b/uefi/src/table/system.rs
index ab7cb3305..87ed22a55 100644
--- a/uefi/src/table/system.rs
+++ b/uefi/src/table/system.rs
@@ -265,7 +265,7 @@ impl SystemTable<Boot> {
         // Allocate a byte slice to hold the memory map. If the
         // allocation fails treat it as an unrecoverable error.
         let buf: *mut u8 = match boot_services.allocate_pool(memory_type, buf_size) {
-            Ok(buf) => buf,
+            Ok(buf) => buf.as_ptr(),
             Err(err) => reset(err.status()),
         };