[dice] Support multi key algorithm switch in pvmfw

Bug: 357008987
Bug: 341630707
Test: atest MicrodroidHostTests
Test: atest libpvmfw.dice.test
Change-Id: Id8803e071004e5622ca537d35cb532203a60758a
diff --git a/guest/pvmfw/Android.bp b/guest/pvmfw/Android.bp
index 5c767a3..51f7802 100644
--- a/guest/pvmfw/Android.bp
+++ b/guest/pvmfw/Android.bp
@@ -118,6 +118,7 @@
         "libciborium",
         "libdiced_open_dice_nostd",
         "libpvmfw_avb_nostd",
+        "libdiced_sample_inputs_nostd",
         "libzerocopy_nostd",
         "libhex",
     ],
diff --git a/guest/pvmfw/src/dice.rs b/guest/pvmfw/src/dice.rs
index f3a2337..b597309 100644
--- a/guest/pvmfw/src/dice.rs
+++ b/guest/pvmfw/src/dice.rs
@@ -21,7 +21,7 @@
 use ciborium::Value;
 use core::mem::size_of;
 use diced_open_dice::{
-    bcc_handover_main_flow, hash, Config, DiceMode, Hash, InputValues, HIDDEN_SIZE,
+    bcc_handover_main_flow, hash, Config, DiceContext, DiceMode, Hash, InputValues, HIDDEN_SIZE,
 };
 use pvmfw_avb::{Capability, DebugLevel, Digest, VerifiedBootData};
 use zerocopy::AsBytes;
@@ -102,6 +102,7 @@
         instance_hash: Option<Hash>,
         deferred_rollback_protection: bool,
         next_bcc: &mut [u8],
+        context: DiceContext,
     ) -> Result<()> {
         let config = self
             .generate_config_descriptor(instance_hash)
@@ -114,7 +115,7 @@
             self.mode,
             self.make_hidden(salt, deferred_rollback_protection)?,
         );
-        let _ = bcc_handover_main_flow(current_bcc_handover, &dice_inputs, next_bcc)?;
+        let _ = bcc_handover_main_flow(current_bcc_handover, &dice_inputs, next_bcc, context)?;
         Ok(())
     }
 
@@ -177,8 +178,11 @@
     };
     use ciborium::Value;
     use diced_open_dice::DiceArtifacts;
+    use diced_open_dice::DiceContext;
     use diced_open_dice::DiceMode;
+    use diced_open_dice::KeyAlgorithm;
     use diced_open_dice::HIDDEN_SIZE;
+    use diced_sample_inputs::make_sample_bcc_and_cdis;
     use pvmfw_avb::Capability;
     use pvmfw_avb::DebugLevel;
     use pvmfw_avb::Digest;
@@ -298,6 +302,10 @@
         let mut buffer_without_defer = [0; 4096];
         let mut buffer_with_defer = [0; 4096];
         let mut buffer_without_defer_retry = [0; 4096];
+        let context = DiceContext {
+            authority_algorithm: KeyAlgorithm::Ed25519,
+            subject_algorithm: KeyAlgorithm::Ed25519,
+        };
 
         let sample_dice_input: &[u8] = &[
             0xa3, // CDI attest
@@ -321,6 +329,7 @@
                 Some([0u8; 64]),
                 false,
                 &mut buffer_without_defer,
+                context.clone(),
             )
             .unwrap();
         let bcc_handover1 = diced_open_dice::bcc_handover_parse(&buffer_without_defer).unwrap();
@@ -333,6 +342,7 @@
                 Some([0u8; 64]),
                 true,
                 &mut buffer_with_defer,
+                context.clone(),
             )
             .unwrap();
         let bcc_handover2 = diced_open_dice::bcc_handover_parse(&buffer_with_defer).unwrap();
@@ -345,6 +355,7 @@
                 Some([0u8; 64]),
                 false,
                 &mut buffer_without_defer_retry,
+                context.clone(),
             )
             .unwrap();
         let bcc_handover3 =
@@ -353,4 +364,82 @@
         assert_ne!(bcc_handover1.cdi_seal(), bcc_handover2.cdi_seal());
         assert_eq!(bcc_handover1.cdi_seal(), bcc_handover3.cdi_seal());
     }
