paycheck: unit tests + fixes to checker module

This adds missing unit tests for the checker module, bundled with fixes
to some bugs that surfaced due to unit tests. This includes:

* A fake extent (signified by start_block == UINT64_MAX) that
  accompanies a signature data blob bears different requirements than
  previously implemented.  Specifically, the extent sequence must have
  exactly one extent; and the number of blocks is not necessarily one,
  rather it is the correct number that corresponds to the actual length
  of the signature blob.

* REPLACE/REPLACE_BZ operations must contain data.

* MOVE operation validation must ensure that all of the actual message
  extents are being used.

* BSDIFF operation must contain data (the diff).

* Signature pseudo-operation should be a REPLACE.

BUG=chromium-os:34911,chromium-os:33607,chromium-os:7597
TEST=Passes unittests (upcoming); works with actual payloads.

Change-Id: I4d839d1d4da1fbb4a493b208958a139368e2c8ca
Reviewed-on: https://gerrit.chromium.org/gerrit/45429
Tested-by: Gilad Arnold <garnold@chromium.org>
Reviewed-by: Chris Sosa <sosa@chromium.org>
Commit-Queue: Gilad Arnold <garnold@chromium.org>

1 files changed, 340 insertions, 0 deletions

diff --git a/scripts/update_payload/test_utils.py b/scripts/update_payload/test_utils.py
new file mode 100644
index 00000000..d05aafd1
--- /dev/null
+++ b/scripts/update_payload/test_utils.py
@@ -0,0 +1,340 @@
+# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Utilities for unit testing."""
+
+import cStringIO
+import hashlib
+import struct
+import subprocess
+
+import common
+import payload
+import update_metadata_pb2
+
+
+class TestError(Exception):
+  """An error during testing of update payload code."""
+
+
+def _WriteInt(file_obj, size, is_unsigned, val):
+  """Writes a binary-encoded integer to a file.
+
+  It will do the correct conversion based on the reported size and whether or
+  not a signed number is expected. Assumes a network (big-endian) byte
+  ordering.
+
+  Args:
+    file_obj: a file object
+    size: the integer size in bytes (2, 4 or 8)
+    is_unsigned: whether it is signed or not
+    val: integer value to encode
+  Raises:
+    PayloadError if a write error occurred.
+
+  """
+  try:
+    file_obj.write(struct.pack(common.IntPackingFmtStr(size, is_unsigned), val))
+  except IOError, e:
+    raise payload.PayloadError('error writing to file (%s): %s' %
+                               (file_obj.name, e))
+
+
+def _SetMsgField(msg, field_name, val):
+  """Sets or clears a field in a protobuf message."""
+  if val is None:
+    msg.ClearField(field_name)
+  else:
+    setattr(msg, field_name, val)
+
+
+def SignSha256(data, privkey_file_name):
+  """Signs the data's SHA256 hash with an RSA private key.
+
+  Args:
+    data: the data whose SHA256 hash we want to sign
+    privkey_file_name: private key used for signing data
+  Returns:
+    The signature string, prepended with an ASN1 header.
+  Raises:
+    TestError if something goes wrong.
+
+  """
+  # pylint: disable=E1101
+  data_sha256_hash = common.SIG_ASN1_HEADER + hashlib.sha256(data).digest()
+  sign_cmd = ['openssl', 'rsautl', '-sign', '-inkey', privkey_file_name]
+  try:
+    sign_process = subprocess.Popen(sign_cmd, stdin=subprocess.PIPE,
+                                    stdout=subprocess.PIPE)
+    sig, _ = sign_process.communicate(input=data_sha256_hash)
+  except Exception as e:
+    raise TestError('signing subprocess failed: %s' % e)
+
+  return sig
+
+
+class SignaturesGenerator(object):
+  """Generates a payload signatures data block."""
+
+  def __init__(self):
+    self.sigs = update_metadata_pb2.Signatures()
+
+  def AddSig(self, version, data):
+    """Adds a signature to the signature sequence.
+
+    Args:
+      version: signature version (None means do not assign)
+      data: signature binary data (None means do not assign)
+
+    """
+    # Pylint fails to identify a member of the Signatures message.
+    # pylint: disable=E1101
+    sig = self.sigs.signatures.add()
+    if version is not None:
+      sig.version = version
+    if data is not None:
+      sig.data = data
+
+  def ToBinary(self):
+    """Returns the binary representation of the signature block."""
+    return self.sigs.SerializeToString()
+
+
+class PayloadGenerator(object):
+  """Generates an update payload allowing low-level control.
+
+  Attributes:
+    manifest: the protobuf containing the payload manifest
+    version: the payload version identifier
+    block_size: the block size pertaining to update operations
+
+  """
+
+  def __init__(self, version=1):
+    self.manifest = update_metadata_pb2.DeltaArchiveManifest()
+    self.version = version
+    self.block_size = 0
+
+  @staticmethod
+  def _WriteExtent(ex, val):
+    """Returns an Extent message."""
+    start_block, num_blocks = val
+    _SetMsgField(ex, 'start_block', start_block)
+    _SetMsgField(ex, 'num_blocks', num_blocks)
+
+  @staticmethod
+  def _AddValuesToRepeatedField(repeated_field, values, write_func):
+    """Adds values to a repeated message field."""
+    if values:
+      for val in values:
+        new_item = repeated_field.add()
+        write_func(new_item, val)
+
+  @staticmethod
+  def _AddExtents(extents_field, values):
+    """Adds extents to an extents field."""
+    PayloadGenerator._AddValuesToRepeatedField(
+        extents_field, values, PayloadGenerator._WriteExtent)
+
+  def SetBlockSize(self, block_size):
+    """Sets the payload's block size."""
+    self.block_size = block_size
+    _SetMsgField(self.manifest, 'block_size', block_size)
+
+  def SetPartInfo(self, is_kernel, is_new, part_size, part_hash):
+    """Set the partition info entry.
+
+    Args:
+      is_kernel: whether this is kernel partition info
+      is_new: whether to set old (False) or new (True) info
+      part_size: the partition size (in fact, filesystem size)
+      part_hash: the partition hash
+
+    """
+    if is_kernel:
+      # pylint: disable=E1101
+      part_info = (self.manifest.new_kernel_info if is_new
+                   else self.manifest.old_kernel_info)
+    else:
+      # pylint: disable=E1101
+      part_info = (self.manifest.new_rootfs_info if is_new
+                   else self.manifest.old_rootfs_info)
+    _SetMsgField(part_info, 'size', part_size)
+    _SetMsgField(part_info, 'hash', part_hash)
+
+  def AddOperation(self, is_kernel, op_type, data_offset=None,
+                   data_length=None, src_extents=None, src_length=None,
+                   dst_extents=None, dst_length=None, data_sha256_hash=None):
+    """Adds an InstallOperation entry."""
+    # pylint: disable=E1101
+    operations = (self.manifest.kernel_install_operations if is_kernel
+                  else self.manifest.install_operations)
+
+    op = operations.add()
+    op.type = op_type
+
+    _SetMsgField(op, 'data_offset', data_offset)
+    _SetMsgField(op, 'data_length', data_length)
+
+    self._AddExtents(op.src_extents, src_extents)
+    _SetMsgField(op, 'src_length', src_length)
+
+    self._AddExtents(op.dst_extents, dst_extents)
+    _SetMsgField(op, 'dst_length', dst_length)
+
+    _SetMsgField(op, 'data_sha256_hash', data_sha256_hash)
+
+  def SetSignatures(self, sigs_offset, sigs_size):
+    """Set the payload's signature block descriptors."""
+    _SetMsgField(self.manifest, 'signatures_offset', sigs_offset)
+    _SetMsgField(self.manifest, 'signatures_size', sigs_size)
+
+  def _WriteHeaderToFile(self, file_obj, manifest_len):
+    """Writes a payload heaer to a file."""
+    # We need to access protected members in Payload for writing the header.
+    # pylint: disable=W0212
+    file_obj.write(payload.Payload._MAGIC)
+    _WriteInt(file_obj, payload.Payload._VERSION_SIZE, True, self.version)
+    _WriteInt(file_obj, payload.Payload._MANIFEST_LEN_SIZE, True, manifest_len)
+
+  def WriteToFile(self, file_obj, manifest_len=-1, data_blobs=None,
+                  sigs_data=None, padding=None):
+    """Writes the payload content to a file.
+
+    Args:
+      file_obj: a file object open for writing
+      manifest_len: manifest len to dump (otherwise computed automatically)
+      data_blobs: a list of data blobs to be concatenated to the payload
+      sigs_data: a binary Signatures message to be concatenated to the payload
+      padding: stuff to dump past the normal data blobs provided (optional)
+
+    """
+    manifest = self.manifest.SerializeToString()
+    if manifest_len < 0:
+      manifest_len = len(manifest)
+    self._WriteHeaderToFile(file_obj, manifest_len)
+    file_obj.write(manifest)
+    if data_blobs:
+      for data_blob in data_blobs:
+        file_obj.write(data_blob)
+    if sigs_data:
+      file_obj.write(sigs_data)
+    if padding:
+      file_obj.write(padding)
+
+
+class EnhancedPayloadGenerator(PayloadGenerator):
+  """Payload generator with automatic handling of data blobs.
+
+  Attributes:
+    data_blobs: a list of blobs, in the order they were added
+    curr_offset: the currently consumed offset of blobs added to the payload
+
+  """
+
+  def __init__(self):
+    super(EnhancedPayloadGenerator, self).__init__()
+    self.data_blobs = []
+    self.curr_offset = 0
+
+  def AddData(self, data_blob):
+    """Adds a (possibly orphan) data blob."""
+    data_length = len(data_blob)
+    data_offset = self.curr_offset
+    self.curr_offset += data_length
+    self.data_blobs.append(data_blob)
+    return data_length, data_offset
+
+  def AddOperationWithData(self, is_kernel, op_type, src_extents=None,
+                           src_length=None, dst_extents=None, dst_length=None,
+                           data_blob=None, do_hash_data_blob=True):
+    """Adds an install operation and associated data blob.
+
+    This takes care of obtaining a hash of the data blob (if so instructed)
+    and appending it to the internally maintained list of blobs, including the
+    necessary offset/length accounting.
+
+    Args:
+      is_kernel: whether this is a kernel (True) or rootfs (False) operation
+      op_type: one of REPLACE, REPLACE_BZ, MOVE or BSDIFF
+      src_extents: list of (start, length) pairs indicating src block ranges
+      src_length: size of the src data in bytes (needed for BSDIFF)
+      dst_extents: list of (start, length) pairs indicating dst block ranges
+      dst_length: size of the dst data in bytes (needed for BSDIFF)
+      data_blob: a data blob associated with this operation
+      do_hash_data_blob: whether or not to compute and add a data blob hash
+
+    """
+    data_offset = data_length = data_sha256_hash = None
+    if data_blob is not None:
+      if do_hash_data_blob:
+        # pylint: disable=E1101
+        data_sha256_hash = hashlib.sha256(data_blob).digest()
+      data_length, data_offset = self.AddData(data_blob)
+
+    self.AddOperation(is_kernel, op_type, data_offset=data_offset,
+                      data_length=data_length, src_extents=src_extents,
+                      src_length=src_length, dst_extents=dst_extents,
+                      dst_length=dst_length, data_sha256_hash=data_sha256_hash)
+
+  def WriteToFileWithData(self, file_obj, sigs_data=None,
+                          privkey_file_name=None,
+                          do_add_pseudo_operation=False,
+                          is_pseudo_in_kernel=False, padding=None):
+    """Writes the payload content to a file, optionally signing the content.
+
+    Args:
+      file_obj: a file object open for writing
+      sigs_data: signatures blob to be appended to the payload (optional;
+                 payload signature fields assumed to be preset by the caller)
+      privkey_file_name: key used for signing the payload (optional; used only
+                         if explicit signatures blob not provided)
+      do_add_pseudo_operation: whether a pseudo-operation should be added to
+                               account for the signature blob
+      is_pseudo_in_kernel: whether the pseudo-operation should be added to
+                           kernel (True) or rootfs (False) operations
+      padding: stuff to dump past the normal data blobs provided (optional)
+    Raises:
+      TestError: if arguments are inconsistent or something goes wrong.
+
+    """
+    sigs_len = len(sigs_data) if sigs_data else 0
+
+    # Do we need to generate a genuine signatures blob?
+    do_generate_sigs_data = sigs_data is None and privkey_file_name
+
+    if do_generate_sigs_data:
+      # First, sign some arbitrary data to obtain the size of a signature blob.
+      fake_sig = SignSha256('fake-payload-data', privkey_file_name)
+      fake_sigs_gen = SignaturesGenerator()
+      fake_sigs_gen.AddSig(1, fake_sig)
+      sigs_len = len(fake_sigs_gen.ToBinary())
+
+      # Update the payload with proper signature attributes.
+      self.SetSignatures(self.curr_offset, sigs_len)
+
+    # Add a pseudo-operation to account for the signature blob, if requested.
+    if do_add_pseudo_operation:
+      if not self.block_size:
+        raise TestError('cannot add pseudo-operation without knowing the '
+                        'payload block size')
+      self.AddOperation(
+          is_pseudo_in_kernel, common.OpType.REPLACE,
+          data_offset=self.curr_offset, data_length=sigs_len,
+          dst_extents=[(common.PSEUDO_EXTENT_MARKER,
+                        (sigs_len + self.block_size - 1) / self.block_size)])
+
+    if do_generate_sigs_data:
+      # Once all payload fields are updated, dump and sign it.
+      temp_payload_file = cStringIO.StringIO()
+      self.WriteToFile(temp_payload_file, data_blobs=self.data_blobs)
+      sig = SignSha256(temp_payload_file.getvalue(), privkey_file_name)
+      sigs_gen = SignaturesGenerator()
+      sigs_gen.AddSig(1, sig)
+      sigs_data = sigs_gen.ToBinary()
+      assert len(sigs_data) == sigs_len, 'signature blob lengths mismatch'
+
+    # Dump the whole thing, complete with data and signature blob, to a file.
+    self.WriteToFile(file_obj, data_blobs=self.data_blobs, sigs_data=sigs_data,
+                     padding=padding)