From a0a26b745f64e9642c22cf8f1404a0074cb8129d Mon Sep 17 00:00:00 2001
From: Artemis Tosini <me@artem.ist>
Date: Sun, 4 Aug 2024 03:37:32 +0000
Subject: [PATCH] Handle alignment and size in Serialize

---
 src/abi.rs    | 28 ++++++++++++++++++++++------
 src/main.rs   | 15 ++++-----------
 src/object.rs |  2 +-
 3 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/src/abi.rs b/src/abi.rs
index f6e1cf8..2963dfe 100644
--- a/src/abi.rs
+++ b/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..]);
     }
 }
diff --git a/src/main.rs b/src/main.rs
index e646745..12c3abb 100644
--- a/src/main.rs
+++ b/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..]
-}
diff --git a/src/object.rs b/src/object.rs
index fe60cfd..4388ab2 100644
--- a/src/object.rs
+++ b/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]);