Merge pull request #1175 from phip1611/mmap

memory map: improvements to fight pitfalls (MemoryMap now owns the memory)

Author: nicholasbishop

uefi-services/README.md

@@ -1,4 +1,4 @@
 # uefi-services
 
 WARNING: `uefi-services` is deprecated. Functionality was moved to
-`uefi::helpers::init` in `uefi` ´v0.28.0`.
+`uefi::helpers::init` in `uefi@v0.28.0`.

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

@@ -63,45 +63,41 @@ 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();
+    // Ensure that the memory map is freed after each iteration (on drop).
+    // Otherwise, we will have an OOM.
+    for _ in 0..200000 {
+        let mut memory_map = bt
+            .memory_map(MemoryType::LOADER_DATA)
+            .expect("Failed to retrieve UEFI memory map");
 
-    // 2 extra descriptors should be enough.
-    let buf_sz = sizes.map_size + 2 * sizes.entry_size;
+        memory_map.sort();
 
-    // We will use vectors for convenience.
-    let mut buffer = vec![0_u8; buf_sz];
+        // Collect the descriptors into a vector
+        let descriptors = memory_map.entries().copied().collect::<Vec<_>>();
 
-    let mut memory_map = bt
-        .memory_map(&mut buffer)
-        .expect("Failed to retrieve UEFI memory map");
+        // Ensured we have at least one entry.
+        // Real memory maps usually have dozens of entries.
+        assert!(!descriptors.is_empty(), "Memory map is empty");
 
-    memory_map.sort();
+        let mut curr_value = descriptors[0];
 
-    // Collect the descriptors into a vector
-    let descriptors = memory_map.entries().copied().collect::<Vec<_>>();
-
-    // Ensured we have at least one entry.
-    // Real memory maps usually have dozens of entries.
-    assert!(!descriptors.is_empty(), "Memory map is empty");
-
-    let mut curr_value = descriptors[0];
-
-    for value in descriptors.iter().skip(1) {
-        if value.phys_start <= curr_value.phys_start {
-            panic!("memory map sorting failed");
+        for value in descriptors.iter().skip(1) {
+            if value.phys_start <= curr_value.phys_start {
+                panic!("memory map sorting failed");
+            }
+            curr_value = *value;
         }
-        curr_value = *value;
-    }
 
-    // This is pretty much a sanity test to ensure returned memory isn't filled with random values.
-    let first_desc = descriptors[0];
+        // This is pretty much a basic sanity test to ensure returned memory
+        // isn't filled with random values.
+        let first_desc = descriptors[0];
 
-    #[cfg(target_arch = "x86_64")]
-    {
-        let phys_start = first_desc.phys_start;
-        assert_eq!(phys_start, 0, "Memory does not start at address 0");
+        #[cfg(target_arch = "x86_64")]
+        {
+            let phys_start = first_desc.phys_start;
+            assert_eq!(phys_start, 0, "Memory does not start at address 0");
+        }
+        let page_count = first_desc.page_count;
+        assert!(page_count != 0, "Memory map entry has size zero");
     }
-    let page_count = first_desc.page_count;
-    assert!(page_count != 0, "Memory map entry has zero size");
 }

uefi/CHANGELOG.md

@@ -16,14 +16,25 @@
 - Added `ByteConversionError`.
 - Re-exported `CapsuleFlags`.
 - One can now specify in `TimeError` what fields of `Time` are outside its valid
-  range. `Time::is_valid` has been updated accordingly. 
+  range. `Time::is_valid` has been updated accordingly.
+- `MemoryMap::as_raw` which provides raw access to the memory map. This is for
+  example useful if you create your own Multiboot2 bootloader that embeds the
+  EFI mmap in a Multiboot2 boot information structure.
 
 ## Changed
 - `SystemTable::exit_boot_services` is now `unsafe`. See that method's
   documentation for details of obligations for callers.
 - `BootServices::allocate_pool` now returns `NonZero<u8>` instead of
   `*mut u8`.
 - `helpers::system_table` is deprecated, use `table::system_table_boot` instead.
+- `BootServices::memory_map` changed its signature from \
+  `pub fn memory_map<'buf>(&self, buffer: &'buf mut [u8]) -> Result<MemoryMap<'buf>> {` \
+  to \
+  `pub fn memory_map(&self, mt: MemoryType) -> Result<MemoryMap>`
+  - Allocations now happen automatically internally on the UEFI heap. Also, the
+    returned type is automatically freed on the UEFI heap, as long as boot
+    services are not excited. By removing the need for that explicit buffer and
+    the lifetime, the API is simpler.
 
 ## Removed
 - Removed the `panic-on-logger-errors` feature of the `uefi` crate. Logger

uefi/src/lib.rs

@@ -48,7 +48,8 @@
 //!   the Rust standard library. For example, methods that return a
 //!   `Vec` rather than filling a statically-sized array. This requires
 //!   a global allocator; you can use the `global_allocator` feature or
-//!   provide your own.
+//!   provide your own. This is independent of internal direct usages of the
+//!   UEFI boot service allocator which may happen anyway, where necessary.
 //! - `global_allocator`: Set [`allocator::Allocator`] as the global Rust
 //!   allocator. This is a simple allocator that relies on the UEFI pool
 //!   allocator. You can choose to provide your own allocator instead of

uefi/src/mem.rs

@@ -23,10 +23,14 @@ use {core::alloc::Allocator, core::ptr::NonNull};
 ///   success.
 ///
 /// # Feature `unstable` / `allocator_api`
