From d42a3bfef68b274bae986ee550c3b29ae1fb80e7 Mon Sep 17 00:00:00 2001
From: Nicholas Bishop <nbishop@nbishop.net>
Date: Sat, 1 Mar 2025 17:48:45 -0500
Subject: [PATCH] uefi: Improve handling of null-address allocations in
 allocate_pages

The firmware is allowed to return an allocation at address zero. This is not
compatible with Rust, since it's UB to write through a null pointer. If a null
address is returned, retry the allocation. If that second allocation still
fails, return an error rather than panicking.
---
 uefi/CHANGELOG.md |  3 +++
 uefi/src/boot.rs  | 35 ++++++++++++++++++++++++++++++-----
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/uefi/CHANGELOG.md b/uefi/CHANGELOG.md
index 562fa77ce..b609e59d8 100644
--- a/uefi/CHANGELOG.md
+++ b/uefi/CHANGELOG.md
@@ -13,6 +13,9 @@
 - `SimpleNetwork::transmit` now passes the correct buffer size argument.
   Previously it incorrectly added the header size to the buffer length, which
   could cause the firmware to read past the end of the buffer.
+- `boot::allocate_pages` no longer panics if the allocation is at address
+  zero. The allocation is retried instead, and in all failure cases an error is
+  returned rather than panicking.
 
 
 # uefi - 0.34.1 (2025-02-07)
diff --git a/uefi/src/boot.rs b/uefi/src/boot.rs
index 2c3978165..e5de81f53 100644
--- a/uefi/src/boot.rs
+++ b/uefi/src/boot.rs
@@ -134,15 +134,40 @@ pub fn allocate_pages(ty: AllocateType, mem_ty: MemoryType, count: usize) -> Res
     let bt = boot_services_raw_panicking();
     let bt = unsafe { bt.as_ref() };
 
-    let (ty, mut addr) = match ty {
+    let (ty, initial_addr) = match ty {
         AllocateType::AnyPages => (0, 0),
         AllocateType::MaxAddress(addr) => (1, addr),
         AllocateType::Address(addr) => (2, addr),
     };
-    let addr =
-        unsafe { (bt.allocate_pages)(ty, mem_ty, count, &mut addr) }.to_result_with_val(|| addr)?;
-    let ptr = addr as *mut u8;
-    Ok(NonNull::new(ptr).expect("allocate_pages must not return a null pointer if successful"))
+
+    let mut addr1 = initial_addr;
+    unsafe { (bt.allocate_pages)(ty, mem_ty, 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);
+    }
+
+    // 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();
+
+    // Free the original allocation (ignoring errors).
+    let _unused = unsafe { (bt.free_pages)(addr1, count) };
+
+    // Return an error if the second allocation failed, or if it is still at
+    // address zero. Otherwise, return a pointer to the second allocation.
+    r?;
+    if let Some(ptr) = NonNull::new(addr2 as *mut u8) {
+        Ok(ptr)
+    } else {
+        Err(Status::OUT_OF_RESOURCES.into())
+    }
 }
 
 /// Frees memory pages allocated by [`allocate_pages`].