Use add_slot_suffix function in edify script

Whenever a device is retrieved from fstab, wrap it with
add_slot_suffix() if it has slotselect option.

Test: change fstab (changes boot image, which is a static partition),
      change system partition (a dynamic partition),
      generate incremental OTA with --force_non_ab and apply it on
      cuttlefish
Bug: 153581609
Change-Id: Id3f8e4425b65176baf1b0ff1ee07ab3d820a3a7f
(cherry picked from commit ae6e0d5d28f1480dce42d702cf6fe7460bf5dd4c)
Merged-In: Id3f8e4425b65176baf1b0ff1ee07ab3d820a3a7f
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index 9f48f78..96f93a8 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -2686,11 +2686,12 @@
       self.device = 'map_partition("%s")' % partition
     else:
       if OPTIONS.source_info_dict is None:
-        _, device_path = GetTypeAndDevice("/" + partition, OPTIONS.info_dict)
+        _, device_expr = GetTypeAndDeviceExpr("/" + partition,
+                                              OPTIONS.info_dict)
       else:
-        _, device_path = GetTypeAndDevice("/" + partition,
-                                          OPTIONS.source_info_dict)
-      self.device = '"%s"' % device_path
+        _, device_expr = GetTypeAndDeviceExpr("/" + partition,
+                                              OPTIONS.source_info_dict)
+      self.device = device_expr
 
   @property
   def required_cache(self):
@@ -2922,16 +2923,51 @@
     "squashfs": "EMMC"
 }
 
-
-def GetTypeAndDevice(mount_point, info):
+def GetTypeAndDevice(mount_point, info, check_no_slot=True):
+  """
+  Use GetTypeAndDeviceExpr whenever possible. This function is kept for
+  backwards compatibility. It aborts if the fstab entry has slotselect option
+  (unless check_no_slot is explicitly set to False).
+  """
   fstab = info["fstab"]
   if fstab:
+    if check_no_slot:
+      assert not fstab[mount_point].slotselect, \
+             "Use GetTypeAndDeviceExpr instead"
     return (PARTITION_TYPES[fstab[mount_point].fs_type],
             fstab[mount_point].device)
   else:
     raise KeyError
 
 
+def GetTypeAndDeviceExpr(mount_point, info):
+  """
+  Return the filesystem of the partition, and an edify expression that evaluates
+  to the device at runtime.
+  """
+  fstab = info["fstab"]
+  if fstab:
+    p = fstab[mount_point]
+    device_expr = '"%s"' % fstab[mount_point].device
+    if p.slotselect:
+      device_expr = 'add_slot_suffix(%s)' % device_expr
+    return (PARTITION_TYPES[fstab[mount_point].fs_type], device_expr)
+  else:
+    raise KeyError
+
+
+def GetEntryForDevice(fstab, device):
+  """
+  Returns:
+    The first entry in fstab whose device is the given value.
+  """
+  if not fstab:
+    return None
+  for mount_point in fstab:
+    if fstab[mount_point].device == device:
+      return fstab[mount_point]
+  return None
+
 def ParseCertificate(data):
   """Parses and converts a PEM-encoded certificate into DER-encoded.
 
@@ -3056,8 +3092,10 @@
   try:
     # The following GetTypeAndDevice()s need to use the path in the target
     # info_dict instead of source_info_dict.
-    boot_type, boot_device = GetTypeAndDevice("/boot", info_dict)
-    recovery_type, recovery_device = GetTypeAndDevice("/recovery", info_dict)
+    boot_type, boot_device = GetTypeAndDevice("/boot", info_dict,
+                                              check_no_slot=False)
+    recovery_type, recovery_device = GetTypeAndDevice("/recovery", info_dict,
+                                                      check_no_slot=False)
   except KeyError:
     return
 
@@ -3099,8 +3137,8 @@
        'recovery_size': recovery_img.size,
        'recovery_sha1': recovery_img.sha1,
        'boot_type': boot_type,
-       'boot_device': boot_device,
-       'recovery_type': recovery_type,
+       'boot_device': boot_device + '$(getprop ro.boot.slot_suffix)',
+       'recovery_type': recovery_type + '$(getprop ro.boot.slot_suffix)',
        'recovery_device': recovery_device,
        'bonus_args': bonus_args}
 
diff --git a/tools/releasetools/edify_generator.py b/tools/releasetools/edify_generator.py
index 7ed85fe..99e21f1 100644
--- a/tools/releasetools/edify_generator.py
+++ b/tools/releasetools/edify_generator.py
@@ -183,11 +183,30 @@
     It checks the checksums of the given partitions. If none of them matches the
     expected checksum, updater will additionally look for a backup on /cache.
     """
