Merge "vmbase: Handle stack overflows"
diff --git a/pvmfw/src/gpt.rs b/pvmfw/src/gpt.rs
index 6af3047..c3ccb5a 100644
--- a/pvmfw/src/gpt.rs
+++ b/pvmfw/src/gpt.rs
@@ -25,6 +25,7 @@
 use static_assertions::const_assert_eq;
 use uuid::Uuid;
 use virtio_drivers::device::blk::SECTOR_SIZE;
+use zerocopy::FromBytes;
 
 pub enum Error {
     /// VirtIO error during read operation.
@@ -101,8 +102,10 @@
     fn new(mut device: VirtIOBlk) -> Result<Self> {
         let mut blk = [0; Self::LBA_SIZE];
         device.read_block(Header::LBA, &mut blk).map_err(Error::FailedRead)?;
-        let (header_bytes, _) = blk.split_at(size_of::<Header>());
-        let header = Header::from_bytes(header_bytes).ok_or(Error::InvalidHeader)?;
+        let header = Header::read_from_prefix(blk.as_slice()).unwrap();
+        if !header.is_valid() {
+            return Err(Error::InvalidHeader);
+        }
         let entries_count = usize::try_from(header.entries_count()).unwrap();
 
         Ok(Self { device, entries_count })
@@ -151,6 +154,7 @@
 type Lba = u64;
 
 /// Structure as defined in release 2.10 of the UEFI Specification (5.3.2 GPT Header).
+#[derive(FromBytes)]
 #[repr(C, packed)]
 struct Header {
     signature: u64,
@@ -162,7 +166,7 @@
     backup_lba: Lba,
     first_lba: Lba,
     last_lba: Lba,
-    disk_guid: Uuid,
+    disk_guid: u128,
     entries_lba: Lba,
     entries_count: u32,
     entry_size: u32,
@@ -176,20 +180,6 @@
     const LBA: usize = 1;
     const ENTRIES_LBA: usize = 2;
 
-    fn from_bytes(bytes: &[u8]) -> Option<&Self> {
-        let bytes = bytes.get(..size_of::<Self>())?;
-        // SAFETY - We assume that bytes is properly aligned for Header and have verified above
-        // that it holds enough bytes. All potential values of the slice will produce a valid
-        // Header.
-        let header = unsafe { &*bytes.as_ptr().cast::<Self>() };
-
-        if header.is_valid() {
-            Some(header)
-        } else {
-            None
-        }
-    }
-
     fn is_valid(&self) -> bool {
         self.signature() == Self::SIGNATURE
             && self.header_size() == size_of::<Self>().try_into().unwrap()
diff --git a/pvmfw/src/instance.rs b/pvmfw/src/instance.rs
index a974543..95ddefd 100644
--- a/pvmfw/src/instance.rs
+++ b/pvmfw/src/instance.rs
@@ -26,13 +26,14 @@
 use crate::virtio::pci::VirtIOBlkIterator;
 use core::fmt;
 use core::mem::size_of;
-use core::slice;
 use diced_open_dice::DiceMode;
 use diced_open_dice::Hash;
 use diced_open_dice::Hidden;
 use log::trace;
 use uuid::Uuid;
 use virtio_drivers::transport::pci::bus::PciRoot;
+use zerocopy::AsBytes;
+use zerocopy::FromBytes;
 
 pub enum Error {
     /// Unexpected I/O error while accessing the underlying disk.
@@ -121,7 +122,7 @@
             let aead = AeadCtx::new_aes_256_gcm_randnonce(&key).map_err(Error::FailedOpen)?;
             let decrypted = aead.open(&mut entry, payload).map_err(Error::FailedOpen)?;
 
-            let body: &EntryBody = decrypted.as_ref();
+            let body = EntryBody::read_from(decrypted).unwrap();
             if body.code_hash != dice_inputs.code_hash {
                 Err(Error::RecordedCodeHashMismatch)
             } else if body.auth_hash != dice_inputs.auth_hash {
@@ -134,22 +135,21 @@
         }
         PvmfwEntry::New { header_index } => {
             let salt = rand::random_array().map_err(Error::FailedSaltGeneration)?;
-            let entry_body = EntryBody::new(dice_inputs, &salt);
-            let body = entry_body.as_ref();
+            let body = EntryBody::new(dice_inputs, &salt);
 
             let key = key.map_err(Error::FailedSeal)?;
             let aead = AeadCtx::new_aes_256_gcm_randnonce(&key).map_err(Error::FailedSeal)?;
             // We currently only support single-blk entries.
-            assert!(body.len() + aead.aead().unwrap().max_overhead() < blk.len());
-            let encrypted = aead.seal(&mut blk, body).map_err(Error::FailedSeal)?;
+            let plaintext = body.as_bytes();
+            assert!(plaintext.len() + aead.aead().unwrap().max_overhead() < blk.len());
+            let encrypted = aead.seal(&mut blk, plaintext).map_err(Error::FailedSeal)?;
             let payload_size = encrypted.len();
             let payload_index = header_index + 1;
             instance_img.write_block(payload_index, &blk).map_err(Error::FailedIo)?;
 
             let header = EntryHeader::new(PvmfwEntry::UUID, payload_size);
-            let (blk_header, blk_rest) = blk.split_at_mut(size_of::<EntryHeader>());
-            blk_header.copy_from_slice(header.as_ref());
-            blk_rest.fill(0);
+            header.write_to_prefix(blk.as_mut_slice()).unwrap();
+            blk[header.as_bytes().len()..].fill(0);
             instance_img.write_block(header_index, &blk).map_err(Error::FailedIo)?;
 
             Ok((true, salt))
@@ -157,6 +157,7 @@
     }
 }
 
+#[derive(FromBytes)]
 #[repr(C, packed)]
 struct Header {
     magic: [u8; Header::MAGIC.len()],
@@ -174,23 +175,6 @@
     fn version(&self) -> u16 {
         u16::from_le(self.version)
     }
-
-    fn from_bytes(bytes: &[u8]) -> Option<&Self> {
-        let header: &Self = bytes.as_ref();
-
-        if header.is_valid() {
-            Some(header)
-        } else {
-            None
-        }
-    }
-}
-
-impl AsRef<Header> for [u8] {
-    fn as_ref(&self) -> &Header {
-        // SAFETY - Assume that the alignement and size match Header.
-        unsafe { &*self.as_ptr().cast::<Header>() }
-    }
 }
 
 fn find_instance_img(pci_root: &mut PciRoot) -> Result<Partition> {
@@ -223,12 +207,15 @@
     let header_index = indices.next().ok_or(Error::MissingInstanceImageHeader)?;
     partition.read_block(header_index, &mut blk).map_err(Error::FailedIo)?;
     // The instance.img header is only used for discovery/validation.
-    let _ = Header::from_bytes(&blk).ok_or(Error::InvalidInstanceImageHeader)?;
+    let header = Header::read_from_prefix(blk.as_slice()).unwrap();
+    if !header.is_valid() {
+        return Err(Error::InvalidInstanceImageHeader);
+    }
 
     while let Some(header_index) = indices.next() {
         partition.read_block(header_index, &mut blk).map_err(Error::FailedIo)?;
 
-        let header: &EntryHeader = blk[..size_of::<EntryHeader>()].as_ref();
+        let header = EntryHeader::read_from_prefix(blk.as_slice()).unwrap();
         match (header.uuid(), header.payload_size()) {
             (uuid, _) if uuid.is_nil() => return Ok(PvmfwEntry::New { header_index }),
             (PvmfwEntry::UUID, payload_size) => {
@@ -250,7 +237,8 @@
 /// Marks the start of an instance.img entry.
 ///
 /// Note: Virtualization/microdroid_manager/src/instance.rs uses the name "partition".
-#[repr(C)]
+#[derive(AsBytes, FromBytes)]
+#[repr(C, packed)]
 struct EntryHeader {
     uuid: u128,
     payload_size: u64,
@@ -270,22 +258,7 @@
     }
 }
 
-impl AsRef<EntryHeader> for [u8] {
-    fn as_ref(&self) -> &EntryHeader {
-        assert_eq!(self.len(), size_of::<EntryHeader>());
-        // SAFETY - The size of the slice was checked and any value may be considered valid.
-        unsafe { &*self.as_ptr().cast::<EntryHeader>() }
-    }
-}
-
-impl AsRef<[u8]> for EntryHeader {
-    fn as_ref(&self) -> &[u8] {
-        let s = self as *const Self;
-        // SAFETY - Transmute the (valid) bytes into a slice.
-        unsafe { slice::from_raw_parts(s.cast::<u8>(), size_of::<Self>()) }
-    }
-}
-
+#[derive(AsBytes, FromBytes)]
 #[repr(C)]
 struct EntryBody {
     code_hash: Hash,
@@ -320,19 +293,3 @@
         }
     }
 }
-
-impl AsRef<EntryBody> for [u8] {
-    fn as_ref(&self) -> &EntryBody {
-        assert_eq!(self.len(), size_of::<EntryBody>());
-        // SAFETY - The size of the slice was checked and members are validated by accessors.
-        unsafe { &*self.as_ptr().cast::<EntryBody>() }
-    }
-}
-
-impl AsRef<[u8]> for EntryBody {
-    fn as_ref(&self) -> &[u8] {
-        let s = self as *const Self;
-        // SAFETY - Transmute the (valid) bytes into a slice.
-        unsafe { slice::from_raw_parts(s.cast::<u8>(), size_of::<Self>()) }
-    }
-}