releasetools: Create StreamingPropertyFiles class.

This CL breaks down ComputeStreamingMetadata() into mutiple member
functions of StreamingPropertyFiles class, which correspond to the
two-pass logic when generating streaming property files (aka streaming
metadata).

StreamingPropertyFiles.Compute() does the work for the first pass, by
putting placeholders before doing initial signing. Finalize()
corresponds to the second pass, where the placeholders get replaced with
actual data. Verify() can be optionally called to assert the correctness
of the work.

The separation between Compute() and Finalize() is to allow having
multiple StreamingPropertyFiles instances (in coming up CLs). This way
we can call Compute() multiple times for each instance, followed by only
one call to SignOutput(). And similarly for Finalize().

Bug: 74210298
Test: Generate an A/B OTA package. Check the METADATA entry.
Test: python -m unittest test_ota_from_target_files
Change-Id: I45be0372a4863c4405e6d8e20bcb9ccdc29e7e11
(cherry picked from commit ae5e4c30fe)

tools/releasetools/ota_from_target_files.py

@@ -955,55 +955,132 @@ def GetPackageMetadata(target_info, source_info=None):
   return metadata
 
 
-def ComputeStreamingMetadata(zip_file, reserve_space=False,
+class StreamingPropertyFiles(object):
-                             expected_length=None):
+  """Computes the ota-streaming-property-files string for streaming A/B OTA.
-  """Computes the streaming metadata for a given zip.
 
-  When 'reserve_space' is True, we reserve extra space for the offset and
+  Computing the final property-files string requires two passes. Because doing
-  length of the metadata entry itself, although we don't know the final
+  the whole package signing (with signapk.jar) will possibly reorder the ZIP
-  values until the package gets signed. This function will be called again
+  entries, which may in turn invalidate earlier computed ZIP entry offset/size
-  after signing. We then write the actual values and pad the string to the
+  values.
-  length we set earlier. Note that we can't use the actual length of the
+
-  metadata entry in the second run. Otherwise the offsets for other entries
+  This class provides functions to be called for each pass. The general flow is
-  will be changing again.
+  as follows.
+
+    property_files = StreamingPropertyFiles()
+    # The first pass, which writes placeholders before doing initial signing.
+    property_files.Compute()
+    SignOutput()
+
+    # The second pass, by replacing the placeholders with actual data.
+    property_files.Finalize()
+    SignOutput()
+
+  And the caller can additionally verify the final result.
+
+    property_files.Verify()
   """
 
-  def ComputeEntryOffsetSize(name):
+  def __init__(self):
-    """Compute the zip entry offset and size."""
+    self.required = (
-    info = zip_file.getinfo(name)
+        # payload.bin and payload_properties.txt must exist.
-    offset = info.header_offset + len(info.FileHeader())
+        'payload.bin',
-    size = info.file_size
+        'payload_properties.txt',
-    return '%s:%d:%d' % (os.path.basename(name), offset, size)
+    )
+    self.optional = (
+        # care_map.txt is available only if dm-verity is enabled.
+        'care_map.txt',
+        # compatibility.zip is available only if target supports Treble.
+        'compatibility.zip',
+    )
 
-  # payload.bin and payload_properties.txt must exist.
+  def Compute(self, input_zip):
-  offsets = [ComputeEntryOffsetSize('payload.bin'),
+    """Computes and returns a property-files string with placeholders.
-             ComputeEntryOffsetSize('payload_properties.txt')]
 
-  # care_map.txt is available only if dm-verity is enabled.
+    We reserve extra space for the offset and size of the metadata entry itself,
-  if 'care_map.txt' in zip_file.namelist():
+    although we don't know the final values until the package gets signed.
-    offsets.append(ComputeEntryOffsetSize('care_map.txt'))
 
-  if 'compatibility.zip' in zip_file.namelist():
+    Args:
-    offsets.append(ComputeEntryOffsetSize('compatibility.zip'))
+      input_zip: The input ZIP file.
 
