fix panic condition in align_up function

Fixed a panic condition that could cause loading malformed
ELF kernel images containing unaligned values in the
`n_namesz` and `n_descsz` fields of ELFNOTE headers to lead
to the abrupt termination of the vmm process.

Signed-off-by: Alexandra Iordache <aghecen@amazon.com>
Signed-off-by: Andreea Florescu <fandree@amazon.com>

Author: Alexandra Iordache

src/loader/x86_64/elf/mod.rs

@@ -15,6 +15,7 @@
 use std::fmt;
 use std::io::{Read, Seek, SeekFrom};
 use std::mem;
+use std::result;
 
 use vm_memory::{Address, ByteValued, Bytes, GuestAddress, GuestMemory, GuestUsize};
 
@@ -29,9 +30,11 @@ unsafe impl ByteValued for elf::Elf64_Phdr {}
 #[derive(Debug, PartialEq)]
 /// Elf kernel loader errors.
 pub enum Error {
+    /// Invalid alignment.
+    Align,
     /// Loaded big endian binary on a little endian platform.
     BigEndianElfOnLittle,
-    /// Invalid ELF magic number
+    /// Invalid ELF magic number.
     InvalidElfMagicNumber,
     /// Invalid program header size.
     InvalidProgramHeaderSize,
@@ -64,6 +67,7 @@ pub enum Error {
 impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let desc = match self {
+            Error::Align => "Invalid alignment",
             Error::BigEndianElfOnLittle => {
                 "Trying to load big-endian binary on little-endian machine"
             }
@@ -332,8 +336,8 @@ where
 
         // Skip the note header plus the size of its fields (with alignment).
         read_size += nhdr_sz
-            + align_up(u64::from(nhdr.n_namesz), n_align)
-            + align_up(u64::from(nhdr.n_descsz), n_align);
+            + align_up(u64::from(nhdr.n_namesz), n_align)?
+            + align_up(u64::from(nhdr.n_descsz), n_align)?;
 
         kernel_image
             .seek(SeekFrom::Start(phdr.p_offset + read_size as u64))
@@ -350,7 +354,7 @@ where
     // just skip the name field contents.
     kernel_image
         .seek(SeekFrom::Current(
-            align_up(u64::from(nhdr.n_namesz), n_align) as i64 - PVH_NOTE_STR_SZ as i64,
+            align_up(u64::from(nhdr.n_namesz), n_align)? as i64 - PVH_NOTE_STR_SZ as i64,
         ))
         .map_err(|_| Error::SeekNoteHeader)?;
 
@@ -372,18 +376,20 @@ where
     )))
 }
 
-/// Align address upwards. Taken from x86_64 crate:
-/// https://docs.rs/x86_64/latest/x86_64/fn.align_up.html
+/// Align address upwards. Adapted from x86_64 crate:
+/// https://docs.rs/x86_64/latest/x86_64/addr/fn.align_up.html
 ///
-/// Returns the smallest x with alignment `align` so that x >= addr. The alignment must be
-/// a power of 2.
-fn align_up(addr: u64, align: u64) -> usize {
-    assert!(align.is_power_of_two(), "`align` must be a power of two");
+/// Returns the smallest x with alignment `align` so that x >= addr if the alignment is a power of
+/// 2, or an error otherwise.
+fn align_up(addr: u64, align: u64) -> result::Result<usize, Error> {
+    if !align.is_power_of_two() {
+        return Err(Error::Align);
+    }
     let align_mask = align - 1;
     if addr & align_mask == 0 {
-        addr as usize // already aligned
+        Ok(addr as usize) // already aligned
     } else {
-        ((addr | align_mask) + 1) as usize
+        Ok(((addr | align_mask) + 1) as usize)
     }
 }
 
@@ -417,6 +423,10 @@ mod tests {
         include_bytes!("test_badnote.bin").to_vec()
     }
 
+    fn make_bad_align() -> Vec<u8> {
+        include_bytes!("test_align.bin").to_vec()
+    }
+
     #[test]
     fn test_load_elf() {
         let gm = create_guest_mem();
@@ -539,4 +549,14 @@ mod tests {
             Elf::load(&gm, None, &mut Cursor::new(&badnote_image), None).err()
         );
     }
+
+    #[test]
+    fn test_bad_align() {
+        let gm = GuestMemoryMmap::from_ranges(&[(GuestAddress(0x0), (0x10_000_000usize))]).unwrap();
+        let bad_align_image = make_bad_align();
+        assert_eq!(
+            Some(KernelLoaderError::Elf(Error::Align)),
+            Elf::load(&gm, None, &mut Cursor::new(&bad_align_image), None).err()
+        );
+    }
 }