avbtool: Add 'zero_hashtree' command.android-o-mr1-iot-release-1.0.12oreo-mr1-iot-release

This is useful for trading compressed image size for having to
reculculate the hashtree and FEC at runtime. If this is done all of
the hashtree and FEC data is set to zero except for the first eight
bytes which is set to the magic `ZeRoHaSH`. Applications can use this
to detect if recalculation is needed.

Also add a --accept_zeroed_hashtree option to 'avbtool verify_image'
which if given will treat a zeroed hashtree as equivalent to a valid
hashtree.

Bug: 130317158
Test: New unit tests + all tests pass.
Change-Id: I00054145942f1ffacfa77b0e96c787bf0c1a9dba

3 files changed, 250 insertions, 15 deletions

diff --git a/README.md b/README.md
index fbee3f5..0b1ec97 100644
--- a/README.md
+++ b/README.md
@@ -422,6 +422,18 @@ an external file via the `--output_vbmeta_image` option and one can
 also specify that the vbmeta struct and footer not be added to the
 image being operated on.
 
+The hashtree and FEC data in an image can be zeroed out with the following
+command:
+
+    $ avbtool zero_hashtree --image IMAGE
+
+This is useful for trading compressed image size for having to reculculate the
+hashtree and FEC at runtime. If this is done the hashtree and FEC data is set
+to zero except for the first eight bytes which are set to the magic
+`ZeRoHaSH`. Either the hashtree or FEC data or both may be zeroed this way
+so applications should check for the magic both places. Applications can
+use the magic to detect if recalculation is needed.
+
 To calculate the maximum size of an image that will fit in a partition
 of a given size after having used the `avbtool add_hash_footer` or
 `avbtool add_hashtree_footer` commands on it, use the
diff --git a/avbtool b/avbtool
index 610cf19..06d4354 100755
--- a/avbtool
+++ b/avbtool
@@ -1121,7 +1121,7 @@ class AvbDescriptor(object):
     return bytearray(ret)
 
   def verify(self, image_dir, image_ext, expected_chain_partitions_map,
-             image_containing_descriptor):
+             image_containing_descriptor, accept_zeroed_hashtree):
     """Verifies contents of the descriptor - used in verify_image sub-command.
 
     Arguments:
@@ -1130,6 +1130,7 @@ class AvbDescriptor(object):
       expected_chain_partitions_map: A map from partition name to the
         tuple (rollback_index_location, key_blob).
       image_containing_descriptor: The image the descriptor is in.
+      accept_zeroed_hashtree: If True, don't fail if hashtree or FEC data is zeroed out.
 
     Returns:
       True if the descriptor verifies, False otherwise.
@@ -1207,7 +1208,7 @@ class AvbPropertyDescriptor(AvbDescriptor):
     return bytearray(ret)
 
   def verify(self, image_dir, image_ext, expected_chain_partitions_map,
-             image_containing_descriptor):
+             image_containing_descriptor, accept_zeroed_hashtree):
     """Verifies contents of the descriptor - used in verify_image sub-command.
 
     Arguments:
@@ -1216,6 +1217,7 @@ class AvbPropertyDescriptor(AvbDescriptor):
       expected_chain_partitions_map: A map from partition name to the
         tuple (rollback_index_location, key_blob).
       image_containing_descriptor: The image the descriptor is in.
+      accept_zeroed_hashtree: If True, don't fail if hashtree or FEC data is zeroed out.
 
     Returns:
       True if the descriptor verifies, False otherwise.
