[sign_virt_apex] Update vbmeta related bootconfigs

Signing microdroid super images will change vbmeta digest & size, these
are embedded in ramdisk & hence need to be changed too. For this we
detach the bootconfigs from ramdisk, update them & re-attach them.

Other then that, this patch also removes signing of legacy images like
bootloader.

Test: atest
MicrodroidHostTests#testBootSucceedsWhenNonProtectedVmStartsWithImagesSignedWithDifferentKey
Bug: 245277660

Change-Id: Ia1d2ab0a7c76c7ee7435e55bab9a1c9d4f29f202
diff --git a/apex/Android.bp b/apex/Android.bp
index dce8edd..64836ff 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -146,7 +146,10 @@
         },
     },
     required: [
+        // sign_virt_apex should be runnable from outside the source tree,
+        // therefore, any required tool should be listed in build/make/core/Makefile as well.
         "img2simg",
+        "initrd_bootconfig",
         "lpmake",
         "lpunpack",
         "simg2img",
@@ -167,6 +170,7 @@
         // sign_virt_apex
         "avbtool",
         "img2simg",
+        "initrd_bootconfig",
         "lpmake",
         "lpunpack",
         "sign_virt_apex",
diff --git a/apex/sign_virt_apex.py b/apex/sign_virt_apex.py
index 557c8aa..65e8414 100644
--- a/apex/sign_virt_apex.py
+++ b/apex/sign_virt_apex.py
@@ -24,7 +24,7 @@
 
 sign_virt_apex uses external tools which are assumed to be available via PATH.
 - avbtool (--avbtool can override the tool)
-- lpmake, lpunpack, simg2img, img2simg
+- lpmake, lpunpack, simg2img, img2simg, initrd_bootconfig
 """
 import argparse
 import hashlib
@@ -102,6 +102,11 @@
     parser.add_argument(
         'input_dir',
         help='the directory having files to be packaged')
+    parser.add_argument(
+        '--do_not_update_bootconfigs',
+        action='store_true',
+        help='This will NOT update the vbmeta related bootconfigs while signing the apex.\
+            Used for testing only!!')
     args = parser.parse_args(argv)
     # preprocess --key_override into a map
     args.key_overrides = {}
@@ -254,6 +259,82 @@
         RunCommand(args, cmd)
 
 
+def UpdateVbmetaBootconfig(args, initrds, vbmeta_img):
+    # Update the bootconfigs in ramdisk
+    def detach_bootconfigs(initrd_bc, initrd, bc):
+        cmd = ['initrd_bootconfig', 'detach', initrd_bc, initrd, bc]
+        RunCommand(args, cmd)
+
+    def attach_bootconfigs(initrd_bc, initrd, bc):
+        cmd = ['initrd_bootconfig', 'attach',
+               initrd, bc, '--output', initrd_bc]
+        RunCommand(args, cmd)
+
+    # Validate that avb version used while signing the apex is the same as used by build server
+    def validate_avb_version(bootconfigs):
+        cmd = ['avbtool', 'version']
+        stdout, _ = RunCommand(args, cmd)
+        avb_version_curr = stdout.split(" ")[1].strip()
+        avb_version_curr = avb_version_curr[0:avb_version_curr.rfind('.')]
+
+        avb_version_bc = re.search(
+            r"androidboot.vbmeta.avb_version = \"([^\"]*)\"", bootconfigs).group(1)
+        if avb_version_curr != avb_version_bc:
+            raise Exception(f'AVB version mismatch between current & one & \
+                used to build bootconfigs:{avb_version_curr}&{avb_version_bc}')
+
+    def calc_vbmeta_digest():
+        cmd = ['avbtool', 'calculate_vbmeta_digest', '--image',
+               vbmeta_img, '--hash_algorithm', 'sha256']
+        stdout, _ = RunCommand(args, cmd)
+        return stdout.strip()
+
+    def calc_vbmeta_size():
+        cmd = ['avbtool', 'info_image', '--image', vbmeta_img]
+        stdout, _ = RunCommand(args, cmd)
+        size = 0
+        for line in stdout.split("\n"):
+            line = line.split(":")
+            if line[0] in ['Header Block', 'Authentication Block', 'Auxiliary Block']:
+                size += int(line[1].strip()[0:-6])
+        return size
+
+    def update_vbmeta_digest(bootconfigs):
+        # Update androidboot.vbmeta.digest in bootconfigs
+        result = re.search(
+            r"androidboot.vbmeta.digest = \"[^\"]*\"", bootconfigs)
+        if not result:
+            raise ValueError("Failed to find androidboot.vbmeta.digest")
+
+        return bootconfigs.replace(result.group(),
+                                   f'androidboot.vbmeta.digest = "{calc_vbmeta_digest()}"')
+
+    def update_vbmeta_size(bootconfigs):
+        # Update androidboot.vbmeta.size in bootconfigs
+        result = re.search(r"androidboot.vbmeta.size = [0-9]+", bootconfigs)
+        if not result:
+            raise ValueError("Failed to find androidboot.vbmeta.size")
+        return bootconfigs.replace(result.group(),
+                                   f'androidboot.vbmeta.size = {calc_vbmeta_size()}')
+
+    with tempfile.TemporaryDirectory() as work_dir:
+        tmp_initrd = os.path.join(work_dir, 'initrd')
+        tmp_bc = os.path.join(work_dir, 'bc')
+
+        for initrd in initrds:
+            detach_bootconfigs(initrd, tmp_initrd, tmp_bc)
+            bc_file = open(tmp_bc, "rt", encoding="utf-8")
+            bc_data = bc_file.read()
+            validate_avb_version(bc_data)
+            bc_data = update_vbmeta_digest(bc_data)
+            bc_data = update_vbmeta_size(bc_data)
+            bc_file.close()
+            bc_file = open(tmp_bc, "wt", encoding="utf-8")
+            bc_file.write(bc_data)
+            bc_file.flush()
+            attach_bootconfigs(initrd, tmp_initrd, tmp_bc)
+
+
 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)]
@@ -318,43 +399,13 @@
         RunCommand(args, cmd)
 
 
-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()
-
-    # replace bootloader pubkey (overwrite the old one with the new one)
-    ExtractAvbPubkey(args, key, bootloader_pubkey)
-
-    # read new pubkey
-    with open(bootloader_pubkey, 'rb') as f:
-        new_pubkey = f.read()
-
-    assert len(old_pubkey) == len(new_pubkey)
-
-    # replace pubkey embedded in bootloader
-    with open(bootloader, 'r+b') as bl_f:
-        pos = bl_f.read().find(old_pubkey)
-        assert pos != -1
-        bl_f.seek(pos)
-        bl_f.write(new_pubkey)
-
-
 # dict of (key, file) for re-sign/verification. keys are un-versioned for readability.
 virt_apex_files = {
-    'bootloader.pubkey': 'etc/microdroid_bootloader.avbpubkey',
-    'bootloader': 'etc/fs/microdroid_bootloader',
-    'boot.img': 'etc/fs/microdroid_boot.img',
-    'vendor_boot.img': 'etc/fs/microdroid_vendor_boot.img',
-    'init_boot.img': 'etc/fs/microdroid_init_boot.img',
-    'super.img': 'etc/fs/microdroid_super.img',
+    'kernel': 'etc/fs/microdroid_kernel',
     'vbmeta.img': 'etc/fs/microdroid_vbmeta.img',
-    'vbmeta_bootconfig.img': 'etc/fs/microdroid_vbmeta_bootconfig.img',
-    'bootconfig.normal': 'etc/fs/microdroid_bootconfig.normal',
-    'bootconfig.debuggable': 'etc/fs/microdroid_bootconfig.debuggable',
-    'uboot_env.img': 'etc/fs/uboot_env.img'
+    'super.img': 'etc/fs/microdroid_super.img',
+    'initrd_normal.img': 'etc/microdroid_initrd_normal.img',
+    'initrd_debuggable.img': 'etc/microdroid_initrd_debuggable.img',
 }
 
 
@@ -371,17 +422,6 @@
     system_a_img = os.path.join(unpack_dir.name, 'system_a.img')
     vendor_a_img = os.path.join(unpack_dir.name, 'vendor_a.img')
 
-    # Key(pubkey) embedded in bootloader should match with the one used to make VBmeta below
-    # while it's okay to use different keys for other image files.
-    replace_f = Async(ReplaceBootloaderPubkey, args,
-                      key, files['bootloader'], files['bootloader.pubkey'])
-
-    # re-sign bootloader, boot.img, vendor_boot.img, and init_boot.img
-    Async(AddHashFooter, args, key, files['bootloader'], wait=[replace_f])
-    Async(AddHashFooter, args, key, files['boot.img'])
-    Async(AddHashFooter, args, key, files['vendor_boot.img'])
-    Async(AddHashFooter, args, key, files['init_boot.img'])
-
     # re-sign super.img
     # 1. unpack super.img
     # 2. resign system and vendor
@@ -390,27 +430,22 @@
     system_a_f = Async(AddHashTreeFooter, args, key, system_a_img)
     vendor_a_f = Async(AddHashTreeFooter, args, key, vendor_a_img)
     partitions = {"system_a": system_a_img, "vendor_a": vendor_a_img}
-    Async(MakeSuperImage, args, partitions, files['super.img'], wait=[system_a_f, vendor_a_f])
+    Async(MakeSuperImage, args, partitions,
+          files['super.img'], wait=[system_a_f, vendor_a_f])
 
     # re-generate vbmeta from re-signed {system_a, vendor_a}.img
-    Async(MakeVbmetaImage, args, key, files['vbmeta.img'],
-          images=[system_a_img, vendor_a_img],
-          wait=[system_a_f, vendor_a_f])
+    vbmeta_f = Async(MakeVbmetaImage, args, key, files['vbmeta.img'],
+                     images=[system_a_img, vendor_a_img],
+                     wait=[system_a_f, vendor_a_f])
 
-    # Re-sign bootconfigs and the uboot_env with the same key
-    bootconfig_sign_key = key
-    Async(AddHashFooter, args, bootconfig_sign_key, files['bootconfig.normal'])
-    Async(AddHashFooter, args, bootconfig_sign_key, files['bootconfig.debuggable'])
-    Async(AddHashFooter, args, bootconfig_sign_key, files['uboot_env.img'])
+    if not args.do_not_update_bootconfigs:
+        Async(UpdateVbmetaBootconfig, args, [files['initrd_normal.img'],
+                                             files['initrd_debuggable.img']], files['vbmeta.img'],
+              wait=[vbmeta_f])
 
-    # Re-sign vbmeta_bootconfig with chained_partitions to "bootconfig" and
-    # "uboot_env". Note that, for now, `key` and `bootconfig_sign_key` are the
-    # same, but technically they can be different. Vbmeta records pubkeys which
-    # signed chained partitions.
-    Async(MakeVbmetaImage, args, key, files['vbmeta_bootconfig.img'], chained_partitions={
-        'bootconfig': bootconfig_sign_key,
-        'uboot_env': bootconfig_sign_key,
-    })
+    # Re-sign kernel
+    # TODO(b/265382249): Kernel's vbmeta should contain hashes of initrd
+    Async(AddHashFooter, args, key, files['kernel'])
 
 
 def VerifyVirtApex(args):
@@ -430,27 +465,16 @@
             pubkey = f.read()
             pubkey_digest = hashlib.sha1(pubkey).hexdigest()
 
-    def contents(file):
-        with open(file, 'rb') as f:
-            return f.read()
-
-    def check_equals_pubkey(file):
-        assert contents(file) == pubkey, f'pubkey mismatch: {file}'
-
-    def check_contains_pubkey(file):
-        assert contents(file).find(pubkey) != -1, f'pubkey missing: {file}'
-
     def check_avb_pubkey(file):
         info, _ = AvbInfo(args, file)
         assert info is not None, f'no avbinfo: {file}'
         assert info['Public key (sha1)'] == pubkey_digest, f'pubkey mismatch: {file}'
 
     for f in files.values():
-        if f == files['bootloader.pubkey']:
-            Async(check_equals_pubkey, f)
-        elif f == files['bootloader']:
-            Async(check_contains_pubkey, f)
-        elif f == files['super.img']:
+        if f in (files['initrd_normal.img'], files['initrd_debuggable.img']):
+            # TODO(b/245277660): Verify that ramdisks contain the correct vbmeta digest
+            continue
+        if f == files['super.img']:
             Async(check_avb_pubkey, system_a_img)
             Async(check_avb_pubkey, vendor_a_img)
         else:
@@ -467,7 +491,7 @@
             SignVirtApex(args)
         # ensure all tasks are completed without exceptions
         AwaitAll(tasks)
-    except: # pylint: disable=bare-except
+    except:  # pylint: disable=bare-except
         traceback.print_exc()
         sys.exit(1)
 
diff --git a/tests/hostside/Android.bp b/tests/hostside/Android.bp
index 7679c57..6e0cf5a 100644
--- a/tests/hostside/Android.bp
+++ b/tests/hostside/Android.bp
@@ -29,6 +29,7 @@
         // For re-sign test
         "avbtool",
         "img2simg",
+        "initrd_bootconfig",
         "lpmake",
         "lpunpack",
         "mk_payload",
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index 2ee33e6..4186ebb 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -55,7 +55,6 @@
 import org.json.JSONObject;
 import org.junit.After;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestName;
@@ -170,7 +169,11 @@
                 .isSuccess();
     }
 
-    private void resignVirtApex(File virtApexDir, File signingKey, Map<String, File> keyOverrides) {
+    private void resignVirtApex(
+            File virtApexDir,
+            File signingKey,
+            Map<String, File> keyOverrides,
+            boolean updateBootconfigs) {
         File signVirtApex = findTestFile("sign_virt_apex");
 
         RunUtil runUtil = new RunUtil();
@@ -181,6 +184,9 @@
 
         List<String> command = new ArrayList<>();
         command.add(signVirtApex.getAbsolutePath());
+        if (!updateBootconfigs) {
+            command.add("--do_not_update_bootconfigs");
+        }
         keyOverrides.forEach(
                 (filename, keyFile) ->
                         command.add("--key_override " + filename + "=" + keyFile.getPath()));
@@ -268,7 +274,11 @@
     }
 
     private VmInfo runMicrodroidWithResignedImages(
-            File key, Map<String, File> keyOverrides, boolean isProtected) throws Exception {
+            File key,
+            Map<String, File> keyOverrides,
+            boolean isProtected,
+            boolean updateBootconfigs)
+            throws Exception {
         CommandRunner android = new CommandRunner(getDevice());
 
         File virtApexDir = FileUtil.createTempDir("virt_apex");
@@ -281,7 +291,7 @@
         assertWithMessage("Failed to pull " + VIRT_APEX + "etc")
                 .that(getDevice().pullDir(VIRT_APEX + "etc", virtApexEtcDir)).isTrue();
 
-        resignVirtApex(virtApexDir, key, keyOverrides);
+        resignVirtApex(virtApexDir, key, keyOverrides, updateBootconfigs);
 
         // Push back re-signed virt APEX contents and updated microdroid.json
         getDevice().pushDir(virtApexDir, TEST_ROOT);
@@ -450,7 +460,8 @@
 
         // Act
         VmInfo vmInfo =
-                runMicrodroidWithResignedImages(key, /*keyOverrides=*/ Map.of(), protectedVm);
+                runMicrodroidWithResignedImages(
+                        key, /*keyOverrides=*/ Map.of(), protectedVm, /*updateBootconfigs=*/ true);
 
         // Assert
         vmInfo.mProcess.waitFor(5L, TimeUnit.SECONDS);
@@ -461,16 +472,15 @@
         vmInfo.mProcess.destroy();
     }
 
-    // TODO(b/245277660): Resigning the system/vendor image changes the vbmeta hash.
-    // So, unless vbmeta related bootconfigs are updated the following test will fail
     @Test
-    @Ignore("b/245277660")
     @CddTest(requirements = {"9.17/C-2-2", "9.17/C-2-6"})
     public void testBootSucceedsWhenNonProtectedVmStartsWithImagesSignedWithDifferentKey()
             throws Exception {
         File key = findTestFile("test.com.android.virt.pem");
         Map<String, File> keyOverrides = Map.of();
-        VmInfo vmInfo = runMicrodroidWithResignedImages(key, keyOverrides, /*isProtected=*/ false);
+        VmInfo vmInfo =
+                runMicrodroidWithResignedImages(
+                        key, keyOverrides, /*isProtected=*/ false, /*updateBootconfigs=*/ true);
         // Device online means that boot must have succeeded.
         adbConnectToMicrodroid(getDevice(), vmInfo.mCid);
         vmInfo.mProcess.destroy();
@@ -481,10 +491,10 @@
     public void testBootFailsWhenVbMetaDigestDoesNotMatchBootconfig() 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);
         // To be able to stop it, it should be a daemon.
-        VmInfo vmInfo = runMicrodroidWithResignedImages(key, keyOverrides, /*isProtected=*/ false);
+        VmInfo vmInfo =
+                runMicrodroidWithResignedImages(
+                        key, Map.of(), /*isProtected=*/ false, /*updateBootconfigs=*/ false);
         // Wait so that init can print errors to console (time in cuttlefish >> in real device)
         assertThatEventually(
                 100000,