update_payload: Allow check for given metadata size

Allow passing metadata size to check_update_payload so we can verify the
metadata size in omaha equals to the one in the payload.

BUG=chromium:820243
TEST=run paycheck.py with both valid and invalid metadata sizes reports as expected
TEST=unittests

Change-Id: Ib41ce77af77636fffec6752201c363e7fbbf868d
Reviewed-on: https://chromium-review.googlesource.com/955679
Commit-Ready: Amin Hassani <ahassani@chromium.org>
Tested-by: Amin Hassani <ahassani@chromium.org>
Reviewed-by: Ben Chan <benchan@chromium.org>
diff --git a/scripts/paycheck.py b/scripts/paycheck.py
index 96b1032..0050f5b 100755
--- a/scripts/paycheck.py
+++ b/scripts/paycheck.py
@@ -89,6 +89,9 @@
                                 'validation'))
   check_args.add_argument('-m', '--meta-sig', metavar='FILE',
                           help='verify metadata against its signature')
+  check_args.add_argument('-s', '--metadata-size', metavar='NUM', default=0,
+                          help='the metadata size to verify with the one in'
+                          ' payload')
   check_args.add_argument('-p', '--root-part-size', metavar='NUM',
                           default=0, type=int,
                           help='override rootfs partition size auto-inference')
@@ -128,7 +131,8 @@
   args.check = (args.check or args.report or args.assert_type or
                 args.block_size or args.allow_unhashed or
                 args.disabled_tests or args.meta_sig or args.key or
-                args.root_part_size or args.kern_part_size)
+                args.root_part_size or args.kern_part_size or
+                args.metadata_size)
 
   # Check the arguments, enforce payload type accordingly.
   if (args.src_kern is None) != (args.src_root is None):
@@ -202,6 +206,7 @@
           payload.Check(
               pubkey_file_name=args.key,
               metadata_sig_file=metadata_sig_file,
+              metadata_size=int(args.metadata_size),
               report_out_file=report_file,
               assert_type=args.assert_type,
               block_size=int(args.block_size),
diff --git a/scripts/update_payload/checker.py b/scripts/update_payload/checker.py
index e241b0b..48ed0f4 100644
--- a/scripts/update_payload/checker.py
+++ b/scripts/update_payload/checker.py
@@ -1231,13 +1231,14 @@
         raise error.PayloadError('Unknown signature version (%d).' %
                                  sig.version)
 
-  def Run(self, pubkey_file_name=None, metadata_sig_file=None,
+  def Run(self, pubkey_file_name=None, metadata_sig_file=None, metadata_size=0,
           rootfs_part_size=0, kernel_part_size=0, report_out_file=None):
     """Checker entry point, invoking all checks.
 
     Args:
       pubkey_file_name: Public key used for signature verification.
       metadata_sig_file: Metadata signature, if verification is desired.
+      metadata_size: metadata size, if verification is desired
       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
@@ -1258,6 +1259,12 @@
     self.payload.ResetFile()
 
     try:
+      # Check metadata_size (if provided).
+      if metadata_size and self.payload.data_offset != metadata_size:
+        raise error.PayloadError('Invalid payload metadata size in payload(%d) '
+                                 'vs given(%d)' % (self.payload.data_offset,
+                                                   metadata_size))
+
       # Check metadata signature (if provided).
       if metadata_sig_file:
         metadata_sig = base64.b64decode(metadata_sig_file.read())
diff --git a/scripts/update_payload/checker_unittest.py b/scripts/update_payload/checker_unittest.py
index f718234..68f1807 100755
--- a/scripts/update_payload/checker_unittest.py
+++ b/scripts/update_payload/checker_unittest.py
@@ -1127,8 +1127,8 @@
 
   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_mismatched_metadata_size, fail_mismatched_block_size,
+                fail_excess_data, fail_rootfs_part_size_exceeded,
                 fail_kernel_part_size_exceeded):
     """Tests Run()."""
     # Generate a test payload. For this test, we generate a full update that
@@ -1178,6 +1178,11 @@
     else:
       use_block_size = block_size
 
+    # For the unittests 246 is the value that generated for the payload.
+    metadata_size = 246
+    if fail_mismatched_metadata_size:
+      metadata_size += 1
+
     kwargs = {
         'payload_gen_dargs': {
             'privkey_file_name': test_utils._PRIVKEY_FILE_NAME,
@@ -1196,9 +1201,10 @@
 
       kwargs = {'pubkey_file_name': test_utils._PUBKEY_FILE_NAME,
                 'rootfs_part_size': rootfs_part_size,
+                'metadata_size': metadata_size,
                 'kernel_part_size': kernel_part_size}
       should_fail = (fail_wrong_payload_type or fail_mismatched_block_size or
-                     fail_excess_data or
+                     fail_mismatched_metadata_size or fail_excess_data or
                      fail_rootfs_part_size_exceeded or
                      fail_kernel_part_size_exceeded)
       if should_fail:
@@ -1353,6 +1359,7 @@
                       'kernel_part_size_provided': (True, False),
                       'fail_wrong_payload_type': (True, False),
                       'fail_invalid_block_size': (True, False),
+                      'fail_mismatched_metadata_size': (True, False),
                       'fail_mismatched_block_size': (True, False),
                       'fail_excess_data': (True, False),
                       'fail_rootfs_part_size_exceeded': (True, False),
diff --git a/scripts/update_payload/payload.py b/scripts/update_payload/payload.py
index 380d6d0..15f66d0 100644
--- a/scripts/update_payload/payload.py
+++ b/scripts/update_payload/payload.py
@@ -273,14 +273,15 @@
     return not self.IsDelta()
 
   def Check(self, pubkey_file_name=None, metadata_sig_file=None,
-            report_out_file=None, assert_type=None, block_size=0,
-            rootfs_part_size=0, kernel_part_size=0, allow_unhashed=False,
-            disabled_tests=()):
+            metadata_size=0, report_out_file=None, assert_type=None,
+            block_size=0, rootfs_part_size=0, kernel_part_size=0,
+            allow_unhashed=False, disabled_tests=()):
     """Checks the payload integrity.
 
     Args:
       pubkey_file_name: public key used for signature verification
       metadata_sig_file: metadata signature, if verification is desired
+      metadata_size: metadata size, if verification is desired
       report_out_file: file object to dump the report to
       assert_type: assert that payload is either 'full' or 'delta'
       block_size: expected filesystem / payload block size
@@ -300,6 +301,7 @@
         allow_unhashed=allow_unhashed, disabled_tests=disabled_tests)
     helper.Run(pubkey_file_name=pubkey_file_name,
                metadata_sig_file=metadata_sig_file,
+               metadata_size=metadata_size,
                rootfs_part_size=rootfs_part_size,
                kernel_part_size=kernel_part_size,
                report_out_file=report_out_file)