@@ -1369,7 +1371,7 @@ class AvbHashtreeDescriptor(AvbDescriptor):
     return bytearray(ret)
 
   def verify(self, image_dir, image_ext, expected_chain_partitions_map,
-             image_containing_descriptor):
+             image_containing_descriptor, accept_zeroed_hashtree):
     """Verifies contents of the descriptor - used in verify_image sub-command.
 
     Arguments:
@@ -1378,6 +1380,7 @@ class AvbHashtreeDescriptor(AvbDescriptor):
       expected_chain_partitions_map: A map from partition name to the
         tuple (rollback_index_location, key_blob).
       image_containing_descriptor: The image the descriptor is in.
+      accept_zeroed_hashtree: If True, don't fail if hashtree or FEC data is zeroed out.
 
     Returns:
       True if the descriptor verifies, False otherwise.
@@ -1406,17 +1409,22 @@ class AvbHashtreeDescriptor(AvbDescriptor):
     # ... also check that the on-disk hashtree matches
     image.seek(self.tree_offset)
     hash_tree_ondisk = image.read(self.tree_size)
-    if hash_tree != hash_tree_ondisk:
-      sys.stderr.write('hashtree of {} contains invalid data\n'.
+    is_zeroed = (hash_tree_ondisk[0:8] == 'ZeRoHaSH')
+    if is_zeroed and accept_zeroed_hashtree:
+      print ('{}: skipping verification since hashtree is zeroed and --accept_zeroed_hashtree was given'
+             .format(self.partition_name))
+    else:
+      if hash_tree != hash_tree_ondisk:
+        sys.stderr.write('hashtree of {} contains invalid data\n'.
                        format(image_filename))
-      return False
+        return False
+      print ('{}: Successfully verified {} hashtree of {} for image of {} bytes'
+             .format(self.partition_name, self.hash_algorithm, image.filename,
+                     self.image_size))
     # TODO: we could also verify that the FEC stored in the image is
     # correct but this a) currently requires the 'fec' binary; and b)
     # takes a long time; and c) is not strictly needed for
     # verification purposes as we've already verified the root hash.
-    print ('{}: Successfully verified {} hashtree of {} for image of {} bytes'
-           .format(self.partition_name, self.hash_algorithm, image.filename,
-                   self.image_size))
     return True
 
 
@@ -1526,7 +1534,7 @@ class AvbHashDescriptor(AvbDescriptor):
     return bytearray(ret)
 
   def verify(self, image_dir, image_ext, expected_chain_partitions_map,
-             image_containing_descriptor):
+             image_containing_descriptor, accept_zeroed_hashtree):
     """Verifies contents of the descriptor - used in verify_image sub-command.
 
     Arguments:
@@ -1535,6 +1543,7 @@ class AvbHashDescriptor(AvbDescriptor):
       expected_chain_partitions_map: A map from partition name to the
         tuple (rollback_index_location, key_blob).
       image_containing_descriptor: The image the descriptor is in.
+      accept_zeroed_hashtree: If True, don't fail if hashtree or FEC data is zeroed out.
 
     Returns:
       True if the descriptor verifies, False otherwise.
@@ -1636,7 +1645,7 @@ class AvbKernelCmdlineDescriptor(AvbDescriptor):
     return bytearray(ret)
 
   def verify(self, image_dir, image_ext, expected_chain_partitions_map,
-             image_containing_descriptor):
+             image_containing_descriptor, accept_zeroed_hashtree):
     """Verifies contents of the descriptor - used in verify_image sub-command.
 
     Arguments:
@@ -1645,6 +1654,7 @@ class AvbKernelCmdlineDescriptor(AvbDescriptor):
       expected_chain_partitions_map: A map from partition name to the
         tuple (rollback_index_location, key_blob).
       image_containing_descriptor: The image the descriptor is in.
+      accept_zeroed_hashtree: If True, don't fail if hashtree or FEC data is zeroed out.
 
     Returns:
       True if the descriptor verifies, False otherwise.
@@ -1739,7 +1749,7 @@ class AvbChainPartitionDescriptor(AvbDescriptor):
     return bytearray(ret)
 
   def verify(self, image_dir, image_ext, expected_chain_partitions_map,
-             image_containing_descriptor):
+             image_containing_descriptor, accept_zeroed_hashtree):
     """Verifies contents of the descriptor - used in verify_image sub-command.
 
     Arguments:
@@ -1748,6 +1758,7 @@ class AvbChainPartitionDescriptor(AvbDescriptor):
       expected_chain_partitions_map: A map from partition name to the
         tuple (rollback_index_location, key_blob).
       image_containing_descriptor: The image the descriptor is in.
+      accept_zeroed_hashtree: If True, don't fail if hashtree or FEC data is zeroed out.
 
     Returns:
       True if the descriptor verifies, False otherwise.
@@ -2086,6 +2097,63 @@ class Avb(object):
     # And cut...
     image.truncate(new_image_size)
 
