Generate OTA for non-A/B devices with dynamic partitions

Test: sideload full OTA on cuttlefish
Test: sideload incremental OTA on cuttlefish (that grows
      system, shrinks vendor, and move vendor to group foo)

Bug: 111801737

Change-Id: Ie8a267a90b4df9e9e0a2fbcc1b582ab2e353df52
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index dcc083c..be2c108 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -14,6 +14,7 @@
 
 from __future__ import print_function
 
+import collections
 import copy
 import errno
 import getopt
@@ -1523,6 +1524,13 @@
     """Called at the start of full OTA installation."""
     return self._DoCall("FullOTA_InstallBegin")
 
+  def FullOTA_GetBlockDifferences(self):
+    """Called during full OTA installation and verification.
+    Implementation should return a list of BlockDifference objects describing
+    the update on each additional partitions.
+    """
+    return self._DoCall("FullOTA_GetBlockDifferences")
+
   def FullOTA_InstallEnd(self):
     """Called at the end of full OTA installation; typically this is
     used to install the image for the device's baseband processor."""
@@ -1551,6 +1559,13 @@
     verification is complete)."""
     return self._DoCall("IncrementalOTA_InstallBegin")
 
+  def IncrementalOTA_GetBlockDifferences(self):
+    """Called during incremental OTA installation and verification.
+    Implementation should return a list of BlockDifference objects describing
+    the update on each additional partitions.
+    """
+    return self._DoCall("IncrementalOTA_GetBlockDifferences")
+
   def IncrementalOTA_InstallEnd(self):
     """Called at the end of incremental OTA installation; typically
     this is used to install the image for the device's baseband
@@ -1745,11 +1760,29 @@
     self.touched_src_ranges = b.touched_src_ranges
     self.touched_src_sha1 = b.touched_src_sha1
 
-    if src is None:
-      _, self.device = GetTypeAndDevice("/" + partition, OPTIONS.info_dict)
+    # On devices with dynamic partitions, for new partitions,
+    # src is None but OPTIONS.source_info_dict is not.
+    if OPTIONS.source_info_dict is None:
+      is_dynamic_build = OPTIONS.info_dict.get(
+          "use_dynamic_partitions") == "true"
     else:
-      _, self.device = GetTypeAndDevice("/" + partition,
-                                        OPTIONS.source_info_dict)
+      is_dynamic_build = OPTIONS.source_info_dict.get(
+          "use_dynamic_partitions") == "true"
+
+    # For dynamic partitions builds, always check partition list in target build
+    # because new partitions may be added.
+    is_dynamic = is_dynamic_build and partition in shlex.split(
+        OPTIONS.info_dict.get("dynamic_partition_list", "").strip())
+
+    if is_dynamic:
+      self.device = 'map_partition("%s")' % partition
+    else:
+      if OPTIONS.source_info_dict is None:
+        _, device_path = GetTypeAndDevice("/" + partition, OPTIONS.info_dict)
+      else:
+        _, device_path = GetTypeAndDevice("/" + partition,
+                                          OPTIONS.source_info_dict)
+      self.device = '"%s"' % device_path
 
   @property
   def required_cache(self):
@@ -1768,7 +1801,7 @@
     self._WriteUpdate(script, output_zip)
 
     if write_verify_script:
