Merge "Split system and vendor unmounted notice deps"
diff --git a/tools/releasetools/Android.bp b/tools/releasetools/Android.bp
index 8c91470..4fdc707 100644
--- a/tools/releasetools/Android.bp
+++ b/tools/releasetools/Android.bp
@@ -37,6 +37,7 @@
         "releasetools_build_image",
         "releasetools_build_super_image",
         "releasetools_common",
+        "libavbtool",
     ],
     required: [
         "care_map_generator",
diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index b177266..465e1a8 100644
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -46,6 +46,7 @@
 
 from __future__ import print_function
 
+import avbtool
 import datetime
 import logging
 import os
@@ -62,12 +63,12 @@
 import common
 import verity_utils
 import ota_metadata_pb2
-
-from apex_utils import GetApexInfoFromTargetFiles
-from common import ZipDelete, PARTITIONS_WITH_CARE_MAP, ExternalError, RunAndCheckOutput, MakeTempFile, ZipWrite
 import rangelib
 import sparse_img
 
+from apex_utils import GetApexInfoFromTargetFiles
+from common import ZipDelete, PARTITIONS_WITH_CARE_MAP, ExternalError, RunAndCheckOutput, IsSparseImage, MakeTempFile, ZipWrite
+
 if sys.hexversion < 0x02070000:
   print("Python 2.7 or newer is required.", file=sys.stderr)
   sys.exit(1)
@@ -87,6 +88,13 @@
     datetime.datetime.utcfromtimestamp(0)).total_seconds())
 
 
+def ParseAvbFooter(img_path) -> avbtool.AvbFooter:
+  with open(img_path, 'rb') as fp:
+    fp.seek(-avbtool.AvbFooter.SIZE, os.SEEK_END)
+    data = fp.read(avbtool.AvbFooter.SIZE)
+    return avbtool.AvbFooter(data)
+
+
 def GetCareMap(which, imgname):
   """Returns the care_map string for the given partition.
 
@@ -100,15 +108,35 @@
   """
   assert which in PARTITIONS_WITH_CARE_MAP
 
-  # which + "_image_size" contains the size that the actual filesystem image
-  # resides in, which is all that needs to be verified. The additional blocks in
-  # the image file contain verity metadata, by reading which would trigger
-  # invalid reads.
-  image_size = OPTIONS.info_dict.get(which + "_image_size")
-  if not image_size:
+  is_sparse_img = IsSparseImage(imgname)
+  unsparsed_image_size = os.path.getsize(imgname)
+
+  # A verified image contains original image + hash tree data + FEC data
+  # + AVB footer, all concatenated together. The caremap specifies a range
+  # of blocks that update_verifier should read on top of dm-verity device
+  # to verify correctness of OTA updates. When reading off of dm-verity device,
+  # the hashtree and FEC part of image isn't available. So caremap should
+  # only contain the original image blocks.
+  try:
+    avbfooter = None
+    if is_sparse_img:
+      with tempfile.NamedTemporaryFile() as tmpfile:
+        img = sparse_img.SparseImage(imgname)
+        unsparsed_image_size = img.total_blocks * img.blocksize
+        for data in img.ReadBlocks(img.total_blocks - 1, 1):
+          tmpfile.write(data)
+        tmpfile.flush()
+        avbfooter = ParseAvbFooter(tmpfile.name)
+    else:
+      avbfooter = ParseAvbFooter(imgname)
+  except LookupError as e:
+    logger.warning(
+        "Failed to parse avbfooter for partition %s image %s, %s", which, imgname, e)
     return None
 
-  disable_sparse = OPTIONS.info_dict.get(which + "_disable_sparse")
+  image_size = avbfooter.original_image_size
+  assert image_size < unsparsed_image_size, f"AVB footer's original image size {image_size} is larger than or equal to image size on disk {unsparsed_image_size}, this can't happen because a verified image = original image + hash tree data + FEC data + avbfooter."
+  assert image_size > 0
 
   image_blocks = int(image_size) // 4096 - 1
   # It's OK for image_blocks to be 0, because care map ranges are inclusive.
