Prevent AVB signing from using _RESERVED_SIZE

When AVB is enabled with PRODUCT_USE_DYNAMIC_PARTITION_SIZE, AVB
metadata (e.g., hash tree, fec metadata) will consume _RESERVED_SIZE,
resulting in smaller reserved size in file system (e.g., ext4).

Adding additional space for AVB signing and keep the _RESERVED_SIZE only
for file system. This is done by adding a function to binary search an
optimal partition size for a given image size (disk usage + _RESERVED_SIZE).

Bug: 112322265
Test: Build aosp_arm64-userdebug, calculate the running time of
      AVBCalcMinPartitionSize() is about 0.3-0.4 seconds.
Test: python -m unittest test_build_image
Change-Id: I8f0051b57701d6fbba6a9db3756dd194066c74b8

tools/releasetools/build_image.py

@@ -149,13 +149,14 @@ def AVBCalcMaxImageSize(avbtool, footer_type, partition_size, additional_args):
     avbtool: String with path to avbtool.
     footer_type: 'hash' or 'hashtree' for generating footer.
     partition_size: The size of the partition in question.
-    additional_args: Additional arguments to pass to 'avbtool
+    additional_args: Additional arguments to pass to "avbtool add_hash_footer"
-      add_hashtree_image'.
+        or "avbtool add_hashtree_footer".
+
   Returns:
     The maximum image size or 0 if an error occurred.
   """
   cmd = [avbtool, "add_%s_footer" % footer_type,
-         "--partition_size", partition_size, "--calc_max_image_size"]
+         "--partition_size", str(partition_size), "--calc_max_image_size"]
   cmd.extend(shlex.split(additional_args))
 
   (output, exit_code) = RunCommand(cmd)
@@ -165,6 +166,63 @@ def AVBCalcMaxImageSize(avbtool, footer_type, partition_size, additional_args):
     return int(output)
 
 
+def AVBCalcMinPartitionSize(image_size, size_calculator):
+  """Calculates min partition size for a given image size.
+
+  Args:
+    image_size: The size of the image in question.
+    size_calculator: The function to calculate max image size
+        for a given partition size.
+
+  Returns:
+    The minimum partition size required to accommodate the image size.
+  """
+  # Use image size as partition size to approximate final partition size.
+  image_ratio = size_calculator(image_size) / float(image_size)
+
+  # Prepare a binary search for the optimal partition size.
+  lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - BLOCK_SIZE
+
+  # Ensure lo is small enough: max_image_size should <= image_size.
+  delta = BLOCK_SIZE
+  max_image_size = size_calculator(lo)
+  while max_image_size > image_size:
+    image_ratio = max_image_size / float(lo)
+    lo = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE - delta
+    delta *= 2
+    max_image_size = size_calculator(lo)
+
+  hi = lo + BLOCK_SIZE
+
+  # Ensure hi is large enough: max_image_size should >= image_size.
+  delta = BLOCK_SIZE
+  max_image_size = size_calculator(hi)
+  while max_image_size < image_size:
+    image_ratio = max_image_size / float(hi)
+    hi = int(image_size / image_ratio) // BLOCK_SIZE * BLOCK_SIZE + delta
+    delta *= 2
+    max_image_size = size_calculator(hi)
+
+  partition_size = hi
+
+  # Start to binary search.
+  while lo < hi:
+    mid = ((lo + hi) // (2 * BLOCK_SIZE)) * BLOCK_SIZE
+    max_image_size = size_calculator(mid)
+    if max_image_size >= image_size:  # if mid can accommodate image_size
+      if mid < partition_size:  # if a smaller partition size is found
+        partition_size = mid
+      hi = mid
+    else:
+      lo = mid + BLOCK_SIZE
+
+  if OPTIONS.verbose:
+    print("AVBCalcMinPartitionSize({}): partition_size: {}.".format(
+        image_size, partition_size))
+
+  return partition_size
+
+
 def AVBAddFooter(image_path, avbtool, footer_type, partition_size,
                  partition_name, key_path, algorithm, salt,
                  additional_args):
@@ -179,8 +237,8 @@ def AVBAddFooter(image_path, avbtool, footer_type, partition_size,
     key_path: Path to key to use or None.
     algorithm: Name of algorithm to use or None.
     salt: The salt to use (a hexadecimal string) or None.
-    additional_args: Additional arguments to pass to 'avbtool
+    additional_args: Additional arguments to pass to "avbtool add_hash_footer"
-        add_hashtree_image'.
+        or "avbtool add_hashtree_footer".
 
   Returns:
     True if the operation succeeded.
@@ -549,6 +607,17 @@ def BuildImage(in_dir, prop_dict, out_file, target_out=None):
   verity_supported = prop_dict.get("verity") == "true"
   verity_fec_supported = prop_dict.get("verity_fec") == "true"
 
+  avb_footer_type = None
+  if prop_dict.get("avb_hash_enable") == "true":
+    avb_footer_type = "hash"
+  elif prop_dict.get("avb_hashtree_enable") == "true":
+    avb_footer_type = "hashtree"
+
+  if avb_footer_type:
+    avbtool = prop_dict.get("avb_avbtool")
+    avb_signing_args = prop_dict.get(
+        "avb_add_" + avb_footer_type + "_footer_args")
+
   if (prop_dict.get("use_dynamic_partition_size") == "true" and
       "partition_size" not in prop_dict):
     # if partition_size is not defined, use output of `du' + reserved_size
@@ -560,6 +629,13 @@ def BuildImage(in_dir, prop_dict, out_file, target_out=None):
     size += int(prop_dict.get("partition_reserved_size", 0))
     # Round this up to a multiple of 4K so that avbtool works
     size = common.RoundUpTo4K(size)
+    # Adjust partition_size to add more space for AVB footer, to prevent
+    # it from consuming partition_reserved_size.
+    if avb_footer_type:
+      size = AVBCalcMinPartitionSize(
+          size,
+          lambda x: AVBCalcMaxImageSize(
+              avbtool, avb_footer_type, x, avb_signing_args))
     prop_dict["partition_size"] = str(size)
     if OPTIONS.verbose:
       print("Allocating %d MB for %s." % (size // BYTES_IN_MB, out_file))
@@ -576,20 +652,12 @@ def BuildImage(in_dir, prop_dict, out_file, target_out=None):
     prop_dict["image_size"] = str(image_size)
     prop_dict["verity_size"] = str(verity_size)
 
-  avb_footer_type = ''
-  if prop_dict.get("avb_hash_enable") == "true":
-    avb_footer_type = 'hash'
-  elif prop_dict.get("avb_hashtree_enable") == "true":
-    avb_footer_type = 'hashtree'
-
   # Adjust the image size for AVB hash footer or AVB hashtree footer.
   if avb_footer_type:
-    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, additional_args)
+                                         partition_size, avb_signing_args)
     if max_image_size <= 0:
       print("AVBCalcMaxImageSize is <= 0: %d" % max_image_size)
       return False
@@ -722,18 +790,15 @@ def BuildImage(in_dir, prop_dict, out_file, target_out=None):
 
   # Add AVB HASH or HASHTREE footer (metadata).
   if avb_footer_type:
-    avbtool = prop_dict["avb_avbtool"]
     partition_size = prop_dict["partition_size"]
     partition_name = prop_dict["partition_name"]
     # 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")
     salt = prop_dict.get("avb_salt")
-    # 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,
                         partition_size, partition_name, key_path,
-                        algorithm, salt, additional_args):
+                        algorithm, salt, avb_signing_args):
       return False
 
   if run_e2fsck and prop_dict.get("skip_fsck") != "true":

tools/releasetools/test_build_image.py

@@ -15,11 +15,15 @@
 #
 
 import filecmp
+import math
 import os.path
+import random
 import unittest
 
 import common
-from build_image import CheckHeadroom, RunCommand, SetUpInDirAndFsConfig
+from build_image import (
+    AVBCalcMinPartitionSize, BLOCK_SIZE,
+    CheckHeadroom, RunCommand, SetUpInDirAndFsConfig)
 
 
 class BuildImageTest(unittest.TestCase):
@@ -28,6 +32,13 @@ class BuildImageTest(unittest.TestCase):
   EXT4FS_OUTPUT = (
       "Created filesystem with 2777/129024 inodes and 515099/516099 blocks")
 
+  def setUp(self):
+    # To test AVBCalcMinPartitionSize(), by using 200MB to 2GB image size.
+    #   -  51200 = 200MB * 1024 * 1024 / 4096
+    #   - 524288 = 2GB * 1024 * 1024 * 1024 / 4096
+    self._image_sizes = [BLOCK_SIZE * random.randint(51200, 524288) + offset
+                         for offset in range(BLOCK_SIZE)]
+
   def tearDown(self):
     common.Cleanup()
 
@@ -176,3 +187,51 @@ class BuildImageTest(unittest.TestCase):
     self.assertIn('fs-config-system\n', fs_config_data)
     self.assertIn('fs-config-root\n', fs_config_data)
     self.assertEqual('/', prop_dict['mount_point'])
+
+  def test_AVBCalcMinPartitionSize_LinearFooterSize(self):
+    """Tests with footer size which is linear to partition size."""
+    for image_size in self._image_sizes:
+      for ratio in 0.95, 0.56, 0.22:
+        expected_size = common.RoundUpTo4K(int(math.ceil(image_size / ratio)))
+        self.assertEqual(
+            expected_size,
+            AVBCalcMinPartitionSize(image_size, lambda x: int(x * ratio)))
+
+  def test_AVBCalcMinPartitionSize_SlowerGrowthFooterSize(self):
+    """Tests with footer size which grows slower than partition size."""
+
+    def _SizeCalculator(partition_size):
+      """Footer size is the power of 0.95 of partition size."""
+      # Minus footer size to return max image size.
+      return partition_size - int(math.pow(partition_size, 0.95))
+
+    for image_size in self._image_sizes:
+      min_partition_size = AVBCalcMinPartitionSize(image_size, _SizeCalculator)
+      # Checks min_partition_size can accommodate image_size.
+      self.assertGreaterEqual(
+          _SizeCalculator(min_partition_size),
+          image_size)
+      # Checks min_partition_size (round to BLOCK_SIZE) is the minimum.
+      self.assertLess(
+          _SizeCalculator(min_partition_size - BLOCK_SIZE),
+          image_size)
+
+  def test_AVBCalcMinPartitionSize_FasterGrowthFooterSize(self):
+    """Tests with footer size which grows faster than partition size."""
+
+    def _SizeCalculator(partition_size):
+      """Max image size is the power of 0.95 of partition size."""
+      # Max image size grows less than partition size, which means
+      # footer size grows faster than partition size.
+      return int(math.pow(partition_size, 0.95))
+
+    for image_size in self._image_sizes:
+      min_partition_size = AVBCalcMinPartitionSize(image_size, _SizeCalculator)
+      # Checks min_partition_size can accommodate image_size.
+      self.assertGreaterEqual(
+          _SizeCalculator(min_partition_size),
+          image_size)
+      # Checks min_partition_size (round to BLOCK_SIZE) is the minimum.
+      self.assertLess(
+          _SizeCalculator(min_partition_size - BLOCK_SIZE),
+          image_size)