-      self._WritePostInstallVerifyScript(script)
+      self.WritePostInstallVerifyScript(script)
 
   def WriteStrictVerifyScript(self, script):
     """Verify all the blocks in the care_map, including clobbered blocks.
@@ -1782,11 +1815,11 @@
     ranges = self.tgt.care_map
     ranges_str = ranges.to_string_raw()
     script.AppendExtra(
-        'range_sha1("%s", "%s") == "%s" && ui_print("    Verified.") || '
-        'ui_print("\\"%s\\" has unexpected contents.");' % (
+        'range_sha1(%s, "%s") == "%s" && ui_print("    Verified.") || '
+        'ui_print("%s has unexpected contents.");' % (
             self.device, ranges_str,
             self.tgt.TotalSha1(include_clobbered_blocks=True),
-            self.device))
+            self.partition))
     script.AppendExtra("")
 
   def WriteVerifyScript(self, script, touched_blocks_only=False):
@@ -1811,7 +1844,7 @@
 
       ranges_str = ranges.to_string_raw()
       script.AppendExtra(
-          'if (range_sha1("%s", "%s") == "%s" || block_image_verify("%s", '
+          'if (range_sha1(%s, "%s") == "%s" || block_image_verify(%s, '
           'package_extract_file("%s.transfer.list"), "%s.new.dat", '
           '"%s.patch.dat")) then' % (
               self.device, ranges_str, expected_sha1,
@@ -1828,7 +1861,7 @@
         # this check fails, give an explicit log message about the partition
         # having been remounted R/W (the most likely explanation).
         if self.check_first_block:
-          script.AppendExtra('check_first_block("%s");' % (self.device,))
+          script.AppendExtra('check_first_block(%s);' % (self.device,))
 
         # If version >= 4, try block recovery before abort update
         if partition == "system":
@@ -1836,8 +1869,8 @@
         else:
           code = ErrorCode.VENDOR_RECOVER_FAILURE
         script.AppendExtra((
-            'ifelse (block_image_recover("{device}", "{ranges}") && '
-            'block_image_verify("{device}", '
+            'ifelse (block_image_recover({device}, "{ranges}") && '
+            'block_image_verify({device}, '
             'package_extract_file("{partition}.transfer.list"), '
             '"{partition}.new.dat", "{partition}.patch.dat"), '
             'ui_print("{partition} recovered successfully."), '
@@ -1859,14 +1892,14 @@
             'abort("E%d: %s partition has unexpected contents");\n'
             'endif;') % (code, partition))
 
-  def _WritePostInstallVerifyScript(self, script):
+  def WritePostInstallVerifyScript(self, script):
     partition = self.partition
     script.Print('Verifying the updated %s image...' % (partition,))
     # Unlike pre-install verification, clobbered_blocks should not be ignored.
     ranges = self.tgt.care_map
     ranges_str = ranges.to_string_raw()
     script.AppendExtra(
-        'if range_sha1("%s", "%s") == "%s" then' % (
+        'if range_sha1(%s, "%s") == "%s" then' % (
             self.device, ranges_str,
             self.tgt.TotalSha1(include_clobbered_blocks=True)))
 
@@ -1875,7 +1908,7 @@
     if self.tgt.extended:
       ranges_str = self.tgt.extended.to_string_raw()
       script.AppendExtra(
-          'if range_sha1("%s", "%s") == "%s" then' % (
+          'if range_sha1(%s, "%s") == "%s" then' % (
               self.device, ranges_str,
               self._HashZeroBlocks(self.tgt.extended.size())))
       script.Print('Verified the updated %s image.' % (partition,))
@@ -1941,7 +1974,7 @@
     else:
       code = ErrorCode.VENDOR_UPDATE_FAILURE
 
-    call = ('block_image_update("{device}", '
+    call = ('block_image_update({device}, '
             'package_extract_file("{partition}.transfer.list"), '
             '"{new_data_name}", "{partition}.patch.dat") ||\n'
             '  abort("E{code}: Failed to update {partition} image.");'.format(
@@ -2134,3 +2167,209 @@
   logger.info("putting script in %s", sh_location)
 
   output_sink(sh_location, sh)
+
+
+class DynamicPartitionUpdate(object):
+  def __init__(self, src_group=None, tgt_group=None, progress=None,
+               block_difference=None):
+    self.src_group = src_group
+    self.tgt_group = tgt_group
+    self.progress = progress
+    self.block_difference = block_difference
+
+  @property
+  def src_size(self):
+    if not self.block_difference:
+      return 0
+    return DynamicPartitionUpdate._GetSparseImageSize(self.block_difference.src)
+
+  @property
+  def tgt_size(self):
+    if not self.block_difference:
+      return 0
+    return DynamicPartitionUpdate._GetSparseImageSize(self.block_difference.tgt)
+
+  @staticmethod
+  def _GetSparseImageSize(img):
+    if not img:
+      return 0
+    return img.blocksize * img.total_blocks
+
+
+class DynamicGroupUpdate(object):
+  def __init__(self, src_size=None, tgt_size=None):
+    # None: group does not exist. 0: no size limits.
+    self.src_size = src_size
+    self.tgt_size = tgt_size
+
+
+class DynamicPartitionsDifference(object):
+  def __init__(self, info_dict, block_diffs, progress_dict=None,
+               source_info_dict=None):
+    if progress_dict is None:
+      progress_dict = dict()
+
+    self._remove_all_before_apply = False
+    if source_info_dict is None:
+      self._remove_all_before_apply = True
+      source_info_dict = dict()
+
+    block_diff_dict = {e.partition:e for e in block_diffs}
+    assert len(block_diff_dict) == len(block_diffs), \
+        "Duplicated BlockDifference object for {}".format(
+            [partition for partition, count in
+             collections.Counter(e.partition for e in block_diffs).items()
+             if count > 1])
+
+    dynamic_partitions = set(shlex.split(info_dict.get(
+        "dynamic_partition_list", "").strip()))
+    assert set(block_diff_dict.keys()) == dynamic_partitions, \
+        "Dynamic partitions: {}, BlockDifference objects: {}".format(
+            list(dynamic_partitions), list(block_diff_dict.keys()))
+
+    self._partition_updates = dict()
+
+    for p, block_diff in block_diff_dict.items():
+      self._partition_updates[p] = DynamicPartitionUpdate()
+      self._partition_updates[p].block_difference = block_diff
+
+    for p, progress in progress_dict.items():
+      if p in self._partition_updates:
+        self._partition_updates[p].progress = progress
+
+    tgt_groups = shlex.split(info_dict.get(
+        "super_partition_groups", "").strip())
+    src_groups = shlex.split(source_info_dict.get(
+        "super_partition_groups", "").strip())
+
+    for g in tgt_groups:
+      for p in shlex.split(info_dict.get(
+          "super_%s_partition_list" % g, "").strip()):
+        assert p in self._partition_updates, \
+            "{} is in target super_{}_partition_list but no BlockDifference " \
+            "object is provided.".format(p, g)
+        self._partition_updates[p].tgt_group = g
+
+    for g in src_groups:
+      for p in shlex.split(source_info_dict.get(
+          "super_%s_partition_list" % g, "").strip()):
+        assert p in self._partition_updates, \
+            "{} is in source super_{}_partition_list but no BlockDifference " \
+            "object is provided.".format(p, g)
+        self._partition_updates[p].src_group = g
+
+    if self._partition_updates:
+      logger.info("Updating dynamic partitions %s",
+                  self._partition_updates.keys())
+
+    self._group_updates = dict()
+
+    for g in tgt_groups:
+      self._group_updates[g] = DynamicGroupUpdate()
+      self._group_updates[g].tgt_size = int(info_dict.get(
+          "super_%s_group_size" % g, "0").strip())
+
+    for g in src_groups:
+      if g not in self._group_updates:
+        self._group_updates[g] = DynamicGroupUpdate()
+      self._group_updates[g].src_size = int(source_info_dict.get(
+          "super_%s_group_size" % g, "0").strip())
+
+    self._Compute()
+
+  def WriteScript(self, script, output_zip, write_verify_script=False):
+    script.Comment('--- Start patching dynamic partitions ---')
+    for p, u in self._partition_updates.items():
+      if u.src_size and u.tgt_size and u.src_size > u.tgt_size:
+        script.Comment('Patch partition %s' % p)
+        u.block_difference.WriteScript(script, output_zip, progress=u.progress,
+                                       write_verify_script=False)
+
+    op_list_path = MakeTempFile()
+    with open(op_list_path, 'w') as f:
+      for line in self._op_list:
+        f.write('{}\n'.format(line))
+
+    ZipWrite(output_zip, op_list_path, "dynamic_partitions_op_list")
+
+    script.Comment('Update dynamic partition metadata')
+    script.AppendExtra('assert(update_dynamic_partitions('
+                       'package_extract_file("dynamic_partitions_op_list")));')
+
+    if write_verify_script:
+      for p, u in self._partition_updates.items():
+        if u.src_size and u.tgt_size and u.src_size > u.tgt_size:
+          u.block_difference.WritePostInstallVerifyScript(script)
+          script.AppendExtra('unmap_partition("%s");' % p) # ignore errors
+
+    for p, u in self._partition_updates.items():
+      if u.tgt_size and u.src_size <= u.tgt_size:
+        script.Comment('Patch partition %s' % p)
+        u.block_difference.WriteScript(script, output_zip, progress=u.progress,
+                                       write_verify_script=write_verify_script)
+        if write_verify_script:
+          script.AppendExtra('unmap_partition("%s");' % p) # ignore errors
+
+    script.Comment('--- End patching dynamic partitions ---')
+
+  def _Compute(self):
+    self._op_list = list()
+
+    def append(line):
+      self._op_list.append(line)
+
+    def comment(line):
+      self._op_list.append("# %s" % line)
+
+    if self._remove_all_before_apply:
+      comment('Remove all existing dynamic partitions and groups before '
+              'applying full OTA')
+      append('remove_all_groups')
+
+    for p, u in self._partition_updates.items():
+      if u.src_group and not u.tgt_group:
+        append('remove %s' % p)
+
+    for p, u in self._partition_updates.items():
+      if u.src_group and u.tgt_group and u.src_group != u.tgt_group:
+        comment('Move partition %s from %s to default' % (p, u.src_group))
+        append('move %s default' % p)
+
+    for p, u in self._partition_updates.items():
+      if u.src_size and u.tgt_size and u.src_size > u.tgt_size:
+        comment('Shrink partition %s from %d to %d' %
+                (p, u.src_size, u.tgt_size))
+        append('resize %s %s' % (p, u.tgt_size))
+
+    for g, u in self._group_updates.items():
+      if u.src_size is not None and u.tgt_size is None:
+        append('remove_group %s' % g)
+      if (u.src_size is not None and u.tgt_size is not None and
+          u.src_size > u.tgt_size):
+        comment('Shrink group %s from %d to %d' % (g, u.src_size, u.tgt_size))
+        append('resize_group %s %d' % (g, u.tgt_size))
+
+    for g, u in self._group_updates.items():
+      if u.src_size is None and u.tgt_size is not None:
+        comment('Add group %s with maximum size %d' % (g, u.tgt_size))
+        append('add_group %s %d' % (g, u.tgt_size))
+      if (u.src_size is not None and u.tgt_size is not None and
+          u.src_size < u.tgt_size):
+        comment('Grow group %s from %d to %d' % (g, u.src_size, u.tgt_size))
+        append('resize_group %s %d' % (g, u.tgt_size))
+
+    for p, u in self._partition_updates.items():
+      if u.tgt_group and not u.src_group:
+        comment('Add partition %s to group %s' % (p, u.tgt_group))
+        append('add %s %s' % (p, u.tgt_group))
+
+    for p, u in self._partition_updates.items():
+      if u.tgt_size and u.src_size < u.tgt_size:
+        comment('Grow partition %s from %d to %d' % (p, u.src_size, u.tgt_size))
+        append('resize %s %d' % (p, u.tgt_size))
+
+    for p, u in self._partition_updates.items():
+      if u.src_group and u.tgt_group and u.src_group != u.tgt_group:
+        comment('Move partition %s from default to %s' %
+                (p, u.tgt_group))
+        append('move %s %s' % (p, u.tgt_group))