Handle alignment and size in Serialize

src/abi.rs

@@ -3,6 +3,12 @@ use alloc::string::String;
 use alloc::vec::Vec;
 use uefi::proto::console::gop::{GraphicsOutput, PixelFormat};
 
+fn align_slice(unaligned: &mut [u8], alignment: usize) -> &mut [u8] {
+    let base_addr = &unaligned[0] as *const u8 as usize;
+    let offset = base_addr.next_multiple_of(alignment) - base_addr;
+    &mut unaligned[offset..]
+}
+
 /// Serialize object into a buffer for FreeBSD
 pub trait Serialize {
     /// Output buffer must be aligned to this size
@@ -11,8 +17,18 @@ pub trait Serialize {
     /// Calculate output size of object, from base address to final byte written
     fn size(&self) -> usize;
 
-    /// Serliaze data into a buffer. Buffer must be aligned at alignment.
-    fn serialize(&self, out: &mut [u8]);
+    /// Serliaze data into a buffer. Buffer must be aligned at alignment and size of size().
+    fn serialize_raw(&self, out: &mut [u8]);
+
+    /// Serialize to an unaligned buffer, returning the end address
+    #[must_use]
+    fn serialize_unaligned<'a>(&self, out: &'a mut [u8]) -> &'a mut [u8] {
+        let aligned = align_slice(out, Self::alignment());
+        let size = self.size();
+        self.serialize_raw(&mut aligned[..size]);
+
+        &mut aligned[size..]
+    }
 }
 
 #[derive(Debug, Copy, Clone)]
@@ -70,7 +86,7 @@ impl Serialize for EFIFramebufferParams {
         core::mem::size_of::<Self>()
     }
 
-    fn serialize(&self, out: &mut [u8]) {
+    fn serialize_raw(&self, out: &mut [u8]) {
         // we're already repr(C) so just copy
         out.copy_from_slice(unsafe {
             core::slice::from_raw_parts(
@@ -166,7 +182,7 @@ impl Serialize for ModuleInfoValue {
         }
     }
 
-    fn serialize(&self, out: &mut [u8]) {
+    fn serialize_raw(&self, out: &mut [u8]) {
         match self {
             ModuleInfoValue::StaticString(s) => Self::serialize_str(s, out),
             ModuleInfoValue::String(s) => Self::serialize_str(s, out),
@@ -206,9 +222,9 @@ impl Serialize for ModuleInfoItem {
         8 + self.value.size()
     }
 
-    fn serialize(&self, out: &mut [u8]) {
+    fn serialize_raw(&self, out: &mut [u8]) {
         out[..4].copy_from_slice(&u32::to_ne_bytes(self.tag as u32));
         out[4..8].copy_from_slice(&u32::to_ne_bytes(self.value.size().try_into().unwrap()));
-        self.value.serialize(&mut out[8..]);
+        self.value.serialize_raw(&mut out[8..]);
     }
 }

src/main.rs

@@ -6,7 +6,7 @@ mod object;
 mod staging;
 
 use abi::{ModuleInfoItem, Serialize};
-use log::{info};
+use log::info;
 use object::{Kernel, Module};
 use uefi::{fs::FileSystem, prelude::*};
 
@@ -44,11 +44,10 @@ fn main(image_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
     let staging = staging::allocate_staging(system_table.boot_services(), alloc_size);
 
     {
-        let mut aligned = align_slice(staging, Kernel::alignment());
-        kernel.serialize(aligned);
+        let mut next = staging;
+        next = kernel.serialize_unaligned(next);
         for item in kernel_metadata.iter() {
-            aligned = align_slice(aligned, ModuleInfoItem::alignment());
-            item.serialize(aligned);
+            next = item.serialize_unaligned(next);
         }
     }
 
@@ -56,9 +55,3 @@ fn main(image_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
         system_table.boot_services().stall(10_000_000);
     }
 }
-
-fn align_slice(unaligned: &mut [u8], alignment: usize) -> &mut [u8] {
-    let base_addr = &unaligned[0] as *const u8 as usize;
-    let offset = base_addr.next_multiple_of(alignment) - base_addr;
-    &mut unaligned[offset..]
-}

src/object.rs

@@ -146,7 +146,7 @@ impl<'a> Serialize for Kernel<'a> {
         self.size_load_segments() + self.size_symbols()
     }
 
-    fn serialize(&self, out: &mut [u8]) {
+    fn serialize_raw(&self, out: &mut [u8]) {
         // Each section is rounded to sizeof(long)
         let end_load = self.size_load_segments();
         self.serialize_load_segments(&mut out[..end_load]);