Merge "Remove VirtualizationTestCases from CTS"
diff --git a/apex/Android.bp b/apex/Android.bp
index d12b27b..0f30c67 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -152,6 +152,11 @@
     srcs: ["test.com.android.virt.pem"],
 }
 
+filegroup {
+    name: "test2.com.android.virt.pem",
+    srcs: ["test2.com.android.virt.pem"],
+}
+
 // custom tool to replace bytes in a file
 python_binary_host {
     name: "replace_bytes",
diff --git a/apex/sign_virt_apex.py b/apex/sign_virt_apex.py
index 207c938..b659b73 100644
--- a/apex/sign_virt_apex.py
+++ b/apex/sign_virt_apex.py
@@ -51,12 +51,24 @@
         help='the extra signing arguments passed to avbtool.'
     )
     parser.add_argument(
+        '--key_override',
+        metavar="filename=key",
+        action='append',
+        help='Overrides a signing key for a file e.g. microdroid_bootloader=mykey (for testing)')
+    parser.add_argument(
         'key',
         help='path to the private key file.')
     parser.add_argument(
         'input_dir',
         help='the directory having files to be packaged')
-    return parser.parse_args(argv)
+    args = parser.parse_args(argv)
+    # preprocess --key_override into a map
+    args.key_overrides = dict()
+    if args.key_override:
+        for pair in args.key_override:
+            name, key = pair.split('=')
+            args.key_overrides[name] = key
+    return args
 
 
 def RunCommand(args, cmd, env=None, expected_return_values={0}):
@@ -155,6 +167,8 @@
 
 
 def AddHashFooter(args, key, image_path):
+    if os.path.basename(image_path) in args.key_overrides:
+        key = args.key_overrides[os.path.basename(image_path)]
     info, descriptors = AvbInfo(args, image_path)
     if info:
         descriptor = LookUp(descriptors, 'Hash descriptor')
@@ -175,6 +189,8 @@
 
 
 def AddHashTreeFooter(args, key, image_path):
+    if os.path.basename(image_path) in args.key_overrides:
+        key = args.key_overrides[os.path.basename(image_path)]
     info, descriptors = AvbInfo(args, image_path)
     if info:
         descriptor = LookUp(descriptors, 'Hashtree descriptor')
@@ -196,6 +212,8 @@
 
 
 def MakeVbmetaImage(args, key, vbmeta_img, images=None, chained_partitions=None):
+    if os.path.basename(vbmeta_img) in args.key_overrides:
+        key = args.key_overrides[os.path.basename(vbmeta_img)]
     info, descriptors = AvbInfo(args, vbmeta_img)
     if info is None:
         return
@@ -263,6 +281,8 @@
 
 
 def ReplaceBootloaderPubkey(args, key, bootloader, bootloader_pubkey):
+    if os.path.basename(bootloader) in args.key_overrides:
+        key = args.key_overrides[os.path.basename(bootloader)]
     # read old pubkey before replacement
     with open(bootloader_pubkey, 'rb') as f:
         old_pubkey = f.read()
