Test: bootloader checks vbmeta signing key

Bootloader fails when its embedded key doesn't match with vbmeta's
signing key.

Note that VM itself doesn't shut down even when boot fails for now.
It should shut down the VM just like pvmfw does.

Bug: 218934597
Test: atest MicrodroidHostTestCases
Change-Id: I8bfef33d3aafe72f672c035fe796d8636bc95e1c
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/tests/hostside/Android.bp b/tests/hostside/Android.bp
index 67a0e8d..6083fe5 100644
--- a/tests/hostside/Android.bp
+++ b/tests/hostside/Android.bp
@@ -21,6 +21,7 @@
         ":MicrodroidTestApp",
         ":microdroid_general_sepolicy.conf",
         ":test.com.android.virt.pem",
+        ":test2.com.android.virt.pem",
         ":test-payload-metadata",
     ],
     data_native_bins: [
diff --git a/tests/hostside/java/android/virt/test/MicrodroidTestCase.java b/tests/hostside/java/android/virt/test/MicrodroidTestCase.java
index e65459a..579fc18 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,7 +144,7 @@
         boolean writable;
     }
 
-    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();
@@ -152,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(
@@ -167,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");
@@ -179,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);
@@ -278,9 +303,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"));
     }
@@ -288,8 +317,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);