From d4cca4c2f97ce7bb585ee271dfcad3c90c75a404 Mon Sep 17 00:00:00 2001
From: Steven Xu <stevendoesstuffs@protonmail.com>
Date: Sun, 12 Dec 2021 16:05:18 -0600
Subject: [PATCH] Change memory_map_size to return entry size as well

---
 src/table/boot.rs                   | 24 ++++++++++++++++++------
 uefi-test-runner/src/boot/memory.rs | 15 +++++++--------
 uefi-test-runner/src/main.rs        |  7 +++----
 3 files changed, 28 insertions(+), 18 deletions(-)

diff --git a/src/table/boot.rs b/src/table/boot.rs
index a64a731ad..e14d7c387 100644
--- a/src/table/boot.rs
+++ b/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> {
diff --git a/uefi-test-runner/src/boot/memory.rs b/uefi-test-runner/src/boot/memory.rs
index 2f07323e5..f7e705c40 100644
--- a/uefi-test-runner/src/boot/memory.rs
+++ b/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
diff --git a/uefi-test-runner/src/main.rs b/uefi-test-runner/src/main.rs
index 45b21e195..0866fe1f8 100644
--- a/uefi-test-runner/src/main.rs
+++ b/uefi-test-runner/src/main.rs
@@ -9,10 +9,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;
@@ -142,8 +141,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[..])