multiboot2: Get a mutable reference to the memory map

Author: YtvwlD

integration-test/bins/multiboot2_payload/src/verify/chainloader.rs

@@ -1,15 +1,15 @@
 use crate::verify::{print_memory_map, print_module_info};
-use multiboot2::BootInformation;
+use multiboot2::{BootInformation, BootInformationInner};
 
-pub fn run(mbi: &BootInformation) -> anyhow::Result<()> {
+pub fn run<T: AsRef<BootInformationInner>>(mbi: &BootInformation<T>) -> anyhow::Result<()> {
     basic_sanity_checks(mbi)?;
     print_memory_map(mbi)?;
     print_module_info(mbi)?;
     // print_elf_info(mbi)?;
     Ok(())
 }
 
-fn basic_sanity_checks(mbi: &BootInformation) -> anyhow::Result<()> {
+fn basic_sanity_checks<T: AsRef<BootInformationInner>>(mbi: &BootInformation<T>) -> anyhow::Result<()> {
     // Some basic sanity checks
     let bootloader_name = mbi
         .boot_loader_name_tag()

integration-test/bins/multiboot2_payload/src/verify/grub.rs

@@ -1,15 +1,15 @@
 use crate::verify::{print_elf_info, print_memory_map, print_module_info};
-use multiboot2::BootInformation;
+use multiboot2::{BootInformation, BootInformationInner};
 
-pub fn run(mbi: &BootInformation) -> anyhow::Result<()> {
+pub fn run<T: AsRef<BootInformationInner>>(mbi: &BootInformation<T>) -> anyhow::Result<()> {
     basic_sanity_checks(mbi)?;
     print_memory_map(mbi)?;
     print_module_info(mbi)?;
     print_elf_info(mbi)?;
     Ok(())
 }
 
-fn basic_sanity_checks(mbi: &BootInformation) -> anyhow::Result<()> {
+fn basic_sanity_checks<T: AsRef<BootInformationInner>>(mbi: &BootInformation<T>) -> anyhow::Result<()> {
     // Some basic sanity checks
     let bootloader_name = mbi
         .boot_loader_name_tag()

integration-test/bins/multiboot2_payload/src/verify/mod.rs

@@ -3,9 +3,9 @@ mod grub;
 
 use alloc::format;
 use alloc::vec::Vec;
-use multiboot2::BootInformation;
+use multiboot2::{BootInformation, BootInformationInner};
 
-pub fn run(mbi: &BootInformation) -> anyhow::Result<()> {
+pub fn run<T: AsRef<BootInformationInner>>(mbi: &BootInformation<T>) -> anyhow::Result<()> {
     println!("{mbi:#x?}");
     println!();
 
@@ -27,7 +27,7 @@ pub fn run(mbi: &BootInformation) -> anyhow::Result<()> {
     Ok(())
 }
 
-pub(self) fn print_memory_map(mbi: &BootInformation) -> anyhow::Result<()> {
+pub(self) fn print_memory_map<T: AsRef<BootInformationInner>>(mbi: &BootInformation<T>) -> anyhow::Result<()> {
     let memmap = mbi
         .memory_map_tag()
         .ok_or("Should have memory map")
@@ -46,7 +46,7 @@ pub(self) fn print_memory_map(mbi: &BootInformation) -> anyhow::Result<()> {
     Ok(())
 }
 
-pub(self) fn print_elf_info(mbi: &BootInformation) -> anyhow::Result<()> {
+pub(self) fn print_elf_info<T: AsRef<BootInformationInner>>(mbi: &BootInformation<T>) -> anyhow::Result<()> {
     let sections_iter = mbi
         .elf_sections()
         .ok_or("Should have elf sections")
@@ -71,7 +71,7 @@ pub(self) fn print_elf_info(mbi: &BootInformation) -> anyhow::Result<()> {
     Ok(())
 }
 
-pub(self) fn print_module_info(mbi: &BootInformation) -> anyhow::Result<()> {
+pub(self) fn print_module_info<T: AsRef<BootInformationInner>>(mbi: &BootInformation<T>) -> anyhow::Result<()> {
     let modules = mbi.module_tags().collect::<Vec<_>>();
     if modules.len() != 1 {
         Err(anyhow::Error::msg("Should have exactly one boot module"))?

multiboot2/src/lib.rs

@@ -97,7 +97,7 @@ use derive_more::Display;
 #[cfg(feature = "builder")]
 use crate::builder::AsBytes;
 use crate::framebuffer::UnknownFramebufferType;
-use tag::TagIter;
+use tag::{TagIter, TagIterMut};
 
 /// Magic number that a Multiboot2-compliant boot loader will use to identify
 /// the handoff. The location depends on the architecture and the targeted
@@ -150,9 +150,9 @@ impl AsBytes for BootInformationHeader {}
 
 /// This type holds the whole data of the MBI. This helps to better satisfy miri
 /// when it checks for memory issues.
-#[derive(ptr_meta::Pointee)]
+#[derive(ptr_meta::Pointee, Debug)]
 #[repr(C)]
-struct BootInformationInner {
+pub struct BootInformationInner {
     header: BootInformationHeader,
     tags: [u8],
 }
@@ -176,11 +176,17 @@ impl BootInformationInner {
     }
 }
 
+impl AsRef<BootInformationInner> for BootInformationInner {
+    fn as_ref(&self) -> &BootInformationInner {
+        self
+    }
+}
+
 /// A Multiboot 2 Boot Information (MBI) accessor.
 #[repr(transparent)]
-pub struct BootInformation<'a>(&'a BootInformationInner);
+pub struct BootInformation<T: AsRef<BootInformationInner>>(T);
 
-impl<'a> BootInformation<'a> {
+impl BootInformation<&BootInformationInner> {
     /// Loads the [`BootInformation`] from a pointer. The pointer must be valid
     /// and aligned to an 8-byte boundary, as defined by the spec.
     ///
@@ -229,15 +235,50 @@ impl<'a> BootInformation<'a> {
 
         Ok(Self(mbi))
     }
+}
+
+impl BootInformation<&mut BootInformationInner> {
+    /// `BootInformation::load`, but mutably.
+    ///
+    /// # Safety
+    /// The same considerations that apply to `load` also apply here, but the
+    /// memory can be modified (through the `_mut` methods).
+    pub unsafe fn load_mut(ptr: *mut BootInformationHeader) -> Result<Self, MbiLoadError> {
+        // null or not aligned
+        if ptr.is_null() || ptr.align_offset(8) != 0 {
+            return Err(MbiLoadError::IllegalAddress);
+        }
 
+        // mbi: reference to basic header
+        let mbi = &*ptr;
+
+        // Check if total size is not 0 and a multiple of 8.
+        if mbi.total_size == 0 || mbi.total_size & 0b111 != 0 {
+            return Err(MbiLoadError::IllegalTotalSize(mbi.total_size));
+        }
+
+        let slice_size = mbi.total_size as usize - size_of::<BootInformationHeader>();
+        // mbi: reference to full mbi
+        let mbi = ptr_meta::from_raw_parts_mut::<BootInformationInner>(ptr.cast(), slice_size);
+        let mbi = &mut *mbi;
+
+        if !mbi.has_valid_end_tag() {
+            return Err(MbiLoadError::NoEndTag);
+        }
+
+        Ok(Self(mbi))
+    }
+}
+
+impl<T: AsRef<BootInformationInner>> BootInformation<T> {
     /// Get the start address of the boot info.
     pub fn start_address(&self) -> usize {
         self.as_ptr() as usize
     }
 
     /// Get the start address of the boot info as pointer.
     pub fn as_ptr(&self) -> *const () {
-        core::ptr::addr_of!(*self.0).cast()
+        core::ptr::addr_of!(*self.0.as_ref()).cast()
     }
 
     /// Get the end address of the boot info.
@@ -256,7 +297,7 @@ impl<'a> BootInformation<'a> {
 
     /// Get the total size of the boot info struct.
     pub fn total_size(&self) -> usize {
-        self.0.header.total_size as usize
+        self.0.as_ref().header.total_size as usize
     }
 
     // ######################################################
@@ -456,19 +497,44 @@ impl<'a> BootInformation<'a> {
     ///     .unwrap();
     /// assert_eq!(tag.name(), Ok("name"));
     /// ```
-    pub fn get_tag<TagT: TagTrait + ?Sized + 'a>(&'a self) -> Option<&'a TagT> {
+    pub fn get_tag<TagT: TagTrait + ?Sized>(&self) -> Option<&TagT> {
         self.tags()
             .find(|tag| tag.typ == TagT::ID)
             .map(|tag| tag.cast_tag::<TagT>())
     }
 
     /// Returns an iterator over all tags.
     fn tags(&self) -> TagIter {
-        TagIter::new(&self.0.tags)
+        TagIter::new(&self.0.as_ref().tags)
     }
 }
 
-impl fmt::Debug for BootInformation<'_> {
+impl<T: AsRef<BootInformationInner> + AsMut<BootInformationInner>> BootInformation<T> {
+    /// Search for the Memory map tag, return a mutable reference.
+    pub fn memory_map_tag_mut(&mut self) -> Option<&mut MemoryMapTag> {
+        self.get_tag_mut::<MemoryMapTag, _>(TagType::Mmap)
+    }
+
+    fn get_tag_mut<TagT: TagTrait + ?Sized, TagType: Into<TagTypeId>>(
+        &mut self,
+        typ: TagType,
+    ) -> Option<&mut TagT> {
+        let typ = typ.into();
+        self.tags_mut()
+            .find(|tag| tag.typ == typ)
+            .map(|tag| tag.cast_tag_mut::<TagT>())
+    }
+
+    fn tags_mut(&mut self) -> TagIterMut {
+        TagIterMut::new(&mut self.0.as_mut().tags)
+    }
+}
+
+// SAFETY: BootInformation contains a const ptr to memory that is never mutated.
+// Sending this pointer to other threads is sound.
+unsafe impl<T: AsRef<BootInformationInner>> Send for BootInformation<T> {}
+
+impl<T: AsRef<BootInformationInner>> fmt::Debug for BootInformation<T> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         /// Limit how many Elf-Sections should be debug-formatted.
         /// Can be thousands of sections for a Rust binary => this is useless output.
@@ -1234,7 +1300,7 @@ mod tests {
 
     /// Helper for [`grub2`].
     fn test_grub2_boot_info(
-        bi: &BootInformation,
+        bi: &BootInformation<&BootInformationInner>,
         addr: usize,
         string_addr: u64,
         bytes: &[u8],

multiboot2/src/memory_map.rs

@@ -59,6 +59,13 @@ impl MemoryMapTag {
         assert_eq!(self.entry_size as usize, mem::size_of::<MemoryArea>());
         &self.areas
     }
+
+    /// Return a mutable slice with all memory areas.
+    pub fn all_memory_areas_mut(&mut self) -> &mut [MemoryArea] {
+        // If this ever fails, we need to model this differently in this crate.
+        assert_eq!(self.entry_size as usize, mem::size_of::<MemoryArea>());
+        &mut self.areas
+    }
 }
 
 impl TagTrait for MemoryMapTag {

multiboot2/src/tag.rs

@@ -38,6 +38,14 @@ impl Tag {
         unsafe { TagTrait::from_base_tag(self) }
     }
 
+    /// Casts the base tag to the specific tag type, but mutably.
+    pub fn cast_tag_mut<'a, T: TagTrait + ?Sized + 'a>(&'a mut self) -> &'a mut T {
+        assert_eq!(self.typ, T::ID);
+        // Safety: At this point, we trust that "self.size" and the size hint
+        // for DST tags are sane.
+        unsafe { TagTrait::from_base_tag_mut(self) }
+    }
+
     /// Some multiboot2 tags are a DST as they end with a dynamically sized byte
     /// slice. This function parses this slice as [`str`] so that either a valid
     /// UTF-8 Rust string slice without a terminating null byte or an error is
@@ -130,6 +138,55 @@ impl<'a> Iterator for TagIter<'a> {
     }
 }
 
+/// Iterates the MBI's tags from the first tag to the end tag.
+#[derive(Clone, Debug)]
+pub struct TagIterMut<'a> {
+    /// Pointer to the next tag. Updated in each iteration.
+    pub current: *mut Tag,
+    /// The pointer right after the MBI. Used for additional bounds checking.
+    end_ptr_exclusive: *mut u8,
+    /// Lifetime capture of the MBI's memory.
+    _mem: PhantomData<&'a mut ()>,
+}
+
+impl<'a> TagIterMut<'a> {
+    /// Creates a new iterator
+    pub fn new(mem: &'a mut [u8]) -> Self {
+        assert_eq!(mem.as_ptr().align_offset(8), 0);
+        TagIterMut {
+            current: mem.as_mut_ptr().cast(),
+            end_ptr_exclusive: unsafe { mem.as_mut_ptr().add(mem.len()) },
+            _mem: PhantomData,
+        }
+    }
+}
+
+impl<'a> Iterator for TagIterMut<'a> {
+    type Item = &'a mut Tag;
+
+    fn next(&mut self) -> Option<&'a mut Tag> {
+        // This never failed so far. But better be safe.
+        assert!(self.current.cast::<u8>() < self.end_ptr_exclusive);
+
+        let tag = unsafe { &mut *self.current };
+        match tag.typ() {
+            TagType::End => None, // end tag
+            _ => {
+                // We return the tag and update self.current already to the next
+                // tag.
+
+                // next pointer (rounded up to 8-byte alignment)
+                let ptr_offset = (tag.size as usize + 7) & !7;
+
+                // go to next tag
+                self.current = unsafe { self.current.cast::<u8>().add(ptr_offset).cast::<Tag>() };
+
+                Some(tag)
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

multiboot2/src/tag_trait.rs

@@ -56,4 +56,17 @@ pub trait TagTrait: Pointee {
         let ptr = ptr_meta::from_raw_parts(ptr.cast(), Self::dst_size(tag));
         &*ptr
     }
+
+    /// Creates a reference to a (dynamically sized) tag type in a safe way.
+    /// DST tags need to implement a proper [`Self::dst_size`] implementation.
+    ///
+    /// # Safety
+    /// Callers must be sure that the "size" field of the provided [`Tag`] is
+    /// sane and the underlying memory valid. The implementation of this trait
+    /// **must have** a correct [`Self::dst_size`] implementation.
+    unsafe fn from_base_tag_mut<'a>(tag: &mut Tag) -> &'a mut Self {
+        let ptr = core::ptr::addr_of_mut!(*tag);
+        let ptr = ptr_meta::from_raw_parts_mut(ptr.cast(), Self::dst_size(tag));
+        &mut *ptr
+    }
 }