uefi_loader/
uefi_mmap.rs

1//! # UEFI Memory Map Utilities
2//!
3//! Helper functions for dealing with the UEFI memory map after exiting boot services.
4
5use alloc::vec;
6use alloc::vec::Vec;
7use kernel_info::boot::MemoryMapInfo;
8use kernel_qemu::qemu_trace;
9use uefi::boot::MemoryType;
10use uefi::mem::memory_map::MemoryMap;
11use uefi::{Status, boot};
12
13/// Exist the UEFI boot services and retain a copy of the UEFI memory map.
14pub fn exit_boot_services() -> Result<MemoryMapInfo, Status> {
15    uefi::println!("Exiting boot services ...");
16    qemu_trace!("Exiting boot services ...\n");
17
18    // Pre-allocate a buffer while UEFI allocator is still alive.
19    let mut mmap_copy = match allocate_mmap_buffer() {
20        Ok(buf) => buf,
21        Err(status) => {
22            return Err(status);
23        }
24    };
25    let mmap_copy_ptr = mmap_copy.as_mut_ptr();
26
27    // Exit boot services — after this, the UEFI allocator must not be used anymore.
28    let owned_map = unsafe { boot::exit_boot_services(None) };
29
30    // Copy the returned descriptors into our preallocated buffer.
31    let src = owned_map.buffer().as_ptr();
32    let mmap_length = owned_map.buffer().len();
33
34    // Safety: ensure the buffer is large enough (or bail/panic in dev builds).
35    if mmap_length > mmap_copy.len() {
36        qemu_trace!(
37            "Memory map size assertion failed: Expected {}, got {}",
38            mmap_copy.len(),
39            mmap_length
40        );
41        return Err(Status::BUFFER_TOO_SMALL);
42    }
43    unsafe {
44        core::ptr::copy_nonoverlapping(src, mmap_copy_ptr, mmap_length);
45    }
46
47    let mmap = MemoryMapInfo {
48        mmap_ptr: mmap_copy_ptr as u64,
49        mmap_len: mmap_length as u64,
50        mmap_desc_size: owned_map.meta().desc_size as u64,
51        mmap_desc_version: owned_map.meta().desc_version,
52    };
53
54    // Ensure the memory map copy continues to exist.
55    core::mem::forget(mmap_copy);
56
57    qemu_trace!("Boot services exited, we're now flying by instruments.\n");
58    Ok(mmap)
59}
60
61/// Allocate a buffer to hold a copy of the memory map returned from `ExitBootServices`.
62///
63/// This seems to be the opposite of an exact science:
64/// * After boot services were exited, allocation is impossible.
65/// * The number of descriptors changes over time.
66///
67/// As a result, we now overallocate to hopefully have enough headroom
68/// to contain the memory map _after_ exiting.
69fn allocate_mmap_buffer() -> Result<Vec<u8>, Status> {
70    const EXTRA_DESCS: usize = 32;
71
72    // Introspect the memory map.
73    let probe = match boot::memory_map(MemoryType::LOADER_DATA) {
74        Ok(probe) => probe,
75        Err(e) => {
76            uefi::println!("Failed to get memory map: {e:?}");
77            return Err(Status::UNSUPPORTED);
78        }
79    };
80
81    let desc_size = probe.meta().desc_size;
82    let mut needed_size = probe.meta().map_size;
83
84    // We won't use `probe`'s buffer; drop it now to reduce churn.
85    drop(probe);
86
87    // Pre-allocate our own buffer with slack for extra descriptors.
88    // Rule of thumb: + N * desc_size; N=16..64 is usually plenty in QEMU/OVMF.
89    needed_size += EXTRA_DESCS * desc_size;
90
91    // Pre-allocate a buffer while UEFI allocator is still alive.
92    let buf = vec![0u8; needed_size];
93    Ok(buf)
94}