+
+    #[test]
+    fn dice_derivation_with_different_algorithms_is_valid() {
+        let dice_artifacts = make_sample_bcc_and_cdis().unwrap();
+        let bcc_handover0_bytes = to_bcc_handover(&dice_artifacts);
+        let vb_data = VerifiedBootData { debug_level: DebugLevel::Full, ..BASE_VB_DATA };
+        let inputs = PartialInputs::new(&vb_data).unwrap();
+        let mut buffer = [0; 4096];
+
+        inputs
+            .clone()
+            .write_next_bcc(
+                &bcc_handover0_bytes,
+                &[0u8; HIDDEN_SIZE],
+                Some([0u8; 64]),
+                true,
+                &mut buffer,
+                DiceContext {
+                    authority_algorithm: KeyAlgorithm::Ed25519,
+                    subject_algorithm: KeyAlgorithm::EcdsaP256,
+                },
+            )
+            .expect("Failed to derive Ed25519 -> EcdsaP256 BCC");
+        let bcc_handover1 = diced_open_dice::bcc_handover_parse(&buffer).unwrap();
+        let bcc_handover1_bytes = to_bcc_handover(&bcc_handover1);
+        buffer.fill(0);
+
+        inputs
+            .clone()
+            .write_next_bcc(
+                &bcc_handover1_bytes,
+                &[0u8; HIDDEN_SIZE],
+                Some([0u8; 64]),
+                true,
+                &mut buffer,
+                DiceContext {
+                    authority_algorithm: KeyAlgorithm::EcdsaP256,
+                    subject_algorithm: KeyAlgorithm::EcdsaP384,
+                },
+            )
+            .expect("Failed to derive EcdsaP256 -> EcdsaP384 BCC");
+        let bcc_handover2 = diced_open_dice::bcc_handover_parse(&buffer).unwrap();
+        let bcc_handover2_bytes = to_bcc_handover(&bcc_handover2);
+        buffer.fill(0);
+
+        inputs
+            .clone()
+            .write_next_bcc(
+                &bcc_handover2_bytes,
+                &[0u8; HIDDEN_SIZE],
+                Some([0u8; 64]),
+                true,
+                &mut buffer,
+                DiceContext {
+                    authority_algorithm: KeyAlgorithm::EcdsaP384,
+                    subject_algorithm: KeyAlgorithm::Ed25519,
+                },
+            )
+            .expect("Failed to derive EcdsaP384 -> Ed25519 BCC");
+        let _bcc_handover3 = diced_open_dice::bcc_handover_parse(&buffer).unwrap();
+
+        // TODO(b/378813154): Check the DICE chain with `hwtrust` once the profile version
+        // is updated.
+        // The check cannot be done now because parsing the chain causes the following error:
+        // Invalid payload at index 3. Caused by:
+        // 0: opendice.example.p256
+        // 1: unknown profile version
+    }
+
+    fn to_bcc_handover(dice_artifacts: &dyn DiceArtifacts) -> Vec<u8> {
+        let dice_chain = cbor_util::deserialize::<Value>(dice_artifacts.bcc().unwrap()).unwrap();
+        let bcc_handover = Value::Map(vec![
+            (Value::Integer(1.into()), Value::Bytes(dice_artifacts.cdi_attest().to_vec())),
+            (Value::Integer(2.into()), Value::Bytes(dice_artifacts.cdi_seal().to_vec())),
+            (Value::Integer(3.into()), dice_chain),
+        ]);
+        cbor_util::serialize(&bcc_handover).unwrap()
+    }
 }
diff --git a/guest/pvmfw/src/main.rs b/guest/pvmfw/src/main.rs
index aeced51..2c1eba9 100644
--- a/guest/pvmfw/src/main.rs
+++ b/guest/pvmfw/src/main.rs
@@ -45,7 +45,7 @@
 use bssl_avf::Digester;
 use core::ops::Range;
 use cstr::cstr;
