Merge "[Refactor] Remove vmclient::connect()" into main
diff --git a/Android.bp b/Android.bp
index 4fa696f..22581b0 100644
--- a/Android.bp
+++ b/Android.bp
@@ -36,6 +36,7 @@
         "release_avf_enable_dice_changes",
         "release_avf_enable_llpvm_changes",
         "release_avf_enable_multi_tenant_microdroid_vm",
+        "release_avf_enable_remote_attestation",
         "release_avf_enable_vendor_modules",
     ],
     properties: [
@@ -55,6 +56,9 @@
         release_avf_enable_multi_tenant_microdroid_vm: {
             cfgs: ["payload_not_root"],
         },
+        release_avf_enable_remote_attestation: {
+            cfgs: ["remote_attestation"],
+        },
         release_avf_enable_vendor_modules: {
             cfgs: ["vendor_modules"],
         },
diff --git a/apex/sign_virt_apex.py b/apex/sign_virt_apex.py
index a42f5ec..8257aae 100644
--- a/apex/sign_virt_apex.py
+++ b/apex/sign_virt_apex.py
@@ -231,6 +231,9 @@
         if additional_descriptors:
             for image in additional_descriptors:
                 cmd.extend(['--include_descriptors_from_image', image])
+
+        if 'Rollback Index' in info:
+            cmd.extend(['--rollback_index', info['Rollback Index']])
         RunCommand(args, cmd)
 
 
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index 00831dd..bac93a4 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -409,7 +409,19 @@
     },
 }
 
-avb_add_hash_footer {
+soong_config_module_type {
+    name: "flag_aware_avb_add_hash_footer",
+    module_type: "avb_add_hash_footer",
+    config_namespace: "ANDROID",
+    bool_variables: [
+        "release_avf_enable_llpvm_changes",
+    ],
+    properties: [
+        "rollback_index",
+    ],
+}
+
+flag_aware_avb_add_hash_footer {
     name: "microdroid_kernel_signed",
     src: ":empty_file",
     filename: "microdroid_kernel",
@@ -431,6 +443,12 @@
         ":microdroid_initrd_normal_hashdesc",
         ":microdroid_initrd_debug_hashdesc",
     ],
+    // Below are properties that are conditionally set depending on value of build flags.
+    soong_config_variables: {
+        release_avf_enable_llpvm_changes: {
+            rollback_index: 1,
+        },
+    },
 }
 
 prebuilt_etc {
@@ -447,7 +465,7 @@
     },
 }
 