diff --git a/apex/test2.com.android.virt.pem b/apex/test2.com.android.virt.pem
new file mode 100644
index 0000000..7e4614a
--- /dev/null
+++ b/apex/test2.com.android.virt.pem
@@ -0,0 +1,51 @@
+-----BEGIN RSA PRIVATE KEY-----
+MIIJKAIBAAKCAgEAxNrDPSV0GxORgFRVX7mwCVQSc1jxZ0ETudLpICfQzK0MTggm
+tNse5O9K9w7ZdxLHt3llm6wokz4yL/WDVirSMY5W1w4TMuTZJ3iGnBmEpjPXFxhW
+yPCBqfs3ZOQ0ndz6pimI8ZTaKF8cM6Iz/ZkNGUcBjgiUKDJQgL+zwJADK63zrGgJ
+WJdXZfPRHDPLX11kwOAemggobVgFAfnGiJc37n1561Ozq/joYD1OZTIdS8+fKkey
+p5cH9uSZgU4Fpf5xU5EDDEnZvqSYAghZY8rKlfxHuM3GdwLOhrCoUEnSYcycnDFs
+t6mYEvyv13pQRR/Co6udEW7Xha3PV5io79066XBRK0RICY5vb+92TJ7dhp3Wz+4j
+3LMwB+G02GlBWIZ6S/uRLbPpdCUj6waBb3im2+1wg5306peHoU2Pj3CJXtz96pYA
+6zgdM0d78H2KfSmKku+k6tjDKlo6584Xa3p5P6yxLnCSqwIBK6saK0M7j3gPU1P6
+j2G0fVXa20OT0qIWAcm24bBQjCF1kbQzPQnn3nKpTPVJNozH72dXYxgBQd80Sb6s
+yhnxXtiEJp+gd/SET/STOZ+qYvPmEF4TkUbtB1MLdpey3+NkRbGRCfeLeKoXJ6zW
+rS85Wo5U8v5d0nRhRCO8bRIKu3BgbBWdouHmmmUshUn3s6PoatUuDPcaX4cCAwEA
+AQKCAgBc/fbC+OFrNQhD9hLKgJ5fGb8JjFelbkGSQ8bq8MQbHBg2+HTIdMaYYU4p
+fXNPY6jCzG8qZd5ZCEWPEEy1tM7MqC/vsW9yWFcY5T+5l/hoxw5xk4bTr4GhOdJG
+L+OHO2+QdQiDDs0ryxo4bgRfZSCh80ARx42tm28aEvUoHx/QT4FPzWm01vFrcQ00
+ZGdLKoRA9N6f5wCp/q2G7GZT+hNq7w1cFJNIxvGHHQ7ekRjzyiWnRG1p69fQUtcN
+FT1n19XTIyqscGqTO4+vTiBkGtUumxmKfHKnn8TOLb+lBeqDVrQYuORhspTrS0EJ
+6nbm4IUC6jvtk03ukVfkSnJrtTdxYMR74TSgLWvsNHUtZcCDDQ+1W5g4LrMsTsfU
+bUXAwdVcRTy8rlT2EOTWVorMZHOCl9HojorRlK1kjQHStdUqrewTEF8bk1BeBI5M
+ddP2acdjfCbn1MBp4S7CqAXSuEi5TgIuoQ9ZZ8XU0n1/PuLNk9dYEdrP32VXfsM7
+AUznCq/kItrPDCzz+sNzeXXP77HRlBD7snJdtY06932ysYw85AQLEX3F3arqcyJj
+3c0/ddbGR9O4re6S0x2XhOplbaePya22Y+/I9ryWQZZO1OVOZNtTj83toIw/E5LN
+HHBAeMjHSVVyzbIEF0XqCmg59l62CRMFxcNtS67l2pcSIWSCkQKCAQEA7FWr4rUY
+gXRVg5473tM/4VT3yIGwKt83s78b0Z30VjfaZp88wx0gJaNR5HTMsEsBjof4khqu
+AKlTHb4riXBPWNz0Sf1s2PzqlKj7Ke4eowEUOAqrvkAZhZBwIvsiwFPH/mzNTdBD
+PmMDsMFerEOhVAZpDUlTUPiWHCa9ani25TYAKMA6QGqpouAMGw2xN/GtdUD4Dy5b
+0K3B2cg1lhXjWqf8LKftFGh0vsZTR/heAPx8NdrY1EOMid1blIa2KefxhunqudGr
+Owx5wMKpJA5x3sid4qds0+q8X6TWYCQORQJNiw0ggHnnjM8O9cszRIlTbMRJ3j1p
+L4TQR+1JKVu1CQKCAQEA1TwZEkUkx/vGrY358/655Nymqa0oMz/g005DTUHRILRf
+hHLH+bTyx9YSquqvE5uGIbavrJtmRwxFWikPfuBdNf87oxfT2mBESrK73162rNt2
+okgky2AoNeVRu2V93wGCTwESoLzgtruUo2EP5W7Af7KFd5Ri5zSOn0hB9nHwNHQf
+RnWxB35bBpjpAkbATr4xAeJ7t99VN0T/sgUkTVRCAy4aEyyAgz3yXGH5Nn0d76zf
+CzNAUzWuK+nwtWQetkcRcAj4ZncyqzDy7XPaoHeur8X45Bh3e27nNzbblrKm0DFT
+eRD4raxubyYFcge8I64xvt73zHcQOTKuC8kuTfmkDwKCAQBrgmmH3yP/t+Ey16ea
+rPTRV6rEbqKqThLz1Msd50IAerYCmwu0Iqq+FHare6qlw+k4YohkRnjDWkOyMxFx
+G0MtRI5onj2G1D8OU3S2VVlgg5wkBk6sZFJ33QX2E9JyNWq0ReB7NnNwjPBf1wdv
+S/C23ZeqcKHTItJ+iez+410oFhGqeA/Hv/3dVxiKsgbdUTa8MUrm9QrVekXGAXrH
+BLwBQIvJ8LY742zAYE4AXm68+h6zDRQ4M2ZaTPVdMo7pr1bDLeQWldfUK8+zLZpu
+CZgpZY/VTJ6IJK9+vui6oYxQPkTyLY2MhGgeOQ8wJzjyQ5pMz1pfHAaelEd/gOUY
+SFypAoIBAE8uUeEG6/GW/N/VqMuB+2WQyhKXyiW9wq60kSlPF2kdkZqNRNTk7IJo
+a+Yr33dYeSZrwDBIRGJ9nAMu3CIxDmvOq0aUwoaE2NckJ796XDs0A4mfYIpk2omo
+7gC4X1VAKjNMIq6tdIRmg3tnv49i4PiKQiV1ZISWb5+WJWhuRtQziqmPan1t3j9E
+6MF/pEmZNnmMsIRG2k37wTdJ0YElmJ21sNkN3WrexfCoMPKa41LszqZKEcjUVijY
+Zhn1Y7IsEb2YlyT1fkszkgG606RizOtYiGOq8jNTq2hFZqU/EdKdfnGma7GSJi//
+3mXJmYNmW/KUuU+jptKWjyqxOhCacuECggEBAJ4UfR26SKumYjw/HZj7ZNp31Fuw
++kqO9GuHlxieVk0FAk9Wd1L1r1VZReyUfKUah57JdS94iO5XizLA8xcMQ37vw5Ki
+SgKVX6ONVwkAmkHQSVAC5783k74n1PoEKd36DcBurib2SPxwXp/Yl9Y0744K4iaT
+VQSVWl/wd1NDaDY7xrOFw1keqY+hFVL/2zUozui1pypzwYOMvlmyT+UfLUnR3Kdc
+EaQMoRMLK1+ct3lyBr1CmB0tXaF+rm4yMJNrZhym2AUFUi8jOBCf3zUnJ/1O45HW
+iTir6LZxBwHcCSwJn6/HdcoYIEsLqwsVzoTMdFBGIOpB+eWaRA7/cYkRv7A=
+-----END RSA PRIVATE KEY-----
diff --git a/authfs/src/fusefs/mount.rs b/authfs/src/fusefs/mount.rs
index e7f8c94..294c6b1 100644
--- a/authfs/src/fusefs/mount.rs
+++ b/authfs/src/fusefs/mount.rs
@@ -53,8 +53,13 @@
         mount_options.push(MountOption::Extra(value));
     }
 
