paycheck: Properly infer usable target partition size.

The payload checker used to restrict read/write block indexes to the
reported target filesystem size, unless explicitly given a partition
size value to use instead. So far this value was easy for clients (like
paygen) to come up with, because it was constant at 2GB for all known
boards.  However this is no longer the case, and there is no an easy way
for clients to know the actual target partition size after the payload
has been generated (nor is it encoded in the payload). This adds logic
for inferring the usable target partition size into PayloadChecker()
itself, as follows:

1) If a partition size was given, use that.

2) Else, if this is an old delta (minor version < 2), use the
   aforementioned default. This is necessary because older deltas may
   actually read/write data beyond the filesystem size. It is also
   sufficient because any old delta payload we generate should write to
   a 2GB target partition.

3) In all other cases, just use the new filesystem size, as encoded in
   the payload. This is a safe choice for full updates and newer deltas.

The command-line tool is updated accordingly. Note that the usable
kernel partition size inference remains unaffected.

BUG=chromium:508566
TEST=Unit tests (revised)

Change-Id: I987f28fdfe1d82d0f6f565ae9852b7b11bce13e8
Reviewed-on: https://chromium-review.googlesource.com/285447
Reviewed-by: Gilad Arnold <garnold@chromium.org>
Tested-by: Gilad Arnold <garnold@chromium.org>
Commit-Queue: Gilad Arnold <garnold@chromium.org>
diff --git a/host/lib/update_payload/checker.py b/host/lib/update_payload/checker.py
index 5428f66..47461af 100644
--- a/host/lib/update_payload/checker.py
+++ b/host/lib/update_payload/checker.py
@@ -57,6 +57,8 @@
     2: (_TYPE_DELTA,),
 }
 
+_OLD_DELTA_USABLE_PART_SIZE = 2 * 1024 * 1024 * 1024
+
 #
 # Helper functions.
 #
@@ -1152,8 +1154,8 @@
     Args:
       pubkey_file_name: Public key used for signature verification.
       metadata_sig_file: Metadata signature, if verification is desired.
-      rootfs_part_size: The size of rootfs partitions in bytes (default: use
-                        reported filesystem size).
+      rootfs_part_size: The size of rootfs partitions in bytes (default: infer
+                        based on payload type and version).
       kernel_part_size: The size of kernel partitions in bytes (default: use
                         reported filesystem size).
       report_out_file: File object to dump the report to.
@@ -1192,6 +1194,18 @@
       self._CheckManifest(report, rootfs_part_size, kernel_part_size)
       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_rootfs_fs_size
+      if rootfs_part_size:
+        new_rootfs_usable_size = rootfs_part_size
+      elif self.payload_type == _TYPE_DELTA and self.minor_version in (None, 1):
+        new_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
@@ -1200,9 +1214,7 @@
       total_blob_size = self._CheckOperations(
           self.payload.manifest.install_operations, report,
           'install_operations', self.old_rootfs_fs_size,
-          self.new_rootfs_fs_size,
-          rootfs_part_size if rootfs_part_size else self.new_rootfs_fs_size,
-          0, False)
+          self.new_rootfs_fs_size, new_rootfs_usable_size, 0, False)
 
       # Part 4: Examine kernel operations.
       # TODO(garnold)(chromium:243559) as above.
diff --git a/host/lib/update_payload/checker_unittest.py b/host/lib/update_payload/checker_unittest.py
index 871f252..a22373f 100755
--- a/host/lib/update_payload/checker_unittest.py
+++ b/host/lib/update_payload/checker_unittest.py
@@ -1085,8 +1085,11 @@
       self.assertRaises(update_payload.PayloadError,
                         payload_checker._CheckManifestMinorVersion, *args)
 
-  def DoRunTest(self, fail_wrong_payload_type, fail_invalid_block_size,
-                fail_mismatched_block_size, fail_excess_data):
+  def DoRunTest(self, rootfs_part_size_provided, kernel_part_size_provided,
+                fail_wrong_payload_type, fail_invalid_block_size,
+                fail_mismatched_block_size, fail_excess_data,
+                fail_rootfs_part_size_exceeded,
+                fail_kernel_part_size_exceeded):
     # Generate a test payload. For this test, we generate a full update that
     # has sample kernel and rootfs operations. Since most testing is done with
     # internal PayloadChecker methods that are tested elsewhere, here we only
@@ -1096,21 +1099,35 @@
     payload_gen = test_utils.EnhancedPayloadGenerator()
     block_size = test_utils.KiB(4)
     payload_gen.SetBlockSize(block_size)