+    self._CheckSecondTokenNotSlotSuffixed(target, "PatchPartitionExprCheck")
+    self._CheckSecondTokenNotSlotSuffixed(source, "PatchPartitionExprCheck")
+    self.PatchPartitionExprCheck('"%s"' % target, '"%s"' % source)
+
+  def PatchPartitionExprCheck(self, target_expr, source_expr):
+    """Checks whether updater can patch the given partitions.
+
+    It checks the checksums of the given partitions. If none of them matches the
+    expected checksum, updater will additionally look for a backup on /cache.
+
+    Args:
+      target_expr: an Edify expression that serves as the target arg to
+        patch_partition. Must be evaluated to a string in the form of
+        foo:bar:baz:quux
+      source_expr: an Edify expression that serves as the source arg to
+        patch_partition. Must be evaluated to a string in the form of
+        foo:bar:baz:quux
+    """
     self.script.append(self.WordWrap((
-        'patch_partition_check("{target}",\0"{source}") ||\n    abort('
-        '"E{code}: \\"{target}\\" or \\"{source}\\" has unexpected '
-        'contents.");').format(
-            target=target, source=source,
+        'patch_partition_check({target},\0{source}) ||\n    abort('
+        'concat("E{code}: \\"",{target},"\\" or \\"",{source},"\\" has '
+        'unexpected contents."));').format(
+            target=target_expr,
+            source=source_expr,
             code=common.ErrorCode.BAD_PATCH_FILE)))
 
   def CacheFreeSpaceCheck(self, amount):
@@ -218,8 +237,9 @@
       mount_flags = mount_dict.get(p.fs_type, "")
       if p.context is not None:
         mount_flags = p.context + ("," + mount_flags if mount_flags else "")
-      self.script.append('mount("%s", "%s", "%s", "%s", "%s");' % (
-          p.fs_type, common.PARTITION_TYPES[p.fs_type], p.device,
+      self.script.append('mount("%s", "%s", %s, "%s", "%s");' % (
+          p.fs_type, common.PARTITION_TYPES[p.fs_type],
+          self._GetSlotSuffixDeviceForEntry(p),
           p.mount_point, mount_flags))
       self.mounts.add(p.mount_point)
 
@@ -242,8 +262,9 @@
         raise ValueError("Partition %s cannot be tuned\n" % (partition,))
     self.script.append(
         'tune2fs(' + "".join(['"%s", ' % (i,) for i in options]) +
-        '"%s") || abort("E%d: Failed to tune partition %s");' % (
-            p.device, common.ErrorCode.TUNE_PARTITION_FAILURE, partition))
+        '%s) || abort("E%d: Failed to tune partition %s");' % (
+            self._GetSlotSuffixDeviceForEntry(p),
+            common.ErrorCode.TUNE_PARTITION_FAILURE, partition))
 
   def FormatPartition(self, partition):
     """Format the given partition, specified by its mount point (eg,
@@ -252,18 +273,19 @@
     fstab = self.fstab
     if fstab:
       p = fstab[partition]
-      self.script.append('format("%s", "%s", "%s", "%s", "%s");' %
+      self.script.append('format("%s", "%s", %s, "%s", "%s");' %
                          (p.fs_type, common.PARTITION_TYPES[p.fs_type],
-                          p.device, p.length, p.mount_point))
+                          self._GetSlotSuffixDeviceForEntry(p),
+                          p.length, p.mount_point))
 
   def WipeBlockDevice(self, partition):
     if partition not in ("/system", "/vendor"):
       raise ValueError(("WipeBlockDevice doesn't work on %s\n") % (partition,))
     fstab = self.fstab
     size = self.info.get(partition.lstrip("/") + "_size", None)
-    device = fstab[partition].device
+    device = self._GetSlotSuffixDeviceForEntry(fstab[partition])
 
-    self.script.append('wipe_block_device("%s", %s);' % (device, size))
+    self.script.append('wipe_block_device(%s, %s);' % (device, size))
 
   def ApplyPatch(self, srcfile, tgtfile, tgtsize, tgtsha1, *patchpairs):
     """Apply binary patches (in *patchpairs) to the given srcfile to
@@ -296,14 +318,69 @@
     self.PatchPartition(target, source, patch)
 
   def PatchPartition(self, target, source, patch):
-    """Applies the patch to the source partition and writes it to target."""
+    """
+    Applies the patch to the source partition and writes it to target.
+
+    Args:
+      target: the target arg to patch_partition. Must be in the form of
+        foo:bar:baz:quux
+      source: the source arg to patch_partition. Must be in the form of
+        foo:bar:baz:quux
+      patch: the patch arg to patch_partition. Must be an unquoted string.
+    """
+    self._CheckSecondTokenNotSlotSuffixed(target, "PatchPartitionExpr")
+    self._CheckSecondTokenNotSlotSuffixed(source, "PatchPartitionExpr")
+    self.PatchPartitionExpr('"%s"' % target, '"%s"' % source, '"%s"' % patch)
+
+  def PatchPartitionExpr(self, target_expr, source_expr, patch_expr):
+    """
+    Applies the patch to the source partition and writes it to target.
+
+    Args:
+      target_expr: an Edify expression that serves as the target arg to
+        patch_partition. Must be evaluated to a string in the form of
+        foo:bar:baz:quux
+      source_expr: an Edify expression that serves as the source arg to
+        patch_partition. Must be evaluated to a string in the form of
+        foo:bar:baz:quux
+      patch_expr: an Edify expression that serves as the patch arg to
+        patch_partition. Must be evaluated to a string.
+    """
     self.script.append(self.WordWrap((
-        'patch_partition("{target}",\0"{source}",\0'
-        'package_extract_file("{patch}")) ||\n'
-        '    abort("E{code}: Failed to apply patch to {source}");').format(
-            target=target, source=source, patch=patch,
+        'patch_partition({target},\0{source},\0'
+        'package_extract_file({patch})) ||\n'
+        '    abort(concat('
+        '        "E{code}: Failed to apply patch to ",{source}));').format(
+            target=target_expr,
+            source=source_expr,
+            patch=patch_expr,
             code=common.ErrorCode.APPLY_PATCH_FAILURE)))
 
+  def _GetSlotSuffixDeviceForEntry(self, entry=None):
+    """
+    Args:
+      entry: the fstab entry of device "foo"
+    Returns:
+      An edify expression. Caller must not quote result.
+      If foo is slot suffixed, it returns
+        'add_slot_suffix("foo")'
+      Otherwise it returns
+        '"foo"' (quoted)
+    """
+    assert entry is not None
+    if entry.slotselect:
+      return 'add_slot_suffix("%s")' % entry.device
+    return '"%s"' % entry.device
+
+  def _CheckSecondTokenNotSlotSuffixed(self, s, fn):
+    lst = s.split(':')
+    assert(len(s) == 4), "{} does not contain 4 tokens".format(s)
+    if self.fstab:
+      entry = common.GetEntryForDevice(s[1])
+      if entry is not None:
+        assert not entry.slotselect, \
+          "Use %s because %s is slot suffixed" % (fn, s[1])
+
   def WriteRawImage(self, mount_point, fn, mapfn=None):
     """Write the given package file into the partition for the given
     mount point."""
@@ -312,15 +389,16 @@
     if fstab:
       p = fstab[mount_point]
       partition_type = common.PARTITION_TYPES[p.fs_type]
-      args = {'device': p.device, 'fn': fn}
+      device = self._GetSlotSuffixDeviceForEntry(p)
+      args = {'device': device, 'fn': fn}
       if partition_type == "EMMC":
         if mapfn:
           args["map"] = mapfn
           self.script.append(
-              'package_extract_file("%(fn)s", "%(device)s", "%(map)s");' % args)
+              'package_extract_file("%(fn)s", %(device)s, "%(map)s");' % args)
         else:
           self.script.append(
-              'package_extract_file("%(fn)s", "%(device)s");' % args)
+              'package_extract_file("%(fn)s", %(device)s);' % args)
       else:
         raise ValueError(
             "don't know how to write \"%s\" partitions" % p.fs_type)
