uefi: introduce MemoryMapBackingMemory helper type

Author: phip1611

uefi/src/table/boot.rs

@@ -1,6 +1,6 @@
 //! UEFI services available during boot.
 
-use super::Revision;
+use super::{system_table_boot, Revision};
 use crate::data_types::{Align, PhysicalAddress};
 use crate::proto::device_path::DevicePath;
 use crate::proto::loaded_image::LoadedImage;
@@ -1615,6 +1615,105 @@ impl Align for MemoryDescriptor {
 #[repr(C)]
 pub struct MemoryMapKey(usize);
 
+/// The backing memory for the UEFI memory app on the UEFI heap, allocated using
+/// the UEFI boot services allocator. This occupied memory will also be
+/// reflected in the memory map itself.
+///
+/// Although untyped, it is similar to the `Box` type in terms of heap
+/// allocation and deallocation, as well as ownership of the corresponding
+/// memory. Apart from that, this type only has the semantics of a buffer.
+///
+/// The memory is untyped, which is necessary due to the nature of the UEFI
+/// spec. It still ensures a correct alignment to hold [`MemoryDescriptor`]. The
+/// size of the buffer is sufficient to hold the memory map at the point in time
+/// where this is created. Note that due to (not obvious or asynchronous)
+/// allocations/allocations in your environment, this might be outdated at the
+/// time you store the memory map in it.
+///
+/// Note that due to the nature of the UEFI memory app, this buffer might
+/// hold (a few) bytes more than necessary. The `map_size` reported by
+/// `get_memory_map` tells the actual size.
+///
+/// When this type is dropped and boot services are not exited yet, the memory
+/// is freed.
+#[derive(Debug)]
+pub struct MemoryMapBackingMemory(NonNull<[u8]> /* buffer on UEFI heap */);
+
+impl MemoryMapBackingMemory {
+    /// Constructs a new [`MemoryMapBackingMemory`].
+    ///
+    /// # Parameters  
+    /// - `memory_type`: The memory type for the memory map allocation.
+    ///   Typically, [`MemoryType::LOADER_DATA`] for regular UEFI applications.
+    pub(crate) fn new(memory_type: MemoryType) -> Result<Self> {
+        let st = system_table_boot().expect("Should have boot services activated");
+        let bs = st.boot_services();
+
+        let memory_map_size = bs.memory_map_size();
+        let alloc_size = Self::allocation_size_hint(memory_map_size);
+        let ptr = bs.allocate_pool(memory_type, alloc_size)?;
+        assert_eq!(ptr.align_offset(mem::align_of::<MemoryDescriptor>()), 0);
+
+        let ptr = NonNull::new(ptr).expect("UEFI should never return a null ptr. An error should have been reflected via an Err earlier.");
+        let slice = NonNull::slice_from_raw_parts(ptr, alloc_size);
+
+        Ok(Self(slice))
+    }
+
+    /// Returns a best-effort size hint of the memory map size. This is
+    /// especially created with exiting boot services in mind.
+    pub fn allocation_size_hint(mms: MemoryMapSize) -> usize {
+        let MemoryMapSize {
+            desc_size,
+            map_size,
+        } = mms;
+
+        // Allocate space for extra entries beyond the current size of the
+        // memory map. The value of 8 matches the value in the Linux kernel:
+        // https://github.com/torvalds/linux/blob/e544a07438/drivers/firmware/efi/libstub/efistub.h#L173
+        let extra_entries = 8;
+
+        assert!(desc_size > 0);
+        // Although very unlikely, this might fail if the memory descriptor is
+        // extended by a future UEFI revision by a significant amount, we
+        // update the struct, but an old UEFI implementation reports a small
+        // size.
+        assert!(desc_size >= mem::size_of::<MemoryDescriptor>());
+        assert!(map_size > 0);
+
+        // Ensure the mmap size is (somehow) sane.
+        const ONE_GB: usize = 1024 * 1024 * 1024;
+        assert!(map_size <= ONE_GB);
+
+        let extra_size = desc_size * extra_entries;
+        let allocation_size = map_size + extra_size;
+        allocation_size
+    }
+    
+    /// Returns the raw pointer to the beginning of the allocation.
+    pub fn as_ptr_mut(&mut self) -> *mut u8 {
+        self.0.as_ptr().cast()
+    }
+    
+    /// Returns the length of the allocation.
+    pub fn len(&self) -> usize {
+        self.0.len()
+    } 
+}
+
+impl Drop for MemoryMapBackingMemory {
+    fn drop(&mut self) {
+        if let Some(bs) = system_table_boot() {
+            let res = unsafe { bs.boot_services().free_pool(self.0.as_ptr().cast()) };
+            if let Err(e) = res {
+                log::error!("Failed to deallocate memory map: {e:?}");
+            }
+        } else {
+            log::debug!("Boot services are excited. Memory map won't be freed using the UEFI boot services allocator.");
+        }
+    }
+}
+
 /// A structure containing the size of a memory descriptor and the size of the
 /// memory map.
 #[derive(Debug)]