multiboot2: Support setting the EFI memory map tag

Author: YtvwlD

multiboot2/src/builder/information.rs

@@ -1,8 +1,8 @@
 //! Exports item [`Multiboot2InformationBuilder`].
 use crate::builder::traits::StructAsBytes;
 use crate::{
-    BasicMemoryInfoTag, BootInformationInner, BootLoaderNameTag, CommandLineTag, EFISdt32,
-    EFISdt64, ElfSectionsTag, EndTag, FramebufferTag, MemoryMapTag, ModuleTag, RsdpV1Tag,
+    BasicMemoryInfoTag, BootInformationInner, BootLoaderNameTag, CommandLineTag, EFIMemoryMapTag,
+    EFISdt32, EFISdt64, ElfSectionsTag, EndTag, FramebufferTag, MemoryMapTag, ModuleTag, RsdpV1Tag,
     RsdpV2Tag, SmbiosTag,
 };
 
@@ -18,6 +18,7 @@ pub struct Multiboot2InformationBuilder {
     basic_memory_info_tag: Option<BasicMemoryInfoTag>,
     boot_loader_name_tag: Option<Box<BootLoaderNameTag>>,
     command_line_tag: Option<Box<CommandLineTag>>,
+    efi_memory_map_tag: Option<Box<EFIMemoryMapTag>>,
     elf_sections_tag: Option<Box<ElfSectionsTag>>,
     framebuffer_tag: Option<Box<FramebufferTag>>,
     memory_map_tag: Option<Box<MemoryMapTag>>,
@@ -37,6 +38,7 @@ impl Multiboot2InformationBuilder {
             command_line_tag: None,
             efisdt32: None,
             efisdt64: None,
+            efi_memory_map_tag: None,
             elf_sections_tag: None,
             framebuffer_tag: None,
             memory_map_tag: None,
@@ -83,6 +85,9 @@ impl Multiboot2InformationBuilder {
         if let Some(tag) = &self.efisdt64 {
             len += Self::size_or_up_aligned(tag.byte_size())
         }
+        if let Some(tag) = &self.efi_memory_map_tag {
+            len += Self::size_or_up_aligned(tag.byte_size())
+        }
         if let Some(tag) = &self.elf_sections_tag {
             len += Self::size_or_up_aligned(tag.byte_size())
         }
@@ -149,6 +154,9 @@ impl Multiboot2InformationBuilder {
         if let Some(tag) = self.efisdt64.as_ref() {
             Self::build_add_bytes(&mut data, &tag.struct_as_bytes(), false)
         }
+        if let Some(tag) = self.efi_memory_map_tag.as_ref() {
+            Self::build_add_bytes(&mut data, &tag.struct_as_bytes(), false)
+        }
         if let Some(tag) = self.elf_sections_tag.as_ref() {
             Self::build_add_bytes(&mut data, &tag.struct_as_bytes(), false)
         }
@@ -196,6 +204,10 @@ impl Multiboot2InformationBuilder {
         self.efisdt64 = Some(efisdt64);
     }
 
+    pub fn efi_memory_map_tag(&mut self, efi_memory_map_tag: Box<EFIMemoryMapTag>) {
+        self.efi_memory_map_tag = Some(efi_memory_map_tag);
+    }
+
     pub fn elf_sections_tag(&mut self, elf_sections_tag: Box<ElfSectionsTag>) {
         self.elf_sections_tag = Some(elf_sections_tag);
     }

multiboot2/src/lib.rs

@@ -334,6 +334,12 @@ impl BootInformation {
         }
     }
 
+    /// Search for the EFI Memory map tag, return a mutable reference.
+    pub fn efi_memory_map_tag_mut(&mut self) -> Option<&mut EFIMemoryMapTag> {
+        self.get_tag_mut(TagType::EfiMmap)
+            .map(|tag| unsafe { &mut *(tag as *mut Tag as *mut EFIMemoryMapTag) })
+    }
+
     /// Search for the EFI 32-bit image handle pointer.
     pub fn efi_32_ih(&self) -> Option<&EFIImageHandle32> {
         self.get_tag::<EFIImageHandle32, _>(TagType::Efi32Ih)
@@ -1509,16 +1515,6 @@ mod tests {
         assert!(efi_mmap.is_none());
     }
 
-    #[test]
-    /// Compile time test for `EFIMemoryMapTag`.
-    fn efi_memory_map_tag_size() {
-        use super::EFIMemoryMapTag;
-        unsafe {
-            // `EFIMemoryMapTag` is 16 bytes without the 1st entry
-            core::mem::transmute::<[u8; 16], EFIMemoryMapTag>([0u8; 16]);
-        }
-    }
-
     #[test]
     #[cfg(feature = "unstable")]
     /// This test succeeds if it compiles.

multiboot2/src/memory_map.rs

@@ -271,30 +271,72 @@ pub struct EFIMemoryMapTag {
     size: u32,
     desc_size: u32,
     desc_version: u32,
-    first_desc: [EFIMemoryDesc; 0],
+    descs: [EFIMemoryDesc],
 }
 
 impl EFIMemoryMapTag {
+    #[cfg(feature = "builder")]
+    /// Create a new EFI memory map tag with the given memory descriptors.
+    /// Version and size can't be set because you're passing a slice of
+    /// EFIMemoryDescs, not the ones you might have gotten from the firmware.
+    pub fn new(descs: &[EFIMemoryDesc]) -> Box<Self> {
+        // update this when updating EFIMemoryDesc
+        const MEMORY_DESCRIPTOR_VERSION: u32 = 1;
+        let mut bytes = [
+            (mem::size_of::<EFIMemoryDesc>() as u32).to_le_bytes(),
+            MEMORY_DESCRIPTOR_VERSION.to_le_bytes(),
+        ]
+        .concat();
+        for desc in descs {
+            bytes.extend(desc.struct_as_bytes());
+        }
+        let tag = boxed_dst_tag(TagType::EfiMmap.into(), bytes.as_slice());
+        unsafe { Box::from_raw(Box::into_raw(tag) as *mut Self) }
+    }
+
     /// Return an iterator over ALL marked memory areas.
     ///
     /// This differs from `MemoryMapTag` as for UEFI, the OS needs some non-
     /// available memory areas for tables and such.
     pub fn memory_areas(&self) -> EFIMemoryAreaIter {
         let self_ptr = self as *const EFIMemoryMapTag;
-        let start_area = self.first_desc.as_ptr();
+        let start_area = (&self.descs[0]) as *const EFIMemoryDesc;
         EFIMemoryAreaIter {
             current_area: start_area as u64,
             // NOTE: `last_area` is only a bound, it doesn't necessarily point exactly to the last element
-            last_area: (self_ptr as u64
-                + (self.size as u64 - core::mem::size_of::<EFIMemoryMapTag>() as u64)),
+            last_area: (self_ptr as *const () as u64 + self.size as u64),
+            entry_size: self.desc_size,
+            phantom: PhantomData,
+        }
+    }
+
+    #[cfg(feature = "builder")]
+    /// Return an iterator over ALL marked memory areas, mutably.
+    ///
+    /// This differs from `MemoryMapTag` as for UEFI, the OS needs some non-
+    /// available memory areas for tables and such.
+    pub fn memory_areas_mut(&mut self) -> EFIMemoryAreaIterMut {
+        let self_ptr = self as *const EFIMemoryMapTag;
+        let start_area = (&self.descs[0]) as *const EFIMemoryDesc;
+        EFIMemoryAreaIterMut {
+            current_area: start_area as u64,
+            // NOTE: `last_area` is only a bound, it doesn't necessarily point exactly to the last element
+            last_area: (self_ptr as *const () as u64 + self.size as u64),
             entry_size: self.desc_size,
             phantom: PhantomData,
         }
     }
 }
 
+#[cfg(feature = "builder")]
+impl StructAsBytes for EFIMemoryMapTag {
+    fn byte_size(&self) -> usize {
+        self.size.try_into().unwrap()
+    }
+}
+
 /// EFI Boot Memory Map Descriptor
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 #[repr(C)]
 pub struct EFIMemoryDesc {
     typ: u32,
@@ -305,6 +347,13 @@ pub struct EFIMemoryDesc {
     attr: u64,
 }
 
+#[cfg(feature = "builder")]
+impl StructAsBytes for EFIMemoryDesc {
+    fn byte_size(&self) -> usize {
+        mem::size_of::<Self>()
+    }
+}
+
 /// An enum of possible reported region types.
 #[derive(Debug, PartialEq, Eq)]
 pub enum EFIMemoryAreaType {
@@ -359,6 +408,29 @@ pub enum EFIMemoryAreaType {
     EfiUnknown,
 }
 
+impl From<EFIMemoryAreaType> for u32 {
+    fn from(area: EFIMemoryAreaType) -> Self {
+        match area {
+            EFIMemoryAreaType::EfiReservedMemoryType => 0,
+            EFIMemoryAreaType::EfiLoaderCode => 1,
+            EFIMemoryAreaType::EfiLoaderData => 2,
+            EFIMemoryAreaType::EfiBootServicesCode => 3,
+            EFIMemoryAreaType::EfiBootServicesData => 4,
+            EFIMemoryAreaType::EfiRuntimeServicesCode => 5,
+            EFIMemoryAreaType::EfiRuntimeServicesData => 6,
+            EFIMemoryAreaType::EfiConventionalMemory => 7,
+            EFIMemoryAreaType::EfiUnusableMemory => 8,
+            EFIMemoryAreaType::EfiACPIReclaimMemory => 9,
+            EFIMemoryAreaType::EfiACPIMemoryNVS => 10,
+            EFIMemoryAreaType::EfiMemoryMappedIO => 11,
+            EFIMemoryAreaType::EfiMemoryMappedIOPortSpace => 12,
+            EFIMemoryAreaType::EfiPalCode => 13,
+            EFIMemoryAreaType::EfiPersistentMemory => 14,
+            EFIMemoryAreaType::EfiUnknown => panic!("unknown type"),
+        }
+    }
+}
+
 impl EFIMemoryDesc {
     /// The physical address of the memory region.
     pub fn physical_address(&self) -> u64 {
@@ -399,6 +471,19 @@ impl EFIMemoryDesc {
     }
 }
 
+impl Default for EFIMemoryDesc {
+    fn default() -> Self {
+        Self {
+            typ: EFIMemoryAreaType::EfiReservedMemoryType.into(),
+            _padding: 0,
+            phys_addr: 0,
+            virt_addr: 0,
+            num_pages: 0,
+            attr: 0,
+        }
+    }
+}
+
 /// EFI ExitBootServices was not called
 #[derive(Debug)]
 #[repr(C)]
@@ -428,3 +513,25 @@ impl<'a> Iterator for EFIMemoryAreaIter<'a> {
         }
     }
 }
+
+/// An iterator over ALL EFI memory areas, mutably.
+#[derive(Clone, Debug)]
+pub struct EFIMemoryAreaIterMut<'a> {
+    current_area: u64,
+    last_area: u64,
+    entry_size: u32,
+    phantom: PhantomData<&'a mut EFIMemoryDesc>,
+}
+
+impl<'a> Iterator for EFIMemoryAreaIterMut<'a> {
+    type Item = &'a mut EFIMemoryDesc;
+    fn next(&mut self) -> Option<&'a mut EFIMemoryDesc> {
+        if self.current_area > self.last_area {
+            None
+        } else {
+            let area = unsafe { &mut *(self.current_area as *mut EFIMemoryDesc) };
+            self.current_area += self.entry_size as u64;
+            Some(area)
+        }
+    }
+}