update_payload: Implement applying for major version 2 payloads

This commit adds payload major version 2 support to paycheck.py
applying.

BUG=b:794404
TEST=no errors during run_unittests and paycheck.py <major version 2
  payload> --part_names boot system --out_dst_part_paths /tmp/boot_part
  /tmp/system_part (./test_paycheck.sh does not pass for major version 2
  payloads since it currently does not detect version 2 payloads, and
  specifies rootfs/kernel as the partitions to paycheck.py instead of
  system/boot; no regressions when running on major version 1 payloads)

Change-Id: Ic411607cee6f032851d1fa9545bed68fe2d3da77
Reviewed-on: https://chromium-review.googlesource.com/1106656
Commit-Ready: Tudor Brindus <tbrindus@chromium.org>
Tested-by: Tudor Brindus <tbrindus@chromium.org>
Reviewed-by: Amin Hassani <ahassani@chromium.org>

1 files changed, 44 insertions, 29 deletions

diff --git a/scripts/update_payload/applier.py b/scripts/update_payload/applier.py
index dad5ba30..49db8aee 100644
--- a/scripts/update_payload/applier.py
+++ b/scripts/update_payload/applier.py
@@ -632,39 +632,54 @@ class PayloadApplier(object):
     Raises:
       PayloadError if payload application failed.
     """
+    if old_parts is None:
+      old_parts = {}
+
     self.payload.ResetFile()
 
-    # TODO(tbrindus): make payload applying work for major version 2 partitions
-    new_kernel_part = new_parts[common.KERNEL]
-    new_rootfs_part = new_parts[common.ROOTFS]
-    old_kernel_part = old_parts.get(common.KERNEL, None) if old_parts else None
-    old_rootfs_part = old_parts.get(common.ROOTFS, None) if old_parts else None
+    new_part_info = {}
+    old_part_info = {}
+    install_operations = []
+
+    manifest = self.payload.manifest
+    if self.payload.header.version == 1:
+      for name in common.CROS_PARTITIONS:
+        new_part_info[name] = getattr(manifest, 'new_%s_info' % name)
+        old_part_info[name] = getattr(manifest, 'old_%s_info' % name)
+
+        install_operations.append((common.ROOTFS, manifest.install_operations))
+        install_operations.append((common.KERNEL,
+                                   manifest.kernel_install_operations))
+    else:
+      for part in manifest.partitions:
+        name = part.partition_name
+        new_part_info[name] = part.new_partition_info
+        old_part_info[name] = part.old_partition_info
+        install_operations.append((name, part.operations))
+
+    part_names = set(new_part_info.keys())  # Equivalently, old_part_info.keys()
 
     # Make sure the arguments are sane and match the payload.
-    if not (new_kernel_part and new_rootfs_part):
-      raise PayloadError('missing dst {kernel,rootfs} partitions')
-
-    if not (old_kernel_part or old_rootfs_part):
-      if not self.payload.IsFull():
-        raise PayloadError('trying to apply a non-full update without src '
-                           '{kernel,rootfs} partitions')
-    elif old_kernel_part and old_rootfs_part:
-      if not self.payload.IsDelta():
-        raise PayloadError('trying to apply a non-delta update onto src '
-                           '{kernel,rootfs} partitions')
+    new_part_names = set(new_parts.keys())
+    if new_part_names != part_names:
+      raise PayloadError('missing dst partition(s) %s' %
+                         ', '.join(part_names - new_part_names))
+
+    old_part_names = set(old_parts.keys())
+    if part_names - old_part_names:
+      if self.payload.IsDelta():
+        raise PayloadError('trying to apply a delta update without src '
+                           'partition(s) %s' %
+                           ', '.join(part_names - old_part_names))
+    elif old_part_names == part_names:
+      if self.payload.IsFull():
+        raise PayloadError('trying to apply a full update onto src partitions')
     else:
       raise PayloadError('not all src partitions provided')
 
-    # Apply update to rootfs.
-    self._ApplyToPartition(
-        self.payload.manifest.install_operations, 'rootfs',
-        'install_operations', new_rootfs_part,
-        self.payload.manifest.new_rootfs_info, old_rootfs_part,
-        self.payload.manifest.old_rootfs_info)
-
-    # Apply update to kernel update.
-    self._ApplyToPartition(
-        self.payload.manifest.kernel_install_operations, 'kernel',
-        'kernel_install_operations', new_kernel_part,
-        self.payload.manifest.new_kernel_info, old_kernel_part,
-        self.payload.manifest.old_kernel_info)
+    for name, operations in install_operations:
+      # Apply update to partition.
+      self._ApplyToPartition(
+          operations, name, '%s_install_operations' % name, new_parts[name],
+          new_part_info[name], old_parts.get(name, None), old_part_info[name])
+