1use 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
13pub fn exit_boot_services() -> Result<MemoryMapInfo, Status> {
15    uefi::println!("Exiting boot services ...");
16    qemu_trace!("Exiting boot services ...\n");
17
18    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    let owned_map = unsafe { boot::exit_boot_services(None) };
29
30    let src = owned_map.buffer().as_ptr();
32    let mmap_length = owned_map.buffer().len();
33
34    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    core::mem::forget(mmap_copy);
56
57    qemu_trace!("Boot services exited, we're now flying by instruments.\n");
58    Ok(mmap)
59}
60
61fn allocate_mmap_buffer() -> Result<Vec<u8>, Status> {
70    const EXTRA_DESCS: usize = 32;
71
72    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    drop(probe);
86
87    needed_size += EXTRA_DESCS * desc_size;
90
91    let buf = vec![0u8; needed_size];
93    Ok(buf)
94}