+  def zero_hashtree(self, image_filename):
+    """Implements the 'zero_hashtree' command.
+
+    Arguments:
+      image_filename: File to zero hashtree and FEC data from.
+
+    Raises:
+      AvbError: If there's no footer in the image.
+    """
+
+    image = ImageHandler(image_filename)
+
+    (footer, _, descriptors, _) = self._parse_image(image)
+
+    if not footer:
+      raise AvbError('Given image does not have a footer.')
+
+    # Search for a hashtree descriptor to figure out the location and
+    # size of the hashtree and FEC.
+    ht_desc = None
+    for desc in descriptors:
+      if isinstance(desc, AvbHashtreeDescriptor):
+        ht_desc = desc
+        break
+
+    if not ht_desc:
+      raise AvbError('No hashtree descriptor was found.')
+
+    zero_ht_start_offset = ht_desc.tree_offset
+    zero_ht_num_bytes = ht_desc.tree_size
+    zero_fec_start_offset = None
+    zero_fec_num_bytes = 0
+    if ht_desc.fec_offset > 0:
+      if ht_desc.fec_offset != ht_desc.tree_offset + ht_desc.tree_size:
+        raise AvbError('Hash-tree and FEC data must be adjacent.')
+      zero_fec_start_offset = ht_desc.fec_offset
+      zero_fec_num_bytes = ht_desc.fec_size
+    zero_end_offset = zero_ht_start_offset + zero_ht_num_bytes + zero_fec_num_bytes
+    image.seek(zero_end_offset)
+    data = image.read(image.image_size - zero_end_offset)
+
+    # Write zeroes all over hashtree and FEC, except for the first eight bytes
+    # where a magic marker - ZeroHaSH - is placed. Place these markers in the
+    # beginning of both hashtree and FEC. (That way, in the future we can add
+    # options to 'avbtool zero_hashtree' so as to zero out only either/or.)
+    #
+    # Applications can use these markers to detect that the hashtree and/or
+    # FEC needs to be recomputed.
+    image.truncate(zero_ht_start_offset)
+    data_zeroed_firstblock = 'ZeRoHaSH' + '\0'*(image.block_size - 8)
+    image.append_raw(data_zeroed_firstblock)
+    image.append_fill('\0\0\0\0', zero_ht_num_bytes - image.block_size)
+    if zero_fec_start_offset:
+      image.append_raw(data_zeroed_firstblock)
+      image.append_fill('\0\0\0\0', zero_fec_num_bytes - image.block_size)
+    image.append_raw(data)
+
   def resize_image(self, image_filename, partition_size):
     """Implements the 'resize_image' command.
 
@@ -2220,7 +2288,8 @@ class Avb(object):
     if num_printed == 0:
       o.write('    (none)\n')
 
-  def verify_image(self, image_filename, key_path, expected_chain_partitions, follow_chain_partitions):
+  def verify_image(self, image_filename, key_path, expected_chain_partitions, follow_chain_partitions,
+                   accept_zeroed_hashtree):
     """Implements the 'verify_image' command.
 
     Arguments:
@@ -2229,6 +2298,7 @@ class Avb(object):
       expected_chain_partitions: List of chain partitions to check or None.
       follow_chain_partitions: If True, will follows chain partitions even when not
                                specified with the --expected_chain_partition option
