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
diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index bdb34b8..23ae29f 100755
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -391,28 +391,6 @@
   img.Write()
 
 
-def AppendVBMetaArgsForPartition(cmd, partition, image):
-  """Appends the VBMeta arguments for partition.
-
-  It sets up the VBMeta argument by including the partition descriptor from the
-  given 'image', or by configuring the partition as a chained partition.
-
-  Args:
-    cmd: A list of command args that will be used to generate the vbmeta image.
-        The argument for the partition will be appended to the list.
-    partition: The name of the partition (e.g. "system").
-    image: The path to the partition image.
-  """
-  # Check if chain partition is used.
-  key_path = OPTIONS.info_dict.get("avb_" + partition + "_key_path")
-  if key_path:
-    chained_partition_arg = common.GetAvbChainedPartitionArg(
-        partition, OPTIONS.info_dict)
-    cmd.extend(["--chain_partition", chained_partition_arg])
-  else:
-    cmd.extend(["--include_descriptors_from_image", image])
-
-
 def AddVBMeta(output_zip, partitions, name, needed_partitions):
   """Creates a VBMeta image and stores it in output_zip.
 
@@ -442,45 +420,7 @@
     logger.info("%s.img already exists; not rebuilding...", name)
     return img.name
 
-  avbtool = OPTIONS.info_dict["avb_avbtool"]
-  cmd = [avbtool, "make_vbmeta_image", "--output", img.name]
-  common.AppendAVBSigningArgs(cmd, name)
-
-  for partition, path in partitions.items():
-    if partition not in needed_partitions:
-      continue
-    assert (partition in common.AVB_PARTITIONS or
-            partition in common.AVB_VBMETA_PARTITIONS), \
-        'Unknown partition: {}'.format(partition)
-    assert os.path.exists(path), \
-        'Failed to find {} for {}'.format(path, partition)
-    AppendVBMetaArgsForPartition(cmd, partition, path)
-
-  args = OPTIONS.info_dict.get("avb_{}_args".format(name))
-  if args and args.strip():
-    split_args = shlex.split(args)
-    for index, arg in enumerate(split_args[:-1]):
-      # Sanity check that the image file exists. Some images might be defined
-      # as a path relative to source tree, which may not be available at the
-      # same location when running this script (we have the input target_files
-      # zip only). For such cases, we additionally scan other locations (e.g.
-      # IMAGES/, RADIO/, etc) before bailing out.
-      if arg == '--include_descriptors_from_image':
-        image_path = split_args[index + 1]
-        if os.path.exists(image_path):
-          continue
-        found = False
-        for dir_name in ['IMAGES', 'RADIO', 'PREBUILT_IMAGES']:
-          alt_path = os.path.join(
-              OPTIONS.input_tmp, dir_name, os.path.basename(image_path))
-          if os.path.exists(alt_path):
-            split_args[index + 1] = alt_path
-            found = True
-            break
-        assert found, 'Failed to find {}'.format(image_path)
-    cmd.extend(split_args)
-
-  common.RunAndCheckOutput(cmd)
+  common.BuildVBMeta(img.name, partitions, name, needed_partitions)
   img.Write()
   return img.name
 
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index 1175688..2401e46 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -625,6 +625,33 @@
     cmd.extend(["--salt", avb_salt])
 
 
+def GetAvbPartitionArg(partition, image, info_dict = None):
+  """Returns the VBMeta arguments for partition.
+
+  It sets up the VBMeta argument by including the partition descriptor from the
+  given 'image', or by configuring the partition as a chained partition.
+
+  Args:
+    partition: The name of the partition (e.g. "system").
+    image: The path to the partition image.
+    info_dict: A dict returned by common.LoadInfoDict(). Will use
+        OPTIONS.info_dict if None has been given.
+
+  Returns:
+    A list of VBMeta arguments.
+  """
+  if info_dict is None:
+    info_dict = OPTIONS.info_dict
+
+  # Check if chain partition is used.
+  key_path = info_dict.get("avb_" + partition + "_key_path")
+  if key_path:
+    chained_partition_arg = GetAvbChainedPartitionArg(partition, info_dict)
+    return ["--chain_partition", chained_partition_arg]
+  else:
+    return ["--include_descriptors_from_image", image]
+
+
 def GetAvbChainedPartitionArg(partition, info_dict, key=None):
   """Constructs and returns the arg to build or verify a chained partition.
 
