Merge "Short fix for changing return type of constructClipboardHeader" into main
diff --git a/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java b/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java
index b3c0746..a2d8ee6 100644
--- a/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java
+++ b/java/framework/src/android/system/virtualmachine/VirtualMachineConfig.java
@@ -687,12 +687,28 @@
         for (int i = 0; i < config.disks.length; i++) {
             config.disks[i] = new DiskImage();
             config.disks[i].writable = customImageConfig.getDisks()[i].isWritable();
+            String diskImagePath = customImageConfig.getDisks()[i].getImagePath();
+            if (diskImagePath != null) {
+                config.disks[i].image =
+                        ParcelFileDescriptor.open(
+                                new File(diskImagePath),
+                                config.disks[i].writable ? MODE_READ_WRITE : MODE_READ_ONLY);
+            }
 
-            config.disks[i].image =
-                    ParcelFileDescriptor.open(
-                            new File(customImageConfig.getDisks()[i].getImagePath()),
-                            config.disks[i].writable ? MODE_READ_WRITE : MODE_READ_ONLY);
-            config.disks[i].partitions = new Partition[0];
+            List<Partition> partitions = new ArrayList<>();
+            for (VirtualMachineCustomImageConfig.Partition p :
+                    customImageConfig.getDisks()[i].getPartitions()) {
+                Partition part = new Partition();
+                part.label = p.name;
+                part.image =
+                        ParcelFileDescriptor.open(
+                                new File(p.imagePath),
+                                p.writable ? MODE_READ_WRITE : MODE_READ_ONLY);
+                part.writable = p.writable;
+                part.guid = TextUtils.isEmpty(p.guid) ? null : p.guid;
+                partitions.add(part);
+            }
+            config.disks[i].partitions = partitions.toArray(new Partition[0]);
         }
 
         config.displayConfig =
diff --git a/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java b/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
index 125e01c..2a571ff 100644
--- a/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
+++ b/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
@@ -32,6 +32,10 @@
     private static final String KEY_PARAMS = "params";
     private static final String KEY_DISK_WRITABLES = "disk_writables";
     private static final String KEY_DISK_IMAGES = "disk_images";
+    private static final String KEY_PARTITION_LABELS = "partition_labels_";
+    private static final String KEY_PARTITION_IMAGES = "partition_images_";
+    private static final String KEY_PARTITION_WRITABLES = "partition_writables_";
+    private static final String KEY_PARTITION_GUIDS = "partition_guids_";
     private static final String KEY_DISPLAY_CONFIG = "display_config";
     private static final String KEY_TOUCH = "touch";
     private static final String KEY_KEYBOARD = "keyboard";
@@ -155,8 +159,23 @@
         if (writables != null && diskImages != null) {
             if (writables.length == diskImages.length) {
                 for (int i = 0; i < writables.length; i++) {
-                    builder.addDisk(
-                            writables[i] ? Disk.RWDisk(diskImages[i]) : Disk.RODisk(diskImages[i]));
+                    String diskImage = diskImages[i];
+                    diskImage = diskImage.equals("") ? null : diskImage;
+                    Disk disk = writables[i] ? Disk.RWDisk(diskImage) : Disk.RODisk(diskImage);
+                    String[] labels =
+                            customImageConfigBundle.getStringArray(KEY_PARTITION_LABELS + i);
+                    String[] images =
+                            customImageConfigBundle.getStringArray(KEY_PARTITION_IMAGES + i);
+                    boolean[] partitionWritables =
+                            customImageConfigBundle.getBooleanArray(KEY_PARTITION_WRITABLES + i);
+                    String[] guids =
+                            customImageConfigBundle.getStringArray(KEY_PARTITION_GUIDS + i);
+                    for (int j = 0; j < labels.length; j++) {
+                        disk.addPartition(
+                                new Partition(
+                                        labels[j], images[j], partitionWritables[j], guids[j]));
+                    }
+                    builder.addDisk(disk);
                 }
             }
         }