-/// By default, this function works with Rust's default allocation mechanism. If you activate the
-/// `unstable`-feature, it uses the `allocator_api` instead. In that case, the function takes an
-/// additional parameter describing the specific [`Allocator`]. You can use [`alloc::alloc::Global`]
-/// as default.
+/// By default, this function works with the allocator that is set as
+/// `#[global_allocator]`. This might be UEFI allocator but depends on your
+/// use case and how you set up the environment.
+///
+/// If you activate the `unstable`-feature, all allocations uses the provided
+/// allocator (via `allocator_api`) instead. In that case, the function takes an
+/// additional parameter describing the specific [`Allocator`]. You can use
+/// [`alloc::alloc::Global`] which defaults to the `#[global_allocator]`.
 ///
 /// [`Allocator`]: https://doc.rust-lang.org/alloc/alloc/trait.Allocator.html
 /// [`alloc::alloc::Global`]: https://doc.rust-lang.org/alloc/alloc/struct.Global.html

uefi/src/table/boot.rs

@@ -2,6 +2,7 @@ use core::ffi::c_void;
 use core::marker::PhantomData;
 use core::ptr::NonNull;
 use core::slice;
+use uefi::table::boot::{MemoryMapBackingMemory, MemoryMapMeta};
 
 use crate::proto::console::text;
 use crate::{CStr16, Result, Status, StatusExt};
@@ -151,41 +152,23 @@ impl SystemTable<Boot> {
         unsafe { &*(*self.table).boot_services.cast_const().cast() }
     }
 
-    /// Get the size in bytes of the buffer to allocate for storing the memory
-    /// map in `exit_boot_services`.
-    ///
-    /// This map contains some extra room to avoid needing to allocate more than
-    /// once.
-    ///
-    /// Returns `None` on overflow.
-    fn memory_map_size_for_exit_boot_services(&self) -> Option<usize> {
-        // 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;
-
-        let memory_map_size = self.boot_services().memory_map_size();
-        let extra_size = memory_map_size.entry_size.checked_mul(extra_entries)?;
-        memory_map_size.map_size.checked_add(extra_size)
-    }
-
     /// Get the current memory map and exit boot services.
     unsafe fn get_memory_map_and_exit_boot_services(
         &self,
-        buf: &'static mut [u8],
-    ) -> Result<MemoryMap<'static>> {
+        buf: &mut [u8],
+    ) -> Result<MemoryMapMeta> {
         let boot_services = self.boot_services();
 
         // Get the memory map.
-        let memory_map = boot_services.memory_map(buf)?;
+        let memory_map = boot_services.get_memory_map(buf)?;
 
         // Try to exit boot services using the memory map key. Note that after
         // the first call to `exit_boot_services`, there are restrictions on
         // what boot services functions can be called. In UEFI 2.8 and earlier,
         // only `get_memory_map` and `exit_boot_services` are allowed. Starting
         // in UEFI 2.9 other memory allocation functions may also be called.
         boot_services
-            .exit_boot_services(boot_services.image_handle(), memory_map.key())
+            .exit_boot_services(boot_services.image_handle(), memory_map.map_key)
             .map(move |()| memory_map)
     }
 
@@ -247,44 +230,37 @@ impl SystemTable<Boot> {
     pub unsafe fn exit_boot_services(
         self,
         memory_type: MemoryType,
-    ) -> (SystemTable<Runtime>, MemoryMap<'static>) {
+    ) -> (SystemTable<Runtime>, MemoryMap) {
         crate::helpers::exit();
 
-        let boot_services = self.boot_services();
-
         // Reboot the device.
-        let reset = |status| -> ! { self.runtime_services().reset(ResetType::COLD, status, None) };
-
-        // Get the size of the buffer to allocate. If that calculation
-        // overflows treat it as an unrecoverable error.
-        let buf_size = match self.memory_map_size_for_exit_boot_services() {
-            Some(buf_size) => buf_size,
-            None => reset(Status::ABORTED),
+        let reset = |status| -> ! {
+            {
+                log::warn!("Resetting the machine");
+                self.runtime_services().reset(ResetType::COLD, status, None)
+            }
         };
 
-        // Allocate a byte slice to hold the memory map. If the
-        // allocation fails treat it as an unrecoverable error.
-        let buf: *mut u8 = match boot_services.allocate_pool(memory_type, buf_size) {
-            Ok(buf) => buf.as_ptr(),
-            Err(err) => reset(err.status()),
-        };
+        let mut buf = MemoryMapBackingMemory::new(memory_type).expect("Failed to allocate memory");
 
         // Calling `exit_boot_services` can fail if the memory map key is not
         // current. Retry a second time if that occurs. This matches the
         // behavior of the Linux kernel:
         // https://github.com/torvalds/linux/blob/e544a0743/drivers/firmware/efi/libstub/efi-stub-helper.c#L375
         let mut status = Status::ABORTED;
         for _ in 0..2 {
-            let buf: &mut [u8] = unsafe { slice::from_raw_parts_mut(buf, buf_size) };
-            match unsafe { self.get_memory_map_and_exit_boot_services(buf) } {
+            match unsafe { self.get_memory_map_and_exit_boot_services(buf.as_mut_slice()) } {
                 Ok(memory_map) => {
                     let st = SystemTable {
                         table: self.table,
                         _marker: PhantomData,
                     };
-                    return (st, memory_map);
+                    return (st, MemoryMap::from_initialized_mem(buf, memory_map));
+                }
+                Err(err) => {
+                    log::error!("Error retrieving the memory map for exiting the boot services");
+                    status = err.status()
                 }
-                Err(err) => status = err.status(),
             }
         }