multiboot2: more memory safety in TagIter

Author: phip1611

multiboot2/src/lib.rs

@@ -464,10 +464,9 @@ impl BootInformation<'_> {
             .map(|tag| tag.cast_tag::<TagT>())
     }
 
+    /// Returns an iterator over all tags.
     fn tags(&self) -> TagIter {
-        // The first tag starts 8 bytes after the begin of the boot info header
-        let ptr = core::ptr::addr_of!(self.0.tags).cast();
-        TagIter::new(ptr)
+        TagIter::new(&self.0.tags)
     }
 }
 

multiboot2/src/tag_type.rs

@@ -377,17 +377,25 @@ impl StructAsBytes for EndTag {
     }
 }
 
+/// Iterates the MBI's tags from the first tag to the end tag.
 #[derive(Clone, Debug)]
 pub struct TagIter<'a> {
+    /// Pointer to the next tag. Updated in each iteration.
     pub current: *const Tag,
-    phantom: PhantomData<&'a Tag>,
+    /// The pointer right after the MBI. Used for additional bounds checking.
+    end_ptr_exclusive: *const u8,
+    /// Lifetime capture of the MBI's memory.
+    _mem: PhantomData<&'a ()>,
 }
 
 impl<'a> TagIter<'a> {
-    pub fn new(first: *const Tag) -> Self {
+    /// Creates a new iterator
+    pub fn new(mem: &'a [u8]) -> Self {
+        assert_eq!(mem.as_ptr().align_offset(8), 0);
         TagIter {
-            current: first,
-            phantom: PhantomData,
+            current: mem.as_ptr().cast(),
+            end_ptr_exclusive: unsafe { mem.as_ptr().add(mem.len()) },
+            _mem: PhantomData,
         }
     }
 }
@@ -396,17 +404,25 @@ impl<'a> Iterator for TagIter<'a> {
     type Item = &'a Tag;
 
     fn next(&mut self) -> Option<&'a Tag> {
-        match unsafe { &*self.current } {
+        // This never failed so far. But better be safe.
+        assert!(self.current.cast::<u8>() < self.end_ptr_exclusive);
+
+        let tag = unsafe { &*self.current };
+        match tag {
             &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
-                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 *const _;
+                self.current = unsafe { self.current.cast::<u8>().add(ptr_offset).cast::<Tag>() };
 
                 Some(tag)
             }