avbtool: Add 'resize_image' command.

This only works on images with AVB footers. This feature is needed for
some Treble use-cases where a "golden" system.img is used across
devices with varying 'system' partition sizes.

Bug: 36029318
Test: New unit tests and all unit tests pass.
Change-Id: Idc0c31a79157c52249b3ebcd02c1c3bc5228de7f
diff --git a/README.md b/README.md
index 56d792a..8a58b3e 100644
--- a/README.md
+++ b/README.md
@@ -243,6 +243,13 @@
         [--signing_helper /path/to/external/signer]                                \
         [--append_to_release_string STR]
 
+The size of an image with integrity footers can be changed using the
+`resize_image` command:
+
+    $ avbtool resize_image                                                         \
+        --image IMAGE                                                              \
+        --partition_size SIZE
+
 The integrity footer on an image can be removed from an image. The
 hashtree can optionally be kept in place.
 
diff --git a/avbtool b/avbtool
index 574139c..bc121bc 100755
--- a/avbtool
+++ b/avbtool
@@ -1814,6 +1814,56 @@
     # And cut...
     image.truncate(new_image_size)
 
+  def resize_image(self, image_filename, partition_size):
+    """Implements the 'resize_image' command.
+
+    Arguments:
+      image_filename: File with footer to resize.
+      partition_size: The new size of the image.
+
+    Raises:
+      AvbError: If there's no footer in the image.
+    """
+
+    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))
+
+    (footer, vbmeta_header, descriptors, _) = self._parse_image(image)
+
+    if not footer:
+      raise AvbError('Given image does not have a footer.')
+
+    # The vbmeta blob is always at the end of the data so resizing an
+    # image amounts to just moving the footer around.
+
+    vbmeta_end_offset = footer.vbmeta_offset + footer.vbmeta_size
+    if vbmeta_end_offset % image.block_size != 0:
+      vbmeta_end_offset += image.block_size - (vbmeta_end_offset % image.block_size)
+
+    if partition_size < vbmeta_end_offset + 1*image.block_size:
+        raise AvbError('Requested size of {} is too small for an image '
+                       'of size {}.'
+                       .format(partition_size,
+                               vbmeta_end_offset + 1*image.block_size))
+
+    # Cut at the end of the vbmeta blob and insert a DONT_CARE chunk
+    # with enough bytes such that the final Footer block is at the end
+    # of partition_size.
+    image.truncate(vbmeta_end_offset)
+    image.append_dont_care(partition_size - vbmeta_end_offset -
+                           1*image.block_size)
+
+    # Just reuse the same footer - only difference is that we're
+    # writing it in a different place.
+    footer_blob = footer.encode()
+    footer_blob_with_padding = ('\0'*(image.block_size - AvbFooter.SIZE) +
+                                footer_blob)
+    image.append_raw(footer_blob_with_padding)
+
   def set_ab_metadata(self, misc_image, slot_data):
     """Implements the 'set_ab_metadata' command.
 
@@ -3255,6 +3305,17 @@
                             action='store_true')
     sub_parser.set_defaults(func=self.erase_footer)
 
+    sub_parser = subparsers.add_parser('resize_image',
+                                       help='Resize image with a footer.')
+    sub_parser.add_argument('--image',
+                            help='Image with a footer',
+                            type=argparse.FileType('rwb+'),
+                            required=True)
+    sub_parser.add_argument('--partition_size',
+                            help='New partition size',
+                            type=parse_number)
+    sub_parser.set_defaults(func=self.resize_image)
+
     sub_parser = subparsers.add_parser(
         'info_image',
         help='Show information about vbmeta or footer.')
@@ -3437,6 +3498,10 @@
     """Implements the 'erase_footer' sub-command."""
     self.avb.erase_footer(args.image.name, args.keep_hashtree)
 
+  def resize_image(self, args):
+    """Implements the 'resize_image' sub-command."""
+    self.avb.resize_image(args.image.name, args.partition_size)
+
   def set_ab_metadata(self, args):
     """Implements the 'set_ab_metadata' sub-command."""
     self.avb.set_ab_metadata(args.misc_image, args.slot_data)
diff --git a/test/avbtool_unittest.cc b/test/avbtool_unittest.cc
index e3af844..58926db 100644
--- a/test/avbtool_unittest.cc
+++ b/test/avbtool_unittest.cc
@@ -342,9 +342,41 @@
   return true;  // Keep iterating.
 }
 
+static std::string AddHashFooterGetExpectedVBMetaInfo(
+    const bool sparse_image, const uint64_t partition_size) {
+  return base::StringPrintf(
+      "Footer version:           1.0\n"
+      "Image size:               %" PRIu64
+      " bytes\n"
+      "Original image size:      1052672 bytes\n"
+      "VBMeta offset:            1052672\n"
+      "VBMeta size:              1280 bytes\n"
+      "--\n"
+      "Minimum libavb version:   1.0%s\n"
+      "Header Block:             256 bytes\n"
+      "Authentication Block:     320 bytes\n"
+      "Auxiliary Block:          704 bytes\n"
+      "Algorithm:                SHA256_RSA2048\n"
+      "Rollback Index:           0\n"
+      "Flags:                    0\n"
+      "Release String:           ''\n"
+      "Descriptors:\n"
+      "    Hash descriptor:\n"
+      "      Image Size:            1052672 bytes\n"
+      "      Hash Algorithm:        sha256\n"
+      "      Partition Name:        foobar\n"
+      "      Salt:                  d00df00d\n"
+      "      Digest:                "
+      "9a58cc996d405e08a1e00f96dbfe9104fedf41cb83b1f"
+      "5e4ed357fbcf58d88d9\n",
+      partition_size,
+      sparse_image ? " (Sparse)" : "");
+}
+
 void AvbToolTest::AddHashFooterTest(bool sparse_image) {
   const size_t rootfs_size = 1028 * 1024;
   const size_t partition_size = 1536 * 1024;
+  const size_t resized_partition_size = 1280 * 1024;
 
   // Generate a 1028 KiB file with known content. Some content have
   // been arranged to ensure FILL_DATA segments in the sparse file.
@@ -393,30 +425,7 @@
                    (int)partition_size,
                    ext_vbmeta_path.value().c_str());
 
-    ASSERT_EQ(base::StringPrintf("Footer version:           1.0\n"
-                                 "Image size:               1572864 bytes\n"
-                                 "Original image size:      1052672 bytes\n"
-                                 "VBMeta offset:            1052672\n"
-                                 "VBMeta size:              1280 bytes\n"
-                                 "--\n"
-                                 "Minimum libavb version:   1.0%s\n"
-                                 "Header Block:             256 bytes\n"
-                                 "Authentication Block:     320 bytes\n"
-                                 "Auxiliary Block:          704 bytes\n"
-                                 "Algorithm:                SHA256_RSA2048\n"
-                                 "Rollback Index:           0\n"
-                                 "Flags:                    0\n"
-                                 "Release String:           ''\n"
-                                 "Descriptors:\n"
-                                 "    Hash descriptor:\n"
-                                 "      Image Size:            1052672 bytes\n"
-                                 "      Hash Algorithm:        sha256\n"
-                                 "      Partition Name:        foobar\n"
-                                 "      Salt:                  d00df00d\n"
-                                 "      Digest:                "
-                                 "9a58cc996d405e08a1e00f96dbfe9104fedf41cb83b1f"
-                                 "5e4ed357fbcf58d88d9\n",
-                                 sparse_image ? " (Sparse)" : ""),
+    ASSERT_EQ(AddHashFooterGetExpectedVBMetaInfo(sparse_image, partition_size),
               InfoImage(rootfs_path));
 
     ASSERT_EQ(
@@ -440,6 +449,25 @@
         InfoImage(ext_vbmeta_path));
   }
 
+  // Resize the image and check that the only thing that has changed
+  // is where the footer is. First check that resizing to a smaller
+  // size than the original rootfs fails. Then resize to something
+  // larger than the original rootfs but smaller than the current
+  // partition size.
+  EXPECT_COMMAND(1,
+                 "./avbtool resize_image --image %s "
+                 "--partition_size %d",
+                 rootfs_path.value().c_str(),
+                 (int)(rootfs_size - 16 * 1024));
+  EXPECT_COMMAND(0,
+                 "./avbtool resize_image --image %s "
+                 "--partition_size %d",
+                 rootfs_path.value().c_str(),
+                 (int)resized_partition_size);
+  ASSERT_EQ(
+      AddHashFooterGetExpectedVBMetaInfo(sparse_image, resized_partition_size),
+      InfoImage(rootfs_path));
+
   if (sparse_image) {
     EXPECT_COMMAND(0,
                    "mv %s %s.sparse",