-    kernel_part_size = test_utils.KiB(16)
-    rootfs_part_size = test_utils.MiB(2)
-    payload_gen.SetPartInfo(False, True, rootfs_part_size,
+    kernel_filesystem_size = test_utils.KiB(16)
+    rootfs_filesystem_size = test_utils.MiB(2)
+    payload_gen.SetPartInfo(False, True, rootfs_filesystem_size,
                             hashlib.sha256('fake-new-rootfs-content').digest())
-    payload_gen.SetPartInfo(True, True, kernel_part_size,
+    payload_gen.SetPartInfo(True, True, kernel_filesystem_size,
                             hashlib.sha256('fake-new-kernel-content').digest())
     payload_gen.SetMinorVersion(0)
+
+    rootfs_part_size = 0
+    if rootfs_part_size_provided:
+      rootfs_part_size = rootfs_filesystem_size + block_size
+    rootfs_op_size = rootfs_part_size or rootfs_filesystem_size
+    if fail_rootfs_part_size_exceeded:
+      rootfs_op_size += block_size
     payload_gen.AddOperationWithData(
         False, common.OpType.REPLACE,
-        dst_extents=[(0, rootfs_part_size / block_size)],
-        data_blob=os.urandom(rootfs_part_size))
+        dst_extents=[(0, rootfs_op_size / block_size)],
+        data_blob=os.urandom(rootfs_op_size))
+
+    kernel_part_size = 0
+    if kernel_part_size_provided:
+      kernel_part_size = kernel_filesystem_size + block_size
+    kernel_op_size = kernel_part_size or kernel_filesystem_size
+    if fail_kernel_part_size_exceeded:
+      kernel_op_size += block_size
     payload_gen.AddOperationWithData(
         True, common.OpType.REPLACE,
-        dst_extents=[(0, kernel_part_size / block_size)],
-        data_blob=os.urandom(kernel_part_size))
+        dst_extents=[(0, kernel_op_size / block_size)],
+        data_blob=os.urandom(kernel_op_size))
 
     # Generate payload (complete w/ signature) and create the test object.
     if fail_invalid_block_size:
@@ -1135,9 +1152,14 @@
     else:
       payload_checker = _GetPayloadChecker(payload_gen.WriteToFileWithData,
                                            **kwargs)
-      kwargs = {'pubkey_file_name': test_utils._PUBKEY_FILE_NAME}
+
+      kwargs = {'pubkey_file_name': test_utils._PUBKEY_FILE_NAME,
+                'rootfs_part_size': rootfs_part_size,
+                'kernel_part_size': kernel_part_size}
       should_fail = (fail_wrong_payload_type or fail_mismatched_block_size or
-                     fail_excess_data)
+                     fail_excess_data or
+                     fail_rootfs_part_size_exceeded or
+                     fail_kernel_part_size_exceeded)
       if should_fail:
         self.assertRaises(update_payload.PayloadError, payload_checker.Run,
                           **kwargs)
@@ -1283,10 +1305,14 @@
 
   # Add all Run() test cases.
   AddParametricTests('Run',
-                     {'fail_wrong_payload_type': (True, False),
+                     {'rootfs_part_size_provided': (True, False),
+                      'kernel_part_size_provided': (True, False),
+                      'fail_wrong_payload_type': (True, False),
                       'fail_invalid_block_size': (True, False),
                       'fail_mismatched_block_size': (True, False),
-                      'fail_excess_data': (True, False)})
+                      'fail_excess_data': (True, False),
+                      'fail_rootfs_part_size_exceeded': (True, False),
+                      'fail_kernel_part_size_exceeded': (True, False)})
 
 
 if __name__ == '__main__':
diff --git a/host/paycheck.py b/host/paycheck.py
index 36edc59..5290e9d 100755
--- a/host/paycheck.py
+++ b/host/paycheck.py
@@ -19,14 +19,6 @@
 import update_payload
 
 
-# The default sizes of partitions, based on current partitioning practice.
-# TODO(garnold)(chromium:243559) we should stop using these values once
-# partition sizes are encoded in newly generated payloads; that said, we should
-# allow users to specify partition sizes on the command-line, so as to be able
-# to check older payloads.
-_DEFAULT_ROOTFS_PART_SIZE = 2 * 1024 * 1024 * 1024
-_DEFAULT_KERNEL_PART_SIZE = 16 * 1024 * 1024
-
 _TYPE_FULL = 'full'
 _TYPE_DELTA = 'delta'
 
@@ -85,13 +77,11 @@
   check_opts.add_option('-m', '--meta-sig', metavar='FILE',
                         help='verify metadata against its signature')
   check_opts.add_option('-p', '--root-part-size', metavar='NUM',
-                        default=_DEFAULT_ROOTFS_PART_SIZE, type='int',
-                        help=('override default (%default) rootfs partition '
-                              'size'))
+                        default=0, type='int',
+                        help=('override rootfs partition size auto-inference'))
   check_opts.add_option('-P', '--kern-part-size', metavar='NUM',
-                        default=_DEFAULT_KERNEL_PART_SIZE, type='int',
-                        help=('override default (%default) kernel partition '
-                              'size'))
+                        default=0, type='int',
+                        help=('override kernel partition size auto-inference'))
   parser.add_option_group(check_opts)
 
   trace_opts = optparse.OptionGroup(parser, 'Applying payload')
@@ -135,8 +125,7 @@
   opts.check = (opts.check or opts.report or opts.assert_type or
                 opts.block_size or opts.allow_unhashed or
                 opts.disabled_tests or opts.meta_sig or opts.key or
-                opts.root_part_size != _DEFAULT_ROOTFS_PART_SIZE or
-                opts.kern_part_size != _DEFAULT_KERNEL_PART_SIZE)
+                opts.root_part_size or opts.kern_part_size)
 
   # Check number of arguments, enforce payload type accordingly.
   if len(args) == 3: