From 9626dfb4e1123d75c3efd1bc8cd085e2130bea76 Mon Sep 17 00:00:00 2001
From: Philipp Schuster <phip1611@gmail.com>
Date: Sun, 6 Apr 2025 19:56:20 +0200
Subject: [PATCH 1/2] uefi/doc: improve documentation of exit_boot_services

TL;DR: Document why exit_boot_services needs to do more than just
       call the UEFI function.

Exiting UEFI boot services requires a non-trivial sequence of steps,
including safe retrieval and finalization of the memory map. To clarify
why our exit_boot_services function is more than a thin wrapper, the
documentation has been updated accordingly.

This complexity is also reflected in the Linux source code [0].
In fact, the sequence of steps uefi-rs uses is inspired by the Linux
source code.

[0] https://github.com/torvalds/linux/blob/e544a0743/drivers/firmware/efi/libstub/efi-stub-helper.c#L375
---
 uefi/src/boot.rs | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/uefi/src/boot.rs b/uefi/src/boot.rs
index 4c44bef55..1f0f7b283 100644
--- a/uefi/src/boot.rs
+++ b/uefi/src/boot.rs
@@ -1258,17 +1258,23 @@ unsafe fn get_memory_map_and_exit_boot_services(buf: &mut [u8]) -> Result<Memory
         .to_result_with_val(|| memory_map)
 }
 
-/// Exit UEFI boot services.
+/// Convenient wrapper to exit UEFI boot services along with corresponding
+/// essential steps to get the memory map.
+///
+/// Exiting UEFI boot services requires a **non-trivial sequence of steps**,
+/// including safe retrieval and finalization of the memory map. This wrapper
+/// ensures a safe and spec-compliant transition from UEFI boot services to
+/// runtime services by retrieving the system memory map and invoking
+/// `ExitBootServices()` with the correct memory map key, retrying if necessary.
 ///
 /// After this function completes, UEFI hands over control of the hardware
 /// to the executing OS loader, which implies that the UEFI boot services
 /// are shut down and cannot be used anymore. Only UEFI configuration tables
-/// and run-time services can be used.
+/// and runtime services can be used.
 ///
-/// The memory map at the time of exiting boot services returned. The map is
-/// backed by a pool allocation of the given `memory_type`. Since the boot
-/// services function to free that memory is no longer available after calling
-/// `exit_boot_services`, the allocation will not be freed on drop.
+/// Since the boot services function to free memory is no longer available after
+/// this function returns, the allocation will not be freed on drop. It will
+/// however be reflected by the memory map itself (self-contained).
 ///
 /// Note that once the boot services are exited, associated loggers and
 /// allocators can't use the boot services anymore. For the corresponding

From 4752b4eecca659375a573ed15158c6c696b17104 Mon Sep 17 00:00:00 2001
From: Philipp Schuster <phip1611@gmail.com>
Date: Sun, 6 Apr 2025 20:08:25 +0200
Subject: [PATCH 2/2] uefi: change signature of exit_boot_services

People probably want to use 99% of the time LOADER_DATA here.
Therefore, we indicate in our function signature, that people should
use the default (`None`) if they are unsure what to do here.

LOADER_DATA is also what Linux uses [0].

[0] https://elixir.bootlin.com/linux/v6.13.7/source/drivers/firmware/efi/libstub/mem.c#L24
---
 uefi-test-runner/src/main.rs |  4 ++--
 uefi/CHANGELOG.md            |  2 ++
 uefi/src/boot.rs             | 12 +++++++++++-
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/uefi-test-runner/src/main.rs b/uefi-test-runner/src/main.rs
index bcbbb0f8e..2e45d8438 100644
--- a/uefi-test-runner/src/main.rs
+++ b/uefi-test-runner/src/main.rs
@@ -10,7 +10,7 @@ extern crate alloc;
 
 use alloc::string::ToString;
 use alloc::vec::Vec;
-use uefi::mem::memory_map::{MemoryMap, MemoryType};
+use uefi::mem::memory_map::MemoryMap;
 use uefi::prelude::*;
 use uefi::proto::console::serial::Serial;
 use uefi::proto::device_path::build::{self, DevicePathBuilder};
@@ -209,7 +209,7 @@ fn shutdown() -> ! {
     info!("Testing complete, exiting boot services...");
 
     // Exit boot services as a proof that it works :)
-    let mmap = unsafe { uefi::boot::exit_boot_services(MemoryType::LOADER_DATA) };
+    let mmap = unsafe { uefi::boot::exit_boot_services(None) };
 
     info!("Memory Map:");
     for desc in mmap.entries() {
diff --git a/uefi/CHANGELOG.md b/uefi/CHANGELOG.md
index dfb07a075..0d059b9d9 100644
--- a/uefi/CHANGELOG.md
+++ b/uefi/CHANGELOG.md
@@ -16,6 +16,8 @@
   `&self`.
 - **Breaking:** The `pxe::Mode` struct is now opaque. Use method calls to access
   mode data instead of direct field access.
+- **Breaking:** `exit_boot_services` now consumes a `Option<MemoryType>` which
+  defaults to the recommended value of `MemoryType::LOADER_DATA`.
 - `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 --git a/uefi/src/boot.rs b/uefi/src/boot.rs
index 1f0f7b283..d36e59dae 100644
--- a/uefi/src/boot.rs
+++ b/uefi/src/boot.rs
@@ -1283,6 +1283,13 @@ unsafe fn get_memory_map_and_exit_boot_services(buf: &mut [u8]) -> Result<Memory
 /// `global_allocator` feature is enabled, attempting to use the allocator
 /// after exiting boot services will panic.
 ///
+/// # Arguments
+/// - `custom_memory_type`: The [`MemoryType`] for the UEFI allocation that will
+///   store the final memory map. If you pass `None`, this defaults to the
+///   recommended default value of [`MemoryType::LOADER_DATA`]. If you want a
+///   specific memory region for the memory map, you can pass the desired
+///   [`MemoryType`].
+///
 /// # Safety
 ///
 /// The caller is responsible for ensuring that no references to
@@ -1313,7 +1320,10 @@ unsafe fn get_memory_map_and_exit_boot_services(buf: &mut [u8]) -> Result<Memory
 /// [`Output`]: crate::proto::console::text::Output
 /// [`PoolString`]: crate::proto::device_path::text::PoolString
 #[must_use]
-pub unsafe fn exit_boot_services(memory_type: MemoryType) -> MemoryMapOwned {
+pub unsafe fn exit_boot_services(custom_memory_type: Option<MemoryType>) -> MemoryMapOwned {
+    // LOADER_DATA is the default and also used by the Linux kernel:
+    // https://elixir.bootlin.com/linux/v6.13.7/source/drivers/firmware/efi/libstub/mem.c#L24
+    let memory_type = custom_memory_type.unwrap_or(MemoryType::LOADER_DATA);
     crate::helpers::exit();
 
     let mut buf = MemoryMapBackingMemory::new(memory_type).expect("Failed to allocate memory");