-    fuse::mount(mountpoint, "authfs", libc::MS_NOSUID | libc::MS_NODEV, &mount_options)
-        .expect("Failed to mount fuse");
+    fuse::mount(
+        mountpoint,
+        "authfs",
+        libc::MS_NOSUID | libc::MS_NODEV | libc::MS_NOEXEC,
+        &mount_options,
+    )
+    .expect("Failed to mount fuse");
 
     fuse::worker::start_message_loop(dev_fuse, MAX_WRITE_BYTES, MAX_READ_BYTES, authfs)
 }
diff --git a/compos/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl b/compos/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl
index 8156265..dde75e1 100644
--- a/compos/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl
+++ b/compos/composd/aidl/android/system/composd/IIsolatedCompilationService.aidl
@@ -19,6 +19,13 @@
 import android.system.composd.ICompilationTaskCallback;
 
 interface IIsolatedCompilationService {
+    enum ApexSource {
+        /** Only use the activated APEXes */
+        NoStaged,
+        /** Prefer any staged APEXes, otherwise use the activated ones */
+        PreferStaged,
+    }
+
     /**
      * Compile BCP extensions and system server, using any staged APEXes that are present in
      * preference to active APEXes, writing the results to the pending artifacts directory to be
@@ -41,5 +48,5 @@
      * callback, unless the returned ICompilationTask is cancelled. The caller should maintain
      * a reference to the ICompilationTask until compilation completes or is cancelled.
      */
-    ICompilationTask startTestCompile(ICompilationTaskCallback callback);
+    ICompilationTask startTestCompile(ApexSource apexSource, ICompilationTaskCallback callback);
 }
diff --git a/compos/composd/src/instance_manager.rs b/compos/composd/src/instance_manager.rs
index 848fc8c..f1289e8 100644
--- a/compos/composd/src/instance_manager.rs
+++ b/compos/composd/src/instance_manager.rs
@@ -48,9 +48,12 @@
         self.start_instance(CURRENT_INSTANCE_DIR, vm_parameters)
     }
 
