update_payload: Implement checking for major version 2 payloads

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

For consistency, report messages for affected fields are kept the same across
both major version 1 and 2 checks, even if the particular field name does not
exist in one version.

BUG=b:794404
TEST=no errors during run_unittests and paycheck.py --check <major version 2
  payload> (./test_paycheck.sh does not pass for major version 2 payloads
  since applying is not implemented yet; no regressions when running on major
  version 1 payloads)

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

2 files changed, 141 insertions, 66 deletions

diff --git a/scripts/update_payload/checker.py b/scripts/update_payload/checker.py
index ec8810d9..d3230cc2 100644
--- a/scripts/update_payload/checker.py
+++ b/scripts/update_payload/checker.py
@@ -336,10 +336,7 @@ class PayloadChecker(object):
     self.new_fs_sizes = collections.defaultdict(int)
     self.old_fs_sizes = collections.defaultdict(int)
     self.minor_version = None
-    # TODO(*): When fixing crbug.com/794404, the major version should be
-    # correctly handled in update_payload scripts. So stop forcing
-    # major_verions=1 here and set it to the correct value.
-    self.major_version = 1
+    self.major_version = None
 
   @staticmethod
   def _CheckElem(msg, name, report, is_mandatory, is_submsg, convert=str,
@@ -389,6 +386,38 @@ class PayloadChecker(object):
       return element_result(value, None)
 
   @staticmethod
+  def _CheckRepeatedElemNotPresent(msg, field_name, msg_name):
+    """Checks that a repeated element is not specified in the message.
+
+    Args:
+      msg: The message containing the element.
+      field_name: The name of the element.
+      msg_name: The name of the message object (for error reporting).
+
+    Raises:
+      error.PayloadError if the repeated element is present or non-empty.
+    """
+    if getattr(msg, field_name, None):
+      raise error.PayloadError('%sfield %r not empty.' %
+                               (msg_name + ' ' if msg_name else '', field_name))
+
+  @staticmethod
+  def _CheckElemNotPresent(msg, field_name, msg_name):
+    """Checks that an element is not specified in the message.
+
+    Args:
+      msg: The message containing the element.
+      field_name: The name of the element.
+      msg_name: The name of the message object (for error reporting).
+
+    Raises:
+      error.PayloadError if the repeated element is present.
+    """
+    if msg.HasField(field_name):
+      raise error.PayloadError('%sfield %r exists.' %
+                               (msg_name + ' ' if msg_name else '', field_name))
+
+  @staticmethod
   def _CheckMandatoryField(msg, field_name, report, msg_name, convert=str,
                            linebreak=False, indent=0):
     """Adds a mandatory field; returning first component from _CheckElem."""
@@ -436,6 +465,22 @@ class PayloadChecker(object):
                                 ' in ' + obj_name if obj_name else ''))
 
   @staticmethod
+  def _CheckPresentIffMany(vals, name, obj_name):
+    """Checks that a set of vals and names imply every other element.
+
+    Args:
+      vals: The set of values to be compared.
+      name: The name of the objects holding the corresponding value.
+      obj_name: Name of the object containing these values.
+
+    Raises:
+      error.PayloadError if assertion does not hold.
+    """
+    if any(vals) and not all(vals):
+      raise error.PayloadError('%r is not present in all values%s.' %
+                               (name, ' in ' + obj_name if obj_name else ''))
+
+  @staticmethod
   def _Run(cmd, send_data=None):
     """Runs a subprocess, returns its output.
 
@@ -561,8 +606,9 @@ class PayloadChecker(object):
     Raises:
       error.PayloadError if any of the checks fail.
     """
-    part_sizes = collections.defaultdict(int, part_sizes)
+    self.major_version = self.payload.header.version
 
+    part_sizes = collections.defaultdict(int, part_sizes)
     manifest = self.payload.manifest
     report.AddSection('manifest')
 
@@ -581,31 +627,50 @@ class PayloadChecker(object):
     self._CheckPresentIff(self.sigs_offset, self.sigs_size,
                           'signatures_offset', 'signatures_size', 'manifest')
 
-    for part in common.CROS_PARTITIONS:
-      self.old_part_info[part] = self._CheckOptionalSubMsg(
-          manifest, 'old_%s_info' % part, report)
-      self.new_part_info[part] = self._CheckMandatorySubMsg(
-          manifest, 'new_%s_info' % part, report, 'manifest')
+    if self.major_version == 1:
+      for part in common.CROS_PARTITIONS:
+        self.old_part_info[part] = self._CheckOptionalSubMsg(
+            manifest, 'old_%s_info' % part, report)
+        self.new_part_info[part] = self._CheckMandatorySubMsg(
+            manifest, 'new_%s_info' % part, report, 'manifest')
+
+      # Check: old_kernel_info <==> old_rootfs_info.
+      self._CheckPresentIff(self.old_part_info[common.KERNEL].msg,
+                            self.old_part_info[common.ROOTFS].msg,
+                            'old_kernel_info', 'old_rootfs_info', 'manifest')
+    else:
+      for part in manifest.partitions:
+        name = part.partition_name
+        self.old_part_info[name] = self._CheckOptionalSubMsg(
+            part, 'old_partition_info', report)
+        self.new_part_info[name] = self._CheckMandatorySubMsg(
+            part, 'new_partition_info', report, 'manifest.partitions')
+
+      # Check: Old-style partition infos should not be specified.
+      for part in common.CROS_PARTITIONS:
+        self._CheckElemNotPresent(manifest, 'old_%s_info' % part, 'manifest')
+        self._CheckElemNotPresent(manifest, 'new_%s_info' % part, 'manifest')
 