+      accept_zeroed_hashtree: If True, don't fail if hashtree or FEC data is zeroed out.
     """
     expected_chain_partitions_map = {}
     if expected_chain_partitions:
@@ -2294,7 +2364,8 @@ class Avb(object):
               .format(desc.partition_name, desc.rollback_index_location,
                       hashlib.sha1(desc.public_key).hexdigest()))
       else:
-        if not desc.verify(image_dir, image_ext, expected_chain_partitions_map, image):
+        if not desc.verify(image_dir, image_ext, expected_chain_partitions_map, image,
+                           accept_zeroed_hashtree):
           raise AvbError('Error verifying descriptor.')
       # Honor --follow_chain_partitions - add '--' to make the output more readable.
       if isinstance(desc, AvbChainPartitionDescriptor) and follow_chain_partitions:
@@ -4026,6 +4097,14 @@ class AvbTool(object):
                             action='store_true')
     sub_parser.set_defaults(func=self.erase_footer)
 
+    sub_parser = subparsers.add_parser('zero_hashtree',
+                                       help='Zero out hashtree and FEC data.')
+    sub_parser.add_argument('--image',
+                            help='Image with a footer',
+                            type=argparse.FileType('rwb+'),
+                            required=True)
+    sub_parser.set_defaults(func=self.zero_hashtree)
+
     sub_parser = subparsers.add_parser('extract_vbmeta_image',
                                        help='Extracts vbmeta from an image with a footer.')
     sub_parser.add_argument('--image',
@@ -4086,6 +4165,9 @@ class AvbTool(object):
                             help=('Follows chain partitions even when not '
                                   'specified with the --expected_chain_partition option'),
                             action='store_true')
+    sub_parser.add_argument('--accept_zeroed_hashtree',
+                            help=('Accept images where the hashtree or FEC data is zeroed out'),
+                            action='store_true')
     sub_parser.set_defaults(func=self.verify_image)
 
     sub_parser = subparsers.add_parser(
@@ -4347,6 +4429,10 @@ class AvbTool(object):
     """Implements the 'erase_footer' sub-command."""
     self.avb.erase_footer(args.image.name, args.keep_hashtree)
 
+  def zero_hashtree(self, args):
+    """Implements the 'zero_hashtree' sub-command."""
+    self.avb.zero_hashtree(args.image.name)
+
   def extract_vbmeta_image(self, args):
     """Implements the 'extract_vbmeta_image' sub-command."""
     self.avb.extract_vbmeta_image(args.output, args.image.name,
@@ -4368,7 +4454,8 @@ class AvbTool(object):
     """Implements the 'verify_image' sub-command."""
     self.avb.verify_image(args.image.name, args.key,
                           args.expected_chain_partition,
-                          args.follow_chain_partitions)
+                          args.follow_chain_partitions,
+                          args.accept_zeroed_hashtree)
 
   def calculate_vbmeta_digest(self, args):
     """Implements the 'calculate_vbmeta_digest' sub-command."""
diff --git a/test/avbtool_unittest.cc b/test/avbtool_unittest.cc
index 1dfdf86..d87f494 100644
--- a/test/avbtool_unittest.cc
+++ b/test/avbtool_unittest.cc
@@ -1003,6 +1003,15 @@ void AvbToolTest::AddHashtreeFooterTest(bool sparse_image) {
                    extracted_vbmeta_path.value().c_str());
   }
 
+  /* Zero the hashtree on a copy of the image. */
+  EXPECT_COMMAND(0,
+                 "cp %s %s.zht",
+                 rootfs_path.value().c_str(),
+                 rootfs_path.value().c_str());
+  EXPECT_COMMAND(0,
+                 "./avbtool zero_hashtree --image %s.zht ",
+                 rootfs_path.value().c_str());
+
   if (sparse_image) {
     EXPECT_COMMAND(0,
                    "mv %s %s.sparse",
@@ -1013,6 +1022,16 @@ void AvbToolTest::AddHashtreeFooterTest(bool sparse_image) {
                    rootfs_path.value().c_str(),
                    rootfs_path.value().c_str());
     EXPECT_COMMAND(0, "rm -f %s.sparse", rootfs_path.value().c_str());
+
+    EXPECT_COMMAND(0,
+                   "mv %s.zht %s.zht.sparse",
+                   rootfs_path.value().c_str(),
+                   rootfs_path.value().c_str());
+    EXPECT_COMMAND(0,
+                   "simg2img %s.zht.sparse %s.zht",
+                   rootfs_path.value().c_str(),
+                   rootfs_path.value().c_str());
+    EXPECT_COMMAND(0, "rm -f %s.zht.sparse", rootfs_path.value().c_str());
   }
 
   // To check that we generate the correct hashtree we can use
@@ -1038,6 +1057,11 @@ void AvbToolTest::AddHashtreeFooterTest(bool sparse_image) {
   std::string part_data;
   ASSERT_TRUE(base::ReadFileToString(rootfs_path, &part_data));
 
+  // Also read the zeroed hash-tree version.
+  std::string zht_part_data;
+  ASSERT_TRUE(base::ReadFileToString(
+      base::FilePath(rootfs_path.value() + ".zht"), &zht_part_data));
+
   // Check footer contains correct data.
   AvbFooter f;
   EXPECT_NE(0,
@@ -1095,6 +1119,33 @@ void AvbToolTest::AddHashtreeFooterTest(bool sparse_image) {
   EXPECT_EQ("e811611467dcd6e8dc4324e45f706c2bdd51db67",
             mem_to_hexstring(desc_end + o, d.root_digest_len));
 
+  // Check that the zeroed hashtree version differ only by the hashtree + fec
+  // being zeroed out.
+  EXPECT_EQ(part_data.size(), zht_part_data.size());
+  size_t zht_ht_begin = d.tree_offset;
+  size_t zht_ht_end = zht_ht_begin + d.tree_size;
+  size_t zht_fec_begin = zht_ht_end;
+  size_t zht_fec_end = zht_fec_begin + d.fec_size;
+  EXPECT_EQ(0, memcmp(part_data.data(), zht_part_data.data(), zht_ht_begin));
+  EXPECT_NE(0,
+            memcmp(part_data.data() + zht_ht_begin,
+                   zht_part_data.data() + zht_ht_begin,
+                   zht_fec_end - zht_ht_begin));
+  EXPECT_EQ(0,
+            memcmp(part_data.data() + zht_fec_end,
+                   zht_part_data.data() + zht_fec_end,
+                   zht_part_data.size() - zht_fec_end));
+  EXPECT_EQ(0, strncmp(zht_part_data.data() + zht_ht_begin, "ZeRoHaSH", 8));
+  for (size_t n = zht_ht_begin + 8; n < zht_ht_end; n++) {
+    EXPECT_EQ(0, zht_part_data.data()[n]);
+  }
+  if (d.fec_size > 0) {
+    EXPECT_EQ(0, strncmp(zht_part_data.data() + zht_fec_begin, "ZeRoHaSH", 8));
+    for (size_t n = zht_fec_begin + 8; n < zht_fec_end; n++) {
+      EXPECT_EQ(0, zht_part_data.data()[n]);
+    }
+  }
+
   // Check that we correctly generate dm-verity kernel cmdline
   // snippets, if requested.
   base::FilePath vbmeta_dmv_path = testdir_.Append("vbmeta_dm_verity_desc.bin");
@@ -1247,6 +1298,15 @@ void AvbToolTest::AddHashtreeFooterFECTest(bool sparse_image) {
               InfoImage(rootfs_path));
   }
 
+  /* Zero the hashtree and FEC on a copy of the image. */
+  EXPECT_COMMAND(0,
+                 "cp %s %s.zht",
+                 rootfs_path.value().c_str(),
+                 rootfs_path.value().c_str());
+  EXPECT_COMMAND(0,
+                 "./avbtool zero_hashtree --image %s.zht ",
+                 rootfs_path.value().c_str());
+
   if (sparse_image) {
     EXPECT_COMMAND(0,
                    "mv %s %s.sparse",
@@ -1257,6 +1317,16 @@ void AvbToolTest::AddHashtreeFooterFECTest(bool sparse_image) {
                    rootfs_path.value().c_str(),
                    rootfs_path.value().c_str());
     EXPECT_COMMAND(0, "rm -f %s.sparse", rootfs_path.value().c_str());
+
+    EXPECT_COMMAND(0,
+                   "mv %s.zht %s.zht.sparse",
+                   rootfs_path.value().c_str(),
+                   rootfs_path.value().c_str());
+    EXPECT_COMMAND(0,
+                   "simg2img %s.zht.sparse %s.zht",
+                   rootfs_path.value().c_str(),
+                   rootfs_path.value().c_str());
+    EXPECT_COMMAND(0, "rm -f %s.zht.sparse", rootfs_path.value().c_str());
   }
 
   /* TODO: would be nice to verify that the FEC data is correct. */
@@ -1265,6 +1335,11 @@ void AvbToolTest::AddHashtreeFooterFECTest(bool sparse_image) {
   std::string part_data;
   ASSERT_TRUE(base::ReadFileToString(rootfs_path, &part_data));
 
+  // Also read the zeroed hash-tree version.
+  std::string zht_part_data;
+  ASSERT_TRUE(base::ReadFileToString(
+      base::FilePath(rootfs_path.value() + ".zht"), &zht_part_data));
+
   // Check footer contains correct data.
   AvbFooter f;
   EXPECT_NE(0,
@@ -1324,6 +1399,33 @@ void AvbToolTest::AddHashtreeFooterFECTest(bool sparse_image) {
   EXPECT_EQ("e811611467dcd6e8dc4324e45f706c2bdd51db67",
             mem_to_hexstring(desc_end + o, d.root_digest_len));
 
+  // Check that the zeroed hashtree version differ only by the hashtree + fec
+  // being zeroed out.
+  EXPECT_EQ(part_data.size(), zht_part_data.size());
+  size_t zht_ht_begin = d.tree_offset;
+  size_t zht_ht_end = zht_ht_begin + d.tree_size;
+  size_t zht_fec_begin = zht_ht_end;
+  size_t zht_fec_end = zht_fec_begin + d.fec_size;
+  EXPECT_EQ(0, memcmp(part_data.data(), zht_part_data.data(), zht_ht_begin));
+  EXPECT_NE(0,
+            memcmp(part_data.data() + zht_ht_begin,
+                   zht_part_data.data() + zht_ht_begin,
+                   zht_fec_end - zht_ht_begin));
+  EXPECT_EQ(0,
+            memcmp(part_data.data() + zht_fec_end,
+                   zht_part_data.data() + zht_fec_end,
+                   zht_part_data.size() - zht_fec_end));
+  EXPECT_EQ(0, strncmp(zht_part_data.data() + zht_ht_begin, "ZeRoHaSH", 8));
+  for (size_t n = zht_ht_begin + 8; n < zht_ht_end; n++) {
+    EXPECT_EQ(0, zht_part_data.data()[n]);
+  }
+  if (d.fec_size > 0) {
+    EXPECT_EQ(0, strncmp(zht_part_data.data() + zht_fec_begin, "ZeRoHaSH", 8));
+    for (size_t n = zht_fec_begin + 8; n < zht_fec_end; n++) {
+      EXPECT_EQ(0, zht_part_data.data()[n]);
+    }
+  }
+
   // Check that we correctly generate dm-verity kernel cmdline
   // snippets, if requested.
   base::FilePath vbmeta_dmv_path = testdir_.Append("vbmeta_dm_verity_desc.bin");
@@ -2392,6 +2494,40 @@ TEST_F(AvbToolTest, VerifyImageWithHashAndHashtree) {
   }
 }
 
+TEST_F(AvbToolTest, VerifyImageWithHashAndZeroedHashtree) {
+  const size_t system_partition_size = 10 * 1024 * 1024;
+  const size_t system_image_size = 8 * 1024 * 1024;
+  base::FilePath system_path = GenerateImage("system.img", system_image_size);
+  EXPECT_COMMAND(0,
+                 "./avbtool add_hashtree_footer --salt d00df00d --image %s "
+                 "--partition_size %zd --partition_name system "
+                 "--internal_release_string \"\" ",
+                 system_path.value().c_str(),
+                 system_partition_size);
+
+  GenerateVBMetaImage("vbmeta.img",
+                      "SHA256_RSA2048",
+                      0,
+                      base::FilePath("test/data/testkey_rsa2048.pem"),
+                      base::StringPrintf("--include_descriptors_from_image %s ",
+                                         system_path.value().c_str()));
+
+  EXPECT_COMMAND(0,
+                 "./avbtool verify_image --image %s --accept_zeroed_hashtree",
+                 vbmeta_image_path_.value().c_str());
+
+  EXPECT_COMMAND(
+      0, "./avbtool zero_hashtree --image %s", system_path.value().c_str());
+
+  EXPECT_COMMAND(1,
+                 "./avbtool verify_image --image %s",
+                 vbmeta_image_path_.value().c_str());
+
+  EXPECT_COMMAND(0,
+                 "./avbtool verify_image --image %s --accept_zeroed_hashtree",
+                 vbmeta_image_path_.value().c_str());
+}
+
 TEST_F(AvbToolTest, VerifyImageWithHashAndHashtreeCorruptHash) {
   GenerateImageWithHashAndHashtreeSetup();