diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index 1d8090a..aab082b 100755
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -192,14 +192,14 @@
   shutil.copy(dtbo_prebuilt_path, img.name)
 
   # AVB-sign the image as needed.
-  if OPTIONS.info_dict.get("board_avb_enable") == "true":
+  if OPTIONS.info_dict.get("avb_enable") == "true":
     avbtool = os.getenv('AVBTOOL') or OPTIONS.info_dict["avb_avbtool"]
     part_size = OPTIONS.info_dict["dtbo_size"]
     # The AVB hash footer will be replaced if already present.
     cmd = [avbtool, "add_hash_footer", "--image", img.name,
            "--partition_size", str(part_size), "--partition_name", "dtbo"]
-    cmd.extend(shlex.split(OPTIONS.info_dict["avb_signing_args"]))
-    args = OPTIONS.info_dict.get("board_avb_dtbo_add_hash_footer_args")
+    common.AppendAVBSigningArgs(cmd, "dtbo")
+    args = OPTIONS.info_dict.get("avb_dtbo_add_hash_footer_args")
     if args and args.strip():
       cmd.extend(shlex.split(args))
     p = common.Run(cmd, stdout=subprocess.PIPE)
@@ -271,7 +271,7 @@
   # by the avb tool.
   is_verity_partition = "verity_block_device" in image_props
   verity_supported = (image_props.get("verity") == "true" or
-                      image_props.get("board_avb_enable") == "true")
+                      image_props.get("avb_enable") == "true")
   is_avb_enable = image_props.get("avb_hashtree_enable") == "true"
   if verity_supported and (is_verity_partition or is_avb_enable):
     adjusted_blocks_value = image_props.get("partition_size")
@@ -334,25 +334,51 @@
   img.Write()
 
 
+def AppendVBMetaArgsForPartition(cmd, partition, img_path, public_key_dir):
+  if not img_path:
+    return
+
+  # Check if chain partition is used.
+  key_path = OPTIONS.info_dict.get("avb_" + partition + "_key_path")
+  if key_path:
+    # extract public key in AVB format to be included in vbmeta.img
+    avbtool = os.getenv('AVBTOOL') or OPTIONS.info_dict["avb_avbtool"]
+    public_key_path = os.path.join(public_key_dir, "%s.avbpubkey" % partition)
+    p = common.Run([avbtool, "extract_public_key", "--key", key_path,
+                    "--output", public_key_path],
+                   stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    p.communicate()
+    assert p.returncode == 0, \
+        "avbtool extract_public_key fail for partition: %r" % partition
+
+    rollback_index_location = OPTIONS.info_dict[
+        "avb_" + partition + "_rollback_index_location"]
+    cmd.extend(["--chain_partition", "%s:%s:%s" % (
+        partition, rollback_index_location, public_key_path)])
+  else:
+    cmd.extend(["--include_descriptors_from_image", img_path])
+
+
 def AddVBMeta(output_zip, boot_img_path, system_img_path, vendor_img_path,
               dtbo_img_path, prefix="IMAGES/"):
   """Create a VBMeta image and store it in output_zip."""
   img = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "vbmeta.img")
   avbtool = os.getenv('AVBTOOL') or OPTIONS.info_dict["avb_avbtool"]
-  cmd = [avbtool, "make_vbmeta_image",
-         "--output", img.name,
-         "--include_descriptors_from_image", boot_img_path,
-         "--include_descriptors_from_image", system_img_path]
-  if vendor_img_path is not None:
-    cmd.extend(["--include_descriptors_from_image", vendor_img_path])
-  if dtbo_img_path is not None:
-    cmd.extend(["--include_descriptors_from_image", dtbo_img_path])
-  if OPTIONS.info_dict.get("system_root_image") == "true":
-    cmd.extend(["--setup_rootfs_from_kernel", system_img_path])
-  cmd.extend(shlex.split(OPTIONS.info_dict["avb_signing_args"]))
-  args = OPTIONS.info_dict.get("board_avb_make_vbmeta_image_args")
+  cmd = [avbtool, "make_vbmeta_image", "--output", img.name]
+  common.AppendAVBSigningArgs(cmd, "vbmeta")
+
+  public_key_dir = tempfile.mkdtemp(prefix="avbpubkey-")
+  OPTIONS.tempfiles.append(public_key_dir)
+
+  AppendVBMetaArgsForPartition(cmd, "boot", boot_img_path, public_key_dir)
+  AppendVBMetaArgsForPartition(cmd, "system", system_img_path, public_key_dir)
+  AppendVBMetaArgsForPartition(cmd, "vendor", vendor_img_path, public_key_dir)
+  AppendVBMetaArgsForPartition(cmd, "dtbo", dtbo_img_path, public_key_dir)
+
+  args = OPTIONS.info_dict.get("avb_vbmeta_args")
   if args and args.strip():
     cmd.extend(shlex.split(args))