@@ -189,7 +208,26 @@
             String[] images = new String[disks.length];
             for (int i = 0; i < disks.length; i++) {
                 writables[i] = disks[i].writable;
-                images[i] = disks[i].imagePath;
+                String imagePath = disks[i].imagePath;
+                images[i] = imagePath == null ? "" : imagePath;
+
+                int numPartitions = disks[i].getPartitions().size();
+                String[] partitionLabels = new String[numPartitions];
+                String[] partitionImages = new String[numPartitions];
+                boolean[] partitionWritables = new boolean[numPartitions];
+                String[] partitionGuids = new String[numPartitions];
+
+                for (int j = 0; j < numPartitions; j++) {
+                    Partition p = disks[i].getPartitions().get(j);
+                    partitionLabels[j] = p.name;
+                    partitionImages[j] = p.imagePath;
+                    partitionWritables[j] = p.writable;
+                    partitionGuids[j] = p.guid == null ? "" : p.guid;
+                }
+                pb.putStringArray(KEY_PARTITION_LABELS + i, partitionLabels);
+                pb.putStringArray(KEY_PARTITION_IMAGES + i, partitionImages);
+                pb.putBooleanArray(KEY_PARTITION_WRITABLES + i, partitionWritables);
+                pb.putStringArray(KEY_PARTITION_GUIDS + i, partitionGuids);
             }
             pb.putBooleanArray(KEY_DISK_WRITABLES, writables);
             pb.putStringArray(KEY_DISK_IMAGES, images);
@@ -232,10 +270,12 @@
     public static final class Disk {
         private final boolean writable;
         private final String imagePath;
+        private final List<Partition> partitions;
 
         private Disk(boolean writable, String imagePath) {
             this.writable = writable;
             this.imagePath = imagePath;
+            this.partitions = new ArrayList<>();
         }
 
         /** @hide */
@@ -257,6 +297,32 @@
         public String getImagePath() {
             return imagePath;
         }
+
+        /** @hide */
+        public Disk addPartition(Partition p) {
+            this.partitions.add(p);
+            return this;
+        }
+
+        /** @hide */
+        public List<Partition> getPartitions() {
+            return partitions;
+        }
+    }
+
+    /** @hide */
+    public static final class Partition {
+        public final String name;
+        public final String imagePath;
+        public final boolean writable;
+        public final String guid;
+
+        public Partition(String name, String imagePath, boolean writable, String guid) {
+            this.name = name;
+            this.imagePath = imagePath;
+            this.writable = writable;
+            this.guid = guid;
+        }
     }
 
     /** @hide */
diff --git a/libs/vmconfig/Android.bp b/libs/vmconfig/Android.bp
index 728033c..657e604 100644
--- a/libs/vmconfig/Android.bp
+++ b/libs/vmconfig/Android.bp
@@ -14,6 +14,7 @@
         "libsemver",
         "libserde",
         "libserde_json",
+        "libuuid",
     ],
     apex_available: [
         "com.android.virt",
diff --git a/libs/vmconfig/src/lib.rs b/libs/vmconfig/src/lib.rs
index 1413b51..ff115f3 100644
--- a/libs/vmconfig/src/lib.rs
+++ b/libs/vmconfig/src/lib.rs
@@ -32,6 +32,7 @@
 use std::io::BufReader;
 use std::num::NonZeroU32;
 use std::path::{Path, PathBuf};
+use uuid::Uuid;
 
 /// Configuration for a particular VM to be started.
 #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
@@ -176,6 +177,9 @@
     /// Whether the partition should be writable.
     #[serde(default)]
     pub writable: bool,
+    /// GUID of this partition.
+    #[serde(default)]
+    pub guid: Option<Uuid>,
 }
 
 impl Partition {
@@ -184,6 +188,7 @@
             image: Some(open_parcel_file(&self.path, self.writable)?),
             writable: self.writable,
             label: self.label.to_owned(),
+            guid: None,
         })
     }
 }
diff --git a/service_vm/manager/src/lib.rs b/service_vm/manager/src/lib.rs
index 987325d..78ed85b 100644
--- a/service_vm/manager/src/lib.rs
+++ b/service_vm/manager/src/lib.rs
@@ -220,6 +220,7 @@
         label: "vm-instance".to_owned(),
         image: Some(instance_img),
         writable: true,
+        guid: None,
     }];
     let rialto = File::open(RIALTO_PATH).context("Failed to open Rialto kernel binary")?;
     let instance_id_file = Path::new(VIRT_DATA_DIR).join(INSTANCE_ID_FILENAME);
diff --git a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
index b646ea1..ec1a553 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -27,12 +27,14 @@
 import static com.google.common.truth.Truth.assertWithMessage;
 import static com.google.common.truth.TruthJUnit.assume;
 