@@ -647,6 +674,65 @@
   return "{}:{}:{}".format(partition, rollback_index_location, pubkey_path)
 
 
+def BuildVBMeta(image_path, partitions, name, needed_partitions):
+  """Creates a VBMeta image.
+
+  It generates the requested VBMeta image. The requested image could be for
+  top-level or chained VBMeta image, which is determined based on the name.
+
+  Args:
+    image_path: The output path for the new VBMeta image.
+    partitions: A dict that's keyed by partition names with image paths as
+        values. Only valid partition names are accepted, as listed in
+        common.AVB_PARTITIONS.
+    name: Name of the VBMeta partition, e.g. 'vbmeta', 'vbmeta_system'.
+    needed_partitions: Partitions whose descriptors should be included into the
+        generated VBMeta image.
+
+  Raises:
+    AssertionError: On invalid input args.
+  """
+  avbtool = OPTIONS.info_dict["avb_avbtool"]
+  cmd = [avbtool, "make_vbmeta_image", "--output", image_path]
+  AppendAVBSigningArgs(cmd, name)
+
+  for partition, path in partitions.items():
+    if partition not in needed_partitions:
+      continue
+    assert (partition in AVB_PARTITIONS or
+            partition in AVB_VBMETA_PARTITIONS), \
+        'Unknown partition: {}'.format(partition)
+    assert os.path.exists(path), \
+        'Failed to find {} for {}'.format(path, partition)
+    cmd.extend(GetAvbPartitionArg(partition, path))
+
+  args = OPTIONS.info_dict.get("avb_{}_args".format(name))
+  if args and args.strip():
+    split_args = shlex.split(args)
+    for index, arg in enumerate(split_args[:-1]):
+      # Sanity check that the image file exists. Some images might be defined
+      # as a path relative to source tree, which may not be available at the
+      # same location when running this script (we have the input target_files
+      # zip only). For such cases, we additionally scan other locations (e.g.
+      # IMAGES/, RADIO/, etc) before bailing out.
+      if arg == '--include_descriptors_from_image':
+        image_path = split_args[index + 1]
+        if os.path.exists(image_path):
+          continue
+        found = False
+        for dir_name in ['IMAGES', 'RADIO', 'PREBUILT_IMAGES']:
+          alt_path = os.path.join(
+              OPTIONS.input_tmp, dir_name, os.path.basename(image_path))
+          if os.path.exists(alt_path):
+            split_args[index + 1] = alt_path
+            found = True
+            break
+        assert found, 'Failed to find {}'.format(image_path)
+    cmd.extend(split_args)
+
+  RunAndCheckOutput(cmd)
+
+
 def _BuildBootableImage(sourcedir, fs_config_file, info_dict=None,
                         has_ramdisk=False, two_step_image=False):
   """Build a bootable image from the specified sourcedir.
diff --git a/tools/releasetools/merge_builds.py b/tools/releasetools/merge_builds.py
index 7724d6f..ca348cf 100644
--- a/tools/releasetools/merge_builds.py
+++ b/tools/releasetools/merge_builds.py
@@ -24,9 +24,8 @@
 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 @@
 
   --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.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 @@
   # 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 @@
     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 @@
       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 @@
           "framework_images=",
           "product_out_framework=",
           "product_out_vendor=",
+          "build_vbmeta",
+          "framework_misc_info_keys=",
       ],
       extra_option_handler=option_handler)
 
diff --git a/tools/releasetools/test_add_img_to_target_files.py b/tools/releasetools/test_add_img_to_target_files.py
index 08e0190..3d0766f 100644
--- a/tools/releasetools/test_add_img_to_target_files.py
+++ b/tools/releasetools/test_add_img_to_target_files.py
@@ -21,7 +21,7 @@
 import common
 import test_utils
 from add_img_to_target_files import (
-    AddCareMapForAbOta, AddPackRadioImages, AppendVBMetaArgsForPartition,
+    AddCareMapForAbOta, AddPackRadioImages,
     CheckAbOtaImages, GetCareMap)
 from rangelib import RangeSet
 
@@ -379,32 +379,6 @@
     # The existing entry should be scheduled to be replaced.
     self.assertIn('META/care_map.pb', OPTIONS.replace_updated_files_list)
 
-  def test_AppendVBMetaArgsForPartition(self):
-    OPTIONS.info_dict = {}
-    cmd = []
-    AppendVBMetaArgsForPartition(cmd, 'system', '/path/to/system.img')
-    self.assertEqual(
-        ['--include_descriptors_from_image', '/path/to/system.img'], cmd)
-
-  @test_utils.SkipIfExternalToolsUnavailable()
-  def test_AppendVBMetaArgsForPartition_vendorAsChainedPartition(self):
-    testdata_dir = test_utils.get_testdata_dir()
-    pubkey = os.path.join(testdata_dir, 'testkey.pubkey.pem')
-    OPTIONS.info_dict = {
-        'avb_avbtool': 'avbtool',
-        'avb_vendor_key_path': pubkey,
-        'avb_vendor_rollback_index_location': 5,
-    }
-    cmd = []
-    AppendVBMetaArgsForPartition(cmd, 'vendor', '/path/to/vendor.img')
-    self.assertEqual(2, len(cmd))
-    self.assertEqual('--chain_partition', cmd[0])
-    chained_partition_args = cmd[1].split(':')
-    self.assertEqual(3, len(chained_partition_args))
-    self.assertEqual('vendor', chained_partition_args[0])
-    self.assertEqual('5', chained_partition_args[1])
-    self.assertTrue(os.path.exists(chained_partition_args[2]))
-
   def test_GetCareMap(self):
     sparse_image = test_utils.construct_sparse_image([
         (0xCAC1, 6),
diff --git a/tools/releasetools/test_common.py b/tools/releasetools/test_common.py
index bcfb1c1..ceb023f 100644
--- a/tools/releasetools/test_common.py
+++ b/tools/releasetools/test_common.py
@@ -1137,6 +1137,30 @@
     }
     self.assertEqual(merged_dict, expected_merged_dict)
 
+  def test_GetAvbPartitionArg(self):
+    info_dict = {}
+    cmd = common.GetAvbPartitionArg('system', '/path/to/system.img', info_dict)
+    self.assertEqual(
+        ['--include_descriptors_from_image', '/path/to/system.img'], cmd)
+
+  @test_utils.SkipIfExternalToolsUnavailable()
+  def test_AppendVBMetaArgsForPartition_vendorAsChainedPartition(self):
+    testdata_dir = test_utils.get_testdata_dir()
+    pubkey = os.path.join(testdata_dir, 'testkey.pubkey.pem')
+    info_dict = {
+        'avb_avbtool': 'avbtool',
+        'avb_vendor_key_path': pubkey,
+        'avb_vendor_rollback_index_location': 5,
+    }
+    cmd = common.GetAvbPartitionArg('vendor', '/path/to/vendor.img', info_dict)
+    self.assertEqual(2, len(cmd))
+    self.assertEqual('--chain_partition', cmd[0])
+    chained_partition_args = cmd[1].split(':')
+    self.assertEqual(3, len(chained_partition_args))
+    self.assertEqual('vendor', chained_partition_args[0])
+    self.assertEqual('5', chained_partition_args[1])
+    self.assertTrue(os.path.exists(chained_partition_args[2]))
+
 
 class InstallRecoveryScriptFormatTest(test_utils.ReleaseToolsTestCase):
   """Checks the format of install-recovery.sh.