uefi: MemoryMap: more realistic unit tests

Author: phip1611

Diff for: uefi/src/table/boot.rs

@@ -1999,21 +1999,27 @@ pub struct ProtocolSearchKey(NonNull<c_void>);
 
 #[cfg(test)]
 mod tests {
-    use core::mem::{size_of, size_of_val};
-    use uefi::table::boot::MemoryMapKey;
-
-    use crate::table::boot::{MemoryAttribute, MemoryMap, MemoryType};
-
-    use super::{MemoryDescriptor, MemoryMapIter};
-
-    fn buffer_to_map(buffer: &mut [MemoryDescriptor]) -> MemoryMap {
-        let byte_buffer = {
-            unsafe {
-                core::slice::from_raw_parts_mut(buffer.as_mut_ptr() as *mut u8, size_of_val(buffer))
-            }
-        };
-
-        MemoryMap::from_raw(byte_buffer.len(), byte_buffer, size_of::<MemoryDescriptor>(), MemoryMapKey(0))
+    use super::*;
+
+    /// Builds an EFI memory map buffer from the provider [`MemoryDescriptor`]s
+    /// that is close to a real world structure. This means that the reference
+    /// for an entry is not `size_of::<MemoryDescriptor>` but `desc_size`,
+    /// which is usually a few byte larger.
+    fn descs_to_buffer(descs: &[MemoryDescriptor]) -> (Vec<u8>, usize) {
+        let desc_size = unsafe { MemoryDescriptor::uefi_desc_size() };
+        let diff = desc_size - mem::size_of::<MemoryDescriptor>();
+
+        let mut buf = Vec::with_capacity(descs.len() * desc_size);
+
+        for desc in descs {
+            let bytes = unsafe {
+                let ptr = ptr::addr_of!(*desc);
+                slice::from_raw_parts(ptr.cast::<u8>(), mem::size_of::<MemoryDescriptor>())
+            };
+            buf.extend(bytes);
+            buf.extend([0].repeat(diff));
+        }
+        (buf, desc_size)
     }
 
     #[test]
@@ -2029,7 +2035,7 @@ mod tests {
             att: MemoryAttribute::empty(),
         };
 
-        let mut buffer = [
+        let descs = [
             MemoryDescriptor {
                 phys_start: 2000,
                 ..BASE
@@ -2045,7 +2051,8 @@ mod tests {
             },
         ];
 
-        let mut mem_map = buffer_to_map(&mut buffer);
+        let (mut buffer, desc_size) = descs_to_buffer(&descs);
+        let mut mem_map = MemoryMap::from_raw(&mut buffer, desc_size, MemoryMapKey(0));
 
         mem_map.sort();
 
@@ -2067,7 +2074,7 @@ mod tests {
             att: MemoryAttribute::empty(),
         };
 
-        const BUFFER: [MemoryDescriptor; 4] = [
+        let descs = [
             MemoryDescriptor {
                 phys_start: 2000,
                 ..BASE
@@ -2083,12 +2090,11 @@ mod tests {
             },
         ];
 
-        let mut buffer = BUFFER;
-
-        let mut mem_map = buffer_to_map(&mut buffer);
+        let (mut buffer, desc_size) = descs_to_buffer(&descs);
+        let mut mem_map = MemoryMap::from_raw(&mut buffer, desc_size, MemoryMapKey(0));
 
         for index in 0..3 {
-            assert_eq!(mem_map.get(index), BUFFER.get(index))
+            assert_eq!(mem_map.get(index), descs.get(index))
         }
 
         let mut_desc = mem_map.get_mut(2).unwrap();
@@ -2097,7 +2103,7 @@ mod tests {
 
         let desc = mem_map.get(2).unwrap();
 
-        assert_ne!(*desc, BUFFER[2]);
+        assert_ne!(*desc, descs[2]);
     }
 
     // Added for debug purposes on test failure