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,9 +159,22 @@ 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(mut self) -> HeaderBytes {
+        // We allocate more than necessary so that we can ensure an correct
+        // alignment within this data.
+        let mut data = Vec::<u8>::with_capacity(self.expected_len() + 7);
+        // Pointer to check that no relocation happened.
+        let alloc_ptr = data.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 = data.as_ptr().align_offset(8);
+        data.extend([0].repeat(offset));
 
         Self::build_add_bytes(
             &mut data,
@@ -176,7 +220,28 @@ impl HeaderBuilder {
 
         Self::build_add_bytes(&mut data, &EndHeaderTag::new().struct_as_bytes(), true);
 
-        data
+        assert_eq!(
+            alloc_ptr,
+            data.as_ptr(),
+            "Vector was reallocated. Alignment of header probably broken!"
+        );
+        assert_eq!(
+            data[0..offset].iter().sum::<u8>(),
+            0,
+            "The offset to alignment area should be zero."
+        );
+
+        let boxed = data.into_boxed_slice();
+        assert_eq!(
+            boxed.as_ptr(),
+            alloc_ptr,
+            "Memory was reallocated. Alignment of header probably broken!"
+        );
+
+        HeaderBytes {
+            offset,
+            bytes: boxed,
+        }
     }
 
     // clippy thinks this can be a const fn but the compiler denies it