Merge "Test: bootloader checks vbmeta signing key"
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);