-avb_add_hash_footer {
+flag_aware_avb_add_hash_footer {
     name: "microdroid_kernel_with_modules_signed",
     src: ":empty_file",
     filename: "microdroid_kernel_with_modules",
@@ -465,6 +483,12 @@
         ":microdroid_initrd_normal_hashdesc",
         ":microdroid_initrd_debug_hashdesc",
     ],
+    // Below are properties that are conditionally set depending on value of build flags.
+    soong_config_variables: {
+        release_avf_enable_llpvm_changes: {
+            rollback_index: 1,
+        },
+    },
 }
 
 prebuilt_etc {
diff --git a/pvmfw/src/fdt.rs b/pvmfw/src/fdt.rs
index 244b192..61de423 100644
--- a/pvmfw/src/fdt.rs
+++ b/pvmfw/src/fdt.rs
@@ -732,6 +732,7 @@
     strict_boot: bool,
     debug_policy: Option<&mut [u8]>,
     debuggable: bool,
+    kaslr_seed: u64,
 ) -> libfdt::Result<()> {
     if let Some(debug_policy) = debug_policy {
         let backup = Vec::from(fdt.as_slice());
@@ -753,6 +754,7 @@
     if let Some(mut chosen) = fdt.chosen_mut()? {
         empty_or_delete_prop(&mut chosen, cstr!("avf,strict-boot"), strict_boot)?;
         empty_or_delete_prop(&mut chosen, cstr!("avf,new-instance"), new_instance)?;
+        chosen.setprop_inplace(cstr!("kaslr-seed"), &kaslr_seed.to_be_bytes())?;
     };
     if !debuggable {
         if let Some(bootargs) = read_bootargs_from(fdt)? {
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index c6aa309..d39d51c 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -52,6 +52,7 @@
 use vmbase::heap;
 use vmbase::memory::flush;
 use vmbase::memory::MEMORY;
+use vmbase::rand;
 use vmbase::virtio::pci;
 
 const NEXT_BCC_SIZE: usize = GUEST_PAGE_SIZE;
@@ -154,12 +155,24 @@
     })?;
     flush(next_bcc);
 
+    let kaslr_seed = u64::from_ne_bytes(rand::random_array().map_err(|e| {
+        error!("Failed to generated guest KASLR seed: {e}");
+        RebootReason::InternalError
+    })?);
     let strict_boot = true;
-    modify_for_next_stage(fdt, next_bcc, new_instance, strict_boot, debug_policy, debuggable)
-        .map_err(|e| {
-            error!("Failed to configure device tree: {e}");
-            RebootReason::InternalError
-        })?;
+    modify_for_next_stage(
+        fdt,
+        next_bcc,
+        new_instance,
+        strict_boot,
+        debug_policy,
+        debuggable,
+        kaslr_seed,
+    )
+    .map_err(|e| {
+        error!("Failed to configure device tree: {e}");
+        RebootReason::InternalError
+    })?;
 
     info!("Starting payload...");
 
diff --git a/tests/benchmark/Android.bp b/tests/benchmark/Android.bp
index 90ba575..80fdff7 100644
--- a/tests/benchmark/Android.bp
+++ b/tests/benchmark/Android.bp
@@ -28,6 +28,7 @@
     compile_multilib: "64",
     required: ["perf-setup"],
     host_required: ["MicrodroidTestPreparer"],
+    data: [":test_microdroid_vendor_image"],
 }
 
 cc_library_shared {
diff --git a/tests/benchmark/AndroidTest.xml b/tests/benchmark/AndroidTest.xml
index 8c8bfbe..11b17f4 100644
--- a/tests/benchmark/AndroidTest.xml
+++ b/tests/benchmark/AndroidTest.xml
@@ -30,6 +30,14 @@
         <option name="post-push" value="chmod 755 /data/local/tmp/perf-setup.sh;/data/local/tmp/perf-setup.sh" />
         <option name="cleanup" value="true" />
     </target_preparer>
+    <target_preparer class="com.android.tradefed.targetprep.RunCommandTargetPreparer">
+        <option name="run-command" value="mkdir -p /data/local/tmp/microdroid-bench" />
+        <option name="teardown-command" value="rm -rf /data/local/tmp/microdroid-bench" />
+    </target_preparer>
+    <target_preparer class="com.android.compatibility.common.tradefed.targetprep.FilePusher">
+        <option name="cleanup" value="true" />
+        <option name="push" value="test_microdroid_vendor_image.img->/data/local/tmp/microdroid-bench/microdroid_vendor_image.img" />
+    </target_preparer>
     <target_preparer class="com.android.microdroid.test.preparer.DisableMicrodroidDebugPolicyPreparer" />
     <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
         <option name="package" value="com.android.microdroid.benchmark" />
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 625f26a..1917654 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/MicrodroidBenchmarks.java
@@ -38,6 +38,7 @@
 import android.system.virtualmachine.VirtualMachineConfig;
 import android.system.virtualmachine.VirtualMachineException;
 import android.system.Os;
+import android.system.virtualmachine.VirtualMachineManager;
 import android.util.Log;
 
 import com.android.microdroid.test.common.MetricsProcessor;
@@ -260,6 +261,26 @@
     }
 
     @Test
+    public void testMicrodroidDebugBootTime_withVendorPartition() throws Exception {
+        assumeFeatureEnabled(VirtualMachineManager.FEATURE_VENDOR_MODULES);
+
+        File vendorDiskImage =
+                new File("/data/local/tmp/microdroid-bench/microdroid_vendor_image.img");
+        BootTimeStats stats =
+                runBootTimeTest(
+                        "test_vm_boot_time_debug_with_vendor_partition",
+                        (builder) ->
+                                builder.setDebugLevel(DEBUG_LEVEL_FULL)
+                                        .setVmOutputCaptured(true)
+                                        .setVendorDiskImage(vendorDiskImage));
+        reportMetrics(stats.get(BootTimeMetric.TOTAL), "boot_time", "ms");
+        reportMetrics(stats.get(BootTimeMetric.VM_START), "vm_starting_time", "ms");
+        reportMetrics(stats.get(BootTimeMetric.BOOTLOADER), "bootloader_time", "ms");
+        reportMetrics(stats.get(BootTimeMetric.KERNEL), "kernel_boot_time", "ms");
+        reportMetrics(stats.get(BootTimeMetric.USERSPACE), "userspace_boot_time", "ms");
+    }
+
+    @Test
     public void testMicrodroidImageSize() throws IOException {
         Bundle bundle = new Bundle();
         for (File file : new File(APEX_ETC_FS).listFiles()) {
diff --git a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
index 9f03ab7..e2795ec 100644
--- a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
@@ -20,6 +20,8 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.TruthJUnit.assume;
 
+import static org.junit.Assume.assumeTrue;
+
 import android.app.Instrumentation;
 import android.app.UiAutomation;
 import android.content.Context;
@@ -545,4 +547,12 @@
     protected interface RunTestsAgainstTestService {
         void runTests(ITestService testService, TestResults testResults) throws Exception;
     }
+
+    protected void assumeFeatureEnabled(String featureName) throws Exception {
+        assumeTrue(featureName + " not enabled", isFeatureEnabled(featureName));
+    }
+
+    protected boolean isFeatureEnabled(String featureName) throws Exception {
+        return getVirtualMachineManager().isFeatureEnabled(featureName);
+    }
 }
diff --git a/tests/hostside/Android.bp b/tests/hostside/Android.bp
index 6301a57..e3d9cbe 100644
--- a/tests/hostside/Android.bp
+++ b/tests/hostside/Android.bp
@@ -16,6 +16,7 @@
     static_libs: [
         "MicrodroidHostTestHelper",
         "compatibility-host-util",
+        "cts-host-utils",
         "cts-statsd-atom-host-test-utils",
         "microdroid_payload_metadata",
     ],
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index 21960b4..d861761 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -30,6 +30,8 @@
 
 import static java.util.stream.Collectors.toList;
 
+import android.cts.host.utils.DeviceJUnit4ClassRunnerWithParameters;
+import android.cts.host.utils.DeviceJUnit4Parameterized;
 import android.cts.statsdatom.lib.ConfigUtils;
 import android.cts.statsdatom.lib.ReportUtils;
 
@@ -42,7 +44,6 @@
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.device.TestDevice;
-import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner.TestMetrics;
 import com.android.tradefed.util.CommandResult;
 import com.android.tradefed.util.FileUtil;
@@ -58,6 +59,8 @@
 import org.junit.Test;
 import org.junit.rules.TestName;
 import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.UseParametersRunnerFactory;
 import org.xml.sax.Attributes;
 import org.xml.sax.helpers.DefaultHandler;
 
@@ -67,6 +70,7 @@
 import java.io.PipedOutputStream;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -77,7 +81,8 @@
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
-@RunWith(DeviceJUnit4ClassRunner.class)
+@RunWith(DeviceJUnit4Parameterized.class)
+@UseParametersRunnerFactory(DeviceJUnit4ClassRunnerWithParameters.RunnerFactory.class)
 public class MicrodroidHostTests extends MicrodroidHostTestCaseBase {
     private static final String APK_NAME = "MicrodroidTestApp.apk";
     private static final String PACKAGE_NAME = "com.android.microdroid.test";
@@ -97,6 +102,14 @@
         }
     }
 
+    @Parameterized.Parameters(name = "protectedVm={0}")
+    public static Collection<Object[]> params() {
+        return List.of(new Object[] {true}, new Object[] {false});
+    }
+
+    @Parameterized.Parameter(0)
+    public boolean mProtectedVm;
+
     @Rule public TestLogData mTestLogs = new TestLogData();
     @Rule public TestName mTestName = new TestName();
     @Rule public TestMetrics mMetrics = new TestMetrics();
@@ -259,9 +272,11 @@
         File virtApexEtcDir = new File(virtApexDir, "etc");
         // We need only etc/ directory for images
         assertWithMessage("Failed to mkdir " + virtApexEtcDir)
-                .that(virtApexEtcDir.mkdirs()).isTrue();
+                .that(virtApexEtcDir.mkdirs())
+                .isTrue();
         assertWithMessage("Failed to pull " + VIRT_APEX + "etc")
-                .that(getDevice().pullDir(VIRT_APEX + "etc", virtApexEtcDir)).isTrue();
+                .that(getDevice().pullDir(VIRT_APEX + "etc", virtApexEtcDir))
+                .isTrue();
 
         resignVirtApex(virtApexDir, key, keyOverrides, updateBootconfigs);
 
@@ -320,10 +335,12 @@
         final String initrdPath = TEST_ROOT + "etc/microdroid_initrd_debuggable.img";
         config.put("initrd", initrdPath);
         // Add instance image as a partition in disks[1]
-        disks.put(new JSONObject()
-                .put("writable", true)
-                .put("partitions",
-                        new JSONArray().put(newPartition("vm-instance", instanceImgPath))));
+        disks.put(
+                new JSONObject()
+                        .put("writable", true)
+                        .put(
+                                "partitions",
+                                new JSONArray().put(newPartition("vm-instance", instanceImgPath))));
         // Add payload image disk with partitions:
         // - payload-metadata
         // - apexes: com.android.os.statsd, com.android.adbd, [sharedlib apex](optional)
@@ -373,7 +390,7 @@
     @CddTest(requirements = {"9.17/C-2-1", "9.17/C-2-2", "9.17/C-2-6"})
     public void protectedVmRunsPvmfw() throws Exception {
         // Arrange
-        assumeProtectedVmSupported();
+        assumeProtectedVm();
         final String configPath = "assets/vm_config_apex.json";
 
         // Act
@@ -401,16 +418,16 @@
     @CddTest(requirements = {"9.17/C-2-1", "9.17/C-2-2", "9.17/C-2-6"})
     public void protectedVmWithImageSignedWithDifferentKeyRunsPvmfw() throws Exception {
         // Arrange
-        assumeProtectedVmSupported();
+        assumeProtectedVm();
         File key = findTestFile("test.com.android.virt.pem");
 
         // Act
         VmInfo vmInfo =
                 runMicrodroidWithResignedImages(
                         key,
-                        /*keyOverrides=*/ Map.of(),
-                        /*isProtected=*/ true,
-                        /*updateBootconfigs=*/ true);
+                        /* keyOverrides= */ Map.of(),
+                        /* isProtected= */ true,
+                        /* updateBootconfigs= */ true);
 
         // Assert
         vmInfo.mProcess.waitFor(5L, TimeUnit.SECONDS);
@@ -425,12 +442,12 @@
     @CddTest(requirements = {"9.17/C-2-2", "9.17/C-2-6"})
     public void testBootSucceedsWhenNonProtectedVmStartsWithImagesSignedWithDifferentKey()
             throws Exception {
-        assumeNonProtectedVmSupported();
+        assumeNonProtectedVm();
         File key = findTestFile("test.com.android.virt.pem");
         Map<String, File> keyOverrides = Map.of();
         VmInfo vmInfo =
                 runMicrodroidWithResignedImages(
-                        key, keyOverrides, /*isProtected=*/ false, /*updateBootconfigs=*/ true);
+                        key, keyOverrides, /* isProtected= */ false, /* updateBootconfigs= */ true);
         assertThatEventually(
                 100000,
                 () -> getDevice().pullFileContents(LOG_PATH),
@@ -443,13 +460,13 @@
     @CddTest(requirements = {"9.17/C-2-2", "9.17/C-2-6"})
     public void testBootFailsWhenVbMetaDigestDoesNotMatchBootconfig() throws Exception {
         // protectedVmWithImageSignedWithDifferentKeyRunsPvmfw() is the protected case.
-        assumeNonProtectedVmSupported();
+        assumeNonProtectedVm();
         // Sign everything with key1 except vbmeta
         File key = findTestFile("test.com.android.virt.pem");
         // To be able to stop it, it should be a daemon.
         VmInfo vmInfo =
                 runMicrodroidWithResignedImages(
-                        key, Map.of(), /*isProtected=*/ false, /*updateBootconfigs=*/ false);
+                        key, Map.of(), /* isProtected= */ false, /* updateBootconfigs= */ false);
         // Wait so that init can print errors to console (time in cuttlefish >> in real device)
         assertThatEventually(
                 100000,
@@ -521,26 +538,12 @@
     }
 
     @Test
-    public void testTombstonesAreGeneratedUponUserspaceCrashOnNonPvm() throws Exception {
-        assumeNonProtectedVmSupported();
+    public void testTombstonesAreGeneratedUponUserspaceCrash() throws Exception {
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        testTombstonesAreGeneratedUponUserspaceCrash(false);
-    }
-
-    @Test
-    public void testTombstonesAreGeneratedUponUserspaceCrashOnPvm() throws Exception {
-        assumeProtectedVmSupported();
-        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
-        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        testTombstonesAreGeneratedUponUserspaceCrash(true);
-    }
-
-    private void testTombstonesAreGeneratedUponUserspaceCrash(boolean protectedVm)
-            throws Exception {
         assertThat(
                         isTombstoneGeneratedWithCmd(
-                                protectedVm,
+                                mProtectedVm,
                                 "assets/vm_config.json",
                                 "kill",
                                 "-SIGSEGV",
@@ -549,28 +552,12 @@
     }
 
     @Test
-    public void testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrashOnNonPvm()
-            throws Exception {
-        assumeNonProtectedVmSupported();
+    public void testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrash() throws Exception {
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrash(false);
-    }
-
-    @Test
-    public void testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrashOnPvm()
-            throws Exception {
-        assumeProtectedVmSupported();
-        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
-        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrash(true);
-    }
-
-    private void testTombstonesAreNotGeneratedIfNotExportedUponUserspaceCrash(boolean protectedVm)
-            throws Exception {
         assertThat(
                         isTombstoneGeneratedWithCmd(
-                                protectedVm,
+                                mProtectedVm,
                                 "assets/vm_config_no_tombstone.json",
                                 "kill",
                                 "-SIGSEGV",
@@ -578,12 +565,13 @@
                 .isFalse();
     }
 
-    private void testTombstonesAreGeneratedUponKernelCrash(boolean protectedVm) throws Exception {
+    @Test
+    public void testTombstonesAreGeneratedUponKernelCrash() throws Exception {
         assumeFalse("Cuttlefish is not supported", isCuttlefish());
         assumeFalse("Skipping test because ramdump is disabled on user build", isUserBuild());
         assertThat(
                         isTombstoneGeneratedWithCmd(
-                                protectedVm,
+                                mProtectedVm,
                                 "assets/vm_config.json",
                                 "echo",
                                 "c",
@@ -592,18 +580,6 @@
                 .isTrue();
     }
 
-    @Test
-    public void testTombstonesAreGeneratedUponKernelCrashOnNonPvm() throws Exception {
-        assumeNonProtectedVmSupported();
-        testTombstonesAreGeneratedUponKernelCrash(/* protectedVm=*/ false);
-    }
-
-    @Test
-    public void testTombstonesAreGeneratedUponKernelCrashOnPvm() throws Exception {
-        assumeProtectedVmSupported();
-        testTombstonesAreGeneratedUponKernelCrash(/* protectedVm=*/ true);
-    }
-
     private boolean isTombstoneGeneratedWithVmRunApp(
             boolean protectedVm, boolean debuggable, String... additionalArgs) throws Exception {
         // we can't use microdroid builder as it wants ADB connection (debuggable)
@@ -642,48 +618,18 @@
     }
 
     @Test
-    public void testTombstonesAreGeneratedWithCrashPayloadOnPvm() throws Exception {
-        assumeProtectedVmSupported();
+    public void testTombstonesAreGeneratedWithCrashPayload() throws Exception {
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assertThat(
-                        isTombstoneGeneratedWithCrashPayload(
-                                /*protectedVm=*/ true, /*debuggable=*/ true))
+        assertThat(isTombstoneGeneratedWithCrashPayload(mProtectedVm, /* debuggable= */ true))
                 .isTrue();
     }
 
     @Test
-    public void testTombstonesAreGeneratedWithCrashPayloadOnNonPvm() throws Exception {
-        assumeNonProtectedVmSupported();
+    public void testTombstonesAreNotGeneratedWithCrashPayloadWhenNonDebuggable() throws Exception {
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assertThat(
-                        isTombstoneGeneratedWithCrashPayload(
-                                /*protectedVm=*/ false, /*debuggable=*/ true))
-                .isTrue();
-    }
-
-    @Test
-    public void testTombstonesAreNotGeneratedWithCrashPayloadWhenNonDebuggableOnPvm()
-            throws Exception {
-        assumeProtectedVmSupported();
-        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
-        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assertThat(
-                        isTombstoneGeneratedWithCrashPayload(
-                                /*protectedVm=*/ true, /*debuggable=*/ false))
-                .isFalse();
-    }
-
-    @Test
-    public void testTombstonesAreNotGeneratedWithCrashPayloadWhenNonDebuggableOnNonPvm()
-            throws Exception {
-        assumeNonProtectedVmSupported();
-        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
-        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assertThat(
-                        isTombstoneGeneratedWithCrashPayload(
-                                /*protectedVm=*/ false, /*debuggable=*/ false))
+        assertThat(isTombstoneGeneratedWithCrashPayload(mProtectedVm, /* debuggable= */ false))
                 .isFalse();
     }
 
@@ -694,62 +640,23 @@
     }
 
     @Test
-    public void testTombstonesAreGeneratedWithCrashConfigOnPvm() throws Exception {
-        assumeProtectedVmSupported();
+    public void testTombstonesAreGeneratedWithCrashConfig() throws Exception {
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assertThat(isTombstoneGeneratedWithCrashConfig(/*protectedVm=*/ true, /*debuggable=*/ true))
+        assertThat(isTombstoneGeneratedWithCrashConfig(mProtectedVm, /* debuggable= */ true))
                 .isTrue();
     }
 
     @Test
-    public void testTombstonesAreGeneratedWithCrashConfigOnNonPvm() throws Exception {
-        assumeNonProtectedVmSupported();
+    public void testTombstonesAreNotGeneratedWithCrashConfigWhenNonDebuggable() throws Exception {
         // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
         assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assertThat(
-                        isTombstoneGeneratedWithCrashConfig(
-                                /*protectedVm=*/ false, /*debuggable=*/ true))
-                .isTrue();
-    }
-
-    @Test
-    public void testTombstonesAreNotGeneratedWithCrashConfigWhenNonDebuggableOnPvm()
-            throws Exception {
-        assumeProtectedVmSupported();
-        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
-        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assertThat(
-                        isTombstoneGeneratedWithCrashConfig(
-                                /*protectedVm=*/ true, /*debuggable=*/ false))
+        assertThat(isTombstoneGeneratedWithCrashConfig(mProtectedVm, /* debuggable= */ false))
                 .isFalse();
     }
 
     @Test
-    public void testTombstonesAreNotGeneratedWithCrashConfigWhenNonDebuggableOnNonPvm()
-            throws Exception {
-        assumeNonProtectedVmSupported();
-        // TODO(b/291867858): tombstones are failing in HWASAN enabled Microdroid.
-        assumeFalse("tombstones are failing in HWASAN enabled Microdroid.", isHwasan());
-        assertThat(
-                        isTombstoneGeneratedWithCrashConfig(
-                                /*protectedVm=*/ false, /*debuggable=*/ false))
-                .isFalse();
-    }
-
-    @Test
-    public void testTelemetryPushedAtomsOnNonPvm() throws Exception {
-        assumeNonProtectedVmSupported();
-        testTelemetryPushedAtoms(false);
-    }
-
-    @Test
-    public void testTelemetryPushedAtomsOnPvm() throws Exception {
-        assumeProtectedVmSupported();
-        testTelemetryPushedAtoms(true);
-    }
-
-    private void testTelemetryPushedAtoms(boolean protectedVm) throws Exception {
+    public void testTelemetryPushedAtoms() throws Exception {
         // Reset statsd config and report before the test
         ConfigUtils.removeConfig(getDevice());
         ReportUtils.clearReports(getDevice());
@@ -770,7 +677,7 @@
                         .debugLevel("full")
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
-                        .protectedVm(protectedVm)
+                        .protectedVm(mProtectedVm)
                         .build(device);
         microdroid.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
         device.shutdownMicrodroid(microdroid);
@@ -797,7 +704,7 @@
                 data.get(0).getAtom().getVmCreationRequested();
         assertThat(atomVmCreationRequested.getHypervisor())
                 .isEqualTo(AtomsProto.VmCreationRequested.Hypervisor.PKVM);
-        assertThat(atomVmCreationRequested.getIsProtected()).isEqualTo(protectedVm);
+        assertThat(atomVmCreationRequested.getIsProtected()).isEqualTo(mProtectedVm);
         assertThat(atomVmCreationRequested.getCreationSucceeded()).isTrue();
         assertThat(atomVmCreationRequested.getBinderExceptionCode()).isEqualTo(0);
         assertThat(atomVmCreationRequested.getVmIdentifier()).isEqualTo("VmRunApp");
@@ -826,19 +733,7 @@
 
     @Test
     @CddTest(requirements = {"9.17/C-1-1", "9.17/C-1-2", "9.17/C/1-3"})
-    public void testMicrodroidBootsOnPvm() throws Exception {
-        assumeProtectedVmSupported();
-        testMicrodroidBoots(true);
-    }
-
-    @Test
-    @CddTest(requirements = {"9.17/C-1-1", "9.17/C-1-2", "9.17/C/1-3"})
-    public void testMicrodroidBootsOnNonPvm() throws Exception {
-        assumeNonProtectedVmSupported();
-        testMicrodroidBoots(false);
-    }
-
-    private void testMicrodroidBoots(boolean protectedVm) throws Exception {
+    public void testMicrodroidBoots() throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
 
         final String configPath = "assets/vm_config.json"; // path inside the APK
@@ -847,7 +742,7 @@
                         .debugLevel("full")
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
-                        .protectedVm(protectedVm)
+                        .protectedVm(mProtectedVm)
                         .build(getAndroidDevice());
         mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
         CommandRunner microdroid = new CommandRunner(mMicrodroidDevice);
@@ -906,25 +801,14 @@
     }
 
     @Test
-    public void testMicrodroidRamUsageOnPvm() throws Exception {
-        assumeProtectedVmSupported();
-        testMicrodroidRamUsage(true);
-    }
-
-    @Test
-    public void testMicrodroidRamUsageOnNonPvm() throws Exception {
-        assumeNonProtectedVmSupported();
-        testMicrodroidRamUsage(false);
-    }
-
-    private void testMicrodroidRamUsage(boolean protectedVm) throws Exception {
+    public void testMicrodroidRamUsage() throws Exception {
         final String configPath = "assets/vm_config.json";
         mMicrodroidDevice =
                 MicrodroidBuilder.fromDevicePath(getPathForPackage(PACKAGE_NAME), configPath)
                         .debugLevel("full")
                         .memoryMib(minMemorySize())
                         .cpuTopology("match_host")
-                        .protectedVm(protectedVm)
+                        .protectedVm(mProtectedVm)
                         .build(getAndroidDevice());
         mMicrodroidDevice.waitForBootComplete(BOOT_COMPLETE_TIMEOUT);
         mMicrodroidDevice.enableAdbRoot();
@@ -949,8 +833,7 @@
         for (Map.Entry<Integer, String> proc :
                 ProcessUtil.getProcessMap(microdroidExec).entrySet()) {
             for (Map.Entry<String, Long> stat :
-                    ProcessUtil.getProcessSmapsRollup(proc.getKey(), microdroidExec)
-                            .entrySet()) {
+                    ProcessUtil.getProcessSmapsRollup(proc.getKey(), microdroidExec).entrySet()) {
                 String name = stat.getKey().toLowerCase();
                 mMetrics.addTestMetric(
                         mMetricPrefix + "smaps/" + name + "/" + proc.getValue(),
@@ -1066,8 +949,8 @@
     }
 
     @Test
-    public void testDevcieAssignment() throws Exception {
-        assumeProtectedVmSupported();
+    public void testDeviceAssignment() throws Exception {
+        assumeProtectedVm();
         assumeVfioPlatformSupported();
 
         List<String> devices = getAssignableDevices();
@@ -1095,6 +978,11 @@
         prepareVirtualizationTestSetup(getDevice());
 
         getDevice().installPackage(findTestFile(APK_NAME), /* reinstall */ false);
+
+        // Skip test if given device doesn't support protected or non-protected VM.
+        assumeTrue(
+                "Microdroid is not supported for specific VM protection type",
+                getAndroidDevice().supportsMicrodroid(mProtectedVm));
     }
 
     @After
@@ -1119,16 +1007,12 @@
                         "android.permission.USE_CUSTOM_VIRTUAL_MACHINE");
     }
 
-    private void assumeProtectedVmSupported() throws Exception {
-        assumeTrue(
-                "Test skipped because protected VMs are not supported",
-                getAndroidDevice().supportsMicrodroid(true));
+    private void assumeProtectedVm() throws Exception {
+        assumeTrue("This test is only for protected VM.", mProtectedVm);
     }
 
-    private void assumeNonProtectedVmSupported() throws Exception {
-        assumeTrue(
-                "Test skipped because non-protected VMs are not supported",
-                getAndroidDevice().supportsMicrodroid(false));
+    private void assumeNonProtectedVm() throws Exception {
+        assumeFalse("This test is only for non-protected VM.", mProtectedVm);
     }
 
     private void assumeVfioPlatformSupported() throws Exception {
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 9fe5614..4b9f803 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -2254,12 +2254,4 @@
                 .that(KERNEL_VERSION)
                 .isNotEqualTo("5.4");
     }
-
-    private void assumeFeatureEnabled(String featureName) throws Exception {
-        assumeTrue(featureName + " not enabled", isFeatureEnabled(featureName));
-    }
-
-    private boolean isFeatureEnabled(String featureName) throws Exception {
-        return getVirtualMachineManager().isFeatureEnabled(featureName);
-    }
 }
diff --git a/virtualizationservice/src/aidl.rs b/virtualizationservice/src/aidl.rs
index 6f5a487..645a82b 100644
--- a/virtualizationservice/src/aidl.rs
+++ b/virtualizationservice/src/aidl.rs
@@ -160,10 +160,20 @@
     fn requestCertificate(&self, csr: &[u8]) -> binder::Result<Vec<u8>> {
         check_manage_access()?;
         info!("Received csr. Getting certificate...");
-        request_certificate(csr)
-            .context("Failed to get certificate")
+        if cfg!(remote_attestation) {
+            request_certificate(csr)
+                .context("Failed to get certificate")
+                .with_log()
+                .or_service_specific_exception(-1)
+        } else {
+            Err(Status::new_exception_str(
+                ExceptionCode::UNSUPPORTED_OPERATION,
+                Some(
+                    "requestCertificate is not supported with the remote_attestation feature disabled",
+                ),
+            ))
             .with_log()
-            .or_service_specific_exception(-1)
+        }
     }
 
     fn getAssignableDevices(&self) -> binder::Result<Vec<AssignableDevice>> {