diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index fb5626d..1abf5a5 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -665,14 +665,9 @@
   makeint("userdata_size")
   makeint("cache_size")
   makeint("recovery_size")
+  makeint("boot_size")
   makeint("fstab_version")
 
-  boot_images = "boot.img"
-  if "boot_images" in d:
-    boot_images = d["boot_images"]
-  for b in boot_images.split():
-    makeint(b.replace(".img","_size"))
-
   # Load recovery fstab if applicable.
   d["fstab"] = _FindAndLoadRecoveryFstab(d, input_file, read_helper)
 
@@ -1226,10 +1221,7 @@
   # AVB: if enabled, calculate and add hash to boot.img or recovery.img.
   if info_dict.get("avb_enable") == "true":
     avbtool = info_dict["avb_avbtool"]
-    if partition_name == "recovery":
-      part_size = info_dict["recovery_size"]
-    else:
-      part_size = info_dict[image_name.replace(".img","_size")]
+    part_size = info_dict[partition_name + "_size"]
     cmd = [avbtool, "add_hash_footer", "--image", img.name,
            "--partition_size", str(part_size), "--partition_name",
            partition_name]