-use diced_open_dice::{bcc_handover_parse, DiceArtifacts, Hidden};
+use diced_open_dice::{bcc_handover_parse, DiceArtifacts, DiceContext, Hidden, VM_KEY_ALGORITHM};
 use libfdt::{Fdt, FdtNode};
 use log::{debug, error, info, trace, warn};
 use pvmfw_avb::verify_payload;
@@ -202,6 +202,13 @@
 
     trace!("BCC leaf subject public key algorithm: {:?}", bcc.leaf_subject_pubkey().cose_alg);
 
+    let dice_context = DiceContext {
+        authority_algorithm: bcc.leaf_subject_pubkey().cose_alg.try_into().map_err(|e| {
+            error!("{e}");
+            RebootReason::InternalError
+        })?,
+        subject_algorithm: VM_KEY_ALGORITHM,
+    };
     dice_inputs
         .write_next_bcc(
             new_bcc_handover.as_ref(),
@@ -209,6 +216,7 @@
             instance_hash,
             defer_rollback_protection,
             next_bcc,
+            dice_context,
         )
         .map_err(|e| {
             error!("Failed to derive next-stage DICE secrets: {e:?}");
diff --git a/libs/dice/open_dice/Android.bp b/libs/dice/open_dice/Android.bp
index 3c5b6ea..4241c47 100644
--- a/libs/dice/open_dice/Android.bp
+++ b/libs/dice/open_dice/Android.bp
@@ -19,6 +19,9 @@
         "libopen_dice_cbor_bindgen_nostd",
         "libzeroize_nostd",
     ],
+    features: [
+        "multialg",
+    ],
     whole_static_libs: [
         "libcrypto_baremetal",
     ],
@@ -172,7 +175,11 @@
         "libopen_dice_cbor_bindgen.rust_defaults",
         "libopen_dice_bindgen_nostd.rust_defaults",
     ],