-    # Check: old_kernel_info <==> old_rootfs_info.
-    self._CheckPresentIff(self.old_part_info[common.KERNEL].msg,
-                          self.old_part_info[common.ROOTFS].msg,
-                          'old_kernel_info', 'old_rootfs_info', 'manifest')
+      # Check: If old_partition_info is specified anywhere, it must be
+      # specified everywhere.
+      old_part_msgs = [part.msg for part in self.old_part_info.values() if part]
+      self._CheckPresentIffMany(old_part_msgs, 'old_partition_info',
+                                'manifest.partitions')
 
-    if self.old_part_info[common.KERNEL].msg:  # equivalently, rootfs msg
+    is_delta = any(part and part.msg for part in self.old_part_info.values())
+    if is_delta:
       # Assert/mark delta payload.
       if self.payload_type == _TYPE_FULL:
         raise error.PayloadError(
             'Apparent full payload contains old_{kernel,rootfs}_info.')
       self.payload_type = _TYPE_DELTA
 
-      for part in common.CROS_PARTITIONS:
+      for part, (msg, part_report) in self.old_part_info.iteritems():
         # Check: {size, hash} present in old_{kernel,rootfs}_info.
         field = 'old_%s_info' % part
-        msg, report = self.old_part_info[part]
         self.old_fs_sizes[part] = self._CheckMandatoryField(msg, 'size',
-                                                            report, field)
-        self._CheckMandatoryField(msg, 'hash', report, field,
+                                                            part_report, field)
+        self._CheckMandatoryField(msg, 'hash', part_report, field,
                                   convert=common.FormatSha256)
 
         # Check: old_{kernel,rootfs} size must fit in respective partition.
@@ -621,13 +686,11 @@ class PayloadChecker(object):
       self.payload_type = _TYPE_FULL
 
     # Check: new_{kernel,rootfs}_info present; contains {size, hash}.
-    for part in common.CROS_PARTITIONS:
+    for part, (msg, part_report) in self.new_part_info.iteritems():
       field = 'new_%s_info' % part
