uefi: MemoryMap now owns the memory

This is an attempt to simplify the overall complex handling of obtaining the UEFI memory
map. We have the following pre-requisites and use-cases all to keep in mind when
designing the functions and associated helper types:
- the memory map itself needs memory; typically on the UEFI heap
- acquiring that memory and storing the memory map inside it are two distinct steps
- the memory map is not just a slice of [MemoryDescriptor] (desc_size is bigger than
  size_of)
- the required size can be obtained by another boot service function
- the needed memory might change due to hidden or asynchronous allocations
  between the allocation of a buffer and storing the memory map inside it
- when boot services are excited, best practise has shown (looking at linux code)
  that one should use the same buffer (with some extra capacity) and call
  exit_boot_services with that buffer at most two times in a loop

This makes it hard to come up with an ergonomic solution such as using a Box
or any other high-level Rust type.

The main simplification of my design is that the MemoryMap type now doesn't
has a reference to memory anymore but actually owns it. This also models the
real world use case where one typically obtains the memory map once when
boot services are exited. A &'static [u8] on the MemoryMap just creates more
confusion that it brings any benefit.

The MemoryMap now knows whether boot services are still active and frees
that memory, or it doesn't if the boot services are exited. This means
less fiddling with life-times and less cognitive overhead when
- reading the code
- calling BootService::memory_map independently of exit_boot_services

Author: phip1611

uefi-test-runner/src/boot/memory.rs