+import android.app.Application;
 import android.app.Instrumentation;
+import android.content.ComponentCallbacks2;
+import android.content.Context;
 import android.os.Bundle;
 import android.os.ParcelFileDescriptor;
 import android.os.ParcelFileDescriptor.AutoCloseInputStream;
 import android.os.ParcelFileDescriptor.AutoCloseOutputStream;
-import android.os.Process;
 import android.os.RemoteException;
 import android.system.Os;
 import android.system.virtualmachine.VirtualMachine;
@@ -619,7 +621,8 @@
                         .setMemoryBytes(256 * ONE_MEBI)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine(vmName, config);
-        MemoryReclaimListener listener = new MemoryReclaimListener(this::executeCommand);
+        MemoryReclaimListener listener =
+                new MemoryReclaimListener(this::executeCommand, getContext());
         BenchmarkVmListener.create(listener).runToFinish(TAG, vm);
         assertWithMessage("VM failed to start").that(listener.mPreCrosvm).isNotNull();
         assertWithMessage("Post trim stats not available").that(listener.mPostCrosvm).isNotNull();
@@ -654,11 +657,13 @@
     }
 
     private static class MemoryReclaimListener implements BenchmarkVmListener.InnerListener {
-        MemoryReclaimListener(Function<String, String> shellExecutor) {
+        MemoryReclaimListener(Function<String, String> shellExecutor, Context applicationCtx) {
             mShellExecutor = shellExecutor;
+            mApplication = (Application) applicationCtx;
         }
 
         public final Function<String, String> mShellExecutor;
+        private final Application mApplication;
 
         public CrosvmStats mPreCrosvm;
         public CrosvmStats mPostCrosvm;
@@ -674,7 +679,7 @@
             service.allocAnonMemory(256);
             mPreCrosvm = new CrosvmStats(vmPid, mShellExecutor);
             // Send a memory trim hint to cause memory reclaim.
-            mShellExecutor.apply("am send-trim-memory " + Process.myPid() + " RUNNING_CRITICAL");
+            mApplication.onTrimMemory(ComponentCallbacks2.TRIM_MEMORY_RUNNING_CRITICAL);
             // Give time for the memory reclaim to do its work.
             try {
                 Thread.sleep(isCuttlefish() ? 10000 : 5000);
diff --git a/tests/testapk/Android.bp b/tests/testapk/Android.bp
index e32ff88..8314f43 100644
--- a/tests/testapk/Android.bp
+++ b/tests/testapk/Android.bp
@@ -57,6 +57,9 @@
         "MicrodroidCrashNativeLib",
         "libmicrodroid_testlib_rust",
         "libvm_attestation_test_payload",
+
+        // Non-VM payload libraries
+        "libhwtrust_jni",
     ],
     min_sdk_version: "33",
 }
@@ -186,3 +189,20 @@
         "libvm_payload_rs",
     ],
 }
+
+rust_ffi_shared {
+    name: "libhwtrust_jni",
+    crate_name: "hwtrust_jni",
+    srcs: ["src/native/hwtrust_jni.rs"],
+    prefer_rlib: true,
+    rustlibs: [
+        "libandroid_logger",
+        "libanyhow",
+        "liblog_rust",
+        "libhwtrust",
+        "libjni",
+    ],
+    shared_libs: [
+        "libcrypto",
+    ],
+}
diff --git a/tests/testapk/src/java/com/android/microdroid/test/HwTrustJni.java b/tests/testapk/src/java/com/android/microdroid/test/HwTrustJni.java
new file mode 100644
index 0000000..3b237aa
--- /dev/null
+++ b/tests/testapk/src/java/com/android/microdroid/test/HwTrustJni.java
@@ -0,0 +1,31 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.microdroid.test;
+
+class HwTrustJni {
+    static {
+        System.loadLibrary("hwtrust_jni");
+    }
+
+    /**
+     * Validates a DICE chain.
+     *
+     * @param diceChain The dice chain to validate.
+     * @return true if the dice chain is valid, false otherwise.
+     */
+    public static native boolean validateDiceChain(byte[] diceChain);
+}
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 fd67659..658b1bb 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -1308,6 +1308,36 @@
     }
 
     @Test
