builder: use new abstraction to guarantee alignment

Author: phip1611

multiboot2-header/Changelog.md

@@ -7,6 +7,7 @@
 - **BREAKING** renamed `MULTIBOOT2_HEADER_MAGIC` to `MAGIC`
 - **BREAKING** renamed `Multiboot2HeaderBuilder` to `HeaderBuilder`
 - **BREAKING** renamed `from_addr` to `load`. The function now consumes a ptr.
+- **BREAKING** `HeaderBuilder::build` now returns a value of type `HeaderBytes`
 - added the optional `unstable` feature (requires nightly)
   - implement `core::error::Error` for `LoadError`
 

multiboot2-header/src/builder/header.rs

@@ -8,8 +8,40 @@ use crate::{
     EntryAddressHeaderTag, EntryEfi32HeaderTag, EntryEfi64HeaderTag, FramebufferHeaderTag,
     ModuleAlignHeaderTag, Multiboot2BasicHeader, RelocatableHeaderTag,
 };
+use alloc::boxed::Box;
 use alloc::vec::Vec;
 use core::mem::size_of;
+use core::ops::Deref;
+
+/// Holds the raw bytes of a boot information built with [`HeaderBuilder`]
+/// on the heap. The bytes returned by [`HeaderBytes::as_bytes`] are
+/// guaranteed to be properly aligned.
+#[derive(Clone, Debug)]
+pub struct HeaderBytes {
+    // Offset into the bytes where the header starts. This is necessary to
+    // guarantee alignment at the moment.
+    offset: usize,
+    bytes: Box<[u8]>,
+}
+
+impl HeaderBytes {
+    /// Returns the bytes. They are guaranteed to be correctly aligned.
+    pub fn as_bytes(&self) -> &[u8] {
+        let slice = &self.bytes[self.offset..];
+        // At this point, the alignment is guaranteed. If not, something is
+        // broken fundamentally.
+        assert_eq!(slice.as_ptr().align_offset(8), 0);
+        slice
+    }
+}
+
+impl Deref for HeaderBytes {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        self.as_bytes()
+    }
+}
 
 /// Builder to construct a valid Multiboot2 header dynamically at runtime.
 /// The tags will appear in the order of their corresponding enumeration,
