optimize TAP decoding and verification

Signed-off-by: Wojciech Ozga <woz@zurich.ibm.com>

Author: wojciechozga

confidential-vms/linux_vm/hypervisor_rootfs/run_linux_vm_qemu.sh

@@ -12,7 +12,7 @@ TAP=/root/linux_vm/cove_tap_qemu
 HOST_PORT="$((3000 + RANDOM % 3000))"
 INTERACTIVE="-nographic"
 SMP=2
-MEMORY=2G
+MEMORY=1G
 
 for i in "$@"; do
   case $i in

security-monitor/rust-crates/riscv_cove_tap/src/lib.rs

@@ -17,6 +17,17 @@ pub use parser::AttestationPayloadParser;
 #[cfg(feature = "serializer")]
 pub use serializer::AttestationPayloadSerializer;
 
+#[macro_export]
+macro_rules! ensure {
+    ($cond:expr, $error:expr) => {
+        if !$cond {
+            Err($error)
+        } else {
+            Ok(())
+        }
+    };
+}
+
 pub use spec::*;
 pub use error::*;
 pub mod spec;

security-monitor/rust-crates/riscv_cove_tap/src/parser.rs

@@ -5,6 +5,7 @@ pub use crate::error::TapError;
 use alloc::vec::Vec;
 use crate::spec::*;
 use alloc::vec;
