multiboot2: Get a mutable reference to the memory map

Author: YtvwlD

multiboot2/src/lib.rs

@@ -63,8 +63,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,
@@ -276,6 +276,11 @@ impl BootInformation {
         self.get_tag::<MemoryMapTag, _>(TagType::Mmap)
     }
 
+    /// 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)
+    }
+
     /// Get an iterator of all module tags.
     pub fn module_tags(&self) -> ModuleIter {
         module::module_iter(self.tags())
@@ -427,9 +432,23 @@ impl BootInformation {
             .map(|tag| tag.cast_tag::<TagT>())
     }
 
+    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(&self) -> TagIter {
         TagIter::new(unsafe { self.inner.offset(1) } as *const _)
     }
+
+    fn tags_mut(&mut self) -> TagIterMut {
+        TagIterMut::new(unsafe { self.inner.offset(1) } as *mut _)
+    }
 }
 
 impl BootInformationInner {
@@ -567,6 +586,19 @@ pub trait TagTrait: Pointee {
         let ptr = ptr_meta::from_raw_parts(ptr, 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 = tag as *mut _ as *mut ();
+        let ptr = ptr_meta::from_raw_parts_mut(ptr, Self::dst_size(tag));
+        &mut *ptr
+    }
 }
 
 // All sized tags automatically have a Pointee implementation where

multiboot2/src/memory_map.rs

@@ -61,6 +61,18 @@ impl MemoryMapTag {
             phantom: PhantomData,
         }
     }
+
+    /// Return a mutable iterator over all marked memory areas.
+    pub fn all_memory_areas_mut(&mut self) -> impl Iterator<Item = &mut MemoryArea> {
+        let self_ptr = self as *mut MemoryMapTag;
+        let start_area = (&mut self.areas[0]) as *mut MemoryArea;
+        MemoryAreaIterMut {
+            current_area: start_area as u64,
+            last_area: (self_ptr as *const () as u64 + (self.size - self.entry_size) as u64),
+            entry_size: self.entry_size,
+            phantom: PhantomData,
+        }
+    }
 }
 
 impl TagTrait for MemoryMapTag {
@@ -172,6 +184,28 @@ impl<'a> Iterator for MemoryAreaIter<'a> {
     }
 }
 
+/// A mutable iterator over all memory areas
+#[derive(Clone, Debug)]
+pub struct MemoryAreaIterMut<'a> {
+    current_area: u64,
+    last_area: u64,
+    entry_size: u32,
+    phantom: PhantomData<&'a MemoryArea>,
+}
+
+impl<'a> Iterator for MemoryAreaIterMut<'a> {
+    type Item = &'a mut MemoryArea;
+    fn next(&mut self) -> Option<&'a mut MemoryArea> {
+        if self.current_area > self.last_area {
+            None
+        } else {
+            let area = unsafe { &mut *(self.current_area as *mut MemoryArea) };
+            self.current_area += self.entry_size as u64;
+            Some(area)
+        }
+    }
+}
+
 /// Basic memory info
 ///
 /// This tag includes "basic memory information".

multiboot2/src/tag_type.rs

@@ -310,6 +310,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
@@ -414,6 +421,43 @@ impl<'a> Iterator for TagIter<'a> {
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct TagIterMut<'a> {
+    pub current: *mut Tag,
+    phantom: PhantomData<&'a mut Tag>,
+}
+
+impl<'a> TagIterMut<'a> {
+    pub fn new(first: *mut Tag) -> Self {
+        TagIterMut {
+            current: first,
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl<'a> Iterator for TagIterMut<'a> {
+    type Item = &'a mut Tag;
+
+    fn next(&mut self) -> Option<&'a mut Tag> {
+        match unsafe { &mut *self.current } {
+            &mut Tag {
+                // END-Tag
+                typ: TagTypeId(0),
+                size: 8,
+            } => None, // end tag
+            tag => {
+                // go to next tag
+                let mut tag_addr = self.current as usize;
+                tag_addr += ((tag.size + 7) & !7) as usize; //align at 8 byte
+                self.current = tag_addr as *mut _;
+
+                Some(tag)
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;