+    @VsrTest(requirements = {"VSR-7.1-001.005"})
+    public void protectedVmHasValidDiceChain() throws Exception {
+        // This test validates two things regarding the pVM DICE chain:
+        // 1. The DICE chain is well-formed that all the entries conform to the DICE spec.
+        // 2. Each entry in the DICE chain is signed by the previous entry's subject public key.
+        assumeSupportedDevice();
+        assumeProtectedVM();
+        assumeVsrCompliant();
+        assumeTrue("Vendor API must be at least 202404", getVendorApiLevel() >= 202404);
+
+        grantPermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
+        VirtualMachineConfig config =
+                newVmConfigBuilderWithPayloadConfig("assets/vm_config.json")
+                        .setDebugLevel(DEBUG_LEVEL_FULL)
+                        .build();
+        VirtualMachine vm = forceCreateNewVirtualMachine("bcc_vm_for_vsr", config);
+        TestResults testResults =
+                runVmTestService(
+                        TAG,
+                        vm,
+                        (service, results) -> {
+                            results.mBcc = service.getBcc();
+                        });
+        testResults.assertNoException();
+        byte[] bccBytes = testResults.mBcc;
+        assertThat(bccBytes).isNotNull();
+        assertThat(HwTrustJni.validateDiceChain(bccBytes)).isTrue();
+    }
+
+    @Test
     @CddTest(requirements = {
             "9.17/C-1-1",
             "9.17/C-1-2"
diff --git a/tests/testapk/src/native/hwtrust_jni.rs b/tests/testapk/src/native/hwtrust_jni.rs
new file mode 100644
index 0000000..3b00364
--- /dev/null
+++ b/tests/testapk/src/native/hwtrust_jni.rs
@@ -0,0 +1,57 @@
+/*
+ * Copyright (c) 2024, The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//! JNI bindings to call into `hwtrust` from Java.
+
+use anyhow::Result;
+use hwtrust::{dice, session::Session};
+use jni::objects::{JByteArray, JClass};
+use jni::sys::jboolean;
+use jni::JNIEnv;
+use log::{debug, error, info};
+
+/// Validates the given DICE chain.
+#[no_mangle]
+pub extern "system" fn Java_com_android_microdroid_test_HwTrustJni_validateDiceChain(
+    env: JNIEnv,
+    _class: JClass,
+    dice_chain: JByteArray,
+) -> jboolean {
+    android_logger::init_once(
+        android_logger::Config::default()
+            .with_tag("hwtrust_jni")
+            .with_max_level(log::LevelFilter::Debug),
+    );
+    debug!("Starting the DICE chain validation ...");
+    match validate_dice_chain(env, dice_chain) {
+        Ok(_) => {
+            info!("DICE chain validated successfully");
+            true
+        }
+        Err(e) => {
+            error!("Failed to validate DICE chain: {:?}", e);
+            false
+        }
+    }
+    .into()
+}
+
+fn validate_dice_chain(env: JNIEnv, jdice_chain: JByteArray) -> Result<()> {
+    let dice_chain = env.convert_byte_array(jdice_chain)?;
+    let session = Session::default();
+    let _chain = dice::Chain::from_cbor(&session, &dice_chain)?;
+    Ok(())
+}
diff --git a/virtualizationmanager/Android.bp b/virtualizationmanager/Android.bp
index ada66dd..ae85934 100644
--- a/virtualizationmanager/Android.bp
+++ b/virtualizationmanager/Android.bp
@@ -71,6 +71,7 @@
         "liblibfdt",
         "libfsfdt",
         "libhypervisor_props",
+        "libuuid",
         // TODO(b/202115393) stabilize the interface
         "packagemanager_aidl-rust",
     ],
diff --git a/virtualizationmanager/src/composite.rs b/virtualizationmanager/src/composite.rs
index a4b7eae..681ec59 100644
--- a/virtualizationmanager/src/composite.rs
+++ b/virtualizationmanager/src/composite.rs
@@ -23,6 +23,8 @@
 use std::os::unix::io::AsRawFd;
 use std::path::{Path, PathBuf};
 
+use uuid::Uuid;
+
 /// Constructs a composite disk image for the given list of partitions, and opens it ready to use.
 ///
 /// Returns the composite disk image file, and a list of files whose file descriptors must be passed
@@ -105,6 +107,7 @@
                 partition_type: ImagePartitionType::LinuxFilesystem,
                 writable: partition.writable,
                 size,
+                part_guid: partition.guid.as_deref().map(Uuid::parse_str).transpose()?,
             })
         })
         .collect::<Result<_, Error>>()?;
diff --git a/virtualizationmanager/src/payload.rs b/virtualizationmanager/src/payload.rs
index 9d0c7d6..82d9ba0 100644
--- a/virtualizationmanager/src/payload.rs
+++ b/virtualizationmanager/src/payload.rs
@@ -289,6 +289,7 @@
         label: "payload-metadata".to_owned(),
         image: Some(metadata_file),
         writable: false,
+        guid: None,
     }];
 
     for (i, apex_info) in apex_infos.iter().enumerate() {
@@ -297,17 +298,20 @@
             label: format!("microdroid-apex-{}", i),
             image: Some(apex_file),
             writable: false,
+            guid: None,
         });
     }
     partitions.push(Partition {
         label: "microdroid-apk".to_owned(),
         image: Some(ParcelFileDescriptor::new(apk_file)),
         writable: false,
+        guid: None,
     });
     partitions.push(Partition {
         label: "microdroid-apk-idsig".to_owned(),
         image: Some(ParcelFileDescriptor::new(idsig_file)),
         writable: false,
+        guid: None,
     });
 
     // we've already checked that extra_apks and extraIdsigs are in the same size.
@@ -319,6 +323,7 @@
             label: format!("extra-apk-{i}"),
             image: Some(ParcelFileDescriptor::new(extra_apk_file)),
             writable: false,
+            guid: None,
         });
 
         partitions.push(Partition {
@@ -330,6 +335,7 @@
                     .with_context(|| format!("Failed to clone the extra idsig #{i}"))?,
             )),
             writable: false,
+            guid: None,
         });
     }
 
@@ -416,6 +422,7 @@
             label: "microdroid-vendor".to_owned(),
             image: Some(ParcelFileDescriptor::new(vendor_image)),
             writable: false,
+            guid: None,
         }],
     })
 }
@@ -439,6 +446,7 @@
         label: "vm-instance".to_owned(),
         image: Some(ParcelFileDescriptor::new(instance_file)),
         writable: true,
+        guid: None,
     }];
 
     if let Some(file) = storage_image {
@@ -446,6 +454,7 @@
             label: "encryptedstore".to_owned(),
             image: Some(ParcelFileDescriptor::new(file)),
             writable: true,
+            guid: None,
         });
     }
 
diff --git a/virtualizationservice/aidl/android/system/virtualizationservice/Partition.aidl b/virtualizationservice/aidl/android/system/virtualizationservice/Partition.aidl
index 825c3da..11a2115 100644
--- a/virtualizationservice/aidl/android/system/virtualizationservice/Partition.aidl
+++ b/virtualizationservice/aidl/android/system/virtualizationservice/Partition.aidl
@@ -25,4 +25,7 @@
 
     /** Whether the partition should be writable by the VM. */
     boolean writable;
+
+    /** GUID of the partition. If not set, automatically created */
+    @nullable String guid;
 }
diff --git a/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java b/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
index cf83bff..ef94be5 100644
--- a/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
+++ b/vmlauncher_app/java/com/android/virtualization/vmlauncher/MainActivity.java
@@ -131,6 +131,26 @@
                                     VirtualMachineCustomImageConfig.Disk.RODisk(
                                             item.getString("image")));
                         }
+                    } else if (item.has("partitions")) {
+                        boolean diskWritable = item.optBoolean("writable", false);
+                        VirtualMachineCustomImageConfig.Disk disk =
+                                diskWritable
+                                        ? VirtualMachineCustomImageConfig.Disk.RWDisk(null)
+                                        : VirtualMachineCustomImageConfig.Disk.RODisk(null);
+                        JSONArray partitions = item.getJSONArray("partitions");
+                        for (int j = 0; j < partitions.length(); j++) {
+                            JSONObject partition = partitions.getJSONObject(j);
+                            String label = partition.getString("label");
+                            String path = partition.getString("path");
+                            boolean partitionWritable =
+                                    diskWritable && partition.optBoolean("writable", false);
+                            String guid = partition.optString("guid");
+                            VirtualMachineCustomImageConfig.Partition p =
+                                    new VirtualMachineCustomImageConfig.Partition(
+                                            label, path, partitionWritable, guid);
+                            disk.addPartition(p);
+                        }
+                        customImageConfigBuilder.addDisk(disk);
                     }
                 }
             }