diff --git a/uefi-test-runner/src/boot/memory.rs b/uefi-test-runner/src/boot/memory.rs index 1a2dcfaf3..5134f5d24 100644 --- a/uefi-test-runner/src/boot/memory.rs +++ b/uefi-test-runner/src/boot/memory.rs @@ -91,17 +91,28 @@ fn memory_map(bt: &BootServices) { // We will use vectors for convenience. let mut buffer = vec![0_u8; buf_sz]; - let (_key, desc_iter) = bt + let mut memory_map = bt .memory_map(&mut buffer) .expect("Failed to retrieve UEFI memory map"); + memory_map.sort(); + // Collect the descriptors into a vector - let descriptors = desc_iter.copied().collect::>(); + let descriptors = memory_map.entries().copied().collect::>(); // 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"); + } + 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]; diff --git a/uefi/src/table/boot.rs b/uefi/src/table/boot.rs index d55984d29..309740263 100644 --- a/uefi/src/table/boot.rs +++ b/uefi/src/table/boot.rs @@ -425,10 +425,7 @@ impl BootServices { /// /// * [`uefi::Status::BUFFER_TOO_SMALL`] /// * [`uefi::Status::INVALID_PARAMETER`] - pub fn memory_map<'buf>( - &self, - buffer: &'buf mut [u8], - ) -> Result<(MemoryMapKey, MemoryMapIter<'buf>)> { + pub fn memory_map<'buf>(&self, buffer: &'buf mut [u8]) -> Result> { let mut map_size = buffer.len(); MemoryDescriptor::assert_aligned(buffer); let map_buffer = buffer.as_mut_ptr().cast::(); @@ -453,13 +450,13 @@ impl BootServices { } .into_with_val(move || { let len = map_size / entry_size; - let iter = MemoryMapIter { - buffer, + + MemoryMap { + key: map_key, + buf: buffer, entry_size, - index: 0, len, - }; - (map_key, iter) + } }) } @@ -1993,6 +1990,108 @@ pub struct MemoryMapSize { pub map_size: usize, } +/// An iterator of [`MemoryDescriptor`] that is always associated with the +/// unique [`MemoryMapKey`] contained in the struct. +/// +/// To iterate over the entries, call [`MemoryMap::entries`]. To get a sorted +/// map, you manually have to call [`MemoryMap::sort`] first. +pub struct MemoryMap<'buf> { + key: MemoryMapKey, + buf: &'buf mut [u8], + entry_size: usize, + len: usize, +} + +impl<'buf> MemoryMap<'buf> { + #[must_use] + /// Returns the unique [`MemoryMapKey`] associated with the memory map. + pub fn key(&self) -> MemoryMapKey { + self.key + } + + /// Sorts the memory map by physical address in place. + /// This operation is optional and should be invoked only once. + pub fn sort(&mut self) { + unsafe { + self.qsort(0, self.len - 1); + } + } + + /// Hoare partition scheme for quicksort. + /// Must be called with `low` and `high` being indices within bounds. + unsafe fn qsort(&mut self, low: usize, high: usize) { + if low >= high { + return; + } + + let p = self.partition(low, high); + self.qsort(low, p); + self.qsort(p + 1, high); + } + + unsafe fn partition(&mut self, low: usize, high: usize) -> usize { + let pivot = self.get_element_phys_addr(low + (high - low) / 2); + + let mut left_index = low.wrapping_sub(1); + let mut right_index = high.wrapping_add(1); + + loop { + while { + left_index = left_index.wrapping_add(1); + + self.get_element_phys_addr(left_index) < pivot + } {} + + while { + right_index = right_index.wrapping_sub(1); + + self.get_element_phys_addr(right_index) > pivot + } {} + + if left_index >= right_index { + return right_index; + } + + self.swap(left_index, right_index); + } + } + + /// Indices must be smaller than len. + unsafe fn swap(&mut self, index1: usize, index2: usize) { + if index1 == index2 { + return; + } + + let base = self.buf.as_mut_ptr(); + + unsafe { + ptr::swap_nonoverlapping( + base.add(index1 * self.entry_size), + base.add(index2 * self.entry_size), + self.entry_size, + ); + } + } + + fn get_element_phys_addr(&self, index: usize) -> PhysicalAddress { + let offset = index.checked_mul(self.entry_size).unwrap(); + let elem = unsafe { &*self.buf.as_ptr().add(offset).cast::() }; + elem.phys_start + } + + /// Returns an iterator over the contained memory map. To get a sorted map, + /// call [`MemoryMap::sort`] first. + #[must_use] + pub fn entries(&self) -> MemoryMapIter { + MemoryMapIter { + buffer: self.buf, + entry_size: self.entry_size, + index: 0, + len: self.len, + } + } +} + /// An iterator of [`MemoryDescriptor`]. The underlying memory map is always /// associated with a unique [`MemoryMapKey`]. #[derive(Debug, Clone)] @@ -2014,12 +2113,16 @@ impl<'buf> Iterator for MemoryMapIter<'buf> { fn next(&mut self) -> Option { if self.index < self.len { - let ptr = self.buffer.as_ptr() as usize + self.entry_size * self.index; + let descriptor = unsafe { + &*self + .buffer + .as_ptr() + .add(self.entry_size * self.index) + .cast::() + }; self.index += 1; - let descriptor = unsafe { &*(ptr as *const MemoryDescriptor) }; - Some(descriptor) } else { None @@ -2197,3 +2300,93 @@ pub enum InterfaceType: i32 => { #[derive(Debug, Clone, Copy)] #[repr(transparent)] pub struct ProtocolSearchKey(NonNull); + +#[cfg(test)] +mod tests { + use core::mem::size_of; + + use crate::table::boot::{MemoryAttribute, MemoryMap, MemoryMapKey, MemoryType}; + + use super::{MemoryDescriptor, MemoryMapIter}; + + #[test] + fn mem_map_sorting() { + // Doesn't matter what type it is. + const TY: MemoryType = MemoryType::RESERVED; + + const BASE: MemoryDescriptor = MemoryDescriptor { + ty: TY, + phys_start: 0, + virt_start: 0, + page_count: 0, + att: MemoryAttribute::empty(), + }; + + let mut buffer = [ + MemoryDescriptor { + phys_start: 2000, + ..BASE + }, + MemoryDescriptor { + phys_start: 3000, + ..BASE + }, + BASE, + MemoryDescriptor { + phys_start: 1000, + ..BASE + }, + ]; + + let desc_count = buffer.len(); + + let byte_buffer = { + let size = desc_count * size_of::(); + unsafe { core::slice::from_raw_parts_mut(buffer.as_mut_ptr() as *mut u8, size) } + }; + + let mut mem_map = MemoryMap { + // Key doesn't matter + key: MemoryMapKey(0), + len: desc_count, + buf: byte_buffer, + entry_size: size_of::(), + }; + + mem_map.sort(); + + if !is_sorted(&mem_map.entries()) { + panic!("mem_map is not sorted: {}", mem_map); + } + } + + // Added for debug purposes on test failure + impl core::fmt::Display for MemoryMap<'_> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + writeln!(f)?; + for desc in self.entries() { + writeln!(f, "{:?}", desc)?; + } + Ok(()) + } + } + + fn is_sorted(iter: &MemoryMapIter) -> bool { + let mut iter = iter.clone(); + let mut curr_start; + + if let Some(val) = iter.next() { + curr_start = val.phys_start; + } else { + return true; + } + + for desc in iter { + if desc.phys_start <= curr_start { + return false; + } + curr_start = desc.phys_start + } + true + } +} diff --git a/uefi/src/table/system.rs b/uefi/src/table/system.rs index e78b38035..03c294029 100644 --- a/uefi/src/table/system.rs +++ b/uefi/src/table/system.rs @@ -7,7 +7,7 @@ use core::{ptr, slice}; use crate::proto::console::text; use crate::{CStr16, Char16, Handle, Result, Status}; -use super::boot::{BootServices, MemoryDescriptor, MemoryMapIter, MemoryType}; +use super::boot::{BootServices, MemoryDescriptor, MemoryMap, MemoryType}; use super::runtime::{ResetType, RuntimeServices}; use super::{cfg, Header, Revision}; @@ -158,11 +158,11 @@ impl SystemTable { unsafe fn get_memory_map_and_exit_boot_services( &self, buf: &'static mut [u8], - ) -> Result> { + ) -> Result> { let boot_services = self.boot_services(); // Get the memory map. - let (memory_map_key, memory_map_iter) = boot_services.memory_map(buf)?; + let memory_map = boot_services.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 @@ -170,8 +170,8 @@ impl SystemTable { // 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) - .map(move |()| memory_map_iter) + .exit_boot_services(boot_services.image_handle(), memory_map.key()) + .map(move |()| memory_map) } /// Exit the UEFI boot services. @@ -212,7 +212,7 @@ impl SystemTable { /// [`Logger::disable`]: crate::logger::Logger::disable /// [`uefi_services::init`]: https://docs.rs/uefi-services/latest/uefi_services/fn.init.html #[must_use] - pub fn exit_boot_services(self) -> (SystemTable, MemoryMapIter<'static>) { + pub fn exit_boot_services(self) -> (SystemTable, MemoryMap<'static>) { let boot_services = self.boot_services(); // Reboot the device.