@@ -61,16 +93,11 @@ impl HeaderBuilder {
     /// easily calculate the size of a Multiboot2 header, where
     /// all the tags are 8-byte aligned.
     const fn size_or_up_aligned(size: usize) -> usize {
-        let remainder = size % 8;
-        if remainder == 0 {
-            size
-        } else {
-            size + 8 - remainder
-        }
+        (size + 7) & !7
     }
 
-    /// Returns the expected length of the Multiboot2 header,
-    /// when the `build()`-method gets called.
+    /// Returns the expected length of the Multiboot2 header, when the
+    /// [`Self::build`]-method gets called.
     pub fn expected_len(&self) -> usize {
         let base_len = size_of::<Multiboot2BasicHeader>();
         // size_or_up_aligned not required, because basic header length is 16 and the
@@ -115,8 +142,12 @@ impl HeaderBuilder {
     }
 
     /// Adds the bytes of a tag to the final Multiboot2 header byte vector.
-    /// Align should be true for all tags except the end tag.
     fn build_add_bytes(dest: &mut Vec<u8>, source: &[u8], is_end_tag: bool) {
+        let vec_next_write_ptr = unsafe { dest.as_ptr().add(dest.len()) };
+        // At this point, the alignment is guaranteed. If not, something is
+        // broken fundamentally.
+        assert_eq!(vec_next_write_ptr.align_offset(8), 0);
+
         dest.extend(source);
         if !is_end_tag {
             let size = source.len();
@@ -128,55 +159,89 @@ impl HeaderBuilder {
     }
 
     /// Constructs the bytes for a valid Multiboot2 header with the given properties.
-    /// The bytes can be casted to a Multiboot2 structure.
-    pub fn build(mut self) -> Vec<u8> {
-        let mut data = Vec::new();
+    pub fn build(self) -> HeaderBytes {
+        const ALIGN: usize = 8;
+
+        // We allocate more than necessary so that we can ensure an correct
+        // alignment within this data.
+        let alloc_len = self.expected_len() + 7;
+        let mut bytes = Vec::<u8>::with_capacity(alloc_len);
+        // Pointer to check that no relocation happened.
+        let alloc_ptr = bytes.as_ptr();
+
+        // As long as there is no nice way in stable Rust to guarantee the
+        // alignment of a vector, I add zero bytes at the beginning and the
+        // header might not start at the start of the allocation.
+        //
+        // Unfortunately, it is not possible to reliably test this in a unit
+        // test as long as the allocator_api feature is not stable.
+        // Due to my manual testing, however, it works.
+        let offset = bytes.as_ptr().align_offset(ALIGN);
+        bytes.extend([0].repeat(offset));
 
         Self::build_add_bytes(
-            &mut data,
+            &mut bytes,
             // important that we write the correct expected length into the header!
             &Multiboot2BasicHeader::new(self.arch, self.expected_len() as u32).struct_as_bytes(),
             false,
         );
 
-        if self.information_request_tag.is_some() {
-            Self::build_add_bytes(
-                &mut data,
-                &self.information_request_tag.take().unwrap().build(),
-                false,
-            )
+        if let Some(irs) = self.information_request_tag {
+            Self::build_add_bytes(&mut bytes, &irs.build(), false)
         }
         if let Some(tag) = self.address_tag.as_ref() {
-            Self::build_add_bytes(&mut data, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.entry_tag.as_ref() {
-            Self::build_add_bytes(&mut data, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.console_tag.as_ref() {
-            Self::build_add_bytes(&mut data, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.framebuffer_tag.as_ref() {
-            Self::build_add_bytes(&mut data, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.module_align_tag.as_ref() {
-            Self::build_add_bytes(&mut data, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.efi_bs_tag.as_ref() {
-            Self::build_add_bytes(&mut data, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.efi_32_tag.as_ref() {
-            Self::build_add_bytes(&mut data, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.efi_64_tag.as_ref() {
-            Self::build_add_bytes(&mut data, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
         }
         if let Some(tag) = self.relocatable_tag.as_ref() {
-            Self::build_add_bytes(&mut data, &tag.struct_as_bytes(), false)
+            Self::build_add_bytes(&mut bytes, &tag.struct_as_bytes(), false)
         }
+        Self::build_add_bytes(&mut bytes, &EndHeaderTag::new().struct_as_bytes(), true);
+
+        // Ensure that the vector has the same length as it's capacity. This is
+        // important so that miri doesn't complain that the boxed memory is
+        // smaller than the original allocation.
+        bytes.extend([0].repeat(bytes.capacity() - bytes.len()));
+
+        assert_eq!(
+            alloc_ptr,
+            bytes.as_ptr(),
+            "Vector was reallocated. Alignment of header probably broken!"
+        );
+        assert_eq!(
+            bytes[0..offset].iter().sum::<u8>(),
+            0,
+            "The offset to alignment area should be zero."
+        );
+
+        // Construct a box from a vec without `into_boxed_slice`. The latter
+        // calls `shrink` on the allocator, which might reallocate this memory.
+        // We don't want that!
+        let bytes = unsafe { Box::from_raw(bytes.leak()) };
 
-        Self::build_add_bytes(&mut data, &EndHeaderTag::new().struct_as_bytes(), true);
+        assert_eq!(bytes.len(), alloc_len);
 
-        data
+        HeaderBytes { offset, bytes }
     }
 
     // clippy thinks this can be a const fn but the compiler denies it
@@ -245,62 +310,69 @@ mod tests {
 
     #[test]
     fn test_builder() {
-        let builder = HeaderBuilder::new(HeaderTagISA::I386);
-        // Multiboot2 basic header + end tag
-        let mut expected_len = 16 + 8;
-        assert_eq!(builder.expected_len(), expected_len);
+        // Step 1/2: Build Header
+        let mb2_hdr_data = {
+            let builder = HeaderBuilder::new(HeaderTagISA::I386);
+            // Multiboot2 basic header + end tag
+            let mut expected_len = 16 + 8;
+            assert_eq!(builder.expected_len(), expected_len);
 
-        // add information request tag
-        let ifr_builder =
-            InformationRequestHeaderTagBuilder::new(HeaderTagFlag::Required).add_irs(&[
-                MbiTagType::EfiMmap,
-                MbiTagType::Cmdline,
-                MbiTagType::ElfSections,
-            ]);
-        let ifr_tag_size_with_padding = ifr_builder.expected_len() + 4;
-        assert_eq!(
-            ifr_tag_size_with_padding % 8,
-            0,
-            "the length of the IFR tag with padding must be a multiple of 8"
-        );
-        expected_len += ifr_tag_size_with_padding;
-        let builder = builder.information_request_tag(ifr_builder);
-        assert_eq!(builder.expected_len(), expected_len);
+            // add information request tag
+            let ifr_builder = InformationRequestHeaderTagBuilder::new(HeaderTagFlag::Required)
+                .add_irs(&[
+                    MbiTagType::EfiMmap,
+                    MbiTagType::Cmdline,
+                    MbiTagType::ElfSections,
+                ]);
+            let ifr_tag_size_with_padding = ifr_builder.expected_len() + 4;
+            assert_eq!(
+                ifr_tag_size_with_padding % 8,
+                0,
+                "the length of the IFR tag with padding must be a multiple of 8"
+            );
+            expected_len += ifr_tag_size_with_padding;
+            let builder = builder.information_request_tag(ifr_builder);
+            assert_eq!(builder.expected_len(), expected_len);
 
-        let builder = builder.relocatable_tag(RelocatableHeaderTag::new(
-            HeaderTagFlag::Required,
-            0x1337,
-            0xdeadbeef,
-            4096,
-            RelocatableHeaderTagPreference::None,
-        ));
-        expected_len += 0x18;
-        assert_eq!(builder.expected_len(), expected_len);
+            let builder = builder.relocatable_tag(RelocatableHeaderTag::new(
+                HeaderTagFlag::Required,
+                0x1337,
+                0xdeadbeef,
+                4096,
+                RelocatableHeaderTagPreference::None,
+            ));
+            expected_len += 0x18;
+            assert_eq!(builder.expected_len(), expected_len);
 
-        println!("builder: {:#?}", builder);
-        println!("expected_len: {} bytes", builder.expected_len());
+            println!("builder: {:#?}", builder);
+            println!("expected_len: {} bytes", builder.expected_len());
 
-        let mb2_hdr_data = builder.build();
-        let mb2_hdr = mb2_hdr_data.as_ptr().cast();
-        let mb2_hdr = unsafe { Multiboot2Header::load(mb2_hdr) }
-            .expect("the generated header to be loadable");
-        println!("{:#?}", mb2_hdr);
-        assert_eq!(
-            mb2_hdr.relocatable_tag().unwrap().flags(),
-            HeaderTagFlag::Required
-        );
-        assert_eq!(mb2_hdr.relocatable_tag().unwrap().min_addr(), 0x1337);
-        assert_eq!(mb2_hdr.relocatable_tag().unwrap().max_addr(), 0xdeadbeef);
-        assert_eq!(mb2_hdr.relocatable_tag().unwrap().align(), 4096);
-        assert_eq!(
-            mb2_hdr.relocatable_tag().unwrap().preference(),
-            RelocatableHeaderTagPreference::None
-        );
+            builder.build()
+        };
+
+        // Step 2/2: Test the built Header
+        {
+            let mb2_hdr = mb2_hdr_data.as_ptr().cast();
+            let mb2_hdr = unsafe { Multiboot2Header::load(mb2_hdr) }
+                .expect("the generated header to be loadable");
+            println!("{:#?}", mb2_hdr);
+            assert_eq!(
+                mb2_hdr.relocatable_tag().unwrap().flags(),
+                HeaderTagFlag::Required
+            );
+            assert_eq!(mb2_hdr.relocatable_tag().unwrap().min_addr(), 0x1337);
+            assert_eq!(mb2_hdr.relocatable_tag().unwrap().max_addr(), 0xdeadbeef);
+            assert_eq!(mb2_hdr.relocatable_tag().unwrap().align(), 4096);
+            assert_eq!(
+                mb2_hdr.relocatable_tag().unwrap().preference(),
+                RelocatableHeaderTagPreference::None
+            );
 
-        /* you can write the binary to a file and a tool such as crate "bootinfo"
-           will be able to fully parse the MB2 header
-        let mut file = std::file::File::create("mb2_hdr.bin").unwrap();
-        use std::io::Write;
-        file.write_all(mb2_hdr_data.as_slice()).unwrap();*/
+            /* you can write the binary to a file and a tool such as crate "bootinfo"
+               will be able to fully parse the MB2 header
+            let mut file = std::file::File::create("mb2_hdr.bin").unwrap();
+            use std::io::Write;
+            file.write_all(mb2_hdr_data.as_slice()).unwrap();*/
+        }
     }
 }