-      msg, report = self.new_part_info[part]
-
-      self.new_fs_sizes[part] = self._CheckMandatoryField(msg, 'size', report,
-                                                          field)
-      self._CheckMandatoryField(msg, 'hash', report, field,
+      self.new_fs_sizes[part] = self._CheckMandatoryField(msg, 'size',
+                                                          part_report, field)
+      self._CheckMandatoryField(msg, 'hash', part_report, field,
                                 convert=common.FormatSha256)
 
       # Check: new_{kernel,rootfs} size must fit in respective partition.
@@ -1266,7 +1329,7 @@ class PayloadChecker(object):
       # Part 1: Check the file header.
       report.AddSection('header')
       # Check: Payload version is valid.
-      if self.payload.header.version != 1:
+      if self.payload.header.version not in (1, 2):
         raise error.PayloadError('Unknown payload version (%d).' %
                                  self.payload.header.version)
       report.AddField('version', self.payload.header.version)
@@ -1276,44 +1339,54 @@ class PayloadChecker(object):
       self._CheckManifest(report, part_sizes)
       assert self.payload_type, 'payload type should be known by now'
 
-      # Infer the usable partition size when validating rootfs operations:
-      # - If rootfs partition size was provided, use that.
-      # - Otherwise, if this is an older delta (minor version < 2), stick with
-      #   a known constant size. This is necessary because older deltas may
-      #   exceed the filesystem size when moving data blocks around.
-      # - Otherwise, use the encoded filesystem size.
-      new_rootfs_usable_size = self.new_fs_sizes[common.ROOTFS]
-      old_rootfs_usable_size = self.old_fs_sizes[common.ROOTFS]
-      if part_sizes.get(common.ROOTFS, 0):
-        new_rootfs_usable_size = part_sizes[common.ROOTFS]
-        old_rootfs_usable_size = part_sizes[common.ROOTFS]
-      elif self.payload_type == _TYPE_DELTA and self.minor_version in (None, 1):
-        new_rootfs_usable_size = _OLD_DELTA_USABLE_PART_SIZE
-        old_rootfs_usable_size = _OLD_DELTA_USABLE_PART_SIZE
-
-      # Part 3: Examine rootfs operations.
-      # TODO(garnold)(chromium:243559) only default to the filesystem size if
-      # no explicit size provided *and* the partition size is not embedded in
-      # the payload; see issue for more details.
-      report.AddSection('rootfs operations')
-      total_blob_size = self._CheckOperations(
-          self.payload.manifest.install_operations, report,
-          'install_operations', self.old_fs_sizes[common.ROOTFS],
-          self.new_fs_sizes[common.ROOTFS], old_rootfs_usable_size,
-          new_rootfs_usable_size, 0, False)
-
-      # Part 4: Examine kernel operations.
-      # TODO(garnold)(chromium:243559) as above.
-      report.AddSection('kernel operations')
-      old_kernel_fs_size = self.old_fs_sizes[common.KERNEL]
-      new_kernel_fs_size = self.new_fs_sizes[common.KERNEL]
-      kernel_part_size = part_sizes.get(common.KERNEL, None)
-      total_blob_size += self._CheckOperations(
-          self.payload.manifest.kernel_install_operations, report,
-          'kernel_install_operations', old_kernel_fs_size, new_kernel_fs_size,
-          kernel_part_size if kernel_part_size else old_kernel_fs_size,
-          kernel_part_size if kernel_part_size else new_kernel_fs_size,
-          total_blob_size, True)
+      manifest = self.payload.manifest
+
+      # Part 3: Examine partition operations.
+      install_operations = []
+      if self.major_version == 1:
+        # partitions field should not ever exist in major version 1 payloads
+        self._CheckRepeatedElemNotPresent(manifest, 'partitions', 'manifest')
+
+        install_operations.append((common.ROOTFS, manifest.install_operations))
+        install_operations.append((common.KERNEL,
+                                   manifest.kernel_install_operations))
+
+      else:
+        self._CheckRepeatedElemNotPresent(manifest, 'install_operations',
+                                          'manifest')
+        self._CheckRepeatedElemNotPresent(manifest, 'kernel_install_operations',
+                                          'manifest')
+
+        for update in manifest.partitions:
+          install_operations.append((update.partition_name, update.operations))
+
+      total_blob_size = 0
+      for part, operations in install_operations:
+        report.AddSection('%s operations' % part)
+
+        new_fs_usable_size = self.new_fs_sizes[part]
+        old_fs_usable_size = self.old_fs_sizes[part]
+
+        if part_sizes.get(part, None):
+          new_fs_usable_size = old_fs_usable_size = part_sizes[part]
+        # Infer the usable partition size when validating rootfs operations:
+        # - If rootfs partition size was provided, use that.
+        # - Otherwise, if this is an older delta (minor version < 2), stick with
+        #   a known constant size. This is necessary because older deltas may
+        #   exceed the filesystem size when moving data blocks around.
+        # - Otherwise, use the encoded filesystem size.
+        elif self.payload_type == _TYPE_DELTA and part == common.ROOTFS and \
+            self.minor_version in (None, 1):
+          new_fs_usable_size = old_fs_usable_size = _OLD_DELTA_USABLE_PART_SIZE
+
+        # TODO(garnold)(chromium:243559) only default to the filesystem size if
+        # no explicit size provided *and* the partition size is not embedded in
+        # the payload; see issue for more details.
+        total_blob_size += self._CheckOperations(
+            operations, report, '%s_install_operations' % part,
+            self.old_fs_sizes[part], self.new_fs_sizes[part],
+            old_fs_usable_size, new_fs_usable_size, total_blob_size,
+            self.major_version == 1 and part == common.KERNEL)
 
       # Check: Operations data reach the end of the payload file.
       used_payload_size = self.payload.data_offset + total_blob_size
@@ -1322,11 +1395,11 @@ class PayloadChecker(object):
             'Used payload size (%d) different from actual file size (%d).' %
             (used_payload_size, payload_file_size))
 
-      # Part 5: Handle payload signatures message.
+      # Part 4: Handle payload signatures message.
       if self.check_payload_sig and self.sigs_size:
         self._CheckSignatures(report, pubkey_file_name)
 
-      # Part 6: Summary.
+      # Part 5: Summary.
       report.AddSection('summary')
       report.AddField('update type', self.payload_type)
 
diff --git a/scripts/update_payload/checker_unittest.py b/scripts/update_payload/checker_unittest.py
index ed5ee80d..98bf6127 100755
--- a/scripts/update_payload/checker_unittest.py
+++ b/scripts/update_payload/checker_unittest.py
@@ -891,6 +891,8 @@ class PayloadCheckerTest(mox.MoxTestBase):
                           self.NewExtentList((1, 16)))
         total_src_blocks = 16
 
+    # TODO(tbrindus): add major version 2 tests.
+    payload_checker.major_version = 1
     if op_type in (common.OpType.REPLACE, common.OpType.REPLACE_BZ):
       payload_checker.minor_version = 0
     elif op_type in (common.OpType.MOVE, common.OpType.BSDIFF):