multiboot2: Get a mutable reference to the memory map

Author: YtvwlD

multiboot2/src/lib.rs

@@ -67,8 +67,8 @@ pub use memory_map::{
 pub use module::{ModuleIter, ModuleTag};
 pub use rsdp::{RsdpV1Tag, RsdpV2Tag};
 pub use smbios::SmbiosTag;
-use tag_type::TagIter;
 pub use tag_type::{EndTag, Tag, TagType, TagTypeId};
+use tag_type::{TagIter, TagIterMut};
 pub use vbe_info::{
     VBECapabilities, VBEControlInfo, VBEDirectColorAttributes, VBEField, VBEInfoTag,
     VBEMemoryModel, VBEModeAttributes, VBEModeInfo, VBEWindowAttributes,
@@ -105,8 +105,10 @@ pub const MAGIC: u32 = 0x36d76289;
 /// # Safety
 /// Deprecated. Please use BootInformation::load() instead.
 #[deprecated = "Please use BootInformation::load() instead."]
-pub unsafe fn load<'a>(address: usize) -> Result<BootInformation<'a>, MbiLoadError> {
-    let ptr = address as *const BootInformationHeader;
+pub unsafe fn load<'a>(
+    address: usize,
+) -> Result<BootInformation<&'a BootInformationInner>, MbiLoadError> {
+    let ptr = address as *mut BootInformationHeader;
     BootInformation::load(ptr)
 }
 
@@ -116,8 +118,8 @@ pub unsafe fn load<'a>(address: usize) -> Result<BootInformation<'a>, MbiLoadErr
 pub unsafe fn load_with_offset<'a>(
     address: usize,
     offset: usize,
-) -> Result<BootInformation<'a>, MbiLoadError> {
-    let ptr = address as *const u8;
+) -> Result<BootInformation<&'a BootInformationInner>, MbiLoadError> {
+    let ptr = address as *mut u8;
     let ptr = ptr.add(offset);
     BootInformation::load(ptr.cast())
 }
@@ -172,9 +174,9 @@ impl StructAsBytes 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],
 }
@@ -198,11 +200,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.
     ///
@@ -251,15 +259,46 @@ impl<'a> BootInformation<'a> {
 
         Ok(Self(mbi))
     }
+}
+
+impl BootInformation<&mut BootInformationInner> {
+    /// [`load`], but mutably.
+    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.
@@ -276,7 +315,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
     }
 
     /// Search for the basic memory info tag.
@@ -454,10 +493,10 @@ impl<'a> BootInformation<'a> {
     ///     .unwrap();
     /// assert_eq!(tag.name(), Ok("name"));
     /// ```
-    pub fn get_tag<TagT: TagTrait + ?Sized + 'a, TagType: Into<TagTypeId>>(
-        &'a self,
+    pub fn get_tag<TagT: TagTrait + ?Sized, TagType: Into<TagTypeId>>(
+        &self,
         typ: TagType,
-    ) -> Option<&'a TagT> {
+    ) -> Option<&TagT> {
         let typ = typ.into();
         self.tags()
             .find(|tag| tag.typ == typ)
@@ -466,15 +505,36 @@ impl<'a> BootInformation<'a> {
 
     /// Returns an iterator over all tags.
     fn tags(&self) -> TagIter {
-        TagIter::new(&self.0.tags)
+        TagIter::new(&self.0.as_ref().tags)
+    }
+}
+
+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 Send for BootInformation<'_> {}
+unsafe impl<T: AsRef<BootInformationInner>> Send for BootInformation<T> {}
 
-impl fmt::Debug for BootInformation<'_> {
+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.
@@ -557,6 +617,19 @@ pub trait TagTrait: Pointee {
         let ptr = ptr_meta::from_raw_parts(ptr.cast(), Self::dst_size(tag));
         &*ptr
     }
+
+    /// Creates a mutable 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
+    }
 }
 
 // All sized tags automatically have a Pointee implementation where
@@ -1280,7 +1353,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

@@ -60,6 +60,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_type.rs

@@ -311,6 +311,13 @@ impl Tag {
         unsafe { TagTrait::from_base_tag(self) }
     }
 
+    /// Casts the base tag to the specific tag type, but mutably.
+    pub(crate) fn cast_tag_mut<'a, T: TagTrait + ?Sized>(&mut self) -> &'a mut T {
+        // 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
@@ -431,6 +438,59 @@ impl<'a> Iterator for TagIter<'a> {
     }
 }
 
+/// Iterates the MBI's tags from the first tag to the end tag, mutably.
+#[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 ()>,
+}
+
+impl<'a> TagIterMut<'a> {
+    /// Creates a new iterator
+    pub fn new(mem: &'a mut [u8]) -> Self {
+        assert_eq!(mem.as_mut_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 {
+            &mut Tag {
+                // END-Tag
+                typ: TagTypeId(0),
+                size: 8,
+            } => None, // end tag
+            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::*;