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

Author: phip1611

Diff for: 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 for: 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 for: uefi/src/boot.rs

@@ -1260,7 +1260,7 @@ unsafe fn get_memory_map_and_exit_boot_services(buf: &mut [u8]) -> Result<Memory
 
 /// Convenient wrapper to exit UEFI boot services along with corresponding
 /// essential steps to get the memory map.
-/// 
+///
 /// This wrapper ensures a safe and spec-compliant transition from UEFI boot
 /// services phase to the runtime services phase by retrieving the system
 /// memory map and invoking `ExitBootServices()` with the correct memory map
@@ -1282,6 +1282,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
@@ -1312,7 +1319,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");