-    pub fn start_test_instance(&self) -> Result<Arc<CompOsInstance>> {
+    pub fn start_test_instance(&self, prefer_staged: bool) -> Result<Arc<CompOsInstance>> {
         let mut vm_parameters = new_vm_parameters()?;
         vm_parameters.debug_mode = true;
+        if prefer_staged {
+            vm_parameters.config_path = Some(PREFER_STAGED_VM_CONFIG_PATH.to_owned());
+        }
         self.start_instance(TEST_INSTANCE_DIR, vm_parameters)
     }
 
diff --git a/compos/composd/src/service.rs b/compos/composd/src/service.rs
index f4798d7..a9b8202 100644
--- a/compos/composd/src/service.rs
+++ b/compos/composd/src/service.rs
@@ -22,7 +22,9 @@
 use android_system_composd::aidl::android::system::composd::{
     ICompilationTask::{BnCompilationTask, ICompilationTask},
     ICompilationTaskCallback::ICompilationTaskCallback,
-    IIsolatedCompilationService::{BnIsolatedCompilationService, IIsolatedCompilationService},
+    IIsolatedCompilationService::{
+        ApexSource::ApexSource, BnIsolatedCompilationService, IIsolatedCompilationService,
+    },
 };
 use android_system_composd::binder::{
     self, BinderFeatures, ExceptionCode, Interface, Status, Strong, ThreadState,
@@ -58,10 +60,16 @@
 
     fn startTestCompile(
         &self,
+        apex_source: ApexSource,
         callback: &Strong<dyn ICompilationTaskCallback>,
     ) -> binder::Result<Strong<dyn ICompilationTask>> {
         check_permissions()?;
-        to_binder_result(self.do_start_test_compile(callback))
+        let prefer_staged = match apex_source {
+            ApexSource::NoStaged => false,
+            ApexSource::PreferStaged => true,
+            _ => unreachable!("Invalid ApexSource {:?}", apex_source),
+        };
+        to_binder_result(self.do_start_test_compile(prefer_staged, callback))
     }
 }
 
@@ -85,9 +93,11 @@
 
     fn do_start_test_compile(
         &self,
+        prefer_staged: bool,
         callback: &Strong<dyn ICompilationTaskCallback>,
     ) -> Result<Strong<dyn ICompilationTask>> {
-        let comp_os = self.instance_manager.start_test_instance().context("Starting CompOS")?;
+        let comp_os =
+            self.instance_manager.start_test_instance(prefer_staged).context("Starting CompOS")?;
 
         let target_dir_name = TEST_ARTIFACTS_SUBDIR.to_owned();
         let task = OdrefreshTask::start(
diff --git a/compos/composd_cmd/composd_cmd.rs b/compos/composd_cmd/composd_cmd.rs
index 546c4af..9f535d5 100644
--- a/compos/composd_cmd/composd_cmd.rs
+++ b/compos/composd_cmd/composd_cmd.rs
@@ -20,6 +20,7 @@
     aidl::android::system::composd::{
         ICompilationTask::ICompilationTask,
         ICompilationTaskCallback::{BnCompilationTaskCallback, ICompilationTaskCallback},
+        IIsolatedCompilationService::ApexSource::ApexSource,
         IIsolatedCompilationService::IIsolatedCompilationService,
     },
     binder::{
@@ -33,22 +34,25 @@
 use std::time::Duration;
 
 fn main() -> Result<()> {
-    let app = clap::App::new("composd_cmd").arg(
-        clap::Arg::with_name("command")
-            .index(1)
-            .takes_value(true)
-            .required(true)
-            .possible_values(&["staged-apex-compile", "test-compile"]),
-    );
+    #[rustfmt::skip]
+    let app = clap::App::new("composd_cmd")
+        .subcommand(
+            clap::SubCommand::with_name("staged-apex-compile"))
+        .subcommand(
+            clap::SubCommand::with_name("test-compile")
+                .arg(clap::Arg::with_name("prefer-staged").long("prefer-staged")),
+        );
     let args = app.get_matches();
-    let command = args.value_of("command").unwrap();
 
     ProcessState::start_thread_pool();
 
-    match command {
-        "staged-apex-compile" => run_staged_apex_compile()?,
-        "test-compile" => run_test_compile()?,
-        _ => panic!("Unexpected command {}", command),
+    match args.subcommand() {
+        ("staged-apex-compile", _) => run_staged_apex_compile()?,
+        ("test-compile", Some(sub_matches)) => {
+            let prefer_staged = sub_matches.is_present("prefer-staged");
+            run_test_compile(prefer_staged)?;
+        }
+        _ => panic!("Unrecognized subcommand"),
     }
 
     println!("All Ok!");
@@ -108,8 +112,9 @@
     run_async_compilation(|service, callback| service.startStagedApexCompile(callback))
 }
 
-fn run_test_compile() -> Result<()> {
-    run_async_compilation(|service, callback| service.startTestCompile(callback))
+fn run_test_compile(prefer_staged: bool) -> Result<()> {
+    let apex_source = if prefer_staged { ApexSource::PreferStaged } else { ApexSource::NoStaged };
+    run_async_compilation(|service, callback| service.startTestCompile(apex_source, callback))
 }
 
 fn run_async_compilation<F>(start_compile_fn: F) -> Result<()>
diff --git a/compos/service/java/com/android/server/compos/IsolatedCompilationJobService.java b/compos/service/java/com/android/server/compos/IsolatedCompilationJobService.java
index f801a8d..75f5334 100644
--- a/compos/service/java/com/android/server/compos/IsolatedCompilationJobService.java
+++ b/compos/service/java/com/android/server/compos/IsolatedCompilationJobService.java
@@ -176,7 +176,8 @@
             try {
                 ICompilationTask composTask;
                 if (jobId == DAILY_JOB_ID) {
-                    composTask = composd.startTestCompile(this);
+                    composTask = composd.startTestCompile(
+                            IIsolatedCompilationService.ApexSource.NoStaged, this);
                 } else {
                     composTask = composd.startStagedApexCompile(this);
                 }
diff --git a/compos/verify/verify.rs b/compos/verify/verify.rs
index 55fe1bd..184b9ff 100644
--- a/compos/verify/verify.rs
+++ b/compos/verify/verify.rs
@@ -17,6 +17,7 @@
 //! A tool to verify a CompOS signature. It starts a CompOS VM as part of this to retrieve the
 //!  public key. The tool is intended to be run by odsign during boot.
 
+use android_logger::LogId;
 use anyhow::{bail, Context, Result};
 use compos_aidl_interface::binder::ProcessState;
 use compos_common::compos_client::{VmInstance, VmParameters};
@@ -39,7 +40,8 @@
     android_logger::init_once(
         android_logger::Config::default()
             .with_tag("compos_verify")
-            .with_min_level(log::Level::Info),
+            .with_min_level(log::Level::Info)
+            .with_log_id(LogId::System), // Needed to log successfully early in boot
     );
 
     // Redirect panic messages to logcat.
diff --git a/tests/hostside/Android.bp b/tests/hostside/Android.bp
index 67a0e8d..fa9da9d 100644
--- a/tests/hostside/Android.bp
+++ b/tests/hostside/Android.bp
@@ -21,7 +21,10 @@
         ":MicrodroidTestApp",
         ":microdroid_general_sepolicy.conf",
         ":test.com.android.virt.pem",
+        ":test2.com.android.virt.pem",
         ":test-payload-metadata",
+        ":com.android.adbd{.apex}",
+        ":com.android.os.statsd{.apex}",
     ],
     data_native_bins: [
         "sepolicy-analyze",
diff --git a/tests/hostside/helper/Android.bp b/tests/hostside/helper/Android.bp
index aa748ab..4ca0bf0 100644
--- a/tests/hostside/helper/Android.bp
+++ b/tests/hostside/helper/Android.bp
@@ -2,12 +2,9 @@
     default_applicable_licenses: ["Android-Apache-2.0"],
 }
 
-java_test_helper_library {
+java_library_host {
     name: "VirtualizationTestHelper",
-    host_supported: true,
-    device_supported: false,
     srcs: ["java/**/*.java"],
-    test_suites: ["general-tests"],
     libs: [
         "tradefed",
         "compatibility-tradefed",
diff --git a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
index e45e524..0f6204c 100644
--- a/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
+++ b/tests/hostside/helper/java/android/virt/test/VirtualizationTestCaseBase.java
@@ -194,6 +194,22 @@
         }
     }
 
+    public String getPathForPackage(String packageName)
+            throws DeviceNotAvailableException {
+        return getPathForPackage(getDevice(), packageName);
+    }
+
+    // Get the path to the installed apk. Note that
+    // getDevice().getAppPackageInfo(...).getCodePath() doesn't work due to the incorrect
+    // parsing of the "=" character. (b/190975227). So we use the `pm path` command directly.
+    private static String getPathForPackage(ITestDevice device, String packageName)
+            throws DeviceNotAvailableException {
+        CommandRunner android = new CommandRunner(device);
+        String pathLine = android.run("pm", "path", packageName);
+        assertTrue("package not found", pathLine.startsWith("package:"));
+        return pathLine.substring("package:".length());
+    }
+
     public static String startMicrodroid(
             ITestDevice androidDevice,
             IBuildInfo buildInfo,
@@ -247,13 +263,8 @@
             androidDevice.installPackage(apkFile, /* reinstall */ true);
         }
 
-        // Get the path to the installed apk. Note that
-        // getDevice().getAppPackageInfo(...).getCodePath() doesn't work due to the incorrect
-        // parsing of the "=" character. (b/190975227). So we use the `pm path` command directly.
         if (apkPath == null) {
-            apkPath = android.run("pm", "path", packageName);
-            assertTrue(apkPath.startsWith("package:"));
-            apkPath = apkPath.substring("package:".length());
+            apkPath = getPathForPackage(androidDevice, packageName);
         }
 
         android.run("mkdir", "-p", TEST_ROOT);
diff --git a/tests/hostside/java/android/virt/test/MicrodroidTestCase.java b/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
index d0bc91a..1218b68 100644
--- a/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
+++ b/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
@@ -50,6 +50,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.Callable;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -143,15 +144,7 @@
         boolean writable;
     }
 
-    private String getPathForPackage(String packageName)
-            throws DeviceNotAvailableException {
-        CommandRunner android = new CommandRunner(getDevice());
-        String pathLine = android.run("pm", "path", packageName);
-        assertTrue("package not found", pathLine.startsWith("package:"));
-        return pathLine.substring("package:".length());
-    }
-
-    private void resignVirtApex(File virtApexDir, File signingKey) {
+    private void resignVirtApex(File virtApexDir, File signingKey, Map<String, File> keyOverrides) {
         File signVirtApex = findTestFile("sign_virt_apex");
 
         RunUtil runUtil = new RunUtil();
@@ -160,14 +153,21 @@
         String path = signVirtApex.getParentFile().getPath() + separator + System.getenv("PATH");
         runUtil.setEnvVariable("PATH", path);
 
-        String resignCommand = String.format("sign_virt_apex %s %s",
-                                        signingKey.getPath(),
-                                        virtApexDir.getPath());
+        List<String> command = new ArrayList<String>();
+        command.add("sign_virt_apex");
+        for (Map.Entry<String, File> entry : keyOverrides.entrySet()) {
+            String filename = entry.getKey();
+            File overridingKey = entry.getValue();
+            command.add("--key_override " + filename + "=" + overridingKey.getPath());
+        }
+        command.add(signingKey.getPath());
+        command.add(virtApexDir.getPath());
+
         CommandResult result = runUtil.runTimedCmd(
                                     20 * 1000,
                                     "/bin/bash",
                                     "-c",
-                                    resignCommand);
+                                    String.join(" ", command));
         String out = result.getStdout();
         String err = result.getStderr();
         assertEquals(
@@ -175,8 +175,26 @@
                 CommandStatus.SUCCESS, result.getStatus());
     }
 
-    private String runMicrodroidWithResignedImages(boolean isProtected, boolean daemonize,
-            String consolePath) throws DeviceNotAvailableException, IOException {
+    private static <T> void assertThatEventually(long timeoutMillis, Callable<T> callable,
+            org.hamcrest.Matcher<T> matcher) throws Exception {
+        long start = System.currentTimeMillis();
+        while (true) {
+            try {
+                assertThat(callable.call(), matcher);
+                return;
+            } catch (Throwable e) {
+                if (System.currentTimeMillis() - start < timeoutMillis) {
+                    Thread.sleep(500);
+                } else {
+                    throw e;
+                }
+            }
+        }
+    }
+
+    private String runMicrodroidWithResignedImages(File key, Map<String, File> keyOverrides,
+            boolean isProtected, boolean daemonize, String consolePath)
+            throws DeviceNotAvailableException, IOException {
         CommandRunner android = new CommandRunner(getDevice());
 
         File virtApexDir = FileUtil.createTempDir("virt_apex");
@@ -187,8 +205,7 @@
         assertTrue(virtApexEtcDir.mkdirs());
         assertTrue(getDevice().pullDir(VIRT_APEX + "etc", virtApexEtcDir));
 
-        File testKey = findTestFile("test.com.android.virt.pem");
-        resignVirtApex(virtApexDir, testKey);
+        resignVirtApex(virtApexDir, key, keyOverrides);
 
         // Push back re-signed virt APEX contents and updated microdroid.json
         getDevice().pushDir(virtApexDir, TEST_ROOT);
@@ -207,6 +224,12 @@
         final String payloadMetadataPath = TEST_ROOT + "payload-metadata.img";
         getDevice().pushFile(findTestFile("test-payload-metadata.img"), payloadMetadataPath);
 
+        // push APEXes required for the VM.
+        final String statsdApexPath = TEST_ROOT + "com.android.os.statsd.apex";
+        final String adbdApexPath = TEST_ROOT + "com.android.adbd.apex";
+        getDevice().pushFile(findTestFile("com.android.os.statsd.apex"), statsdApexPath);
+        getDevice().pushFile(findTestFile("com.android.adbd.apex"), adbdApexPath);
+
         // Since Java APP can't start a VM with a custom image, here, we start a VM using `vm run`
         // command with a VM Raw config which is equiv. to what virtualizationservice creates with
         // a VM App config.
@@ -247,17 +270,8 @@
         // - apk and idsig
         Disk payloadDisk = new Disk();
         payloadDisk.addPartition("payload-metadata", payloadMetadataPath);
-        String[] apexes = {"com.android.os.statsd", "com.android.adbd"};
-        for (int i = 0; i < apexes.length; i++) {
-            String apexPath = getPathForPackage(apexes[i]);
-            String filename = apexes[i] + ".apex";
-            File localApexFile = new File(virtApexDir, filename);
-            String remoteApexFile = TEST_ROOT + filename;
-            // Since `adb shell vm` can't access apex_data_file, we `adb pull/push` apex files.
-            getDevice().pullFile(apexPath, localApexFile);
-            getDevice().pushFile(localApexFile, remoteApexFile);
-            payloadDisk.addPartition("microdroid-apex-" + i, remoteApexFile);
-        }
+        payloadDisk.addPartition("microdroid-apex-0", statsdApexPath);
+        payloadDisk.addPartition("microdroid-apex-1", adbdApexPath);
         payloadDisk.addPartition("microdroid-apk", apkPath);
         payloadDisk.addPartition("microdroid-apk-idsig", idSigPath);
         config.disks.add(payloadDisk);
@@ -286,9 +300,13 @@
     public void testBootFailsWhenProtectedVmStartsWithImagesSignedWithDifferentKey()
             throws Exception {
         assumeTrue(isProtectedVmSupported());
+
+        File key = findTestFile("test.com.android.virt.pem");
+        Map<String, File> keyOverrides = Map.of();
+        boolean isProtected = true;
+        boolean daemonize = false;  // VM should shut down due to boot failure.
         String consolePath = TEST_ROOT + "console";
-        // Run VM without --daemonize. It will shut down due to boot failure.
-        runMicrodroidWithResignedImages(/*protected=*/true, /*daemonize=*/false, consolePath);
+        runMicrodroidWithResignedImages(key, keyOverrides, isProtected, daemonize, consolePath);
         assertThat(getDevice().pullFileContents(consolePath),
                 containsString("pvmfw boot failed"));
     }
@@ -296,8 +314,54 @@
     @Test
     public void testBootSucceedsWhenNonProtectedVmStartsWithImagesSignedWithDifferentKey()
             throws Exception {
-        String cid = runMicrodroidWithResignedImages(/*protected=*/false,
-                /*daemonize=*/true, /*consolePath=*/null);
+        File key = findTestFile("test.com.android.virt.pem");
+        Map<String, File> keyOverrides = Map.of();
+        boolean isProtected = false;
+        boolean daemonize = true;
+        String consolePath = TEST_ROOT + "console";
+        String cid = runMicrodroidWithResignedImages(key, keyOverrides, isProtected, daemonize,
+                consolePath);
+        // Adb connection to the microdroid means that boot succeeded.
+        adbConnectToMicrodroid(getDevice(), cid);
+        shutdownMicrodroid(getDevice(), cid);
+    }
+
+    @Test
+    public void testBootFailsWhenBootloaderAndVbMetaAreSignedWithDifferentKeys()
+            throws Exception {
+        // Sign everything with key1 except vbmeta
+        File key = findTestFile("test.com.android.virt.pem");
+        File key2 = findTestFile("test2.com.android.virt.pem");
+        Map<String, File> keyOverrides = Map.of(
+                "microdroid_vbmeta.img", key2);
+        boolean isProtected = false;  // Not interested in pvwfw
+        boolean daemonize = true;  // Bootloader fails and enters prompts.
+                                   // To be able to stop it, it should be a daemon.
+        String consolePath = TEST_ROOT + "console";
+        String cid = runMicrodroidWithResignedImages(key, keyOverrides, isProtected, daemonize,
+                consolePath);
+        // Wail for a while so that bootloader prints errors to console
+        assertThatEventually(10000, () -> getDevice().pullFileContents(consolePath),
+                containsString("Public key was rejected"));
+        shutdownMicrodroid(getDevice(), cid);
+    }
+
+    @Test
+    public void testBootSucceedsWhenBootloaderAndVbmetaHaveSameSigningKeys()
+            throws Exception {
+        // Sign everything with key1 except bootloader and vbmeta
+        File key = findTestFile("test.com.android.virt.pem");
+        File key2 = findTestFile("test2.com.android.virt.pem");
+        Map<String, File> keyOverrides = Map.of(
+                "microdroid_bootloader", key2,
+                "microdroid_vbmeta.img", key2,
+                "microdroid_vbmeta_bootconfig.img", key2);
+        boolean isProtected = false;  // Not interested in pvwfw
+        boolean daemonize = true;  // Bootloader should succeed.
+                                   // To be able to stop it, it should be a daemon.
+        String consolePath = TEST_ROOT + "console";
+        String cid = runMicrodroidWithResignedImages(key, keyOverrides, isProtected, daemonize,
+                consolePath);
         // Adb connection to the microdroid means that boot succeeded.
         adbConnectToMicrodroid(getDevice(), cid);
         shutdownMicrodroid(getDevice(), cid);
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 cd9f284..27e1846 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -73,9 +73,20 @@
     private static final String KERNEL_VERSION = SystemProperties.get("ro.kernel.version");
 
     private static class Inner {
+        public boolean mProtectedVm;
         public Context mContext;
         public VirtualMachineManager mVmm;
         public VirtualMachine mVm;
+
+        Inner(boolean protectedVm) {
+            mProtectedVm = protectedVm;
+        }
+
+        /** Create a new VirtualMachineConfig.Builder with the parameterized protection mode. */
+        public VirtualMachineConfig.Builder newVmConfigBuilder(String payloadConfigPath) {
+            return new VirtualMachineConfig.Builder(mContext, payloadConfigPath)
+                            .protectedVm(mProtectedVm);
+        }
     }
 
     @Parameterized.Parameters(name = "protectedVm={0}")
@@ -112,7 +123,7 @@
                 .that(HypervisorProperties.hypervisor_vm_supported().orElse(false))
                 .isTrue();
         }
-        mInner = new Inner();
+        mInner = new Inner(mProtectedVm);
         mInner.mContext = ApplicationProvider.getApplicationContext();
         mInner.mVmm = VirtualMachineManager.getInstance(mInner.mContext);
     }
@@ -181,8 +192,7 @@
             .isNotEqualTo("5.4");
 
         VirtualMachineConfig.Builder builder =
-                new VirtualMachineConfig.Builder(mInner.mContext, "assets/vm_config_extra_apk.json")
-                        .protectedVm(mProtectedVm);
+                mInner.newVmConfigBuilder("assets/vm_config_extra_apk.json");
         if (Build.SUPPORTED_ABIS.length > 0) {
             String primaryAbi = Build.SUPPORTED_ABIS[0];
             switch(primaryAbi) {
@@ -261,9 +271,7 @@
             .that(KERNEL_VERSION)
             .isNotEqualTo("5.4");
 
-        VirtualMachineConfig.Builder builder =
-                new VirtualMachineConfig.Builder(mInner.mContext, "assets/vm_config.json")
-                        .protectedVm(mProtectedVm);
+        VirtualMachineConfig.Builder builder = mInner.newVmConfigBuilder("assets/vm_config.json");
         VirtualMachineConfig normalConfig = builder.debugLevel(DebugLevel.NONE).build();
         mInner.mVm = mInner.mVmm.getOrCreate("test_vm", normalConfig);
         VmEventListener listener =
@@ -306,10 +314,9 @@
 
     private VmCdis launchVmAndGetCdis(String instanceName)
             throws VirtualMachineException, InterruptedException {
-        VirtualMachineConfig.Builder builder =
-                new VirtualMachineConfig.Builder(mInner.mContext, "assets/vm_config.json")
-                        .protectedVm(mProtectedVm);
-        VirtualMachineConfig normalConfig = builder.debugLevel(DebugLevel.NONE).build();
+        VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder("assets/vm_config.json")
+                .debugLevel(DebugLevel.NONE)
+                .build();
         mInner.mVm = mInner.mVmm.getOrCreate(instanceName, normalConfig);
         final VmCdis vmCdis = new VmCdis();
         final CompletableFuture<Exception> exception = new CompletableFuture<>();
@@ -391,10 +398,9 @@
             .that(KERNEL_VERSION)
             .isNotEqualTo("5.4");
 
-        VirtualMachineConfig.Builder builder =
-                new VirtualMachineConfig.Builder(mInner.mContext, "assets/vm_config.json")
-                        .protectedVm(mProtectedVm);
-        VirtualMachineConfig normalConfig = builder.debugLevel(DebugLevel.NONE).build();
+        VirtualMachineConfig normalConfig = mInner.newVmConfigBuilder("assets/vm_config.json")
+                .debugLevel(DebugLevel.NONE)
+                .build();
         mInner.mVm = mInner.mVmm.getOrCreate("bcc_vm", normalConfig);
         final VmCdis vmCdis = new VmCdis();
         final CompletableFuture<byte[]> bcc = new CompletableFuture<>();
@@ -433,6 +439,10 @@
 
     private static final UUID MICRODROID_PARTITION_UUID =
             UUID.fromString("cf9afe9a-0662-11ec-a329-c32663a09d75");
+    private static final UUID U_BOOT_AVB_PARTITION_UUID =
+            UUID.fromString("7e8221e7-03e6-4969-948b-73a4c809a4f2");
+    private static final UUID U_BOOT_ENV_PARTITION_UUID =
+            UUID.fromString("0ab72d30-86ae-4d05-81b2-c1760be2b1f9");
     private static final long BLOCK_SIZE = 512;
 
     // Find the starting offset which holds the data of a partition having UUID.
@@ -460,53 +470,11 @@
         file.writeByte(b ^ 1);
     }
 
-    @Test
-    public void bootFailsWhenInstanceDiskIsCompromised()
-            throws VirtualMachineException, InterruptedException, IOException {
-        assume().withMessage("Skip on Cuttlefish. b/195765441")
-                .that(android.os.Build.DEVICE)
-                .isNotEqualTo("vsoc_x86_64");
-
-        VirtualMachineConfig config =
-                new VirtualMachineConfig.Builder(mInner.mContext, "assets/vm_config.json")
-                        .protectedVm(mProtectedVm)
-                        .debugLevel(DebugLevel.NONE)
-                        .build();
-
-        // Remove any existing VM so we can start from scratch
-        VirtualMachine oldVm = mInner.mVmm.getOrCreate("test_vm_integrity", config);
-        oldVm.delete();
-
-        mInner.mVm = mInner.mVmm.getOrCreate("test_vm_integrity", config);
-
-        final CompletableFuture<Boolean> payloadReady = new CompletableFuture<>();
-        VmEventListener listener =
-                new VmEventListener() {
-                    @Override
-                    public void onPayloadReady(VirtualMachine vm) {
-                        payloadReady.complete(true);
-                        forceStop(vm);
-                    }
-                };
-        listener.runToFinish(mInner.mVm);
-        assertThat(payloadReady.getNow(false)).isTrue();
-
-        // Launch the same VM after flipping a bit of the instance image.
-        // Flip actual data, as flipping trivial bits like the magic string isn't interesting.
-        File vmRoot = new File(mInner.mContext.getFilesDir(), "vm");
-        File vmDir = new File(vmRoot, "test_vm_integrity");
-        File instanceImgPath = new File(vmDir, "instance.img");
-        RandomAccessFile instanceFile = new RandomAccessFile(instanceImgPath, "rw");
-
-        // microdroid data partition must exist.
-        OptionalLong microdroidPartitionOffset =
-                findPartitionDataOffset(instanceFile, MICRODROID_PARTITION_UUID);
-        assertThat(microdroidPartitionOffset.isPresent()).isTrue();
-
-        flipBit(instanceFile, microdroidPartitionOffset.getAsLong());
-        mInner.mVm = mInner.mVmm.get("test_vm_integrity"); // re-load the vm with new instance disk
+    private boolean tryBootVm(String vmName)
+            throws VirtualMachineException, InterruptedException {
+        mInner.mVm = mInner.mVmm.get(vmName); // re-load the vm before running tests
         final CompletableFuture<Boolean> payloadStarted = new CompletableFuture<>();
-        listener =
+        VmEventListener listener =
                 new VmEventListener() {
                     @Override
                     public void onPayloadStarted(VirtualMachine vm, ParcelFileDescriptor stream) {
@@ -515,7 +483,45 @@
                     }
                 };
         listener.runToFinish(mInner.mVm);
-        assertThat(payloadStarted.getNow(false)).isFalse();
-        flipBit(instanceFile, microdroidPartitionOffset.getAsLong());
+        return payloadStarted.getNow(false);
+    }
+
+    @Test
+    public void bootFailsWhenInstanceDiskIsCompromised()
+            throws VirtualMachineException, InterruptedException, IOException {
+        assume().withMessage("Skip on Cuttlefish. b/195765441")
+                .that(android.os.Build.DEVICE)
+                .isNotEqualTo("vsoc_x86_64");
+
+        VirtualMachineConfig config = mInner.newVmConfigBuilder("assets/vm_config.json")
+                .debugLevel(DebugLevel.NONE)
+                .build();
+
+        // Remove any existing VM so we can start from scratch
+        VirtualMachine oldVm = mInner.mVmm.getOrCreate("test_vm_integrity", config);
+        oldVm.delete();
+        mInner.mVmm.getOrCreate("test_vm_integrity", config);
+
+        assertThat(tryBootVm("test_vm_integrity")).isTrue();
+
+        // Launch the same VM after flipping a bit of the instance image.
+        // Flip actual data, as flipping trivial bits like the magic string isn't interesting.
+        File vmRoot = new File(mInner.mContext.getFilesDir(), "vm");
+        File vmDir = new File(vmRoot, "test_vm_integrity");
+        File instanceImgPath = new File(vmDir, "instance.img");
+        RandomAccessFile instanceFile = new RandomAccessFile(instanceImgPath, "rw");
+
+        // partitions may or may not exist
+        for (UUID uuid :
+                new UUID[] {
+                    MICRODROID_PARTITION_UUID, U_BOOT_AVB_PARTITION_UUID, U_BOOT_ENV_PARTITION_UUID
+                }) {
+            OptionalLong offset = findPartitionDataOffset(instanceFile, uuid);
+            if (!offset.isPresent()) continue;
+
+            flipBit(instanceFile, offset.getAsLong());
+            assertThat(tryBootVm("test_vm_integrity")).isFalse();
+            flipBit(instanceFile, offset.getAsLong());
+        }
     }
 }