add missing functionality in multiboot2-header (finding the header, getting tags)

multiboot2-header/src/header.rs

@@ -1,8 +1,10 @@
 use crate::{
     AddressHeaderTag, ConsoleHeaderTag, EfiBootServiceHeaderTag, EndHeaderTag,
     EntryAddressHeaderTag, EntryEfi32HeaderTag, EntryEfi64HeaderTag, FramebufferHeaderTag,
-    HeaderTag, HeaderTagISA, HeaderTagType, InformationRequestHeaderTag, RelocatableHeaderTag,
+    HeaderTag, HeaderTagISA, HeaderTagType, InformationRequestHeaderTag, ModuleAlignHeaderTag,
+    RelocatableHeaderTag,
 };
+use core::convert::TryInto;
 use core::fmt::{Debug, Formatter};
 use core::mem::size_of;
 
@@ -59,6 +61,28 @@ impl<'a> Multiboot2Header<'a> {
         Self { inner: reference }
     }
 
+    /// Find the header in a given slice.
+    pub fn find_header(buffer: &[u8]) -> Option<(&[u8], u32)> {
+        // the magic is 32 bit aligned and inside the first 8192 bytes
+        assert!(buffer.len() >= 8192);
+        let mut chunks = buffer[0..8192].chunks_exact(4);
+        let magic_index = match chunks.position(|vals| {
+            u32::from_le_bytes(vals.try_into().unwrap()) // yes, there's 4 bytes here
+            == MULTIBOOT2_HEADER_MAGIC
+        }) {
+            Some(idx) => idx * 4,
+            None => return None,
+        };
+        chunks.next(); // arch
+        let header_length: usize = u32::from_le_bytes(chunks.next().unwrap().try_into().unwrap())
+            .try_into()
+            .unwrap();
+        Some((
+            &buffer[magic_index..magic_index + header_length],
+            magic_index as u32,
+        ))
+    }
+
     /// Wrapper around [`Multiboot2BasicHeader::verify_checksum`].
     pub const fn verify_checksum(&self) -> bool {
         self.inner.verify_checksum()
@@ -87,6 +111,66 @@ impl<'a> Multiboot2Header<'a> {
     pub const fn calc_checksum(magic: u32, arch: HeaderTagISA, length: u32) -> u32 {
         Multiboot2BasicHeader::calc_checksum(magic, arch, length)
     }
+
+    /// Search for the address header tag.
+    pub fn address_tag(&self) -> Option<&AddressHeaderTag> {
+        self.get_tag(HeaderTagType::Address)
+            .map(|tag| unsafe { &*(tag as *const HeaderTag as *const AddressHeaderTag) })
+    }
+
+    /// Search for the entry address header tag.
+    pub fn entry_address_tag(&self) -> Option<&EntryAddressHeaderTag> {
+        self.get_tag(HeaderTagType::EntryAddress)
+            .map(|tag| unsafe { &*(tag as *const HeaderTag as *const EntryAddressHeaderTag) })
+    }
+
+    /// Search for the EFI32 entry address header tag.
+    pub fn entry_address_efi32_tag(&self) -> Option<&EntryEfi32HeaderTag> {
+        self.get_tag(HeaderTagType::EntryAddressEFI32)
+            .map(|tag| unsafe { &*(tag as *const HeaderTag as *const EntryEfi32HeaderTag) })
+    }
+
+    /// Search for the EFI64 entry address header tag.
+    pub fn entry_address_efi64_tag(&self) -> Option<&EntryEfi64HeaderTag> {
+        self.get_tag(HeaderTagType::EntryAddressEFI64)
+            .map(|tag| unsafe { &*(tag as *const HeaderTag as *const EntryEfi64HeaderTag) })
+    }
+
+    /// Search for the console flags header tag.
+    pub fn console_flags_tag(&self) -> Option<&ConsoleHeaderTag> {
+        self.get_tag(HeaderTagType::ConsoleFlags)
+            .map(|tag| unsafe { &*(tag as *const HeaderTag as *const ConsoleHeaderTag) })
+    }
+
+    /// Search for the framebuffer header tag.
+    pub fn framebuffer_tag(&self) -> Option<&FramebufferHeaderTag> {
+        self.get_tag(HeaderTagType::Framebuffer)
+            .map(|tag| unsafe { &*(tag as *const HeaderTag as *const FramebufferHeaderTag) })
+    }
+
+    /// Search for the module align header tag.
+    pub fn module_align_tag(&self) -> Option<&ModuleAlignHeaderTag> {
+        self.get_tag(HeaderTagType::ModuleAlign)
+            .map(|tag| unsafe { &*(tag as *const HeaderTag as *const ModuleAlignHeaderTag) })
+    }
+
+    /// Search for the EFI Boot Services header tag.
+    pub fn efi_boot_services_tag(&self) -> Option<&EfiBootServiceHeaderTag> {
+        self.get_tag(HeaderTagType::EfiBS)
+            .map(|tag| unsafe { &*(tag as *const HeaderTag as *const EfiBootServiceHeaderTag) })
+    }
+
+    /// Search for the EFI32 entry address header tag.
+    pub fn relocatable_tag(&self) -> Option<&RelocatableHeaderTag> {
+        self.get_tag(HeaderTagType::Relocatable)
+            .map(|tag| unsafe { &*(tag as *const HeaderTag as *const RelocatableHeaderTag) })
+    }
+
+    fn get_tag(&self, typ: HeaderTagType) -> Option<&HeaderTag> {
+        self.iter()
+            .map(|tag| unsafe { tag.as_ref() }.unwrap())
+            .find(|tag| tag.typ() == typ)
+    }
 }
 
 impl<'a> Debug for Multiboot2Header<'a> {
@@ -313,6 +397,10 @@ impl Debug for Multiboot2HeaderTagIter {
                 let entry = t as *const EntryEfi64HeaderTag;
                 let entry = &*(entry);
                 debug.entry(entry);
+            } else if typ == HeaderTagType::ModuleAlign {
+                let entry = t as *const ModuleAlignHeaderTag;
+                let entry = &*(entry);
+                debug.entry(entry);
             } else if typ == HeaderTagType::Relocatable {
                 let entry = t as *const RelocatableHeaderTag;
                 let entry = &*(entry);