Fix subtle bugs when handling dice data by consolidating logic

oak_dice/src/evidence.rs

@@ -17,13 +17,22 @@
 //! C-like structs for representing DICE data in environments where we don't have protocol buffer
 //! support.
 
-use alloc::{format, string::String};
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec::Vec,
+};
 use strum::{Display, FromRepr};
 use zerocopy::{AsBytes, FromBytes, FromZeroes};
 
 /// The maximum size of the signed attestation report.
 pub const REPORT_SIZE: usize = 2048;
 
+/// The size of the attestation report generated by AMD SEV-SNP.
+///
+/// See Table 21 in <https://www.amd.com/system/files/TechDocs/56860.pdf>.
+const AMD_SEV_SNP_ATTESTATION_REPORT_SIZE: usize = 1184;
+
 /// The maximum size of an ECDSA private key.
 pub const PRIVATE_KEY_SIZE: usize = 64;
 
@@ -74,8 +83,35 @@ pub struct RootLayerEvidence {
 }
 
 impl RootLayerEvidence {
-    pub fn get_tee_platform(&self) -> Option<TeePlatform> {
-        TeePlatform::from_repr(self.tee_platform)
+    pub fn get_tee_platform(&self) -> Result<TeePlatform, String> {
+        TeePlatform::from_repr(self.tee_platform).ok_or("invalid TEE Platform value".to_string())
+    }
+
+    pub fn get_remote_attestation_report(&self) -> Result<&[u8], String> {
+        match self.get_tee_platform()? {
+            TeePlatform::AmdSevSnp => {
+                if self.remote_attestation_report.len() < AMD_SEV_SNP_ATTESTATION_REPORT_SIZE {
+                    Err("attestation report smaller than expected".to_string())
+                } else {
+                    Ok(&self.remote_attestation_report[..AMD_SEV_SNP_ATTESTATION_REPORT_SIZE])
+                }
+            }
+            _ => Ok(&self.remote_attestation_report),
+        }
+    }
+
+    pub fn set_remote_attestation_report(&mut self, src: &[u8]) -> Result<(), String> {
+        crate::utils::padded_copy_from_slice(&mut self.remote_attestation_report, src)?;
+        Ok(())
+    }
+
+    pub fn get_eca_public_key(&self) -> Result<Vec<u8>, String> {
+        crate::utils::cbor_encoded_bytes_to_vec(&self.eca_public_key)
+    }
+
+    pub fn set_eca_public_key(&mut self, src: &[u8]) -> Result<(), String> {
+        crate::utils::padded_copy_from_slice(&mut self.eca_public_key, src)?;
+        Ok(())
     }
 }
 

oak_dice/src/utils.rs

@@ -16,7 +16,11 @@
 
 //! Utilities to handle encoded keys and certificates
 
-use alloc::{format, string::String, vec::Vec};
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec::Vec,
+};
 use core::result::Result;
 
 /// Extracts the bytes used to encode a CBOR object from a slice that might include unused bytes by
@@ -29,3 +33,16 @@ pub fn cbor_encoded_bytes_to_vec(bytes: &[u8]) -> Result<Vec<u8>, String> {
         .map_err(|err| format!("failed to write bytes: {:?}", err))?;
     Ok(result)
 }
+
+/// Like the slice `copy_from_slice` method but does not panic if slices are not
+/// the same length. In the case of a shorter source slice, only overwrites
+/// the beginning of the destination slice. In the case of a longer source slice
+/// it throws an error.
+pub(crate) fn padded_copy_from_slice(dst: &mut [u8], src: &[u8]) -> Result<(), String> {
+    if dst.len() < src.len() {
+        dst[..src.len()].copy_from_slice(src);
+        Ok(())
+    } else {
+        Err("destination slice shorter than source slice".to_string())
+    }
+}

oak_remote_attestation/src/dice.rs

@@ -33,11 +33,6 @@ use oak_dice::{
 };
 use p256::ecdsa::{SigningKey, VerifyingKey};
 
-/// The size of the attestation report generated by AMD SEV-SNP.
-///
-/// See Table 21 in <https://www.amd.com/system/files/TechDocs/56860.pdf>.
-const AMD_SEV_SNP_ATTESTATION_REPORT_SIZE: usize = 1184;
-
 /// Builds the DICE evidence and certificate authority for the next DICE layer.
 pub struct DiceBuilder {
     evidence: Evidence,
@@ -191,24 +186,22 @@ impl TryFrom<DiceData> for DiceBuilder {
 impl TryFrom<Stage0DiceData> for DiceData {
     type Error = anyhow::Error;
     fn try_from(value: Stage0DiceData) -> anyhow::Result<Self> {
-        let (platform, remote_attestation_report) =
-            if value.root_layer_evidence.tee_platform == TeePlatform::AmdSevSnp as u64 {
-                (
-                    TeePlatform::AmdSevSnp as u64,
-                    value.root_layer_evidence.remote_attestation_report
-                        [..AMD_SEV_SNP_ATTESTATION_REPORT_SIZE]
-                        .to_vec(),
-                )
-            } else {
-                (TeePlatform::Unspecified as u64, Vec::new())
-            };
-        let platform = platform.try_into().map_err(anyhow::Error::msg)?;
-        let eca_public_key = oak_dice::utils::cbor_encoded_bytes_to_vec(
-            &value.root_layer_evidence.eca_public_key[..],
-        )
-        .map_err(anyhow::Error::msg)?;
+        let platform: TeePlatform = value
+            .root_layer_evidence
+            .get_tee_platform()
+            .map_err(anyhow::Error::msg)?
+            .try_into()?;
+        let remote_attestation_report = value
+            .root_layer_evidence
+            .get_remote_attestation_report()
+            .map_err(anyhow::Error::msg)?
+            .to_vec();
+        let eca_public_key = value
+            .root_layer_evidence
+            .get_eca_public_key()
+            .map_err(anyhow::Error::msg)?;
         let root_layer = Some(RootLayerEvidence {
-            platform,
+            platform: platform as i32,
             remote_attestation_report,
             eca_public_key,
         });
@@ -236,6 +229,19 @@ impl TryFrom<Stage0DiceData> for DiceData {
     }
 }
 
+impl TryFrom<oak_dice::evidence::TeePlatform> for TeePlatform {
+    type Error = anyhow::Error;
+
+    fn try_from(src: oak_dice::evidence::TeePlatform) -> Result<Self, Self::Error> {
+        match src {
+            oak_dice::evidence::TeePlatform::AmdSevSnp => Ok(TeePlatform::AmdSevSnp),
+            oak_dice::evidence::TeePlatform::IntelTdx => Ok(TeePlatform::IntelTdx),
+            oak_dice::evidence::TeePlatform::Unspecified => Ok(TeePlatform::Unspecified),
+            _ => Err(anyhow::Error::msg("invalid TeePlatform")),
+        }
+    }
+}
+
 impl LayerEvidence {
     /// Extracts the `ClaimsSet` encoded in the ECA certificate of the layer.
     pub fn get_claims(&self) -> anyhow::Result<ClaimsSet> {

stage0/src/dice_attestation.rs

@@ -132,10 +132,14 @@ pub fn generate_dice_data(measurements: &Measurements) -> &'static Stage0DiceDat
 
     result.magic = STAGE0_MAGIC;
     result.root_layer_evidence.tee_platform = TeePlatform::AmdSevSnp as u64;
-    result.root_layer_evidence.remote_attestation_report[..report_bytes.len()]
-        .copy_from_slice(report_bytes);
-    result.root_layer_evidence.eca_public_key[..stage0_eca_verifying_key.len()]
-        .copy_from_slice(&stage0_eca_verifying_key);
+    result
+        .root_layer_evidence
+        .set_remote_attestation_report(report_bytes)
+        .expect("failed to set remote attestation report");
+    result
+        .root_layer_evidence
+        .set_eca_public_key(&stage0_eca_verifying_key)
+        .expect("failed to set eca public key");
     result.layer_1_evidence.eca_certificate[..stage1_eca_cert.len()]
         .copy_from_slice(&stage1_eca_cert);
     result.layer_1_certificate_authority.eca_private_key[..stage1_eca_signing_key.len()]