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))
diff --git a/tools/releasetools/ota_from_target_files.py b/tools/releasetools/ota_from_target_files.py
index 7ec8ad8..3fbcbcf 100755
--- a/tools/releasetools/ota_from_target_files.py
+++ b/tools/releasetools/ota_from_target_files.py
@@ -826,32 +826,51 @@
   # See the notes in WriteBlockIncrementalOTAPackage().
   allow_shared_blocks = target_info.get('ext4_share_dup_blocks') == "true"
 
-  # Full OTA is done as an "incremental" against an empty source image. This
-  # has the effect of writing new data from the package to the entire
-  # partition, but lets us reuse the updater code that writes incrementals to
-  # do it.
-  system_tgt = common.GetSparseImage("system", OPTIONS.input_tmp, input_zip,
-                                     allow_shared_blocks)
-  system_tgt.ResetFileMap()
-  system_diff = common.BlockDifference("system", system_tgt, src=None)
-  system_diff.WriteScript(script, output_zip,
-                          write_verify_script=OPTIONS.verify)
+  def GetBlockDifference(partition):
+    # Full OTA is done as an "incremental" against an empty source image. This
+    # has the effect of writing new data from the package to the entire
+    # partition, but lets us reuse the updater code that writes incrementals to
+    # do it.
+    tgt = common.GetSparseImage(partition, OPTIONS.input_tmp, input_zip,
+                                allow_shared_blocks)
+    tgt.ResetFileMap()
+    diff = common.BlockDifference(partition, tgt, src=None)
+    return diff
 
-  boot_img = common.GetBootableImage(
-      "boot.img", "boot.img", OPTIONS.input_tmp, "BOOT")
+  device_specific_diffs = device_specific.FullOTA_GetBlockDifferences()
+  if device_specific_diffs:
+    assert all(isinstance(diff, common.BlockDifference)
+               for diff in device_specific_diffs), \
+        "FullOTA_GetBlockDifferences is not returning a list of " \
+        "BlockDifference objects"
 
+  progress_dict = dict()
+  block_diffs = [GetBlockDifference("system")]
   if HasVendorPartition(input_zip):
-    script.ShowProgress(0.1, 0)
+    block_diffs.append(GetBlockDifference("vendor"))
+    progress_dict["vendor"] = 0.1
+  if device_specific_diffs:
+    block_diffs += device_specific_diffs
 
-    vendor_tgt = common.GetSparseImage("vendor", OPTIONS.input_tmp, input_zip,
-                                       allow_shared_blocks)
-    vendor_tgt.ResetFileMap()
-    vendor_diff = common.BlockDifference("vendor", vendor_tgt)
-    vendor_diff.WriteScript(script, output_zip,
-                            write_verify_script=OPTIONS.verify)
+  if target_info.get('use_dynamic_partitions') == "true":
+    # Use empty source_info_dict to indicate that all partitions / groups must
+    # be re-added.
+    dynamic_partitions_diff = common.DynamicPartitionsDifference(
+        info_dict=OPTIONS.info_dict,
+        block_diffs=block_diffs,
+        progress_dict=progress_dict)
+    dynamic_partitions_diff.WriteScript(script, output_zip,
+                                        write_verify_script=OPTIONS.verify)
+  else:
+    for block_diff in block_diffs:
+      block_diff.WriteScript(script, output_zip,
+                             progress=progress_dict.get(block_diff.partition),
+                             write_verify_script=OPTIONS.verify)
 
   AddCompatibilityArchiveIfTrebleEnabled(input_zip, output_zip, target_info)
 
+  boot_img = common.GetBootableImage(
+      "boot.img", "boot.img", OPTIONS.input_tmp, "BOOT")
   common.CheckSize(boot_img.data, "boot.img", target_info)
   common.ZipWriteStr(output_zip, "boot.img", boot_img.data)
 
@@ -1571,18 +1590,43 @@
   system_diff.WriteVerifyScript(script, touched_blocks_only=True)
   if vendor_diff:
     vendor_diff.WriteVerifyScript(script, touched_blocks_only=True)
+  device_specific_diffs = device_specific.IncrementalOTA_GetBlockDifferences()
+  if device_specific_diffs:
+    assert all(isinstance(diff, common.BlockDifference)
+               for diff in device_specific_diffs), \
+        "IncrementalOTA_GetBlockDifferences is not returning a list of " \
+        "BlockDifference objects"
+    for diff in device_specific_diffs:
+      diff.WriteVerifyScript(script, touched_blocks_only=True)
 
   script.Comment("---- start making changes here ----")
 
   device_specific.IncrementalOTA_InstallBegin()
 
-  system_diff.WriteScript(script, output_zip,
-                          progress=0.8 if vendor_diff else 0.9,
-                          write_verify_script=OPTIONS.verify)
-
+  block_diffs = [system_diff]
+  progress_dict = {"system": 0.8 if vendor_diff else 0.9}
   if vendor_diff:
-    vendor_diff.WriteScript(script, output_zip, progress=0.1,
-                            write_verify_script=OPTIONS.verify)
+    block_diffs.append(vendor_diff)
+    progress_dict["vendor"] = 0.1
+  if device_specific_diffs:
+    block_diffs += device_specific_diffs
+
+  if OPTIONS.source_info_dict.get("use_dynamic_partitions") == "true":
+    if OPTIONS.target_info_dict.get("use_dynamic_partitions") != "true":
+      raise RuntimeError(
+          "can't generate incremental that disables dynamic partitions")
+    dynamic_partitions_diff = common.DynamicPartitionsDifference(
+        info_dict=OPTIONS.target_info_dict,
+        source_info_dict=OPTIONS.source_info_dict,
+        block_diffs=block_diffs,
+        progress_dict=progress_dict)
+    dynamic_partitions_diff.WriteScript(
+        script, output_zip, write_verify_script=OPTIONS.verify)
+  else:
+    for block_diff in block_diffs:
+      block_diff.WriteScript(script, output_zip,
+                             progress=progress_dict.get(block_diff.partition),
+                             write_verify_script=OPTIONS.verify)
 
   if OPTIONS.two_step:
     common.ZipWriteStr(output_zip, "boot.img", target_boot.data)