Merge pull request rust-osdev#1361 from nicholasbishop/bishop-memmap-…

…update

uefi: Use global system table in MemoryMapBackingMemory

uefi/src/boot.rs

@@ -178,6 +178,51 @@ pub unsafe fn free_pool(ptr: NonNull<u8>) -> Result {
     unsafe { (bt.free_pool)(ptr.as_ptr()) }.to_result()
 }
 
+/// Queries the `get_memory_map` function of UEFI to retrieve the current
+/// size of the map. Returns a [`MemoryMapMeta`].
+///
+/// It is recommended to add a few more bytes for a subsequent allocation
+/// for the memory map, as the memory map itself also needs heap memory,
+/// and other allocations might occur before that call.
+#[must_use]
+pub(crate) fn memory_map_size() -> MemoryMapMeta {
+    let bt = boot_services_raw_panicking();
+    let bt = unsafe { bt.as_ref() };
+
+    let mut map_size = 0;
+    let mut map_key = MemoryMapKey(0);
+    let mut desc_size = 0;
+    let mut desc_version = 0;
+
+    let status = unsafe {
+        (bt.get_memory_map)(
+            &mut map_size,
+            ptr::null_mut(),
+            &mut map_key.0,
+            &mut desc_size,
+            &mut desc_version,
+        )
+    };
+    assert_eq!(status, Status::BUFFER_TOO_SMALL);
+
+    assert_eq!(
+        map_size % desc_size,
+        0,
+        "Memory map must be a multiple of the reported descriptor size."
+    );
+
+    let mmm = MemoryMapMeta {
+        desc_size,
+        map_size,
+        map_key,
+        desc_version,
+    };
+
+    mmm.assert_sanity_checks();
+
+    mmm
+}
+
 /// Stores the current UEFI memory map in an UEFI-heap allocated buffer
 /// and returns a [`MemoryMapOwned`].
 ///

uefi/src/mem/memory_map/impl_.rs

@@ -2,7 +2,7 @@
 //! as well as relevant helper types, such as [`MemoryMapBackingMemory`].
 
 use super::*;
-use crate::table::system_table_boot;
+use crate::boot;
 use core::fmt::{Debug, Display, Formatter};
 use core::ops::{Index, IndexMut};
 use core::ptr::NonNull;
@@ -274,12 +274,9 @@ impl MemoryMapBackingMemory {
     /// - `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) -> crate::Result<Self> {
-        let st = system_table_boot().expect("Should have boot services activated");
-        let bs = st.boot_services();
-
-        let memory_map_meta = bs.memory_map_size();
+        let memory_map_meta = boot::memory_map_size();
         let len = Self::safe_allocation_size_hint(memory_map_meta);
-        let ptr = bs.allocate_pool(memory_type, len)?.as_ptr();
+        let ptr = boot::allocate_pool(memory_type, len)?.as_ptr();
 
         // Should be fine as UEFI always has  allocations with a guaranteed
         // alignment of 8 bytes.
@@ -339,20 +336,15 @@ impl MemoryMapBackingMemory {
 }
 
 // Don't drop when we use this in unit tests.
-#[cfg(not(test))]
 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 boot::are_boot_services_active() {
+            let res = unsafe { boot::free_pool(self.0.cast()) };
             if let Err(e) = res {
                 log::error!("Failed to deallocate memory map: {e:?}");
             }
         } else {
-            #[cfg(test)]
-            log::debug!("Boot services are not available in unit tests.");
-
-            #[cfg(not(test))]
-            log::debug!("Boot services are excited. Memory map won't be freed using the UEFI boot services allocator.");
+            log::debug!("Boot services are exited. Memory map won't be freed using the UEFI boot services allocator.");
         }
     }
 }

uefi/src/table/boot.rs

@@ -168,48 +168,6 @@ impl BootServices {
         unsafe { (self.0.free_pages)(addr, count) }.to_result()
     }
 
-    /// Queries the `get_memory_map` function of UEFI to retrieve the current
-    /// size of the map. Returns a [`MemoryMapMeta`].
-    ///
-    /// It is recommended to add a few more bytes for a subsequent allocation
-    /// for the memory map, as the memory map itself also needs heap memory,
-    /// and other allocations might occur before that call.
-    #[must_use]
-    pub(crate) fn memory_map_size(&self) -> MemoryMapMeta {
-        let mut map_size = 0;
-        let mut map_key = MemoryMapKey(0);
-        let mut desc_size = 0;
-        let mut desc_version = 0;
-
-        let status = unsafe {
-            (self.0.get_memory_map)(
-                &mut map_size,
-                ptr::null_mut(),
-                &mut map_key.0,
-                &mut desc_size,
-                &mut desc_version,
-            )
-        };
-        assert_eq!(status, Status::BUFFER_TOO_SMALL);
-
-        assert_eq!(
-            map_size % desc_size,
-            0,
-            "Memory map must be a multiple of the reported descriptor size."
-        );
-
-        let mmm = MemoryMapMeta {
-            desc_size,
-            map_size,
-            map_key,
-            desc_version,
-        };
-
-        mmm.assert_sanity_checks();
-
-        mmm
-    }
-
     /// Stores the current UEFI memory map in an UEFI-heap allocated buffer
     /// and returns a [`MemoryMapOwned`].
     ///