Parse Payload v2 header.

The update payload v2 contains an extra field in the header with the
size of the metadata signatures and the metadata signatures stored
right after the metadata. This patch parses the new payload format.

BUG=b:22024447
TEST=cros payload show payload-v2.bin; served a payload v2 with devserver.py

Change-Id: I8ce85af1df505f82f62a9d1cd57910cee6921f84
Reviewed-on: https://chromium-review.googlesource.com/306090
Commit-Ready: Alex Deymo <deymo@chromium.org>
Tested-by: Alex Deymo <deymo@chromium.org>
Reviewed-by: Gilad Arnold <garnold@chromium.org>
diff --git a/autoupdate.py b/autoupdate.py
index e20d4c2..a2a293f 100644
--- a/autoupdate.py
+++ b/autoupdate.py
@@ -566,12 +566,15 @@
     Returns:
       The size of the payload metadata, as reported in the payload header.
     """
-    # Handle corner-case where unit tests pass in empty payload files.
-    if os.path.getsize(payload_filename) < 20:
+    try:
+      with open(payload_filename) as payload_file:
+        payload = update_payload.Payload(payload_file)
+        payload.Init()
+        return payload.metadata_size
+    except (IOError, update_payload.PayloadError):
+      # For unit tests we may not have real files, so it's ok to ignore these
+      # errors.
       return 0
-    stream = open(payload_filename, 'rb')
-    stream.seek(16)
-    return struct.unpack('>I', stream.read(4))[0] + 20
 
   def GetLocalPayloadAttrs(self, payload_dir):
     """Returns hashes, size and delta flag of a local update payload.
diff --git a/host/lib/update_payload/common.py b/host/lib/update_payload/common.py
index 9d008a9..38d949b 100644
--- a/host/lib/update_payload/common.py
+++ b/host/lib/update_payload/common.py
@@ -21,6 +21,9 @@
     '\x00\x04\x20'
 )
 
+CHROMEOS_MAJOR_PAYLOAD_VERSION = 1
+BRILLO_MAJOR_PAYLOAD_VERSION = 2
+
 INPLACE_MINOR_PAYLOAD_VERSION = 1
 SOURCE_MINOR_PAYLOAD_VERSION = 2
 
diff --git a/host/lib/update_payload/payload.py b/host/lib/update_payload/payload.py
index b13aa11..ccd3240 100644
--- a/host/lib/update_payload/payload.py
+++ b/host/lib/update_payload/payload.py
@@ -49,14 +49,52 @@
   class _PayloadHeader(object):
     """Update payload header struct."""
 
-    def __init__(self, version, manifest_len):
-      self.version = version
-      self.manifest_len = manifest_len
+    # Header constants; sizes are in bytes.
+    _MAGIC = 'CrAU'
+    _VERSION_SIZE = 8
+    _MANIFEST_LEN_SIZE = 8
+    _METADATA_SIGNATURE_LEN_SIZE = 4
 
-  # Header constants; sizes are in bytes.
-  _MAGIC = 'CrAU'
-  _VERSION_SIZE = 8
-  _MANIFEST_LEN_SIZE = 8
+    def __init__(self):
+      self.version = None
+      self.manifest_len = None
+      self.metadata_signature_len = None
+      self.size = None
+
+    def ReadFromPayload(self, payload_file, hasher=None):
+      """Reads the payload header from a file.
+
+      Reads the payload header from the |payload_file| and updates the |hasher|
+      if one is passed. The parsed header is stored in the _PayloadHeader
+      instance attributes.
+
+      Args:
+        payload_file: a file object
+        hasher: an optional hasher to pass the value through
+      Returns:
+        None.
+      Raises:
+        PayloadError if a read error occurred or the header is invalid.
+      """
+      # Verify magic
+      magic = common.Read(payload_file, len(self._MAGIC), hasher=hasher)
+      if magic != self._MAGIC:
+        raise PayloadError('invalid payload magic: %s' % magic)
+
+      self.version = _ReadInt(payload_file, self._VERSION_SIZE, True,
+                              hasher=hasher)
+      self.manifest_len = _ReadInt(payload_file, self._MANIFEST_LEN_SIZE, True,
+                                   hasher=hasher)
+      self.size = (len(self._MAGIC) + self._VERSION_SIZE +
+                   self._MANIFEST_LEN_SIZE)
+      self.metadata_signature_len = 0
+
+      if self.version == common.BRILLO_MAJOR_PAYLOAD_VERSION:
+        self.size += self._METADATA_SIGNATURE_LEN_SIZE
+        self.metadata_signature_len = _ReadInt(
+            payload_file, self._METADATA_SIGNATURE_LEN_SIZE, True,
+            hasher=hasher)
+
 
   def __init__(self, payload_file):
     """Initialize the payload object.
@@ -70,7 +108,9 @@
     self.is_init = False
     self.header = None
     self.manifest = None
-    self.data_offset = 0
+    self.data_offset = None
+    self.metadata_signature = None
+    self.metadata_size = None
 
   def _ReadHeader(self):
     """Reads and returns the payload header.
@@ -81,17 +121,9 @@
       PayloadError if a read error occurred.
 
     """
-    # Verify magic
-    magic = common.Read(self.payload_file, len(self._MAGIC),
-                        hasher=self.manifest_hasher)
-    if magic != self._MAGIC:
-      raise PayloadError('invalid payload magic: %s' % magic)
-
-    return self._PayloadHeader(
-        _ReadInt(self.payload_file, self._VERSION_SIZE, True,
-                 hasher=self.manifest_hasher),
-        _ReadInt(self.payload_file, self._MANIFEST_LEN_SIZE, True,
-                 hasher=self.manifest_hasher))
+    header = self._PayloadHeader()
+    header.ReadFromPayload(self.payload_file, self.manifest_hasher)
+    return header
 
   def _ReadManifest(self):
     """Reads and returns the payload manifest.
@@ -108,6 +140,23 @@
     return common.Read(self.payload_file, self.header.manifest_len,
                        hasher=self.manifest_hasher)
 
+  def _ReadMetadataSignature(self):
+    """Reads and returns the metadata signatures.
+
+    Returns:
+      A string containing the metadata signatures protobuf in binary form or
+      an empty string if no metadata signature found in the payload.
+    Raises:
+      PayloadError if a read error occurred.
+
+    """
+    if not self.header:
+      raise PayloadError('payload header not present')
+
+    return common.Read(
+        self.payload_file, self.header.metadata_signature_len,
+        offset=self.header.size + self.header.manifest_len)
+
   def ReadDataBlob(self, offset, length):
     """Reads and returns a single data blob from the update payload.
 
@@ -148,9 +197,14 @@
     self.manifest = update_metadata_pb2.DeltaArchiveManifest()
     self.manifest.ParseFromString(manifest_raw)
 
-    # Store data offset.
-    self.data_offset = (len(self._MAGIC) + self._VERSION_SIZE +
-                        self._MANIFEST_LEN_SIZE + self.header.manifest_len)
+    # Read the metadata signature (if any).
+    metadata_signature_raw = self._ReadMetadataSignature()
+    if metadata_signature_raw:
+      self.metadata_signature = update_metadata_pb2.Signatures()
+      self.metadata_signature.ParseFromString(metadata_signature_raw)
+
+    self.metadata_size = self.header.size + self.header.manifest_len
+    self.data_offset = self.metadata_size + self.header.metadata_signature_len
 
     self.is_init = True