+
   p = common.Run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
   p.communicate()
   assert p.returncode == 0, "avbtool make_vbmeta_image failed"
@@ -553,7 +579,7 @@
     banner("dtbo")
     dtbo_img_path = AddDtbo(output_zip)
 
-  if OPTIONS.info_dict.get("board_avb_enable") == "true":
+  if OPTIONS.info_dict.get("avb_enable") == "true":
     banner("vbmeta")
     boot_contents = boot_image.WriteToTemp()
     AddVBMeta(output_zip, boot_contents.name, system_img_path,
diff --git a/tools/releasetools/build_image.py b/tools/releasetools/build_image.py
index de75a6b..6de9763 100755
--- a/tools/releasetools/build_image.py
+++ b/tools/releasetools/build_image.py
@@ -109,18 +109,19 @@
   Returns:
     The maximum image size or 0 if an error occurred.
   """
-  cmdline = "%s add_%s_footer " % (avbtool, footer_type)
-  cmdline += "--partition_size %d " % partition_size
-  cmdline += "--calc_max_image_size "
-  cmdline += additional_args
-  (output, exit_code) = RunCommand(shlex.split(cmdline))
+  cmd =[avbtool, "add_%s_footer" % footer_type,
+        "--partition_size", partition_size, "--calc_max_image_size"]
+  cmd.extend(shlex.split(additional_args))
+
+  (output, exit_code) = RunCommand(cmd)
   if exit_code != 0:
     return 0
   else:
     return int(output)
 
 def AVBAddFooter(image_path, avbtool, footer_type, partition_size,
-                 partition_name, signing_args, additional_args):
+                 partition_name, key_path, algorithm,
+                 additional_args):
   """Adds dm-verity hashtree and AVB metadata to an image.
 
   Args:
@@ -129,19 +130,24 @@
     footer_type: 'hash' or 'hashtree' for generating footer.
     partition_size: The size of the partition in question.
     partition_name: The name of the partition - will be embedded in metadata.
-    signing_args: Arguments for signing the image.
+    key_path: Path to key to use or None.
+    algorithm: Name of algorithm to use or None.
     additional_args: Additional arguments to pass to 'avbtool
       add_hashtree_image'.
   Returns:
     True if the operation succeeded.
   """
-  cmdline = "%s add_%s_footer " % (avbtool, footer_type)
-  cmdline += "--partition_size %d " % partition_size
-  cmdline += "--partition_name %s " % partition_name
-  cmdline += "--image %s " % image_path
-  cmdline += signing_args + " "
-  cmdline += additional_args
-  (_, exit_code) = RunCommand(shlex.split(cmdline))
+  cmd =[avbtool, "add_%s_footer" % footer_type,
+        "--partition_size", partition_size,
+        "--partition_name", partition_name,
+        "--image", image_path]
+
+  if key_path and algorithm:
+    cmd.extend(["--key", key_path, "--algorithm", algorithm])
+
+  cmd.extend(shlex.split(additional_args))
+
+  (_, exit_code) = RunCommand(cmd)
   return exit_code == 0
 
 def AdjustPartitionSizeForVerity(partition_size, fec_supported):
@@ -420,8 +426,8 @@
     avb_footer_type = 'hashtree'
 
   if avb_footer_type:
-    avbtool = prop_dict.get("avb_avbtool")
-    partition_size = int(prop_dict.get("partition_size"))
+    avbtool = prop_dict["avb_avbtool"]
+    partition_size = prop_dict["partition_size"]
     # avb_add_hash_footer_args or avb_add_hashtree_footer_args.
     additional_args = prop_dict["avb_add_" + avb_footer_type + "_footer_args"]
     max_image_size = AVBCalcMaxImageSize(avbtool, avb_footer_type, partition_size,
@@ -429,7 +435,7 @@
     if max_image_size == 0:
       return False
     prop_dict["partition_size"] = str(max_image_size)
-    prop_dict["original_partition_size"] = str(partition_size)
+    prop_dict["original_partition_size"] = partition_size
 
   if fs_type.startswith("ext"):
     build_command = [prop_dict["ext_mkuserimg"]]
@@ -572,14 +578,16 @@
 
   # Add AVB HASH or HASHTREE footer (metadata).
   if avb_footer_type:
-    avbtool = prop_dict.get("avb_avbtool")
-    original_partition_size = int(prop_dict.get("original_partition_size"))
+    avbtool = prop_dict["avb_avbtool"]
+    original_partition_size = prop_dict["original_partition_size"]
     partition_name = prop_dict["partition_name"]
-    signing_args = prop_dict["avb_signing_args"]
+    # key_path and algorithm are only available when chain partition is used.
+    key_path = prop_dict.get("avb_key_path")
+    algorithm = prop_dict.get("avb_algorithm")
     # avb_add_hash_footer_args or avb_add_hashtree_footer_args
     additional_args = prop_dict["avb_add_" + avb_footer_type + "_footer_args"]
     if not AVBAddFooter(out_file, avbtool, avb_footer_type, original_partition_size,
-                      partition_name, signing_args, additional_args):
+                        partition_name, key_path, algorithm, additional_args):
       return False
 
   if run_fsck and prop_dict.get("skip_fsck") != "true":
@@ -624,8 +632,7 @@
       "verity_key",
       "verity_signer_cmd",
       "verity_fec",
-      "board_avb_enable",
-      "avb_signing_args",
+      "avb_enable",
       "avb_avbtool"
       )
   for p in common_props:
@@ -633,6 +640,11 @@
 
   d["mount_point"] = mount_point
   if mount_point == "system":
+    copy_prop("avb_system_hashtree_enable", "avb_hashtree_enable")
+    copy_prop("avb_system_add_hashtree_footer_args",
+              "avb_add_hashtree_footer_args")
+    copy_prop("avb_system_key_path", "avb_key_path")
+    copy_prop("avb_system_algorithm", "avb_algorithm")
     copy_prop("fs_type", "fs_type")
     # Copy the generic system fs type first, override with specific one if
     # available.
@@ -650,13 +662,15 @@
     copy_prop("system_squashfs_block_size", "squashfs_block_size")
     copy_prop("system_squashfs_disable_4k_align", "squashfs_disable_4k_align")
     copy_prop("system_base_fs_file", "base_fs_file")
-    copy_prop("system_avb_hashtree_enable", "avb_hashtree_enable")
-    copy_prop("system_avb_add_hashtree_footer_args",
-              "avb_add_hashtree_footer_args")
     copy_prop("system_extfs_inode_count", "extfs_inode_count")
   elif mount_point == "system_other":
     # We inherit the selinux policies of /system since we contain some of its files.
     d["mount_point"] = "system"
+    copy_prop("avb_system_hashtree_enable", "avb_hashtree_enable")
+    copy_prop("avb_system_add_hashtree_footer_args",
+              "avb_add_hashtree_footer_args")
+    copy_prop("avb_system_key_path", "avb_key_path")
+    copy_prop("avb_system_algorithm", "avb_algorithm")
     copy_prop("fs_type", "fs_type")
     copy_prop("system_fs_type", "fs_type")
     copy_prop("system_size", "partition_size")
@@ -667,9 +681,6 @@
     copy_prop("system_squashfs_compressor_opt", "squashfs_compressor_opt")
     copy_prop("system_squashfs_block_size", "squashfs_block_size")
     copy_prop("system_base_fs_file", "base_fs_file")
-    copy_prop("system_avb_hashtree_enable", "avb_hashtree_enable")
-    copy_prop("system_avb_add_hashtree_footer_args",
-              "avb_add_hashtree_footer_args")
     copy_prop("system_extfs_inode_count", "extfs_inode_count")
   elif mount_point == "data":
     # Copy the generic fs type first, override with specific one if available.
@@ -682,6 +693,11 @@
     copy_prop("cache_fs_type", "fs_type")
     copy_prop("cache_size", "partition_size")
   elif mount_point == "vendor":
+    copy_prop("avb_vendor_hashtree_enable", "avb_hashtree_enable")
+    copy_prop("avb_vendor_add_hashtree_footer_args",
+              "avb_add_hashtree_footer_args")
+    copy_prop("avb_vendor_key_path", "avb_key_path")
+    copy_prop("avb_vendor_algorithm", "avb_algorithm")
     copy_prop("vendor_fs_type", "fs_type")
     copy_prop("vendor_size", "partition_size")
     copy_prop("vendor_journal_size", "journal_size")
@@ -692,9 +708,6 @@
     copy_prop("vendor_squashfs_block_size", "squashfs_block_size")
     copy_prop("vendor_squashfs_disable_4k_align", "squashfs_disable_4k_align")
     copy_prop("vendor_base_fs_file", "base_fs_file")
-    copy_prop("vendor_avb_hashtree_enable", "avb_hashtree_enable")
-    copy_prop("vendor_avb_add_hashtree_footer_args",
-              "avb_add_hashtree_footer_args")
     copy_prop("vendor_extfs_inode_count", "extfs_inode_count")
   elif mount_point == "oem":
     copy_prop("fs_type", "fs_type")
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index 652fadf..2ad628b 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -345,6 +345,15 @@
     print("%-25s = (%s) %s" % (k, type(v).__name__, v))
 
 
+def AppendAVBSigningArgs(cmd, partition):
+  """Append signing arguments for avbtool."""
+  # e.g., "--key path/to/signing_key --algorithm SHA256_RSA4096"
+  key_path = OPTIONS.info_dict.get("avb_" + partition + "_key_path")
+  algorithm = OPTIONS.info_dict.get("avb_" + partition + "_algorithm")
+  if key_path and algorithm:
+    cmd.extend(["--key", key_path, "--algorithm", algorithm])
+
+
 def _BuildBootableImage(sourcedir, fs_config_file, info_dict=None,
                         has_ramdisk=False, two_step_image=False):
   """Build a bootable image from the specified sourcedir.
@@ -480,13 +489,13 @@
     img_keyblock.close()
 
   # AVB: if enabled, calculate and add hash to boot.img.
-  if info_dict.get("board_avb_enable", None) == "true":
+  if info_dict.get("avb_enable") == "true":
     avbtool = os.getenv('AVBTOOL') or info_dict["avb_avbtool"]
     part_size = info_dict["boot_size"]
     cmd = [avbtool, "add_hash_footer", "--image", img.name,
            "--partition_size", str(part_size), "--partition_name", "boot"]
-    cmd.extend(shlex.split(info_dict["avb_signing_args"]))
-    args = info_dict.get("board_avb_boot_add_hash_footer_args")
+    AppendAVBSigningArgs(cmd, "boot")
+    args = info_dict.get("avb_boot_add_hash_footer_args")
     if args and args.strip():
       cmd.extend(shlex.split(args))
     p = Run(cmd, stdout=subprocess.PIPE)
diff --git a/tools/releasetools/ota_from_target_files.py b/tools/releasetools/ota_from_target_files.py
index 1c8fe65..1b0f68b 100755
--- a/tools/releasetools/ota_from_target_files.py
+++ b/tools/releasetools/ota_from_target_files.py
@@ -1168,7 +1168,7 @@
   # into A/B OTA package.
   target_zip = zipfile.ZipFile(target_file, "r")
   if (OPTIONS.info_dict.get("verity") == "true" or
-      OPTIONS.info_dict.get("board_avb_enable") == "true"):
+      OPTIONS.info_dict.get("avb_enable") == "true"):
     care_map_path = "META/care_map.txt"
     namelist = target_zip.namelist()
     if care_map_path in namelist: