multiboot2: Support setting the EFI memory map tag

Author: YtvwlD

multiboot2/src/builder/information.rs

@@ -2,8 +2,8 @@
 use crate::builder::traits::StructAsBytes;
 use crate::{
     BasicMemoryInfoTag, BootInformationInner, BootLoaderNameTag, CommandLineTag,
-    EFISdt32, EFISdt64, ElfSectionsTag, EndTag, FramebufferTag, MemoryMapTag,
-    ModuleTag, RsdpV1Tag, RsdpV2Tag, SmbiosTag,
+    EFISdt32, EFISdt64, EFIMemoryMapTag, ElfSectionsTag, EndTag, FramebufferTag,
+    MemoryMapTag, ModuleTag, RsdpV1Tag, RsdpV2Tag, SmbiosTag,
 };
 
 use alloc::boxed::Box;
@@ -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())
         }
@@ -151,6 +156,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)
         }
@@ -198,6 +206,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

@@ -325,6 +325,13 @@ impl BootInformation {
         }
     }
 
+    /// Search for the EFI Memory map tag, return a mutable reference.
+    pub fn efi_memory_map_tag_mut(&self) -> Option<&mut EFIMemoryMapTag> {
+        self
+            .get_tag(TagType::EfiMmap)
+            .map(|tag| unsafe { &mut *(tag as *const Tag as *mut EFIMemoryMapTag) })
+    }
+
     /// Search for the EFI 32-bit image handle pointer.
     pub fn efi_32_ih(&self) -> Option<&EFIImageHandle32> {
         self.get_tag(TagType::Efi32Ih)
@@ -1430,14 +1437,4 @@ mod tests {
         let efi_mmap = bi.efi_memory_map_tag();
         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 + `EFIMemoryDesc` is 40 bytes.
-            core::mem::transmute::<[u8; 56], EFIMemoryMapTag>([0u8; 56]);
-        }
-    }
 }

multiboot2/src/memory_map.rs

@@ -266,28 +266,69 @@ pub struct EFIMemoryMapTag {
     size: u32,
     desc_size: u32,
     desc_version: u32,
-    first_desc: EFIMemoryDesc,
+    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, 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 *const EFIMemoryDesc;
+        let start_area = (&self.descs[0]) as *const EFIMemoryDesc;
         EFIMemoryAreaIter {
             current_area: start_area as u64,
-            last_area: (self_ptr as u64 + self.size 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,
+            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,
@@ -298,6 +339,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 {
@@ -352,6 +400,29 @@ pub enum EFIMemoryAreaType {
     EfiUnknown,
 }
 
+impl Into<u32> for EFIMemoryAreaType {
+    fn into(self) -> u32 {
+        match self {
+            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 {
@@ -392,6 +463,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)]
@@ -421,3 +505,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)
+        }
+    }
+}