Add --dynamic_partition_size argument to add_hash_footer.

When --dynamic_partition_size is passed, the assumed partition size
is based on the size of the provided image file.

Change-Id: I410f697429cc372731c1dc611c7f34bcef897196
Bug: 142352330

1 files changed, 29 insertions, 10 deletions

diff --git a/avbtool.py b/avbtool.py
index 5ba8dbf..5fb48b8 100755
--- a/avbtool.py
+++ b/avbtool.py
@@ -3329,7 +3329,8 @@ class Avb(object):
       image.truncate(original_image_size)
       raise AvbError('Appending VBMeta image failed: {}.'.format(e)) from e
 
-  def add_hash_footer(self, image_filename, partition_size, partition_name,
+  def add_hash_footer(self, image_filename, partition_size,
+                      dynamic_partition_size, partition_name,
                       hash_algorithm, salt, chain_partitions, algorithm_name,
                       key_path,
                       public_key_metadata_path, rollback_index, flags,
@@ -3347,6 +3348,7 @@ class Avb(object):
     Arguments:
       image_filename: File to add the footer to.
       partition_size: Size of partition.
+      dynamic_partition_size: Calculate partition size based on image size.
       partition_name: Name of partition (without A/B suffix).
       hash_algorithm: Hash algorithm to use.
       salt: Salt to use as a hexadecimal string or None to use /dev/urandom.
@@ -3380,6 +3382,14 @@ class Avb(object):
     Raises:
       AvbError: If an argument is incorrect of if adding of hash_footer failed.
     """
+    if not partition_size and not dynamic_partition_size:
+      raise AvbError('--dynamic_partition_size required when not specifying a '
+                     'partition size')
+
+    if dynamic_partition_size and calc_max_image_size:
+      raise AvbError('--calc_max_image_size not supported with '
+                     '--dynamic_partition_size')
+
     required_libavb_version_minor = 0
     if use_persistent_digest or do_not_use_ab:
       required_libavb_version_minor = 1
@@ -3395,24 +3405,18 @@ class Avb(object):
     # this size + metadata (footer + vbmeta struct) fits in
     # |partition_size|.
     max_metadata_size = self.MAX_VBMETA_SIZE + self.MAX_FOOTER_SIZE
-    if partition_size < max_metadata_size:
+    if not dynamic_partition_size and partition_size < max_metadata_size:
       raise AvbError('Parition size of {} is too small. '
                      'Needs to be at least {}'.format(
                          partition_size, max_metadata_size))
-    max_image_size = partition_size - max_metadata_size
 
     # If we're asked to only calculate the maximum image size, we're done.
     if calc_max_image_size:
-      print('{}'.format(max_image_size))
+      print('{}'.format(partition_size - max_metadata_size))
       return
 
     image = ImageHandler(image_filename)
 
-    if partition_size % image.block_size != 0:
-      raise AvbError('Partition size of {} is not a multiple of the image '
-                     'block size {}.'.format(partition_size,
-                                             image.block_size))
-
     # If there's already a footer, truncate the image to its original
     # size. This way 'avbtool add_hash_footer' is idempotent (modulo
     # salts).
@@ -3429,6 +3433,18 @@ class Avb(object):
       # Image size is too small to possibly contain a footer.
       original_image_size = image.image_size
 
+    if dynamic_partition_size:
+      partition_size = original_image_size + max_metadata_size
+
+      # Round up to block size.
+      partition_size -= partition_size % -image.block_size
+
+    max_image_size = partition_size - max_metadata_size
+    if partition_size % image.block_size != 0:
+      raise AvbError('Partition size of {} is not a multiple of the image '
+                     'block size {}.'.format(partition_size,
+                                             image.block_size))
+
     # If anything goes wrong from here-on, restore the image back to
     # its original size.
     try:
@@ -4331,6 +4347,9 @@ class AvbTool(object):
     sub_parser.add_argument('--partition_size',
                             help='Partition size',
                             type=parse_number)
+    sub_parser.add_argument('--dynamic_partition_size',
+                            help='Calculate partition size based on image size',
+                            action='store_true')
     sub_parser.add_argument('--partition_name',
                             help='Partition name',
                             default=None)
@@ -4759,7 +4778,7 @@ class AvbTool(object):
     """Implements the 'add_hash_footer' sub-command."""
     args = self._fixup_common_args(args)
     self.avb.add_hash_footer(args.image.name if args.image else None,
-                             args.partition_size,
+                             args.partition_size, args.dynamic_partition_size,
                              args.partition_name, args.hash_algorithm,
                              args.salt, args.chain_partition, args.algorithm,
                              args.key,