+use crate::ensure;
 
 pub struct AttestationPayloadParser {
     pub pointer: *const u8,
@@ -13,23 +14,17 @@ pub struct AttestationPayloadParser {
 
 impl AttestationPayloadParser {
     pub fn from_raw_pointer(pointer: *const u8, size: usize) -> Result<Self, TapError> {
-        Ok(Self {
-            pointer, size
-        })
+        Ok(Self { pointer, size })
     }
 
-    pub fn parse_and_verify(&mut self, decapsulation_key: &Vec<u8>) -> Result<AttestationPayload, TapError> {
-        if self.read_u32()? != ACE_MAGIC_TAP_START {
-            return Err(TapError::InvalidMagicStart());
-        }
+    pub fn parse_and_verify(&mut self, decapsulation_key: &[u8]) -> Result<AttestationPayload, TapError> {
+        ensure!(self.read_u32()? == ACE_MAGIC_TAP_START, TapError::InvalidMagicStart())?;
         self.read_u16()?;
         // if self.read_u16()? as usize != self.size {
         //     return Err(TapError::InvalidSize());
         // }
         let number_of_lockboxes = self.read_u16()?;
-        if usize::from(number_of_lockboxes) > MAX_NUMBER_OF_LOCKBOXES {
-            return Err(TapError::InvalidSize());
-        }
+        ensure!(usize::from(number_of_lockboxes) <= MAX_NUMBER_OF_LOCKBOXES, TapError::InvalidSize())?;
 
         let mut symmetric_key = vec![];
         for _ in 0..number_of_lockboxes {
@@ -44,20 +39,10 @@ impl AttestationPayloadParser {
             let tag_size = self.read_u16()? as usize;
             let tag = self.read_exact(tag_size)?;
             let tsk_size = self.read_u16()? as usize;
-            let tsk = self.read_exact(tsk_size)?;
-            match algorithm.decode(decapsulation_key, esk, nonce, tag, tsk) {
-                Ok(mut tsk) => {
-                    symmetric_key.append(&mut tsk);
-                    break;
-                }
-                Err(e) => {
-                    return Err(e)
-                }
-            };
-        }
-        if symmetric_key.is_empty() {
-            return Err(TapError::NoLockboxFound());
+            symmetric_key = self.read_exact(tsk_size)?;
+            algorithm.decode(decapsulation_key, &esk, &nonce, &tag, &mut symmetric_key)?;
         }
+        ensure!(!symmetric_key.is_empty(), TapError::NoLockboxFound())?;
 
         let payload_encryption_algorithm = PayloadEncryptionAlgorithm::from_u16(self.read_u16()?)?;
         match payload_encryption_algorithm {
@@ -66,9 +51,10 @@ impl AttestationPayloadParser {
         }
 
         let number_of_digests = self.read_u16()?;
-        let mut digests = vec![];
+        let mut digests = Vec::with_capacity(number_of_digests.into());
         for _ in 0..number_of_digests {
             let size = self.read_u16()? as usize;
+            ensure!(size >= 4, TapError::InvalidSize())?;
             let pcr_id = self.read_u16()?;
             let algorithm = DigestAlgorithm::from_u16(self.read_u16()?)?;
             let value = self.read_exact(size-4)?;
@@ -83,18 +69,16 @@ impl AttestationPayloadParser {
         let mut secrets = vec![];
         for _ in 0..number_of_secrets {
             let size = self.read_u16()? as usize;
+            ensure!(size >= 10, TapError::InvalidSize())?;
             let name = self.read_u64()? as u64;
             let value = self.read_exact(size-10)?;
             secrets.push(Secret { name, value });
         }
 
-        Ok(AttestationPayload {
-            digests,
-            secrets,
-        })
+        Ok(AttestationPayload { digests, secrets })
     }
 
-    fn decrypt_aes_gcm_256(&mut self, symmetric_key: &Vec<u8>) -> Result<(), TapError> {
+    fn decrypt_aes_gcm_256(&mut self, symmetric_key: &[u8]) -> Result<(), TapError> {
         use aes_gcm::{AeadInPlace, Aes256Gcm, Key, KeyInit, Tag, Nonce};
 
         let nonce_size = self.read_u16()? as usize;
@@ -103,10 +87,8 @@ impl AttestationPayloadParser {
         let tag = self.read_exact(tag_size)?;
         let payload_size = self.read_u16()? as usize;
 
-        if symmetric_key.len() != 32 {
-            return Err(TapError::InvalidTskSize());
-        }
-        let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(symmetric_key.as_slice()));
+        ensure!(symmetric_key.len() == 32, TapError::InvalidTskSize())?;
+        let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(symmetric_key));
         let nonce = Nonce::from_slice(&nonce);
         let tag = Tag::from_slice(&tag);
         let mut data_slice = unsafe{ core::slice::from_raw_parts_mut(self.pointer as *mut u8, payload_size) };
@@ -133,7 +115,7 @@ impl AttestationPayloadParser {
     }
 
     fn read_exact(&mut self, size: usize) -> Result<Vec<u8>, TapError> {
-        let mut result = vec![];
+        let mut result = Vec::with_capacity(size);
         for _ in 0..size {
             let value = unsafe { self.pointer.read_volatile() };
             self.pointer = self.pointer.wrapping_add(1);

security-monitor/rust-crates/riscv_cove_tap/src/spec.rs

@@ -88,29 +88,27 @@ impl LockboxAlgorithm {
     }
 
     #[cfg(feature = "parser")]
-    pub fn decode(&self, decapsulation_key: &Vec<u8>, esk: Vec<u8>, nonce: Vec<u8>, tag: Vec<u8>, mut tsk: Vec<u8>) -> Result<Vec<u8>, TapError> {
+    pub fn decode(&self, decapsulation_key: &[u8], esk: &[u8], nonce: &[u8], tag: &[u8], tsk: &mut Vec<u8>) -> Result<(), TapError> {
         match self {
             LockboxAlgorithm::Debug => {
-                Ok(tsk)
+                Ok(())
             },
             LockboxAlgorithm::MlKem1024Aes256 => {
                 use aes_gcm::{AeadInPlace, Aes256Gcm, Key, KeyInit, Tag, Nonce};
                 use hybrid_array::Array;
                 use ml_kem::{MlKem1024, KemCore, MlKem1024Params, Encoded, EncodedSizeUser,kem::{Decapsulate, DecapsulationKey}};
 
-                let m = Array::try_from(esk.as_slice())?;
-                let dk_bytes = Encoded::<DecapsulationKey::<MlKem1024Params>>::try_from(decapsulation_key.as_slice())?;
+                let m = Array::try_from(esk)?;
+                let dk_bytes = Encoded::<DecapsulationKey::<MlKem1024Params>>::try_from(decapsulation_key)?;
                 let dk = <MlKem1024 as KemCore>::DecapsulationKey::from_bytes(&dk_bytes);
                 let sk = match dk.decapsulate(&m) {
                     Ok(v) => v,
                     Err(_) => return Err(TapError::KemError())
                 };
 
                 let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(sk.as_slice()));
-                let nonce = Nonce::from_slice(&nonce);
-                let tag = Tag::from_slice(&tag);
-                cipher.decrypt_in_place_detached(nonce, b"", &mut tsk, &tag).unwrap();
-                Ok(tsk)
+                cipher.decrypt_in_place_detached(Nonce::from_slice(nonce), b"", tsk, &Tag::from_slice(tag)).unwrap();
+                Ok(())
             }
         }
     }