Merge "Adds support for optionally generating vbmeta.img in merge_builds."

tools/releasetools/add_img_to_target_files.py

@@ -391,28 +391,6 @@ def AddUserdata(output_zip):
   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 @@ def AddVBMeta(output_zip, partitions, name, needed_partitions):
     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
 

tools/releasetools/common.py

@@ -625,6 +625,33 @@ def AppendAVBSigningArgs(cmd, partition):
     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 @@ def GetAvbChainedPartitionArg(partition, info_dict, key=None):
   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.

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)
 

tools/releasetools/test_add_img_to_target_files.py

@@ -21,7 +21,7 @@ import zipfile
 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 @@ class AddImagesToTargetFilesTest(test_utils.ReleaseToolsTestCase):
     # 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),

tools/releasetools/test_common.py

@@ -1137,6 +1137,30 @@ class CommonUtilsTest(test_utils.ReleaseToolsTestCase):
     }
     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.