Adds support for optionally generating vbmeta.img in merge_builds.

Bug: 137853921
Bug: 138671115
Test: python -m unittest test_common
Test: python -m unittest test_add_img_to_target_files
Test: Ran 'merge_builds --build_vbmeta' for two devices, one with the
vbmeta struct on system.img and another with vbmeta_system.img. Flashed
the regenerated vbmeta.img files on devices, devices boot.

Change-Id: I8d7585c7af468be3d242d8aceeed6d27e6fc6d96

tools/releasetools/merge_builds.py

@@ -24,9 +24,8 @@ build, the framework partial build should always be built with DAP enabled. The
 vendor partial build determines whether the merged result supports DAP.
 
 This script does not require builds to be built with 'make dist'.
-This script assumes that images other than super_empty.img do not require
-regeneration, including vbmeta images.
-TODO(b/137853921): Add support for regenerating vbmeta images.
+This script regenerates super_empty.img and vbmeta.img if necessary. Other
+images are assumed to not require regeneration.
 
 Usage: merge_builds.py [args]
 
@@ -39,6 +38,15 @@ Usage: merge_builds.py [args]
 
   --product_out_vendor product_out_vendor_path
       Path to out/target/product/<vendor build>.
+
+  --build_vbmeta
+      If provided, vbmeta.img will be regenerated in out/target/product/<vendor
+      build>.
+
+  --framework_misc_info_keys
+      The optional path to a newline-separated config file containing keys to
+      obtain from the framework instance of misc_info.txt, used for creating
+      vbmeta.img. The remaining keys come from the vendor instance.
 """
 from __future__ import print_function
 
@@ -55,6 +63,8 @@ OPTIONS = common.OPTIONS
 OPTIONS.framework_images = ("system",)
 OPTIONS.product_out_framework = None
 OPTIONS.product_out_vendor = None
+OPTIONS.build_vbmeta = False
+OPTIONS.framework_misc_info_keys = None
 
 
 def CreateImageSymlinks():
@@ -82,6 +92,7 @@ def BuildSuperEmpty():
   # super_empty.img from the framework build.
   if (framework_dict.get("use_dynamic_partitions") == "true") and (
       vendor_dict.get("use_dynamic_partitions") == "true"):
+    logger.info("Building super_empty.img.")
     merged_dict = dict(vendor_dict)
     merged_dict.update(
         common.MergeDynamicPartitionInfoDicts(
@@ -96,10 +107,52 @@ def BuildSuperEmpty():
     build_super_image.BuildSuperImage(merged_dict, output_super_empty_path)
 
 
+def BuildVBMeta():
+  logger.info("Building vbmeta.img.")
+
+  framework_dict = common.LoadDictionaryFromFile(
+      os.path.join(OPTIONS.product_out_framework, "misc_info.txt"))
+  vendor_dict = common.LoadDictionaryFromFile(
+      os.path.join(OPTIONS.product_out_vendor, "misc_info.txt"))
+  merged_dict = dict(vendor_dict)
+  if OPTIONS.framework_misc_info_keys:
+    for key in common.LoadListFromFile(OPTIONS.framework_misc_info_keys):
+      merged_dict[key] = framework_dict[key]
+
+  # Build vbmeta.img using partitions in product_out_vendor.
+  partitions = {}
+  for partition in common.AVB_PARTITIONS:
+    partition_path = os.path.join(OPTIONS.product_out_vendor,
+                                  "%s.img" % partition)
+    if os.path.exists(partition_path):
+      partitions[partition] = partition_path
+
+  # vbmeta_partitions includes the partitions that should be included into
+  # top-level vbmeta.img, which are the ones that are not included in any
+  # chained VBMeta image plus the chained VBMeta images themselves.
+  vbmeta_partitions = common.AVB_PARTITIONS[:]
+  for partition in common.AVB_VBMETA_PARTITIONS:
+    chained_partitions = merged_dict.get("avb_%s" % partition, "").strip()
+    if chained_partitions:
+      partitions[partition] = os.path.join(OPTIONS.product_out_vendor,
+                                           "%s.img" % partition)
+      vbmeta_partitions = [
+          item for item in vbmeta_partitions
+          if item not in chained_partitions.split()
+      ]
+      vbmeta_partitions.append(partition)
+
+  output_vbmeta_path = os.path.join(OPTIONS.product_out_vendor, "vbmeta.img")
+  OPTIONS.info_dict = merged_dict
+  common.BuildVBMeta(output_vbmeta_path, partitions, "vbmeta",
+                     vbmeta_partitions)
+
+
 def MergeBuilds():
   CreateImageSymlinks()
   BuildSuperEmpty()
-  # TODO(b/137853921): Add support for regenerating vbmeta images.
+  if OPTIONS.build_vbmeta:
+    BuildVBMeta()
 
 
 def main():
@@ -112,6 +165,10 @@ def main():
       OPTIONS.product_out_framework = a
     elif o == "--product_out_vendor":
       OPTIONS.product_out_vendor = a
+    elif o == "--build_vbmeta":
+      OPTIONS.build_vbmeta = True
+    elif o == "--framework_misc_info_keys":
+      OPTIONS.framework_misc_info_keys = a
     else:
       return False
     return True
@@ -123,6 +180,8 @@ def main():
           "framework_images=",
           "product_out_framework=",
           "product_out_vendor=",
+          "build_vbmeta",
+          "framework_misc_info_keys=",
       ],
       extra_option_handler=option_handler)