From eaacf348ff1636002b2a8fcb96c0011f90b6c810 Mon Sep 17 00:00:00 2001
From: landhb <landhb@users.noreply.github.com>
Date: Sat, 8 Jul 2023 10:51:41 -0400
Subject: [PATCH] Add alignment check to ElfBinary::new

---
 src/binary.rs | 44 ++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs    | 17 ++++++++++++++++-
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/src/binary.rs b/src/binary.rs
index 1b9f673..32bd3eb 100644
--- a/src/binary.rs
+++ b/src/binary.rs
@@ -34,6 +34,11 @@ impl<'s> fmt::Debug for ElfBinary<'s> {
 impl<'s> ElfBinary<'s> {
     /// Create a new ElfBinary.
     pub fn new(region: &'s [u8]) -> Result<ElfBinary<'s>, ElfLoaderErr> {
+        // Verify that the slice is aligned properly
+        if !crate::is_aligned_to(region.as_ptr() as usize, crate::ALIGNMENT) {
+            return Err(ElfLoaderErr::UnalignedMemory);
+        }
+
         let file = ElfFile::new(region)?;
 
         // Parse relevant parts out of the the .dynamic section
@@ -382,3 +387,42 @@ impl<'s> ElfBinary<'s> {
         self.file.program_iter().filter(select_load)
     }
 }
+
+#[test]
+fn test_load_unaligned() {
+    use core::ops::Deref;
+
+    #[repr(C, align(8))]
+    struct AlignedStackBuffer([u8; 8096]);
+    impl Deref for AlignedStackBuffer {
+        type Target = [u8; 8096];
+        fn deref(&self) -> &Self::Target {
+            &self.0
+        }
+    }
+
+    impl AlignedStackBuffer {
+        fn slice_at_index(&self, index: usize) -> &[u8] {
+            &self.0[index..]
+        }
+
+        fn buffer_from_file(&mut self, path: &str) {
+            let data = std::fs::read(path).unwrap();
+            let max = core::cmp::min(data.len(), self.0.len());
+            self.0[..max].copy_from_slice(&data[..max]);
+        }
+    }
+
+    // Read the file into an aligned buffer
+    let mut aligned = AlignedStackBuffer([0u8; 8096]);
+    aligned.buffer_from_file("test/test.riscv64");
+
+    // Verify aligned version works
+    let result = ElfBinary::new(aligned.deref());
+    assert!(result.is_ok());
+
+    // Verify unaligned version fails with appropriate error
+    let unaligned = aligned.slice_at_index(1);
+    let result = ElfBinary::new(unaligned);
+    assert_eq!(result.err().unwrap(), ElfLoaderErr::UnalignedMemory);
+}
diff --git a/src/lib.rs b/src/lib.rs
index 0de06c0..98bce8f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -21,12 +21,16 @@ use bitflags::bitflags;
 use xmas_elf::dynamic::*;
 use xmas_elf::program::ProgramIter;
 
-pub use xmas_elf::header::Machine;
+pub use xmas_elf::header::{Header, Machine};
 pub use xmas_elf::program::{Flags, ProgramHeader, ProgramHeader64};
 pub use xmas_elf::sections::{Rel, Rela};
 pub use xmas_elf::symbol_table::{Entry, Entry64};
 pub use xmas_elf::{P32, P64};
 
+/// Required alignment for zero-copy reads provided to xmas_elf by the
+/// zero crate.
+pub(crate) const ALIGNMENT: usize = core::mem::align_of::<Header>();
+
 /// An iterator over [`ProgramHeader`] whose type is `LOAD`.
 pub type LoadableHeaders<'a, 'b> = Filter<ProgramIter<'a, 'b>, fn(&ProgramHeader) -> bool>;
 pub type PAddr = u64;
@@ -47,6 +51,7 @@ pub struct RelocationEntry {
 pub enum ElfLoaderErr {
     ElfParser { source: &'static str },
     OutOfMemory,
+    UnalignedMemory,
     SymbolTableNotFound,
     UnsupportedElfFormat,
     UnsupportedElfVersion,
@@ -69,6 +74,7 @@ impl fmt::Display for ElfLoaderErr {
         match self {
             ElfLoaderErr::ElfParser { source } => write!(f, "Error in ELF parser: {}", source),
             ElfLoaderErr::OutOfMemory => write!(f, "Out of memory"),
+            ElfLoaderErr::UnalignedMemory => write!(f, "Data must be aligned to {:?}", ALIGNMENT),
             ElfLoaderErr::SymbolTableNotFound => write!(f, "No symbol table in the ELF file"),
             ElfLoaderErr::UnsupportedElfFormat => write!(f, "ELF format not supported"),
             ElfLoaderErr::UnsupportedElfVersion => write!(f, "ELF version not supported"),
@@ -166,6 +172,15 @@ pub trait ElfLoader {
     }
 }
 
+/// Utility function to verify alignment.
+///
+/// Note: this may be stabilized in the future as:
+///
+/// [core::ptr::is_aligned_to](https://doc.rust-lang.org/core/primitive.pointer.html#method.is_aligned_to)
+pub(crate) fn is_aligned_to(ptr: usize, align: usize) -> bool {
+    ptr & (align - 1) == 0
+}
+
 #[cfg(doctest)]
 mod test_readme {
     macro_rules! external_doc_test {