Stop truncating the BCC

But only if the RELEASE_AVF_ENABLE_DICE_CHANGES flag is enabled.

Bug: 280405545
Bug: 299472719
Bug: 266172411
Test: atest MicrodroidTests (with all flags enabled)
Change-Id: Iff619b89f81e53dc71f0ef27676b0f7c338f3031
diff --git a/Android.bp b/Android.bp
index d1086ba..cde4419 100644
--- a/Android.bp
+++ b/Android.bp
@@ -33,6 +33,7 @@
     module_type: "rust_defaults",
     config_namespace: "ANDROID",
     bool_variables: [
+        "release_avf_enable_dice_changes",
         "release_avf_enable_multi_tenant_microdroid_vm",
         "release_avf_enable_vendor_modules",
     ],
@@ -44,6 +45,9 @@
 avf_flag_aware_rust_defaults {
     name: "avf_build_flags_rust",
     soong_config_variables: {
+        release_avf_enable_dice_changes: {
+            cfgs: ["dice_changes"],
+        },
         release_avf_enable_multi_tenant_microdroid_vm: {
             cfgs: ["payload_not_root"],
         },
diff --git a/javalib/api/test-current.txt b/javalib/api/test-current.txt
index bedb267..7c61712 100644
--- a/javalib/api/test-current.txt
+++ b/javalib/api/test-current.txt
@@ -19,6 +19,7 @@
 
   public class VirtualMachineManager {
     method @RequiresPermission(android.system.virtualmachine.VirtualMachine.MANAGE_VIRTUAL_MACHINE_PERMISSION) public boolean isFeatureEnabled(String) throws android.system.virtualmachine.VirtualMachineException;
+    field public static final String FEATURE_DICE_CHANGES = "com.android.kvm.DICE_CHANGES";
     field public static final String FEATURE_PAYLOAD_NOT_ROOT = "com.android.kvm.PAYLOAD_NON_ROOT";
     field public static final String FEATURE_VENDOR_MODULES = "com.android.kvm.VENDOR_MODULES";
   }
diff --git a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
index 0a79553..260cd58 100644
--- a/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
+++ b/javalib/src/android/system/virtualmachine/VirtualMachineManager.java
@@ -107,10 +107,17 @@
     @Retention(RetentionPolicy.SOURCE)
     @StringDef(
             prefix = "FEATURE_",
-            value = {FEATURE_PAYLOAD_NOT_ROOT, FEATURE_VENDOR_MODULES})
+            value = {FEATURE_DICE_CHANGES, FEATURE_PAYLOAD_NOT_ROOT, FEATURE_VENDOR_MODULES})
     public @interface Features {}
 
     /**
+     * Feature to include new data in the VM DICE chain.
+     *
+     * @hide
+     */
+    @TestApi
+    public static final String FEATURE_DICE_CHANGES = IVirtualizationService.FEATURE_DICE_CHANGES;
+    /**
      * Feature to run payload as non-root user.
      *
      * @hide
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index ba453e7..c6aa309 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -38,6 +38,7 @@
 use crate::fdt::modify_for_next_stage;
 use crate::helpers::GUEST_PAGE_SIZE;
 use crate::instance::get_or_generate_instance_salt;
+use alloc::borrow::Cow;
 use alloc::boxed::Box;
 use core::ops::Range;
 use diced_open_dice::{bcc_handover_parse, DiceArtifacts};
@@ -132,22 +133,25 @@
         })?;
     trace!("Got salt from instance.img: {salt:x?}");
 
-    // 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_bcc_handover = bcc::truncate(bcc_handover).map_err(|e| {
-        error!("{e}");
-        RebootReason::InternalError
-    })?;
+    let new_bcc_handover = if cfg!(dice_changes) {
+        Cow::Borrowed(current_bcc_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_bcc_handover = bcc::truncate(bcc_handover).map_err(|e| {
+            error!("{e}");
+            RebootReason::InternalError
+        })?;
+        Cow::Owned(truncated_bcc_handover)
+    };
 
-    dice_inputs.write_next_bcc(truncated_bcc_handover.as_slice(), &salt, next_bcc).map_err(
-        |e| {
-            error!("Failed to derive next-stage DICE secrets: {e:?}");
-            RebootReason::SecretDerivationError
-        },
-    )?;
+    dice_inputs.write_next_bcc(new_bcc_handover.as_ref(), &salt, next_bcc).map_err(|e| {
+        error!("Failed to derive next-stage DICE secrets: {e:?}");
+        RebootReason::SecretDerivationError
+    })?;
     flush(next_bcc);
 
     let strict_boot = true;
diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index 473a560..9fe5614 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -1103,11 +1103,18 @@
         assertThat(dataItems.size()).isEqualTo(1);
         assertThat(dataItems.get(0).getMajorType()).isEqualTo(MajorType.ARRAY);
         List<DataItem> rootArrayItems = ((Array) dataItems.get(0)).getDataItems();
-        assertThat(rootArrayItems.size()).isAtLeast(2); // Public key and one certificate
+        assertThat(rootArrayItems.size()).isAtLeast(2); // Root public key and one certificate
         if (mProtectedVm) {
-            // pvmfw truncates the DICE chain it gets, so we expect exactly entries for: public key,
-            // Microdroid and app payload.
-            assertThat(rootArrayItems.size()).isEqualTo(3);
+            if (isFeatureEnabled(VirtualMachineManager.FEATURE_DICE_CHANGES)) {
+                // When a true DICE chain is created, we expect the root public key, at least one
+                // entry for the boot before pvmfw, then pvmfw, vm_entry (Microdroid kernel) and
+                // Microdroid payload entries.
+                assertThat(rootArrayItems.size()).isAtLeast(5);
+            } else {
+                // pvmfw truncates the DICE chain it gets, so we expect exactly entries for
+                // public key, vm_entry (Microdroid kernel) and Microdroid payload.
+                assertThat(rootArrayItems.size()).isEqualTo(3);
+            }
         }
     }
 
@@ -2249,7 +2256,10 @@
     }
 
     private void assumeFeatureEnabled(String featureName) throws Exception {
-        VirtualMachineManager vmm = getVirtualMachineManager();
-        assumeTrue(featureName + " not enabled", vmm.isFeatureEnabled(featureName));
+        assumeTrue(featureName + " not enabled", isFeatureEnabled(featureName));
+    }
+
+    private boolean isFeatureEnabled(String featureName) throws Exception {
+        return getVirtualMachineManager().isFeatureEnabled(featureName);
     }
 }
diff --git a/virtualizationmanager/src/aidl.rs b/virtualizationmanager/src/aidl.rs
index 781d83f..fe2f37d 100644
--- a/virtualizationmanager/src/aidl.rs
+++ b/virtualizationmanager/src/aidl.rs
@@ -36,6 +36,7 @@
     IVirtualizationService::IVirtualizationService,
     IVirtualizationService::FEATURE_PAYLOAD_NON_ROOT,
     IVirtualizationService::FEATURE_VENDOR_MODULES,
+    IVirtualizationService::FEATURE_DICE_CHANGES,
     MemoryTrimLevel::MemoryTrimLevel,
     Partition::Partition,
     PartitionType::PartitionType,
@@ -274,10 +275,11 @@
         // This approach is quite cumbersome, but will do the work for the short term.
         // TODO(b/298012279): make this scalable.
         match feature {
+            FEATURE_DICE_CHANGES => Ok(cfg!(dice_changes)),
             FEATURE_PAYLOAD_NON_ROOT => Ok(cfg!(payload_not_root)),
             FEATURE_VENDOR_MODULES => Ok(cfg!(vendor_modules)),
             _ => {
-                warn!("unknown feature {}", feature);
+                warn!("unknown feature {feature}");
                 Ok(false)
             }
         }
@@ -400,8 +402,9 @@
 
         // Check if partition images are labeled incorrectly. This is to prevent random images
         // which are not protected by the Android Verified Boot (e.g. bits downloaded by apps) from
-        // being loaded in a pVM. This applies to everything but the instance image in the raw config,
-        // and everything but the non-executable, generated partitions in the app config.
+        // being loaded in a pVM. This applies to everything but the instance image in the raw
+        // config, and everything but the non-executable, generated partitions in the app
+        // config.
         config
             .disks
             .iter()
@@ -989,10 +992,10 @@
 /// struct.
 #[derive(Debug, Default)]
 struct State {
-    /// The VMs which have been started. When VMs are started a weak reference is added to this list
-    /// while a strong reference is returned to the caller over Binder. Once all copies of the
-    /// Binder client are dropped the weak reference here will become invalid, and will be removed
-    /// from the list opportunistically the next time `add_vm` is called.
+    /// The VMs which have been started. When VMs are started a weak reference is added to this
+    /// list while a strong reference is returned to the caller over Binder. Once all copies of
+    /// the Binder client are dropped the weak reference here will become invalid, and will be
+    /// removed from the list opportunistically the next time `add_vm` is called.
     vms: Vec<Weak<VmInstance>>,
 }
 
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
index fa50d54..9255e1c 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/IVirtualizationService.aidl
@@ -22,6 +22,7 @@
 import android.system.virtualizationservice.VirtualMachineDebugInfo;
 
 interface IVirtualizationService {
+    const String FEATURE_DICE_CHANGES = "com.android.kvm.DICE_CHANGES";
     const String FEATURE_PAYLOAD_NON_ROOT = "com.android.kvm.PAYLOAD_NON_ROOT";
     const String FEATURE_VENDOR_MODULES = "com.android.kvm.VENDOR_MODULES";