@@ -118,7 +146,7 @@
 
   # For sparse images, we will only check the blocks that are listed in the care
   # map, i.e. the ones with meaningful data.
-  if "extfs_sparse_flag" in OPTIONS.info_dict and not disable_sparse:
+  if is_sparse_img:
     simg = sparse_img.SparseImage(imgname)
     care_map_ranges = simg.care_map.intersect(
         rangelib.RangeSet("0-{}".format(image_blocks)))
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index cc75d2c..bbdff6e 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -2865,7 +2865,7 @@
   zipfile.ZIP64_LIMIT = saved_zip64_limit
 
 
-def ZipDelete(zip_filename, entries):
+def ZipDelete(zip_filename, entries, force=False):
   """Deletes entries from a ZIP file.
 
   Since deleting entries from a ZIP file is not supported, it shells out to
@@ -2883,8 +2883,15 @@
   # If list is empty, nothing to do
   if not entries:
     return
-  cmd = ["zip", "-d", zip_filename] + entries
-  RunAndCheckOutput(cmd)
+  if force:
+    cmd = ["zip", "-q", "-d", zip_filename] + entries
+  else:
+    cmd = ["zip", "-d", zip_filename] + entries
+  if force:
+    p = Run(cmd)
+    p.wait()
+  else:
+    RunAndCheckOutput(cmd)
 
 
 def ZipClose(zip_file):
@@ -3979,6 +3986,8 @@
 
 
 def IsSparseImage(filepath):
+  if not os.path.exists(filepath):
+    return False
   with open(filepath, 'rb') as fp:
     # Magic for android sparse image format
     # https://source.android.com/devices/bootloader/images
diff --git a/tools/releasetools/sparse_img.py b/tools/releasetools/sparse_img.py
index e824a64..a2f7e9e 100644
--- a/tools/releasetools/sparse_img.py
+++ b/tools/releasetools/sparse_img.py
@@ -80,7 +80,7 @@
     self.offset_map = offset_map = []
     self.clobbered_blocks = rangelib.RangeSet(data=clobbered_blocks)
 
-    for i in range(total_chunks):
+    for _ in range(total_chunks):
       header_bin = f.read(12)
       header = struct.unpack("<2H2I", header_bin)
       chunk_type = header[0]
@@ -166,6 +166,11 @@
   def ReadRangeSet(self, ranges):
     return [d for d in self._GetRangeData(ranges)]
 
+  def ReadBlocks(self, start=0, num_blocks=None):
+    if num_blocks is None:
+      num_blocks = self.total_blocks
+    return self._GetRangeData([(start, start + num_blocks)])
+
   def TotalSha1(self, include_clobbered_blocks=False):
     """Return the SHA-1 hash of all data in the 'care' regions.
 
diff --git a/tools/releasetools/test_add_img_to_target_files.py b/tools/releasetools/test_add_img_to_target_files.py
index 5077b38..7b5476d 100644
--- a/tools/releasetools/test_add_img_to_target_files.py
+++ b/tools/releasetools/test_add_img_to_target_files.py
@@ -16,6 +16,7 @@
 
 import os
 import os.path
+import tempfile
 import zipfile
 
 import common
