Merge changes I0dfcac1d,I82287a6e,I210dee00 into main

* changes:
  pvmfw: main: Introduce AVB helper function
  pvmfw: main: Introduce DICE helper functions
  pvmfw: Move instance_hash in dice::PartialInputs
diff --git a/guest/pvmfw/src/dice/mod.rs b/guest/pvmfw/src/dice/mod.rs
index 8317e48..dc7b64e 100644
--- a/guest/pvmfw/src/dice/mod.rs
+++ b/guest/pvmfw/src/dice/mod.rs
@@ -87,11 +87,12 @@
     pub mode: DiceMode,
     pub security_version: u64,
     pub rkp_vm_marker: bool,
+    pub instance_hash: Option<Hash>,
     component_name: String,
 }
 
 impl PartialInputs {
-    pub fn new(data: &VerifiedBootData) -> Result<Self> {
+    pub fn new(data: &VerifiedBootData, instance_hash: Option<Hash>) -> Result<Self> {
         let code_hash = to_dice_hash(data)?;
         let auth_hash = hash(data.public_key)?;
         let mode = to_dice_mode(data.debug_level);
@@ -101,20 +102,27 @@
         let rkp_vm_marker = data.has_capability(Capability::RemoteAttest)
             || data.has_capability(Capability::TrustySecurityVm);
 
-        Ok(Self { code_hash, auth_hash, mode, security_version, rkp_vm_marker, component_name })
+        Ok(Self {
+            code_hash,
+            auth_hash,
+            mode,
+            security_version,
+            rkp_vm_marker,
+            instance_hash,
+            component_name,
+        })
     }
 
     pub fn write_next_handover(
         self,
         current_handover: &[u8],
         salt: &[u8; HIDDEN_SIZE],
-        instance_hash: Option<Hash>,
         deferred_rollback_protection: bool,
         next_handover: &mut [u8],
         context: DiceContext,
     ) -> Result<()> {
         let config = self
-            .generate_config_descriptor(instance_hash)
+            .generate_config_descriptor()
             .map_err(|_| diced_open_dice::DiceError::InvalidInput)?;
 
         let dice_inputs = InputValues::new(
@@ -160,14 +168,14 @@
         )
     }
 
-    fn generate_config_descriptor(&self, instance_hash: Option<Hash>) -> Result<Vec<u8>> {
+    fn generate_config_descriptor(&self) -> Result<Vec<u8>> {
         let mut config = Vec::with_capacity(4);
         config.push((cbor!(COMPONENT_NAME_KEY)?, cbor!(self.component_name.as_str())?));
         config.push((cbor!(SECURITY_VERSION_KEY)?, cbor!(self.security_version)?));
         if self.rkp_vm_marker {
             config.push((cbor!(RKP_VM_MARKER_KEY)?, Value::Null))
         }
-        if let Some(instance_hash) = instance_hash {
+        if let Some(instance_hash) = self.instance_hash {
             config.push((cbor!(INSTANCE_HASH_KEY)?, Value::from(instance_hash.as_slice())));
         }
         let config = Value::Map(config);
@@ -217,7 +225,7 @@
     #[test]
     fn base_data_conversion() {
         let vb_data = BASE_VB_DATA;
-        let inputs = PartialInputs::new(&vb_data).unwrap();
+        let inputs = PartialInputs::new(&vb_data, None).unwrap();
 
         assert_eq!(inputs.mode, DiceMode::kDiceModeNormal);
         assert_eq!(inputs.security_version, 42);
@@ -229,7 +237,7 @@
     #[test]
     fn debuggable_conversion() {
         let vb_data = VerifiedBootData { debug_level: DebugLevel::Full, ..BASE_VB_DATA };
-        let inputs = PartialInputs::new(&vb_data).unwrap();
+        let inputs = PartialInputs::new(&vb_data, None).unwrap();
 
         assert_eq!(inputs.mode, DiceMode::kDiceModeDebug);
     }
@@ -238,7 +246,7 @@
     fn rkp_vm_conversion() {
         let vb_data =
             VerifiedBootData { capabilities: vec![Capability::RemoteAttest], ..BASE_VB_DATA };
-        let inputs = PartialInputs::new(&vb_data).unwrap();
+        let inputs = PartialInputs::new(&vb_data, None).unwrap();
 
         assert!(inputs.rkp_vm_marker);
     }
@@ -246,22 +254,23 @@
     #[test]
     fn base_config_descriptor() {
         let vb_data = BASE_VB_DATA;
-        let inputs = PartialInputs::new(&vb_data).unwrap();
-        let config_map = decode_config_descriptor(&inputs, None);
+        let inputs = PartialInputs::new(&vb_data, None).unwrap();
+        let config_map = decode_config_descriptor(&inputs);
 
         assert_eq!(config_map.get(&COMPONENT_NAME_KEY).unwrap().as_text().unwrap(), "vm_entry");
         assert_eq!(config_map.get(&COMPONENT_VERSION_KEY), None);
         assert_eq!(config_map.get(&RESETTABLE_KEY), None);
         assert_eq!(config_map.get(&SECURITY_VERSION_KEY).unwrap().as_integer().unwrap(), 42.into());
         assert_eq!(config_map.get(&RKP_VM_MARKER_KEY), None);
+        assert_eq!(config_map.get(&INSTANCE_HASH_KEY), None);
     }
 
     #[test]
     fn rkp_vm_config_descriptor_has_rkp_vm_marker_and_component_name() {
         let vb_data =
             VerifiedBootData { capabilities: vec![Capability::RemoteAttest], ..BASE_VB_DATA };
-        let inputs = PartialInputs::new(&vb_data).unwrap();
-        let config_map = decode_config_descriptor(&inputs, Some(HASH));
+        let inputs = PartialInputs::new(&vb_data, Some(HASH)).unwrap();
+        let config_map = decode_config_descriptor(&inputs);
 
         assert_eq!(config_map.get(&COMPONENT_NAME_KEY).unwrap().as_text().unwrap(), "vm_entry");
         assert!(config_map.get(&RKP_VM_MARKER_KEY).unwrap().is_null());
@@ -271,8 +280,8 @@
     fn security_vm_config_descriptor_has_rkp_vm_marker() {
         let vb_data =
             VerifiedBootData { capabilities: vec![Capability::TrustySecurityVm], ..BASE_VB_DATA };
-        let inputs = PartialInputs::new(&vb_data).unwrap();
-        let config_map = decode_config_descriptor(&inputs, Some(HASH));
+        let inputs = PartialInputs::new(&vb_data, Some(HASH)).unwrap();
+        let config_map = decode_config_descriptor(&inputs);
 
         assert!(config_map.get(&RKP_VM_MARKER_KEY).unwrap().is_null());
     }
@@ -281,8 +290,8 @@
     fn config_descriptor_with_instance_hash() {
         let vb_data =
             VerifiedBootData { capabilities: vec![Capability::RemoteAttest], ..BASE_VB_DATA };
-        let inputs = PartialInputs::new(&vb_data).unwrap();
-        let config_map = decode_config_descriptor(&inputs, Some(HASH));
+        let inputs = PartialInputs::new(&vb_data, Some(HASH)).unwrap();
+        let config_map = decode_config_descriptor(&inputs);
         assert_eq!(*config_map.get(&INSTANCE_HASH_KEY).unwrap(), Value::from(HASH.as_slice()));
     }
 
@@ -290,16 +299,13 @@
     fn config_descriptor_without_instance_hash() {
         let vb_data =
             VerifiedBootData { capabilities: vec![Capability::RemoteAttest], ..BASE_VB_DATA };
-        let inputs = PartialInputs::new(&vb_data).unwrap();
-        let config_map = decode_config_descriptor(&inputs, None);
+        let inputs = PartialInputs::new(&vb_data, None).unwrap();
+        let config_map = decode_config_descriptor(&inputs);
         assert!(!config_map.contains_key(&INSTANCE_HASH_KEY));
     }
 
-    fn decode_config_descriptor(
-        inputs: &PartialInputs,
-        instance_hash: Option<Hash>,
-    ) -> HashMap<i64, Value> {
-        let config_descriptor = inputs.generate_config_descriptor(instance_hash).unwrap();
+    fn decode_config_descriptor(inputs: &PartialInputs) -> HashMap<i64, Value> {
+        let config_descriptor = inputs.generate_config_descriptor().unwrap();
 
         let cbor_map =
             cbor_util::deserialize::<Value>(&config_descriptor).unwrap().into_map().unwrap();
@@ -313,7 +319,7 @@
     #[test]
     fn changing_deferred_rpb_changes_secrets() {
         let vb_data = VerifiedBootData { debug_level: DebugLevel::Full, ..BASE_VB_DATA };
-        let inputs = PartialInputs::new(&vb_data).unwrap();
+        let inputs = PartialInputs::new(&vb_data, Some([0u8; 64])).unwrap();
         let mut buffer_without_defer = [0; 4096];
         let mut buffer_with_defer = [0; 4096];
         let mut buffer_without_defer_retry = [0; 4096];
@@ -341,7 +347,6 @@
             .write_next_handover(
                 sample_dice_input,
                 &[0u8; HIDDEN_SIZE],
-                Some([0u8; 64]),
                 false,
                 &mut buffer_without_defer,
                 context.clone(),
@@ -354,7 +359,6 @@
             .write_next_handover(
                 sample_dice_input,
                 &[0u8; HIDDEN_SIZE],
-                Some([0u8; 64]),
                 true,
                 &mut buffer_with_defer,
                 context.clone(),
@@ -367,7 +371,6 @@
             .write_next_handover(
                 sample_dice_input,
                 &[0u8; HIDDEN_SIZE],
-                Some([0u8; 64]),
                 false,
                 &mut buffer_without_defer_retry,
                 context.clone(),
@@ -384,7 +387,7 @@
         let dice_artifacts = make_sample_bcc_and_cdis().unwrap();
         let handover0_bytes = to_serialized_handover(&dice_artifacts);
         let vb_data = VerifiedBootData { debug_level: DebugLevel::Full, ..BASE_VB_DATA };
-        let inputs = PartialInputs::new(&vb_data).unwrap();
+        let inputs = PartialInputs::new(&vb_data, Some([0u8; 64])).unwrap();
         let mut buffer = [0; 4096];
 
         inputs
@@ -392,7 +395,6 @@
             .write_next_handover(
                 &handover0_bytes,
                 &[0u8; HIDDEN_SIZE],
-                Some([0u8; 64]),
                 true,
                 &mut buffer,
                 DiceContext {
@@ -410,7 +412,6 @@
             .write_next_handover(
                 &handover1_bytes,
                 &[0u8; HIDDEN_SIZE],
-                Some([0u8; 64]),
                 true,
                 &mut buffer,
                 DiceContext {
@@ -428,7 +429,6 @@
             .write_next_handover(
                 &handover2_bytes,
                 &[0u8; HIDDEN_SIZE],
-                Some([0u8; 64]),
                 true,
                 &mut buffer,
                 DiceContext {
diff --git a/guest/pvmfw/src/main.rs b/guest/pvmfw/src/main.rs
index 9f1b5e6..f505318 100644
--- a/guest/pvmfw/src/main.rs
+++ b/guest/pvmfw/src/main.rs
@@ -37,12 +37,16 @@
 use crate::rollback::perform_rollback_protection;
 use alloc::borrow::Cow;
 use alloc::boxed::Box;
+use alloc::vec::Vec;
 use bssl_avf::Digester;
-use diced_open_dice::{bcc_handover_parse, DiceArtifacts, DiceContext, Hidden, VM_KEY_ALGORITHM};
+use diced_open_dice::{
+    bcc_handover_parse, DiceArtifacts, DiceContext, Hidden, HIDDEN_SIZE, VM_KEY_ALGORITHM,
+};
 use libfdt::Fdt;
 use log::{debug, error, info, trace, warn};
 use pvmfw_avb::verify_payload;
 use pvmfw_avb::DebugLevel;
+use pvmfw_avb::VerifiedBootData;
 use pvmfw_embedded_key::PUBLIC_KEY;
 use vmbase::heap;
 use vmbase::memory::{flush, SIZE_4KB};
@@ -67,103 +71,43 @@
         debug!("Ramdisk: None");
     }
 
-    let dice_handover = bcc_handover_parse(current_dice_handover).map_err(|e| {
-        error!("Invalid DICE Handover: {e:?}");
-        RebootReason::InvalidDiceHandover
-    })?;
-    trace!("DICE handover: {dice_handover:x?}");
-
-    let dice_chain_info = DiceChainInfo::new(dice_handover.bcc()).map_err(|e| {
-        error!("{e}");
-        RebootReason::InvalidDiceHandover
-    })?;
+    let (dice_handover_bytes, dice_cdi_seal, dice_context, dice_debug_mode) =
+        parse_dice_handover(current_dice_handover)?;
 
     // The bootloader should never pass us a debug policy when the boot is secure (the bootloader
     // is locked). If it gets it wrong, disregard it & log it, to avoid it causing problems.
-    if debug_policy.is_some() && !dice_chain_info.is_debug_mode() {
+    if debug_policy.is_some() && !dice_debug_mode {
         warn!("Ignoring debug policy, DICE handover does not indicate Debug mode");
         debug_policy = None;
     }
 
-    let verified_boot_data = verify_payload(signed_kernel, ramdisk, PUBLIC_KEY).map_err(|e| {
-        error!("Failed to verify the payload: {e}");
-        RebootReason::PayloadVerificationError
-    })?;
-    let debuggable = verified_boot_data.debug_level != DebugLevel::None;
-    if debuggable {
-        info!("Successfully verified a debuggable payload.");
-        info!("Please disregard any previous libavb ERROR about initrd_normal.");
-    }
+    let (verified_boot_data, debuggable, guest_page_size) =
+        perform_verified_boot(signed_kernel, ramdisk)?;
 
-    let guest_page_size = verified_boot_data.page_size.unwrap_or(SIZE_4KB);
     let hyp_page_size = hypervisor_backends::get_granule_size();
     let _ =
         sanitize_device_tree(untrusted_fdt, vm_dtbo, vm_ref_dt, guest_page_size, hyp_page_size)?;
     let fdt = untrusted_fdt; // DT has now been sanitized.
 
-    let next_dice_handover_size = guest_page_size;
-    let next_dice_handover = heap::aligned_boxed_slice(next_dice_handover_size, guest_page_size)
-        .ok_or_else(|| {
-            error!("Failed to allocate the next-stage DICE handover");
-            RebootReason::InternalError
-        })?;
-    // By leaking the slice, its content will be left behind for the next stage.
-    let next_dice_handover = Box::leak(next_dice_handover);
-
-    let dice_inputs = PartialInputs::new(&verified_boot_data).map_err(|e| {
+    let instance_hash = salt_from_instance_id(fdt)?;
+    let dice_inputs = PartialInputs::new(&verified_boot_data, instance_hash).map_err(|e| {
         error!("Failed to compute partial DICE inputs: {e:?}");
         RebootReason::InternalError
     })?;
 
-    let instance_hash = salt_from_instance_id(fdt)?;
-    let (new_instance, salt, defer_rollback_protection) = perform_rollback_protection(
-        fdt,
-        &verified_boot_data,
-        &dice_inputs,
-        dice_handover.cdi_seal(),
-        instance_hash,
-    )?;
+    let (new_instance, salt, defer_rollback_protection) =
+        perform_rollback_protection(fdt, &verified_boot_data, &dice_inputs, &dice_cdi_seal)?;
     trace!("Got salt for instance: {salt:x?}");
 
-    let new_dice_handover = if cfg!(dice_changes) {
-        Cow::Borrowed(current_dice_handover)
-    } else {
-        // It is possible that the DICE chain we were given is rooted in the UDS. We do not want to
-        // give such a chain to the payload, or even the associated CDIs. So remove the
-        // entire chain we were given and taint the CDIs. Note that the resulting CDIs are
-        // still deterministically derived from those we received, so will vary iff they do.
-        // TODO(b/280405545): Remove this post Android 14.
-        let truncated_dice_handover = dice::chain::truncate(dice_handover).map_err(|e| {
-            error!("{e}");
-            RebootReason::InternalError
-        })?;
-        Cow::Owned(truncated_dice_handover)
-    };
-
-    let cose_alg = dice_chain_info.leaf_subject_pubkey().cose_alg;
-    trace!("DICE chain leaf subject public key algorithm: {:?}", cose_alg);
-
-    let dice_context = DiceContext {
-        authority_algorithm: cose_alg.try_into().map_err(|e| {
-            error!("{e}");
-            RebootReason::InternalError
-        })?,
-        subject_algorithm: VM_KEY_ALGORITHM,
-    };
-    dice_inputs
-        .write_next_handover(
-            new_dice_handover.as_ref(),
-            &salt,
-            instance_hash,
-            defer_rollback_protection,
-            next_dice_handover,
-            dice_context,
-        )
-        .map_err(|e| {
-            error!("Failed to derive next-stage DICE secrets: {e:?}");
-            RebootReason::SecretDerivationError
-        })?;
-    flush(next_dice_handover);
+    let next_dice_handover = perform_dice_derivation(
+        dice_handover_bytes.as_ref(),
+        dice_context,
+        dice_inputs,
+        &salt,
+        defer_rollback_protection,
+        guest_page_size,
+        guest_page_size,
+    )?;
 
     let kaslr_seed = u64::from_ne_bytes(rand::random_array().map_err(|e| {
         error!("Failed to generated guest KASLR seed: {e}");
@@ -188,6 +132,102 @@
     Ok((next_dice_handover, debuggable))
 }
 
+fn parse_dice_handover(
+    bytes: &[u8],
+) -> Result<(Cow<'_, [u8]>, Vec<u8>, DiceContext, bool), RebootReason> {
+    let dice_handover = bcc_handover_parse(bytes).map_err(|e| {
+        error!("Invalid DICE Handover: {e:?}");
+        RebootReason::InvalidDiceHandover
+    })?;
+    trace!("DICE handover: {dice_handover:x?}");
+
+    let dice_chain_info = DiceChainInfo::new(dice_handover.bcc()).map_err(|e| {
+        error!("{e}");
+        RebootReason::InvalidDiceHandover
+    })?;
+    let is_debug_mode = dice_chain_info.is_debug_mode();
+    let cose_alg = dice_chain_info.leaf_subject_pubkey().cose_alg;
+    trace!("DICE chain leaf subject public key algorithm: {:?}", cose_alg);
+
+    let dice_context = DiceContext {
+        authority_algorithm: cose_alg.try_into().map_err(|e| {
+            error!("{e}");
+            RebootReason::InternalError
+        })?,
+        subject_algorithm: VM_KEY_ALGORITHM,
+    };
+
+    let cdi_seal = dice_handover.cdi_seal().to_vec();
+
+    let bytes_for_next = if cfg!(dice_changes) {
+        Cow::Borrowed(bytes)
+    } else {
+        // It is possible that the DICE chain we were given is rooted in the UDS. We do not want to
+        // give such a chain to the payload, or even the associated CDIs. So remove the
+        // entire chain we were given and taint the CDIs. Note that the resulting CDIs are
+        // still deterministically derived from those we received, so will vary iff they do.
+        // TODO(b/280405545): Remove this post Android 14.
+        let truncated_bytes = dice::chain::truncate(dice_handover).map_err(|e| {
+            error!("{e}");
+            RebootReason::InternalError
+        })?;
+        Cow::Owned(truncated_bytes)
+    };
+
+    Ok((bytes_for_next, cdi_seal, dice_context, is_debug_mode))
+}
+
+fn perform_dice_derivation<'a>(
+    dice_handover_bytes: &[u8],
+    dice_context: DiceContext,
+    dice_inputs: PartialInputs,
+    salt: &[u8; HIDDEN_SIZE],
+    defer_rollback_protection: bool,
+    next_handover_size: usize,
+    next_handover_align: usize,
+) -> Result<&'a [u8], RebootReason> {
+    let next_dice_handover = heap::aligned_boxed_slice(next_handover_size, next_handover_align)
+        .ok_or_else(|| {
+            error!("Failed to allocate the next-stage DICE handover");
+            RebootReason::InternalError
+        })?;
+    // By leaking the slice, its content will be left behind for the next stage.
+    let next_dice_handover = Box::leak(next_dice_handover);
+
+    dice_inputs
+        .write_next_handover(
+            dice_handover_bytes.as_ref(),
+            salt,
+            defer_rollback_protection,
+            next_dice_handover,
+            dice_context,
+        )
+        .map_err(|e| {
+            error!("Failed to derive next-stage DICE secrets: {e:?}");
+            RebootReason::SecretDerivationError
+        })?;
+    flush(next_dice_handover);
+    Ok(next_dice_handover)
+}
+
+fn perform_verified_boot<'a>(
+    signed_kernel: &[u8],
+    ramdisk: Option<&[u8]>,
+) -> Result<(VerifiedBootData<'a>, bool, usize), RebootReason> {
+    let verified_boot_data = verify_payload(signed_kernel, ramdisk, PUBLIC_KEY).map_err(|e| {
+        error!("Failed to verify the payload: {e}");
+        RebootReason::PayloadVerificationError
+    })?;
+    let debuggable = verified_boot_data.debug_level != DebugLevel::None;
+    if debuggable {
+        info!("Successfully verified a debuggable payload.");
+        info!("Please disregard any previous libavb ERROR about initrd_normal.");
+    }
+    let guest_page_size = verified_boot_data.page_size.unwrap_or(SIZE_4KB);
+
+    Ok((verified_boot_data, debuggable, guest_page_size))
+}
+
 // Get the "salt" which is one of the input for DICE derivation.
 // This provides differentiation of secrets for different VM instances with same payloads.
 fn salt_from_instance_id(fdt: &Fdt) -> Result<Option<Hidden>, RebootReason> {
diff --git a/guest/pvmfw/src/rollback.rs b/guest/pvmfw/src/rollback.rs
index e51b6d5..c2848a2 100644
--- a/guest/pvmfw/src/rollback.rs
+++ b/guest/pvmfw/src/rollback.rs
@@ -42,8 +42,8 @@
     verified_boot_data: &VerifiedBootData,
     dice_inputs: &PartialInputs,
     cdi_seal: &[u8],
-    instance_hash: Option<Hidden>,
 ) -> Result<(bool, Hidden, bool), RebootReason> {
+    let instance_hash = dice_inputs.instance_hash;
     if let Some(fixed) = get_fixed_rollback_protection(verified_boot_data) {
         // Prevent attackers from impersonating well-known images.
         perform_fixed_index_rollback_protection(verified_boot_data, fixed)?;