Support building chained VBMeta images.

Bug: 112007947
Test: `m dist` with aosp_taimen-userdebug.
Test: Set up a target that uses chained VBMeta images of
      `vbmeta_mainline` and `vbmeta_vendor`. `m dist` and check the
      build log, as well as outputs from
      `avbtool info_image --image vbmeta.img`,
      `avbtool info_image --image vbmeta_mainline.img`,
      `avbtool info_image --image vbmeta_vendor.img`.
Change-Id: Ib1d4e97f583b65245703eae15d211adcd9e83741
diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index 9e8b698..e9419fe 100755
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -393,32 +393,46 @@
     cmd.extend(["--include_descriptors_from_image", image])
 
 
-def AddVBMeta(output_zip, partitions):
-  """Creates a VBMeta image and store it in output_zip.
+def AddVBMeta(output_zip, partitions, name, needed_partitions):
+  """Creates a VBMeta image and stores it in output_zip.
+
+  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:
     output_zip: The output zip file, which needs to be already open.
     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_mainline'.
+    needed_partitions: Partitions whose descriptors should be included into the
+        generated VBMeta image.
+
+  Raises:
+    AssertionError: On invalid input args.
   """
-  img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "vbmeta.img")
+  assert needed_partitions, "Needed partitions must be specified"
+
+  img = OutputFile(
+      output_zip, OPTIONS.input_tmp, "IMAGES", "{}.img".format(name))
   if os.path.exists(img.input_name):
-    print("vbmeta.img already exists; not rebuilding...")
+    print("{}.img already exists; not rebuilding...".format(name))
     return img.input_name
 
   avbtool = os.getenv('AVBTOOL') or OPTIONS.info_dict["avb_avbtool"]
   cmd = [avbtool, "make_vbmeta_image", "--output", img.name]
-  common.AppendAVBSigningArgs(cmd, "vbmeta")
+  common.AppendAVBSigningArgs(cmd, name)
 
   for partition, path in partitions.items():
+    if partition not in needed_partitions:
+      continue
     assert partition in common.AVB_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_vbmeta_args")
+  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]):
@@ -439,7 +453,7 @@
             split_args[index + 1] = alt_path
             found = True
             break
-        assert found, 'failed to find %s' % (image_path,)
+        assert found, 'Failed to find {}'.format(image_path)
     cmd.extend(split_args)
 
   p = common.Run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
@@ -693,8 +707,8 @@
                 os.path.exists(os.path.join(OPTIONS.input_tmp, "IMAGES",
                                             "vendor.img")))
   has_odm = (os.path.isdir(os.path.join(OPTIONS.input_tmp, "ODM")) or
-                 os.path.exists(os.path.join(OPTIONS.input_tmp, "IMAGES",
-                                             "odm.img")))
+             os.path.exists(os.path.join(OPTIONS.input_tmp, "IMAGES",
+                                         "odm.img")))
   has_product = (os.path.isdir(os.path.join(OPTIONS.input_tmp, "PRODUCT")) or
                  os.path.exists(os.path.join(OPTIONS.input_tmp, "IMAGES",
                                              "product.img")))
@@ -806,8 +820,33 @@
     partitions['dtbo'] = AddDtbo(output_zip)
 
   if OPTIONS.info_dict.get("avb_enable") == "true":
+    # 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[:]
+
+    vbmeta_mainline = OPTIONS.info_dict.get("avb_vbmeta_mainline", "").strip()
+    if vbmeta_mainline:
+      banner("vbmeta_mainline")
+      AddVBMeta(
+          output_zip, partitions, "vbmeta_mainline", vbmeta_mainline.split())
+      vbmeta_partitions = [
+          item for item in vbmeta_partitions
+          if item not in vbmeta_mainline.split()]
+      vbmeta_partitions.append("vbmeta_mainline")
+
+    vbmeta_vendor = OPTIONS.info_dict.get("avb_vbmeta_vendor", "").strip()
+    if vbmeta_vendor:
+      banner("vbmeta_vendor")
+      AddVBMeta(
+          output_zip, partitions, "vbmeta_vendor", vbmeta_vendor.split())
+      vbmeta_partitions = [
+          item for item in vbmeta_partitions
+          if item not in vbmeta_vendor.split()]
+      vbmeta_partitions.append("vbmeta_vendor")
+
     banner("vbmeta")
-    AddVBMeta(output_zip, partitions)
+    AddVBMeta(output_zip, partitions, "vbmeta", vbmeta_partitions)
 
   if OPTIONS.info_dict.get("super_size"):
     banner("super_empty")
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index 8d51df6..ee2c6f4 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -408,7 +408,7 @@
     cmd.extend(["--key", key_path, "--algorithm", algorithm])
   avb_salt = OPTIONS.info_dict.get("avb_salt")
   # make_vbmeta_image doesn't like "--salt" (and it's not needed).
-  if avb_salt and partition != "vbmeta":
+  if avb_salt and not partition.startswith("vbmeta"):
     cmd.extend(["--salt", avb_salt])