@@ -124,9 +125,6 @@
   def _test_AddCareMapForAbOta():
     """Helper function to set up the test for test_AddCareMapForAbOta()."""
     OPTIONS.info_dict = {
-        'extfs_sparse_flag': '-s',
-        'system_image_size': 65536,
-        'vendor_image_size': 40960,
         'system_verity_block_device': '/dev/block/system',
         'vendor_verity_block_device': '/dev/block/vendor',
         'system.build.prop': common.PartitionBuildProps.FromDictionary(
@@ -149,9 +147,9 @@
     system_image = test_utils.construct_sparse_image([
         (0xCAC1, 6),
         (0xCAC3, 4),
-        (0xCAC1, 8)])
+        (0xCAC1, 6)], "system")
     vendor_image = test_utils.construct_sparse_image([
-        (0xCAC2, 12)])
+        (0xCAC2, 10)], "vendor")
 
     image_paths = {
         'system': system_image,
@@ -210,9 +208,6 @@
     """Tests the case for device using AVB."""
     image_paths = self._test_AddCareMapForAbOta()
     OPTIONS.info_dict = {
-        'extfs_sparse_flag': '-s',
-        'system_image_size': 65536,
-        'vendor_image_size': 40960,
         'avb_system_hashtree_enable': 'true',
         'avb_vendor_hashtree_enable': 'true',
         'system.build.prop': common.PartitionBuildProps.FromDictionary(
@@ -244,9 +239,6 @@
     """Tests the case for partitions without fingerprint."""
     image_paths = self._test_AddCareMapForAbOta()
     OPTIONS.info_dict = {
-        'extfs_sparse_flag': '-s',
-        'system_image_size': 65536,
-        'vendor_image_size': 40960,
         'system_verity_block_device': '/dev/block/system',
         'vendor_verity_block_device': '/dev/block/vendor',
     }
@@ -266,9 +258,6 @@
     """Tests the case for partitions with thumbprint."""
     image_paths = self._test_AddCareMapForAbOta()
     OPTIONS.info_dict = {
-        'extfs_sparse_flag': '-s',
-        'system_image_size': 65536,
-        'vendor_image_size': 40960,
         'system_verity_block_device': '/dev/block/system',
         'vendor_verity_block_device': '/dev/block/vendor',
         'system.build.prop': common.PartitionBuildProps.FromDictionary(
@@ -298,9 +287,7 @@
   @test_utils.SkipIfExternalToolsUnavailable()
   def test_AddCareMapForAbOta_skipPartition(self):
     image_paths = self._test_AddCareMapForAbOta()
-
-    # Remove vendor_image_size to invalidate the care_map for vendor.img.
-    del OPTIONS.info_dict['vendor_image_size']
+    test_utils.erase_avb_footer(image_paths["vendor"])
 
     care_map_file = os.path.join(OPTIONS.input_tmp, 'META', 'care_map.pb')
     AddCareMapForAbOta(care_map_file, ['system', 'vendor'], image_paths)
@@ -314,10 +301,8 @@
   @test_utils.SkipIfExternalToolsUnavailable()
   def test_AddCareMapForAbOta_skipAllPartitions(self):
     image_paths = self._test_AddCareMapForAbOta()
-
-    # Remove the image_size properties for all the partitions.
-    del OPTIONS.info_dict['system_image_size']
-    del OPTIONS.info_dict['vendor_image_size']
+    test_utils.erase_avb_footer(image_paths["system"])
+    test_utils.erase_avb_footer(image_paths["vendor"])
 
     care_map_file = os.path.join(OPTIONS.input_tmp, 'META', 'care_map.pb')
     AddCareMapForAbOta(care_map_file, ['system', 'vendor'], image_paths)
@@ -396,35 +381,18 @@
     sparse_image = test_utils.construct_sparse_image([
         (0xCAC1, 6),
         (0xCAC3, 4),
-        (0xCAC1, 6)])
-    OPTIONS.info_dict = {
-        'extfs_sparse_flag': '-s',
-        'system_image_size': 53248,
-    }
+        (0xCAC1, 6)], "system")
     name, care_map = GetCareMap('system', sparse_image)
     self.assertEqual('system', name)
-    self.assertEqual(RangeSet("0-5 10-12").to_string_raw(), care_map)
+    self.assertEqual(RangeSet("0-5 10-15").to_string_raw(), care_map)
 
   def test_GetCareMap_invalidPartition(self):
     self.assertRaises(AssertionError, GetCareMap, 'oem', None)
 
-  def test_GetCareMap_invalidAdjustedPartitionSize(self):
-    sparse_image = test_utils.construct_sparse_image([
-        (0xCAC1, 6),
-        (0xCAC3, 4),
-        (0xCAC1, 6)])
-    OPTIONS.info_dict = {
-        'extfs_sparse_flag': '-s',
-        'system_image_size': -45056,
-    }
-    self.assertRaises(AssertionError, GetCareMap, 'system', sparse_image)
-
   def test_GetCareMap_nonSparseImage(self):
-    OPTIONS.info_dict = {
-        'system_image_size': 53248,
-    }
-    # 'foo' is the image filename, which is expected to be not used by
-    # GetCareMap().
-    name, care_map = GetCareMap('system', 'foo')
-    self.assertEqual('system', name)
-    self.assertEqual(RangeSet("0-12").to_string_raw(), care_map)
+    with tempfile.NamedTemporaryFile() as tmpfile:
+      tmpfile.truncate(4096 * 13)
+      test_utils.append_avb_footer(tmpfile.name, "system")
+      name, care_map = GetCareMap('system', tmpfile.name)
+      self.assertEqual('system', name)
+      self.assertEqual(RangeSet("0-12").to_string_raw(), care_map)
diff --git a/tools/releasetools/test_utils.py b/tools/releasetools/test_utils.py
index e30d2b9..5bbcf7f 100755
--- a/tools/releasetools/test_utils.py
+++ b/tools/releasetools/test_utils.py
@@ -19,6 +19,7 @@
 Utils for running unittests.
 """
 
+import avbtool
 import logging
 import os
 import os.path
@@ -57,12 +58,14 @@
   current_dir = os.path.dirname(os.path.realpath(__file__))
   return os.path.join(current_dir, 'testdata')
 
+
 def get_current_dir():
   """Returns the current dir, relative to the script dir."""
   # The script dir is the one we want, which could be different from pwd.
   current_dir = os.path.dirname(os.path.realpath(__file__))
   return current_dir
 
+
 def get_search_path():
   """Returns the search path that has 'framework/signapk.jar' under."""
 
@@ -83,14 +86,33 @@
       # In relative to 'build/make/tools/releasetools' in the Android source.
       ['..'] * 4 + ['out', 'host', 'linux-x86'],
       # Or running the script unpacked from otatools.zip.
-      ['..']):
+          ['..']):
     full_path = os.path.realpath(os.path.join(current_dir, *path))
     if signapk_exists(full_path):
       return full_path
   return None
 
 
-def construct_sparse_image(chunks):
+def append_avb_footer(file_path: str, partition_name: str = ""):
+  avb = avbtool.AvbTool()
+  try:
+    args = ["avbtool", "add_hashtree_footer", "--image", file_path,
+            "--partition_name", partition_name, "--do_not_generate_fec"]
+    avb.run(args)
+  except SystemExit:
+    raise ValueError(f"Failed to append hashtree footer {args}")
+
+
+def erase_avb_footer(file_path: str):
+  avb = avbtool.AvbTool()
+  try:
+    args = ["avbtool", "erase_footer", "--image", file_path]
+    avb.run(args)
+  except SystemExit:
+    raise ValueError(f"Failed to erase hashtree footer {args}")
+
+
+def construct_sparse_image(chunks, partition_name: str = ""):
   """Returns a sparse image file constructed from the given chunks.
 
   From system/core/libsparse/sparse_format.h.
@@ -151,6 +173,7 @@
       if data_size != 0:
         fp.write(os.urandom(data_size))
 
+  append_avb_footer(sparse_image, partition_name)
   return sparse_image
 
 
@@ -201,6 +224,7 @@
   def tearDown(self):
     common.Cleanup()
 
+
 class PropertyFilesTestCase(ReleaseToolsTestCase):
 
   @staticmethod