-  # 'META-INF/com/android/metadata' is required. We don't know its actual
+    Returns:
-  # offset and length (as well as the values for other entries). So we
+      A string with placeholders for the metadata offset/size info, e.g.
-  # reserve 10-byte as a placeholder, which is to cover the space for metadata
+      "payload.bin:679:343,payload_properties.txt:378:45,metadata:        ".
-  # entry ('xx:xxx', since it's ZIP_STORED which should appear at the
+    """
-  # beginning of the zip), as well as the possible value changes in other
+    return self._GetPropertyFilesString(input_zip, reserve_space=True)
-  # entries.
-  if reserve_space:
-    offsets.append('metadata:' + ' ' * 10)
-  else:
-    offsets.append(ComputeEntryOffsetSize(METADATA_NAME))
 
-  value = ','.join(offsets)
+  def Finalize(self, input_zip, reserved_length):
-  if expected_length is not None:
+    """Finalizes a property-files string with actual METADATA offset/size info.
-    assert len(value) <= expected_length, \
+
-        'Insufficient reserved space: reserved=%d, actual=%d' % (
+    The input ZIP file has been signed, with the ZIP entries in the desired
-            expected_length, len(value))
+    place (signapk.jar will possibly reorder the ZIP entries). Now we compute
-    value += ' ' * (expected_length - len(value))
+    the ZIP entry offsets and construct the property-files string with actual
-  return value
+    data. Note that during this process, we must pad the property-files string
+    to the reserved length, so that the METADATA entry size remains the same.
+    Otherwise the entries' offsets and sizes may change again.
+
+    Args:
+      input_zip: The input ZIP file.
+      reserved_length: The reserved length of the property-files string during
+          the call to Compute(). The final string must be no more than this
+          size.
+
+    Returns:
+      A property-files string including the metadata offset/size info, e.g.
+      "payload.bin:679:343,payload_properties.txt:378:45,metadata:69:379  ".
+
+    Raises:
+      AssertionError: If the reserved length is insufficient to hold the final
+          string.
+    """
+    result = self._GetPropertyFilesString(input_zip, reserve_space=False)
+    assert len(result) <= reserved_length, \
+        'Insufficient reserved space: reserved={}, actual={}'.format(
+            reserved_length, len(result))
+    result += ' ' * (reserved_length - len(result))
+    return result
+
+  def Verify(self, input_zip, expected):
+    """Verifies the input ZIP file contains the expected property-files string.
+
+    Args:
+      input_zip: The input ZIP file.
+      expected: The property-files string that's computed from Finalize().
+
+    Raises:
+      AssertionError: On finding a mismatch.
+    """
+    actual = self._GetPropertyFilesString(input_zip)
+    assert actual == expected, \
+        "Mismatching streaming metadata: {} vs {}.".format(actual, expected)
+
+  def _GetPropertyFilesString(self, zip_file, reserve_space=False):
+    """Constructs the property-files string per request."""
+
+    def ComputeEntryOffsetSize(name):
+      """Computes the zip entry offset and size."""
+      info = zip_file.getinfo(name)
+      offset = info.header_offset + len(info.FileHeader())
+      size = info.file_size
+      return '%s:%d:%d' % (os.path.basename(name), offset, size)
+
+    tokens = []
+    for entry in self.required:
+      tokens.append(ComputeEntryOffsetSize(entry))
+    for entry in self.optional:
+      if entry in zip_file.namelist():
+        tokens.append(ComputeEntryOffsetSize(entry))
+
+    # 'META-INF/com/android/metadata' is required. We don't know its actual
+    # offset and length (as well as the values for other entries). So we reserve
+    # 10-byte as a placeholder, which is to cover the space for metadata entry
+    # ('xx:xxx', since it's ZIP_STORED which should appear at the beginning of
+    # the zip), as well as the possible value changes in other entries.
+    if reserve_space:
+      tokens.append('metadata:' + ' ' * 10)
+    else:
+      tokens.append(ComputeEntryOffsetSize(METADATA_NAME))
+
+    return ','.join(tokens)
 
 
 def FinalizeMetadata(metadata, input_file, output_file):
@@ -1028,9 +1105,10 @@ def FinalizeMetadata(metadata, input_file, output_file):
   output_zip = zipfile.ZipFile(
       input_file, 'a', compression=zipfile.ZIP_DEFLATED)
 
+  property_files = StreamingPropertyFiles()
+
   # Write the current metadata entry with placeholders.
-  metadata['ota-streaming-property-files'] = ComputeStreamingMetadata(
+  metadata['ota-streaming-property-files'] = property_files.Compute(output_zip)
-      output_zip, reserve_space=True)
   WriteMetadata(metadata, output_zip)
   common.ZipClose(output_zip)
 
@@ -1043,11 +1121,10 @@ def FinalizeMetadata(metadata, input_file, output_file):
   SignOutput(input_file, prelim_signing)
 
   # Open the signed zip. Compute the final metadata that's needed for streaming.
-  prelim_signing_zip = zipfile.ZipFile(prelim_signing, 'r')
+  with zipfile.ZipFile(prelim_signing, 'r') as prelim_signing_zip:
-  expected_length = len(metadata['ota-streaming-property-files'])
+    expected_length = len(metadata['ota-streaming-property-files'])
-  metadata['ota-streaming-property-files'] = ComputeStreamingMetadata(
+    metadata['ota-streaming-property-files'] = property_files.Finalize(
-      prelim_signing_zip, reserve_space=False, expected_length=expected_length)
+        prelim_signing_zip, expected_length)
-  common.ZipClose(prelim_signing_zip)
 
   # Replace the METADATA entry.
   common.ZipDelete(prelim_signing, METADATA_NAME)
@@ -1060,12 +1137,9 @@ def FinalizeMetadata(metadata, input_file, output_file):
   SignOutput(prelim_signing, output_file)
 
   # Reopen the final signed zip to double check the streaming metadata.
-  output_zip = zipfile.ZipFile(output_file, 'r')
+  with zipfile.ZipFile(output_file, 'r') as output_zip:
-  actual = metadata['ota-streaming-property-files'].strip()
+    property_files.Verify(
-  expected = ComputeStreamingMetadata(output_zip)
+        output_zip, metadata['ota-streaming-property-files'].strip())
-  assert actual == expected, \
-      "Mismatching streaming metadata: %s vs %s." % (actual, expected)
-  common.ZipClose(output_zip)
 
 
 def WriteBlockIncrementalOTAPackage(target_zip, source_zip, output_zip):

tools/releasetools/test_ota_from_target_files.py

@@ -23,10 +23,10 @@ import zipfile
 import common
 import test_utils
 from ota_from_target_files import (
-    _LoadOemDicts, BuildInfo, ComputeStreamingMetadata, GetPackageMetadata,
+    _LoadOemDicts, BuildInfo, GetPackageMetadata,
     GetTargetFilesZipForSecondaryImages,
     GetTargetFilesZipWithoutPostinstallConfig,
-    Payload, PayloadSigner, POSTINSTALL_CONFIG,
+    Payload, PayloadSigner, POSTINSTALL_CONFIG, StreamingPropertyFiles,
     WriteFingerprintAssertion)
 
 
@@ -589,6 +589,12 @@ class OtaFromTargetFilesTest(unittest.TestCase):
     with zipfile.ZipFile(target_file) as verify_zip:
       self.assertNotIn(POSTINSTALL_CONFIG, verify_zip.namelist())
 
+
+class StreamingPropertyFilesTest(unittest.TestCase):
+
+  def tearDown(self):
+    common.Cleanup()
+
   @staticmethod
   def _construct_zip_package(entries):
     zip_file = common.MakeTempFile(suffix='.zip')
@@ -619,20 +625,21 @@ class OtaFromTargetFilesTest(unittest.TestCase):
           expected = entry.replace('.', '-').upper().encode()
         self.assertEqual(expected, input_fp.read(size))
 
-  def test_ComputeStreamingMetadata_reserveSpace(self):
+  def test_Compute(self):
     entries = (
         'payload.bin',
         'payload_properties.txt',
     )
     zip_file = self._construct_zip_package(entries)
+    property_files = StreamingPropertyFiles()
     with zipfile.ZipFile(zip_file, 'r') as zip_fp:
-      streaming_metadata = ComputeStreamingMetadata(zip_fp, reserve_space=True)
+      streaming_metadata = property_files.Compute(zip_fp)
-    tokens = self._parse_streaming_metadata_string(streaming_metadata)
 
+    tokens = self._parse_streaming_metadata_string(streaming_metadata)
     self.assertEqual(3, len(tokens))
     self._verify_entries(zip_file, tokens, entries)
 
-  def test_ComputeStreamingMetadata_reserveSpace_withCareMapTxtAndCompatibilityZip(self):
+  def test_Compute_withCareMapTxtAndCompatibilityZip(self):
     entries = (
         'payload.bin',
         'payload_properties.txt',
@@ -640,22 +647,26 @@ class OtaFromTargetFilesTest(unittest.TestCase):
         'compatibility.zip',
     )
     zip_file = self._construct_zip_package(entries)
+    property_files = StreamingPropertyFiles()
     with zipfile.ZipFile(zip_file, 'r') as zip_fp:
-      streaming_metadata = ComputeStreamingMetadata(zip_fp, reserve_space=True)
+      streaming_metadata = property_files.Compute(zip_fp)
-    tokens = self._parse_streaming_metadata_string(streaming_metadata)
 
+    tokens = self._parse_streaming_metadata_string(streaming_metadata)
     self.assertEqual(5, len(tokens))
     self._verify_entries(zip_file, tokens, entries)
 
-  def test_ComputeStreamingMetadata(self):
+  def test_Finalize(self):
     entries = [
         'payload.bin',
         'payload_properties.txt',
         'META-INF/com/android/metadata',
     ]
     zip_file = self._construct_zip_package(entries)
+    property_files = StreamingPropertyFiles()
     with zipfile.ZipFile(zip_file, 'r') as zip_fp:
-      streaming_metadata = ComputeStreamingMetadata(zip_fp, reserve_space=False)
+      raw_metadata = property_files._GetPropertyFilesString(
+          zip_fp, reserve_space=False)
+      streaming_metadata = property_files.Finalize(zip_fp, len(raw_metadata))
     tokens = self._parse_streaming_metadata_string(streaming_metadata)
 
     self.assertEqual(3, len(tokens))
@@ -664,7 +675,7 @@ class OtaFromTargetFilesTest(unittest.TestCase):
     entries[2] = 'metadata'
     self._verify_entries(zip_file, tokens, entries)
 
-  def test_ComputeStreamingMetadata_withExpectedLength(self):
+  def test_Finalize_assertReservedLength(self):
     entries = (
         'payload.bin',
         'payload_properties.txt',
@@ -672,36 +683,52 @@ class OtaFromTargetFilesTest(unittest.TestCase):
         'META-INF/com/android/metadata',
     )
     zip_file = self._construct_zip_package(entries)
+    property_files = StreamingPropertyFiles()
     with zipfile.ZipFile(zip_file, 'r') as zip_fp:
       # First get the raw metadata string (i.e. without padding space).
-      raw_metadata = ComputeStreamingMetadata(
+      raw_metadata = property_files._GetPropertyFilesString(
-          zip_fp,
+          zip_fp, reserve_space=False)
-          reserve_space=False)
       raw_length = len(raw_metadata)
 
       # Now pass in the exact expected length.
-      streaming_metadata = ComputeStreamingMetadata(
+      streaming_metadata = property_files.Finalize(zip_fp, raw_length)
-          zip_fp,
-          reserve_space=False,
-          expected_length=raw_length)
       self.assertEqual(raw_length, len(streaming_metadata))
 
       # Or pass in insufficient length.
       self.assertRaises(
           AssertionError,
-          ComputeStreamingMetadata,
+          property_files.Finalize,
           zip_fp,
-          reserve_space=False,
+          raw_length - 1)
-          expected_length=raw_length - 1)
 
       # Or pass in a much larger size.
-      streaming_metadata = ComputeStreamingMetadata(
+      streaming_metadata = property_files.Finalize(
           zip_fp,
-          reserve_space=False,
+          raw_length + 20)
-          expected_length=raw_length + 20)
       self.assertEqual(raw_length + 20, len(streaming_metadata))
       self.assertEqual(' ' * 20, streaming_metadata[raw_length:])
 
+  def test_Verify(self):
+    entries = (
+        'payload.bin',
+        'payload_properties.txt',
+        'care_map.txt',
+        'META-INF/com/android/metadata',
+    )
+    zip_file = self._construct_zip_package(entries)
+    property_files = StreamingPropertyFiles()
+    with zipfile.ZipFile(zip_file, 'r') as zip_fp:
+      # First get the raw metadata string (i.e. without padding space).
+      raw_metadata = property_files._GetPropertyFilesString(
+          zip_fp, reserve_space=False)
+
+      # Should pass the test if verification passes.
+      property_files.Verify(zip_fp, raw_metadata)
+
+      # Or raise on verification failure.
+      self.assertRaises(
+          AssertionError, property_files.Verify, zip_fp, raw_metadata + 'x')
+
 
 class PayloadSignerTest(unittest.TestCase):