-    whole_static_libs: ["libopen_dice_cbor_baremetal"],
+    bindgen_flags: [
+        "--rustified-enum DiceKeyAlgorithm",
+        "--allowlist-type=DiceContext",
+    ],
+    whole_static_libs: ["libopen_dice_cbor_baremetal_multialg"],
 }
 
 rust_defaults {
@@ -228,7 +235,7 @@
     rustlibs: [
         "libopen_dice_cbor_bindgen_nostd",
     ],
-    whole_static_libs: ["libopen_dice_android_baremetal"],
+    whole_static_libs: ["libopen_dice_android_baremetal_multialg"],
 }
 
 rust_test {
diff --git a/libs/dice/open_dice/src/bcc.rs b/libs/dice/open_dice/src/bcc.rs
index 9c9545b..fabd7c7 100644
--- a/libs/dice/open_dice/src/bcc.rs
+++ b/libs/dice/open_dice/src/bcc.rs
@@ -14,13 +14,13 @@
 
 //! This module mirrors the content in open-dice/include/dice/android.h
 
-use crate::dice::{Cdi, CdiValues, DiceArtifacts, InputValues, CDI_SIZE};
+use crate::dice::{context, Cdi, CdiValues, DiceArtifacts, InputValues, CDI_SIZE};
 use crate::error::{check_result, DiceError, Result};
 use open_dice_android_bindgen::{
-    DiceAndroidConfigValues, DiceAndroidFormatConfigDescriptor, DiceAndroidHandoverMainFlow,
-    DiceAndroidHandoverParse, DiceAndroidMainFlow, DICE_ANDROID_CONFIG_COMPONENT_NAME,
-    DICE_ANDROID_CONFIG_COMPONENT_VERSION, DICE_ANDROID_CONFIG_RESETTABLE,
-    DICE_ANDROID_CONFIG_RKP_VM_MARKER, DICE_ANDROID_CONFIG_SECURITY_VERSION,
+    DiceAndroidConfigValues, DiceAndroidFormatConfigDescriptor, DiceAndroidHandoverParse,
+    DiceAndroidMainFlow, DICE_ANDROID_CONFIG_COMPONENT_NAME, DICE_ANDROID_CONFIG_COMPONENT_VERSION,
+    DICE_ANDROID_CONFIG_RESETTABLE, DICE_ANDROID_CONFIG_RKP_VM_MARKER,
+    DICE_ANDROID_CONFIG_SECURITY_VERSION,
 };
 use std::{ffi::CStr, ptr};
 
@@ -101,10 +101,11 @@
         // SAFETY: `DiceAndroidMainFlow` only reads the `current_chain` and CDI values and writes
         // to `next_chain` and next CDI values within its bounds. It also reads `input_values` as a
         // constant input and doesn't store any pointer.
-        // The first argument can be null and is not used in the current implementation.
+        // The first argument is a pointer to a valid |DiceContext_| object for multi-alg open-dice
+        // and a null pointer otherwise.
         unsafe {
             DiceAndroidMainFlow(
-                ptr::null_mut(), // context
+                context(),
                 current_cdi_attest.as_ptr(),
                 current_cdi_seal.as_ptr(),
                 current_chain.as_ptr(),
@@ -127,20 +128,23 @@
 /// A handover combines the DICE chain and CDIs in a single CBOR object.
 /// This function takes the current boot stage's handover bundle and produces a
 /// bundle for the next stage.
+#[cfg(feature = "multialg")]
 pub fn bcc_handover_main_flow(
     current_handover: &[u8],
     input_values: &InputValues,
     next_handover: &mut [u8],
+    ctx: crate::dice::DiceContext,
 ) -> Result<usize> {
     let mut next_handover_size = 0;
+    let mut ctx: open_dice_cbor_bindgen::DiceContext_ = ctx.into();
     check_result(
         // SAFETY: The function only reads `current_handover` and writes to `next_handover`
         // within its bounds,
         // It also reads `input_values` as a constant input and doesn't store any pointer.
-        // The first argument can be null and is not used in the current implementation.
+        // The first argument is a pointer to a valid |DiceContext_| object.
         unsafe {
-            DiceAndroidHandoverMainFlow(
-                ptr::null_mut(), // context
+            open_dice_android_bindgen::DiceAndroidHandoverMainFlow(
+                &mut ctx as *mut _ as *mut std::ffi::c_void,
                 current_handover.as_ptr(),
                 current_handover.len(),
                 input_values.as_ptr(),
diff --git a/libs/dice/open_dice/src/dice.rs b/libs/dice/open_dice/src/dice.rs
index 6404508..206769c 100644
--- a/libs/dice/open_dice/src/dice.rs
+++ b/libs/dice/open_dice/src/dice.rs
@@ -15,7 +15,7 @@
 //! Structs and functions about the types used in DICE.
 //! This module mirrors the content in open-dice/include/dice/dice.h
 
-use crate::error::{check_result, Result};
+use crate::error::{check_result, DiceError, Result};
 use coset::iana;
 pub use open_dice_cbor_bindgen::DiceMode;
 use open_dice_cbor_bindgen::{
@@ -23,9 +23,11 @@
     DiceMainFlow, DICE_CDI_SIZE, DICE_HASH_SIZE, DICE_HIDDEN_SIZE, DICE_ID_SIZE,
     DICE_INLINE_CONFIG_SIZE, DICE_PRIVATE_KEY_SEED_SIZE, DICE_PRIVATE_KEY_SIZE,
 };
+#[cfg(feature = "multialg")]
+use open_dice_cbor_bindgen::{DiceContext_, DiceKeyAlgorithm};
 #[cfg(feature = "serde_derive")]
 use serde_derive::{Deserialize, Serialize};
-use std::{marker::PhantomData, ptr};
+use std::{ffi::c_void, marker::PhantomData, ptr};
 use zeroize::{Zeroize, ZeroizeOnDrop};
 
 /// The size of a DICE hash.
@@ -114,6 +116,69 @@
     }
 }
 
+impl TryFrom<iana::Algorithm> for KeyAlgorithm {
+    type Error = DiceError;
+
+    fn try_from(alg: iana::Algorithm) -> Result<Self> {
+        match alg {
+            iana::Algorithm::EdDSA => Ok(KeyAlgorithm::Ed25519),
+            iana::Algorithm::ES256 => Ok(KeyAlgorithm::EcdsaP256),
+            iana::Algorithm::ES384 => Ok(KeyAlgorithm::EcdsaP384),
+            other => Err(DiceError::UnsupportedKeyAlgorithm(other)),
+        }
+    }
+}
+
+#[cfg(feature = "multialg")]
+impl From<KeyAlgorithm> for DiceKeyAlgorithm {
+    fn from(alg: KeyAlgorithm) -> Self {
+        match alg {
+            KeyAlgorithm::Ed25519 => DiceKeyAlgorithm::kDiceKeyAlgorithmEd25519,
+            KeyAlgorithm::EcdsaP256 => DiceKeyAlgorithm::kDiceKeyAlgorithmP256,
+            KeyAlgorithm::EcdsaP384 => DiceKeyAlgorithm::kDiceKeyAlgorithmP384,
+        }
+    }
+}
+
+/// Represents the context used for DICE operations.
+#[cfg(feature = "multialg")]
+#[derive(Debug, Clone)]
+pub struct DiceContext {
+    /// The algorithm used for the authority key.
+    pub authority_algorithm: KeyAlgorithm,
+    /// The algorithm used for the subject key.
+    pub subject_algorithm: KeyAlgorithm,
+}
+
+#[cfg(feature = "multialg")]
+impl From<DiceContext> for DiceContext_ {
+    fn from(ctx: DiceContext) -> Self {
+        DiceContext_ {
+            authority_algorithm: ctx.authority_algorithm.into(),
+            subject_algorithm: ctx.subject_algorithm.into(),
+        }
+    }
+}
+
+#[cfg(feature = "multialg")]
+const VM_DICE_CONTEXT: DiceContext_ = DiceContext_ {
+    authority_algorithm: DiceKeyAlgorithm::kDiceKeyAlgorithmEd25519,
+    subject_algorithm: DiceKeyAlgorithm::kDiceKeyAlgorithmEd25519,
+};
+
+/// Returns the pointer points to |DiceContext_| for DICE operations when `multialg`
+/// feature is enabled.
+#[cfg(feature = "multialg")]
+pub(crate) fn context() -> *mut c_void {
+    &VM_DICE_CONTEXT as *const DiceContext_ as *mut c_void
+}
+
+/// Returns a null pointer when `multialg` feature is disabled.
+#[cfg(not(feature = "multialg"))]
+pub(crate) fn context() -> *mut c_void {
+    ptr::null_mut()
+}
+
 /// A trait for types that represent Dice artifacts, which include:
 ///
 /// - Attestation CDI
@@ -322,10 +387,9 @@
     check_result(
         // SAFETY: The function only reads the current CDI values and inputs and writes
         // to `next_cdi_certificate` and next CDI values within its bounds.
-        // The first argument can be null and is not used in the current implementation.
         unsafe {
             DiceMainFlow(
-                ptr::null_mut(), // context
+                context(),
                 current_cdi_attest.as_ptr(),
                 current_cdi_seal.as_ptr(),
                 input_values.as_ptr(),
diff --git a/libs/dice/open_dice/src/error.rs b/libs/dice/open_dice/src/error.rs
index bef9a9c..9089432 100644
--- a/libs/dice/open_dice/src/error.rs
+++ b/libs/dice/open_dice/src/error.rs
@@ -29,6 +29,8 @@
     BufferTooSmall(usize),
     /// Platform error.
     PlatformError,
+    /// Unsupported key algorithm.
+    UnsupportedKeyAlgorithm(coset::iana::Algorithm),
 }
 
 /// This makes `DiceError` accepted by anyhow.
@@ -43,6 +45,9 @@
                 write!(f, "Buffer too small; need {buffer_required_size} bytes")
             }
             Self::PlatformError => write!(f, "Platform error"),
+            Self::UnsupportedKeyAlgorithm(algorithm) => {
+                write!(f, "Unsupported key algorithm: {algorithm:?}")
+            }
         }
     }
 }
diff --git a/libs/dice/open_dice/src/lib.rs b/libs/dice/open_dice/src/lib.rs
index a347d46..4d05255 100644
--- a/libs/dice/open_dice/src/lib.rs
+++ b/libs/dice/open_dice/src/lib.rs
@@ -28,10 +28,13 @@
 mod ops;
 mod retry;
 
+#[cfg(feature = "multialg")]
+pub use bcc::bcc_handover_main_flow;
 pub use bcc::{
-    bcc_format_config_descriptor, bcc_handover_main_flow, bcc_handover_parse, bcc_main_flow,
-    BccHandover, DiceConfigValues,
+    bcc_format_config_descriptor, bcc_handover_parse, bcc_main_flow, BccHandover, DiceConfigValues,
 };
+#[cfg(feature = "multialg")]
+pub use dice::DiceContext;
 pub use dice::{
     derive_cdi_certificate_id, derive_cdi_private_key_seed, dice_main_flow, Cdi, CdiValues, Config,
     DiceArtifacts, DiceMode, Hash, Hidden, InlineConfig, InputValues, KeyAlgorithm, PrivateKey,
diff --git a/libs/dice/open_dice/src/ops.rs b/libs/dice/open_dice/src/ops.rs
index 7bc0ee5..41951bf 100644
--- a/libs/dice/open_dice/src/ops.rs
+++ b/libs/dice/open_dice/src/ops.rs
@@ -17,7 +17,7 @@
 //! main DICE functions depend on.
 
 use crate::dice::{
-    derive_cdi_private_key_seed, DiceArtifacts, Hash, InputValues, PrivateKey, HASH_SIZE,
+    context, derive_cdi_private_key_seed, DiceArtifacts, Hash, InputValues, PrivateKey, HASH_SIZE,
     PRIVATE_KEY_SEED_SIZE, PRIVATE_KEY_SIZE, VM_KEY_ALGORITHM,
 };
 use crate::error::{check_result, DiceError, Result};
@@ -83,11 +83,12 @@
     let principal = DicePrincipal::kDicePrincipalSubject;
     check_result(
         // SAFETY: The function writes to the `public_key` and `private_key` within the given
-        // bounds, and only reads the `seed`. The first argument context is not used in this
-        // function.
+        // bounds, and only reads the `seed`.
+        // The first argument is a pointer to a valid |DiceContext_| object for multi-alg open-dice
+        // and a null pointer otherwise.
         unsafe {
             DiceKeypairFromSeed(
-                ptr::null_mut(), // context
+                context(),
                 principal,
                 seed.as_ptr(),
                 public_key.as_mut_ptr(),
@@ -118,11 +119,12 @@
     let mut signature = vec![0u8; VM_KEY_ALGORITHM.signature_size()];
     check_result(
         // SAFETY: The function writes to the `signature` within the given bounds, and only reads
-        // the message and the private key. The first argument context is not used in this
-        // function.
+        // the message and the private key.
+        // The first argument is a pointer to a valid |DiceContext_| object for multi-alg open-dice
+        // and a null pointer otherwise.
         unsafe {
             DiceSign(
-                ptr::null_mut(), // context
+                context(),
                 message.as_ptr(),
                 message.len(),
                 private_key.as_ptr(),
@@ -143,10 +145,11 @@
     }
     check_result(
         // SAFETY: only reads the messages, signature and public key as constant values.
-        // The first argument context is not used in this function.
+        // The first argument is a pointer to a valid |DiceContext_| object for multi-alg open-dice
+        // and a null pointer otherwise.
         unsafe {
             DiceVerify(
-                ptr::null_mut(), // context
+                context(),
                 message.as_ptr(),
                 message.len(),
                 signature.as_ptr(),
@@ -171,11 +174,12 @@
     let mut certificate_actual_size = 0;
     check_result(
         // SAFETY: The function writes to the `certificate` within the given bounds, and only reads
-        // the input values and the key seeds. The first argument context is not used in this
-        // function.
+        // the input values and the key seeds.
+        // The first argument is a pointer to a valid |DiceContext_| object for multi-alg open-dice
+        // and a null pointer otherwise.
         unsafe {
             DiceGenerateCertificate(
-                ptr::null_mut(), // context
+                context(),
                 subject_private_key_seed.as_ptr(),
                 authority_private_key_seed.as_ptr(),
                 input_values.as_ptr(),