sign_virt_apex: replace bootloader pubkey

VBmeta's key should match with pubkey embedded in bootloader. The
bootloader pubkey is added to the APEX so that sign_virt_apex can
replace it with a new pubkey.

Bug: 193504286
Test: sign_apex --sign_tool sign_virt_apex ...
      & install & run a VM
Change-Id: Ic8e5ec9cb45434691c8dce0ca09243e181dc59cc
diff --git a/apex/Android.bp b/apex/Android.bp
index 20a863f..983253e 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -59,6 +59,7 @@
         "microdroid.json",
         "microdroid_uboot_env",
         "microdroid_bootloader",
+        "microdroid_bootloader.avbpubkey",
         "microdroid_bootconfig_normal",
         "microdroid_bootconfig_app_debuggable",
         "microdroid_bootconfig_full_debuggable",
diff --git a/apex/sign_virt_apex.py b/apex/sign_virt_apex.py
index 6ecd07e..77f54c4 100644
--- a/apex/sign_virt_apex.py
+++ b/apex/sign_virt_apex.py
@@ -214,11 +214,34 @@
         RunCommand(args, cmd)
 
 
+def ReplaceBootloaderPubkey(args, key, bootloader, bootloader_pubkey):
+    # read old pubkey before replacement
+    with open(bootloader_pubkey, 'rb') as f:
+        old_pubkey = f.read()
+
+    # replace bootloader pubkey
+    RunCommand(args, ['avbtool', 'extract_public_key', '--key', key, '--output', 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)
+
+
 def SignVirtApex(args):
     key = args.key
     input_dir = args.input_dir
 
     # target files in the Virt APEX
+    bootloader_pubkey = os.path.join(input_dir, 'etc', 'microdroid_bootloader.avbpubkey')
     bootloader = os.path.join(input_dir, 'etc', 'microdroid_bootloader')
     boot_img = os.path.join(input_dir, 'etc', 'fs', 'microdroid_boot-5.10.img')
     vendor_boot_img = os.path.join(
@@ -226,6 +249,10 @@
     super_img = os.path.join(input_dir, 'etc', 'fs', 'microdroid_super.img')
     vbmeta_img = os.path.join(input_dir, 'etc', 'fs', 'microdroid_vbmeta.img')
 
+    # Key(pubkey) for bootloader should match with the one used to make VBmeta below
+    # while it's okay to use different keys for other image files.
+    ReplaceBootloaderPubkey(args, key, bootloader, bootloader_pubkey)
+
     # re-sign bootloader, boot.img, vendor_boot.img
     AddHashFooter(args, key, bootloader)
     AddHashFooter(args, key, boot_img)
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index 4d7c218..274b7ed 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -379,6 +379,7 @@
 // MAX_VBMETA_SIZE=64KB, MAX_FOOTER_SIZE=4KB
 avb_hash_footer_kb = "68"
 
+// TODO(b/193504286) remove this when prebuilt bootloader exposes pubkey as well.
 genrule {
     name: "microdroid_bootloader_gen",
     tools: ["avbtool"],
@@ -405,6 +406,22 @@
 }
 
 prebuilt_etc {
+    name: "microdroid_bootloader.avbpubkey",
+    src: ":microdroid_bootloader_pubkey_gen",
+}
+
+genrule {
+    name: "microdroid_bootloader_pubkey_gen",
+    tools: ["avbtool"],
+    srcs: [
+        ":microdroid_crosvm_bootloader",
+        ":avb_testkey_rsa4096",
+    ],
+    out: ["bootloader-pubkey"],
+    cmd: "$(location avbtool) extract_public_key --key $(location :avb_testkey_rsa4096) --output $(out)",
+}
+
+prebuilt_etc {
     name: "microdroid_uboot_env",
     src: ":microdroid_uboot_env_gen",
     arch: {