Change memory_map_size to return entry size as well

Author: Steven Xu

src/table/boot.rs

@@ -218,12 +218,13 @@ impl BootServices {
  (self.free_pages)(addr, count).into()
  }
 
- /// Retrieves the size, in bytes, of the current memory map.
+ /// Returns struct which contains the size of a single memory descriptor
+ /// as well as the size of the current memory map.
  ///
- /// A buffer of this size will be capable of holding the whole current memory map,
- /// including padding. Note, however, that allocations will increase the size of the
- /// memory map, therefore it is better to allocate some extra space.
- pub fn memory_map_size(&self) -> usize {
+ /// Note that the size of the memory map can increase any time an allocation happens,
+ /// so when creating a buffer to put the memory map into, it's recommended to allocate a few extra
+ /// elements worth of space above the size of the current memory map.
+ pub fn memory_map_size(&self) -> MemoryMapSize {
  let mut map_size = 0;
  let mut map_key = MemoryMapKey(0);
  let mut entry_size = 0;
@@ -240,7 +241,10 @@ impl BootServices {
  };
  assert_eq!(status, Status::BUFFER_TOO_SMALL);
 
- map_size
+ MemoryMapSize {
+ entry_size,
+ map_size,
+ }
  }
 
  /// Retrieves the current memory map.
@@ -1297,6 +1301,14 @@ bitflags! {
 #[repr(C)]
 pub struct MemoryMapKey(usize);
 
+/// A structure containing the size of a memory descriptor and the size of the memory map
+pub struct MemoryMapSize {
+ /// Size of a single memory descriptor in bytes
+ pub entry_size: usize,
+ /// Size of the entire memory map in bytes
+ pub map_size: usize,
+}
+
 /// An iterator of memory descriptors
 #[derive(Debug, Clone)]
 struct MemoryMapIter<'buf> {

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

@@ -1,8 +1,7 @@
 use uefi::prelude::*;
-use uefi::table::boot::{AllocateType, BootServices, MemoryDescriptor, MemoryType};
+use uefi::table::boot::{AllocateType, BootServices, MemoryType};
 
 use crate::alloc::vec::Vec;
-use core::mem;
 
 pub fn test(bt: &BootServices) {
  info!("Testing memory functions");
@@ -26,7 +25,7 @@ fn allocate_pages(bt: &BootServices) {
 
  assert_eq!(pgs % 4096, 0, "Page pointer is not page-aligned");
 
- // Reinterprete the page as an array of bytes
+ // Reinterpret the page as an array of bytes
  let buf = unsafe { &mut *(pgs as *mut [u8; 4096]) };
 
  // If these don't fail then we properly allocated some memory.
@@ -84,13 +83,13 @@ fn memmove(bt: &BootServices) {
 fn memory_map(bt: &BootServices) {
  info!("Testing memory map functions");
 
- // Get an estimate of the memory map size.
- let map_sz = bt.memory_map_size();
+ // Get the memory descriptor size and an estimate of the memory map size
+ let sizes = bt.memory_map_size();
 
- // 8 extra descriptors should be enough.
- let buf_sz = map_sz + 8 * mem::size_of::<MemoryDescriptor>();
+ // 2 extra descriptors should be enough.
+ let buf_sz = sizes.map_size + 2 * sizes.entry_size;
 
- // We will use vectors for convencience.
+ // We will use vectors for convenience.
  let mut buffer = vec![0_u8; buf_sz];
 
  let (_key, desc_iter) = bt

uefi-test-runner/src/main.rs

@@ -8,10 +8,9 @@ extern crate log;
 extern crate alloc;
 
 use alloc::string::String;
-use core::mem;
 use uefi::prelude::*;
 use uefi::proto::console::serial::Serial;
-use uefi::table::boot::{MemoryDescriptor, OpenProtocolAttributes, OpenProtocolParams};
+use uefi::table::boot::{OpenProtocolAttributes, OpenProtocolParams};
 
 mod boot;
 mod proto;
@@ -141,8 +140,8 @@ fn shutdown(image: uefi::Handle, mut st: SystemTable<Boot>) -> ! {
  }
 
  // Exit boot services as a proof that it works :)
- let max_mmap_size =
-  st.boot_services().memory_map_size() + 8 * mem::size_of::<MemoryDescriptor>();
+ let sizes = st.boot_services().memory_map_size();
+ let max_mmap_size = sizes.map_size + 2 * sizes.entry_size;
  let mut mmap_storage = vec![0; max_mmap_size].into_boxed_slice();
  let (st, _iter) = st
  .exit_boot_services(image, &mut mmap_storage[..])