diff --git a/tools/releasetools/ota_from_target_files.py b/tools/releasetools/ota_from_target_files.py
index 16d7485..16b278a 100755
--- a/tools/releasetools/ota_from_target_files.py
+++ b/tools/releasetools/ota_from_target_files.py
@@ -1463,7 +1463,8 @@
   required_cache_sizes = [diff.required_cache for diff in
                           block_diff_dict.values()]
   if updating_boot:
-    boot_type, boot_device = common.GetTypeAndDevice("/boot", source_info)
+    boot_type, boot_device_expr = common.GetTypeAndDeviceExpr("/boot",
+                                                              source_info)
     d = common.Difference(target_boot, source_boot)
     _, _, d = d.ComputePatch()
     if d is None:
@@ -1478,11 +1479,11 @@
 
       common.ZipWriteStr(output_zip, "boot.img.p", d)
 
-      script.PatchPartitionCheck(
-          "{}:{}:{}:{}".format(
-              boot_type, boot_device, target_boot.size, target_boot.sha1),
-          "{}:{}:{}:{}".format(
-              boot_type, boot_device, source_boot.size, source_boot.sha1))
+      target_expr = 'concat("{}:",{},":{}:{}")'.format(
+          boot_type, boot_device_expr, target_boot.size, target_boot.sha1)
+      source_expr = 'concat("{}:",{},":{}:{}")'.format(
+          boot_type, boot_device_expr, source_boot.size, source_boot.sha1)
+      script.PatchPartitionExprCheck(target_expr, source_expr)
 
       required_cache_sizes.append(target_boot.size)
 
@@ -1550,12 +1551,11 @@
         logger.info("boot image changed; including patch.")
         script.Print("Patching boot image...")
         script.ShowProgress(0.1, 10)
-        script.PatchPartition(
-            '{}:{}:{}:{}'.format(
-                boot_type, boot_device, target_boot.size, target_boot.sha1),
-            '{}:{}:{}:{}'.format(
-                boot_type, boot_device, source_boot.size, source_boot.sha1),
-            'boot.img.p')
+        target_expr = 'concat("{}:",{},":{}:{}")'.format(
+            boot_type, boot_device_expr, target_boot.size, target_boot.sha1)
+        source_expr = 'concat("{}:",{},":{}:{}")'.format(
+            boot_type, boot_device_expr, source_boot.size, source_boot.sha1)
+        script.PatchPartitionExpr(target_expr, source_expr, '"boot.img.p"')
     else:
       logger.info("boot image unchanged; skipping.")