@@ -63,17 +63,8 @@ fn alloc_alignment() {
 fn memory_map(bt: &BootServices) {
     info!("Testing memory map functions");
 
-    // Get the memory descriptor size and an estimate of the memory map size
-    let sizes = bt.memory_map_size();
-
-    // 2 extra descriptors should be enough.
-    let buf_sz = sizes.map_size + 2 * sizes.desc_size;
-
-    // We will use vectors for convenience.
-    let mut buffer = vec![0_u8; buf_sz];
-
     let mut memory_map = bt
-        .memory_map(&mut buffer)
+        .memory_map(MemoryType::LOADER_DATA)
         .expect("Failed to retrieve UEFI memory map");
 
     memory_map.sort();

uefi/src/table/boot.rs

@@ -215,37 +215,46 @@ impl BootServices {
         }
     }
 
-    /// Stores the current UEFI memory map in the provided buffer.
-    ///
-    /// The allocated buffer must be at least aligned to a [`MemoryDescriptor`]
-    /// and should be big enough to store the whole map. To estimating how big
-    /// the map will be, you can call [`Self::memory_map_size`].
-    ///
-    /// The memory map contains entries of type [`MemoryDescriptor`]. However,
-    /// the relevant step size is always the reported `desc_size` but never
-    /// `size_of::<MemoryDescriptor>()`.
-    ///
-    /// The returned key is a unique identifier of the current configuration of
-    /// memory. Any allocations or such will change the memory map's key.
-    ///
-    /// If you want to store the resulting memory map without having to keep
-    /// the buffer around, you can use `.copied().collect()` on the iterator.
-    /// Note that this will change the current memory map again, if the UEFI
-    /// allocator is used under the hood.
+    /// Stores the current UEFI memory map in an UEFI-heap allocated buffer
+    /// and returns a [`MemoryMap`].
     ///
     /// # Errors
     ///
-    /// See section `EFI_BOOT_SERVICES.GetMemoryMap()` in the UEFI Specification for more details.
+    /// See section `EFI_BOOT_SERVICES.GetMemoryMap()` in the UEFI Specification
+    /// for more details.
     ///
     /// * [`uefi::Status::BUFFER_TOO_SMALL`]
     /// * [`uefi::Status::INVALID_PARAMETER`]
-    pub fn memory_map<'buf>(&self, buffer: &'buf mut [u8]) -> Result<MemoryMap<'buf>> {
-        let mut map_size = buffer.len();
-        MemoryDescriptor::assert_aligned(buffer);
-        let map_buffer = buffer.as_mut_ptr().cast::<MemoryDescriptor>();
+    pub fn memory_map(&self, mt: MemoryType) -> Result<MemoryMap> {
+        let mut buffer = MemoryMapBackingMemory::new(mt)?;
+
+        let GetMemoryMapResult {
+            map_size,
+            map_key,
+            desc_size,
+            desc_version,
+        } = self.get_memory_map(buffer.as_mut_slice())?;
+
+        let len = map_size / desc_size;
+        assert_eq!(map_size % desc_size, 0);
+        assert_eq!(desc_version, MemoryDescriptor::VERSION);
+        Ok(MemoryMap {
+            key: map_key,
+            buf: buffer,
+            desc_size,
+            len,
+        })
+    }
+
+    /// Calls the underlying `GetMemoryMap` function of UEFI. On success,
+    /// the buffer is mutated and contains the map. The map might be shorter
+    /// than the buffer, which is reflected by the return value.
+    pub(crate) fn get_memory_map(&self, buf: &mut [u8]) -> Result<GetMemoryMapResult> {
+        let mut map_size = buf.len();
+        let map_buffer = buf.as_mut_ptr().cast::<MemoryDescriptor>();
         let mut map_key = MemoryMapKey(0);
         let mut desc_size = 0;
-        let mut entry_version = 0;
+        let mut desc_version = 0;
 
         assert_eq!(
             (map_buffer as usize) % mem::align_of::<MemoryDescriptor>(),
@@ -259,18 +268,14 @@ impl BootServices {
                 map_buffer,
                 &mut map_key.0,
                 &mut desc_size,
-                &mut entry_version,
+                &mut desc_version,
             )
         }
-        .to_result_with_val(move || {
-            let len = map_size / desc_size;
-
-            MemoryMap {
-                key: map_key,
-                buf: buffer,
-                desc_size,
-                len,
-            }
+        .to_result_with_val(|| GetMemoryMapResult {
+            map_size,
+            desc_size,
+            map_key,
+            desc_version,
         })
     }
 
@@ -1637,31 +1642,45 @@ pub struct MemoryMapKey(usize);
 /// 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 */);
+pub struct MemoryMapBackingMemory(NonNull<[u8]>);
 
 impl MemoryMapBackingMemory {
     /// Constructs a new [`MemoryMapBackingMemory`].
     ///
-    /// # Parameters  
+    /// # 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)?;
+        let len = Self::allocation_size_hint(memory_map_size);
+        let ptr = bs.allocate_pool(memory_type, len)?;
+        assert_eq!(ptr.align_offset(mem::align_of::<MemoryDescriptor>()), 0);
+        Ok(Self::from_raw(ptr, len))
+    }
+
+    fn from_raw(ptr: *mut u8, len: usize) -> Self {
         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);
+        let slice = NonNull::slice_from_raw_parts(ptr, len);
 
-        Ok(Self(slice))
+        Self(slice)
+    }
+
+    /// Creates an instance from the provided memory, which is not necessarily
+    /// on the UEFI heap.
+    #[cfg(test)]
+    fn from_slice(buffer: &mut [u8]) -> Self {
+        let len = buffer.len();
+        Self::from_raw(buffer.as_mut_ptr(), len)
     }
 
     /// Returns a best-effort size hint of the memory map size. This is
     /// especially created with exiting boot services in mind.
+    #[must_use]
     pub fn allocation_size_hint(mms: MemoryMapSize) -> usize {
         let MemoryMapSize {
             desc_size,
@@ -1689,16 +1708,34 @@ impl MemoryMapBackingMemory {
         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 {
+
+    /// Returns a raw pointer to the beginning of the allocation.
+    #[must_use]
+    pub fn as_ptr(&self) -> *const u8 {
         self.0.as_ptr().cast()
     }
-    
+
+    /// Returns a mutable raw pointer to the beginning of the allocation.
+    #[must_use]
+    pub fn as_mut_ptr(&mut self) -> *mut u8 {
+        self.0.as_ptr().cast()
+    }
+
     /// Returns the length of the allocation.
+    #[must_use]
     pub fn len(&self) -> usize {
         self.0.len()
-    } 
+    }
+
+    /*#[must_use]
+    pub fn as_slice(&self) -> &[u8] {
+        self.0.as_ref()
+    }*/
+
+    #[must_use]
+    pub fn as_mut_slice(&mut self) -> &mut [u8] {
+        unsafe { self.0.as_mut() }
+    }
 }
 
 impl Drop for MemoryMapBackingMemory {
@@ -1724,6 +1761,14 @@ pub struct MemoryMapSize {
     pub map_size: usize,
 }
 
+#[derive(Copy, Clone, Debug)]
+pub struct GetMemoryMapResult {
+    pub map_size: usize,
+    pub desc_size: usize,
+    pub map_key: MemoryMapKey,
+    pub desc_version: u32,
+}
+
 /// An accessory to the memory map that can be either iterated or
 /// indexed like an array.
 ///
@@ -1739,35 +1784,36 @@ pub struct MemoryMapSize {
 /// always use `entry_size` as step-size when interfacing with the memory map on
 /// a low level.
 ///
-///
-///
 /// [0]: https://github.com/tianocore/edk2/blob/7142e648416ff5d3eac6c6d607874805f5de0ca8/MdeModulePkg/Core/PiSmmCore/Page.c#L1059
 #[derive(Debug)]
-pub struct MemoryMap<'buf> {
+pub struct MemoryMap {
+    /// Backing memory, properly initialized at this point.
+    buf: MemoryMapBackingMemory,
     key: MemoryMapKey,
-    buf: &'buf mut [u8],
     /// Usually bound to the size of a [`MemoryDescriptor`] but can indicate if
     /// this field is ever extended by a new UEFI standard.
     desc_size: usize,
     len: usize,
 }
 
-impl<'buf> MemoryMap<'buf> {
-    /// Creates a [`MemoryMap`] from the given buffer and entry size.
-    /// The entry size is usually bound to the size of a [`MemoryDescriptor`]
-    /// but can indicate if this field is ever extended by a new UEFI standard.
-    ///
-    /// This allows parsing a memory map provided by a kernel after boot
-    /// services have already exited.
-    pub fn from_raw(buf: &'buf mut [u8], desc_size: usize) -> Self {
-        assert!(!buf.is_empty());
-        assert_eq!(
-            buf.len() % desc_size,
-            0,
-            "The buffer length must be a multiple of the desc_size"
-        );
+impl MemoryMap {
+    /*pub fn new() -> Self {
+
+    }*/
+
+    /// Creates a [`MemoryMap`] from the give initialized memory map behind
+    /// the buffer and the reported `desc_size` from UEFI.
+    pub(crate) fn from_initialized_mem(
+        buf: MemoryMapBackingMemory,
+        props: GetMemoryMapResult,
+    ) -> Self {
+        let GetMemoryMapResult {
+            map_size,
+            desc_size,
+            ..
+        } = props;
         assert!(desc_size >= mem::size_of::<MemoryDescriptor>());
-        let len = buf.len() / desc_size;
+        let len = map_size / desc_size;
         MemoryMap {
             key: MemoryMapKey(0),
             buf,
@@ -1776,6 +1822,20 @@ impl<'buf> MemoryMap<'buf> {
         }
     }
 
+    #[cfg(test)]
+    fn from_raw(buf: &mut [u8], desc_size: usize) -> Self {
+        let mem = MemoryMapBackingMemory::from_slice(buf);
+        Self::from_initialized_mem(
+            mem,
+            GetMemoryMapResult {
+                map_size: buf.len(),
+                desc_size,
+                map_key: MemoryMapKey(0),
+                desc_version: MemoryDescriptor::VERSION,
+            },
+        )
+    }
+
     #[must_use]
     /// Returns the unique [`MemoryMapKey`] associated with the memory map.
     pub fn key(&self) -> MemoryMapKey {
@@ -1870,7 +1930,7 @@ impl<'buf> MemoryMap<'buf> {
 
     /// Returns a reference to the [`MemoryDescriptor`] at `index` or `None` if out of bounds.
     #[must_use]
-    pub fn get(&self, index: usize) -> Option<&'buf MemoryDescriptor> {
+    pub fn get(&self, index: usize) -> Option<&MemoryDescriptor> {
         if index >= self.len {
             return None;
         }
@@ -1888,7 +1948,7 @@ impl<'buf> MemoryMap<'buf> {
 
     /// Returns a mut reference to the [`MemoryDescriptor`] at `index` or `None` if out of bounds.
     #[must_use]
-    pub fn get_mut(&mut self, index: usize) -> Option<&'buf mut MemoryDescriptor> {
+    pub fn get_mut(&mut self, index: usize) -> Option<&mut MemoryDescriptor> {
         if index >= self.len {
             return None;
         }
@@ -1905,15 +1965,15 @@ impl<'buf> MemoryMap<'buf> {
     }
 }
 
-impl core::ops::Index<usize> for MemoryMap<'_> {
+impl core::ops::Index<usize> for MemoryMap {
     type Output = MemoryDescriptor;
 
     fn index(&self, index: usize) -> &Self::Output {
         self.get(index).unwrap()
     }
 }
 
-impl core::ops::IndexMut<usize> for MemoryMap<'_> {
+impl core::ops::IndexMut<usize> for MemoryMap {
     fn index_mut(&mut self, index: usize) -> &mut Self::Output {
         self.get_mut(index).unwrap()
     }
@@ -1922,13 +1982,13 @@ impl core::ops::IndexMut<usize> for MemoryMap<'_> {
 /// An iterator of [`MemoryDescriptor`]. The underlying memory map is always
 /// associated with a unique [`MemoryMapKey`].
 #[derive(Debug, Clone)]
-pub struct MemoryMapIter<'buf> {
-    memory_map: &'buf MemoryMap<'buf>,
+pub struct MemoryMapIter<'a> {
+    memory_map: &'a MemoryMap,
     index: usize,
 }
 
-impl<'buf> Iterator for MemoryMapIter<'buf> {
-    type Item = &'buf MemoryDescriptor;
+impl<'a> Iterator for MemoryMapIter<'a> {
+    type Item = &'a MemoryDescriptor;
 
     fn size_hint(&self) -> (usize, Option<usize>) {
         let sz = self.memory_map.len - self.index;
@@ -2179,7 +2239,7 @@ mod tests {
     }
 
     // Added for debug purposes on test failure
-    impl core::fmt::Display for MemoryMap<'_> {
+    impl core::fmt::Display for MemoryMap {
         fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
             writeln!(f)?;
             for desc in self.entries() {