diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index bff6608..a882685 100755
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -346,9 +346,15 @@
     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."""
+def AddVBMeta(output_zip, partitions, prefix="IMAGES/"):
+  """Creates a VBMeta image and store it in output_zip.
+
+  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, which include 'boot',
+        'recovery', 'system', 'vendor', 'dtbo'.
+  """
   img = OutputFile(output_zip, OPTIONS.input_tmp, prefix, "vbmeta.img")
   if os.path.exists(img.input_name):
     print("vbmeta.img already exists in %s; not rebuilding..." % (prefix,))
@@ -361,10 +367,12 @@
   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)
+  for partition, path in partitions.items():
+    assert partition in common.AVB_PARTITIONS, 'Unknown partition: %s' % (
+        partition,)
+    assert os.path.exists(path), 'Failed to find %s for partition %s' % (
+        path, partition)
+    AppendVBMetaArgsForPartition(cmd, partition, path, public_key_dir)
 
   args = OPTIONS.info_dict.get("avb_vbmeta_args")
   if args and args.strip():
@@ -544,6 +552,10 @@
     if fp:
       OPTIONS.info_dict["avb_salt"] = hashlib.sha256(fp).hexdigest()
 
+  # A map between partition names and their paths, which could be used when
+  # generating AVB vbmeta image.
+  partitions = dict()
+
   def banner(s):
     print("\n\n++++ " + s + " ++++\n\n")
 
@@ -553,8 +565,8 @@
       "IMAGES/boot.img", "boot.img", OPTIONS.input_tmp, "BOOT")
   # boot.img may be unavailable in some targets (e.g. aosp_arm64).
   if boot_image:
-    boot_img_path = os.path.join(OPTIONS.input_tmp, "IMAGES", "boot.img")
-    if not os.path.exists(boot_img_path):
+    partitions['boot'] = os.path.join(OPTIONS.input_tmp, "IMAGES", "boot.img")
+    if not os.path.exists(partitions['boot']):
       boot_image.WriteToDir(OPTIONS.input_tmp)
       if output_zip:
         boot_image.AddToZip(output_zip)
@@ -565,9 +577,9 @@
     recovery_image = common.GetBootableImage(
         "IMAGES/recovery.img", "recovery.img", OPTIONS.input_tmp, "RECOVERY")
     assert recovery_image, "Failed to create recovery.img."
-    recovery_img_path = os.path.join(
+    partitions['recovery'] = os.path.join(
         OPTIONS.input_tmp, "IMAGES", "recovery.img")
-    if not os.path.exists(recovery_img_path):
+    if not os.path.exists(partitions['recovery']):
       recovery_image.WriteToDir(OPTIONS.input_tmp)
       if output_zip:
         recovery_image.AddToZip(output_zip)
@@ -586,15 +598,17 @@
           recovery_two_step_image.AddToZip(output_zip)
 
   banner("system")
-  system_img_path = AddSystem(
+  partitions['system'] = system_img_path = AddSystem(
       output_zip, recovery_img=recovery_image, boot_img=boot_image)
-  vendor_img_path = None
+
   if has_vendor:
     banner("vendor")
-    vendor_img_path = AddVendor(output_zip)
+    partitions['vendor'] = vendor_img_path = AddVendor(output_zip)
+
   if has_system_other:
     banner("system_other")
     AddSystemOther(output_zip)
+
   if not OPTIONS.is_signing:
     banner("userdata")
     AddUserdata(output_zip)
@@ -605,15 +619,13 @@
     banner("partition-table")
     AddPartitionTable(output_zip)
 
-  dtbo_img_path = None
   if OPTIONS.info_dict.get("has_dtbo") == "true":
     banner("dtbo")
-    dtbo_img_path = AddDtbo(output_zip)
+    partitions['dtbo'] = AddDtbo(output_zip)
 
   if OPTIONS.info_dict.get("avb_enable") == "true":
     banner("vbmeta")
-    AddVBMeta(output_zip, boot_img_path, system_img_path, vendor_img_path,
-              dtbo_img_path)
+    AddVBMeta(output_zip, partitions)
 
   # For devices using A/B update, copy over images from RADIO/ and/or
   # VENDOR_IMAGES/ to IMAGES/ and make sure we have all the needed
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index f80d0ec..75c86cc 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -452,10 +452,12 @@
   else:
     cmd.extend(["--output", img.name])
 
+  # "boot" or "recovery", without extension.
+  partition_name = os.path.basename(sourcedir).lower()
+
   p = Run(cmd, stdout=subprocess.PIPE)
   p.communicate()
-  assert p.returncode == 0, "mkbootimg of %s image failed" % (
-      os.path.basename(sourcedir),)
+  assert p.returncode == 0, "mkbootimg of %s image failed" % (partition_name,)
 
   if (info_dict.get("boot_signer", None) == "true" and
       info_dict.get("verity_key", None)):
@@ -464,7 +466,7 @@
     if two_step_image:
       path = "/boot"
     else:
-      path = "/" + os.path.basename(sourcedir).lower()
+      path = "/" + partition_name
     cmd = [OPTIONS.boot_signer_path]
     cmd.extend(OPTIONS.boot_signer_args)
     cmd.extend([path, img.name,
@@ -476,7 +478,7 @@
 
   # Sign the image if vboot is non-empty.
   elif info_dict.get("vboot", None):
-    path = "/" + os.path.basename(sourcedir).lower()
+    path = "/" + partition_name
     img_keyblock = tempfile.NamedTemporaryFile()
     # We have switched from the prebuilt futility binary to using the tool
     # (futility-host) built from the source. Override the setting in the old
@@ -503,15 +505,16 @@
     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"]
-    AppendAVBSigningArgs(cmd, "boot")
+           "--partition_size", str(part_size), "--partition_name",
+           partition_name]
+    AppendAVBSigningArgs(cmd, partition_name)
     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)
     p.communicate()
     assert p.returncode == 0, "avbtool add_hash_footer of %s failed" % (
-        os.path.basename(OPTIONS.input_tmp))
+        partition_name,)
 
   img.seek(os.SEEK_SET, 0)
   data = img.read()