[dice] Replace DiceContext with OwnedDiceArtifacts in microdroid

The two structs have the same set of data and OwnedDiceArtifacts
is also more secure as it zeros out the DICE CDI values when it
is dropped.

Bug: 268587826
Test: atest microdroid_manager_test MicrodroidTests
Change-Id: I23a5c759af537e0784b9dfda00c2e8717bca4ec7
diff --git a/microdroid_manager/src/dice.rs b/microdroid_manager/src/dice.rs
index cf5e73e..9a2648f 100644
--- a/microdroid_manager/src/dice.rs
+++ b/microdroid_manager/src/dice.rs
@@ -17,7 +17,7 @@
 use anyhow::{bail, Context, Error, Result};
 use byteorder::{NativeEndian, ReadBytesExt};
 use diced_open_dice::{
-    retry_bcc_main_flow, Config, DiceMode, Hash, Hidden, InputValues, OwnedDiceArtifacts, CDI_SIZE,
+    retry_bcc_main_flow, Cdi, Config, DiceMode, Hash, Hidden, InputValues, OwnedDiceArtifacts,
 };
 use keystore2_crypto::ZVec;
 use libc::{c_void, mmap, munmap, MAP_FAILED, MAP_PRIVATE, PROT_READ};
@@ -29,33 +29,16 @@
 use std::ptr::null_mut;
 use std::slice;
 
-/// Artifacts that are kept in the process address space after the artifacts from the driver have
-/// been consumed.
-/// TODO(b/267575445): Replace with `OwnedDiceArtifacts` from the library `diced_open_dice`.
-pub struct DiceContext {
-    pub cdi_attest: [u8; CDI_SIZE],
-    pub cdi_seal: [u8; CDI_SIZE],
-    pub bcc: Vec<u8>,
-}
-
-impl From<OwnedDiceArtifacts> for DiceContext {
-    fn from(dice_artifacts: OwnedDiceArtifacts) -> Self {
-        Self {
-            cdi_attest: dice_artifacts.cdi_values.cdi_attest,
-            cdi_seal: dice_artifacts.cdi_values.cdi_seal,
-            bcc: dice_artifacts.bcc[..].to_vec(),
-        }
-    }
-}
-
-impl DiceContext {
-    pub fn get_sealing_key(&self, salt: &[u8], identifier: &[u8], keysize: u32) -> Result<ZVec> {
-        // Deterministically derive a key to use for sealing data based on salt. Use different salt
-        // for different keys.
-        let mut key = ZVec::new(keysize as usize)?;
-        hkdf(&mut key, Md::sha256(), &self.cdi_seal, salt, identifier)?;
-        Ok(key)
-    }
+/// Derives a sealing key from the DICE sealing CDI.
+pub fn derive_sealing_key(
+    cdi_seal: &Cdi,
+    salt: &[u8],
+    info: &[u8],
+    keysize: usize,
+) -> Result<ZVec> {
+    let mut key = ZVec::new(keysize)?;
+    hkdf(&mut key, Md::sha256(), cdi_seal, salt, info)?;
+    Ok(key)
 }
 
 /// Artifacts that are mapped into the process address space from the driver.
@@ -64,11 +47,11 @@
         driver_path: PathBuf,
         mmap_addr: *mut c_void,
         mmap_size: usize,
-        cdi_attest: &'a [u8; CDI_SIZE],
-        cdi_seal: &'a [u8; CDI_SIZE],
+        cdi_attest: &'a Cdi,
+        cdi_seal: &'a Cdi,
         bcc: &'a [u8],
     },
-    Fake(DiceContext),
+    Fake(OwnedDiceArtifacts),
 }
 
 impl DiceDriver<'_> {
@@ -81,7 +64,7 @@
             log::warn!("Using sample DICE values");
             let dice_artifacts = diced_sample_inputs::make_sample_bcc_and_cdis()
                 .expect("Failed to create sample dice artifacts.");
-            return Ok(Self::Fake(dice_artifacts.into()));
+            return Ok(Self::Fake(dice_artifacts));
         };
 
         let mut file = fs::File::open(driver_path)
@@ -133,12 +116,10 @@
         // input key material is already cryptographically strong.
         let cdi_seal = match self {
             Self::Real { cdi_seal, .. } => cdi_seal,
-            Self::Fake(fake) => &fake.cdi_seal,
+            Self::Fake(fake) => &fake.cdi_values.cdi_seal,
         };
         let salt = &[];
-        let mut key = ZVec::new(32)?;
-        hkdf(&mut key, Md::sha256(), cdi_seal, salt, identifier)?;
-        Ok(key)
+        derive_sealing_key(cdi_seal, salt, identifier, 32)
     }
 
     pub fn derive(
@@ -148,7 +129,7 @@
         authority_hash: Hash,
         debug: bool,
         hidden: Hidden,
-    ) -> Result<DiceContext> {
+    ) -> Result<OwnedDiceArtifacts> {
         let input_values = InputValues::new(
             code_hash,
             Config::Descriptor(config_desc),
@@ -158,7 +139,9 @@
         );
         let (cdi_attest, cdi_seal, bcc) = match &self {
             Self::Real { cdi_attest, cdi_seal, bcc, .. } => (*cdi_attest, *cdi_seal, *bcc),
-            Self::Fake(fake) => (&fake.cdi_attest, &fake.cdi_seal, fake.bcc.as_slice()),
+            Self::Fake(fake) => {
+                (&fake.cdi_values.cdi_attest, &fake.cdi_values.cdi_seal, fake.bcc.as_slice())
+            }
         };
         let dice_artifacts = retry_bcc_main_flow(cdi_attest, cdi_seal, bcc, &input_values)
             .context("DICE derive from driver")?;
@@ -168,7 +151,7 @@
             fs::write(driver_path, "wipe")
                 .map_err(|err| Error::new(err).context("Wiping driver"))?;
         }
-        Ok(dice_artifacts.into())
+        Ok(dice_artifacts)
     }
 }