releasetools: Support packaging secondary payload.

By default, an A/B OTA package doesn't contain the images for the
secondary slot (e.g. system_other.img). Specifying
"--include_secondary" that's introduced in this CL allows generating
a separate payload that will install secondary slot images. Both
payloads will be added to the generated A/B OTA package.

An example A/B OTA package with secondary payload
  |
  +-- payload.bin
  |
  +-- payload_properties.txt
  |
  +-- secondary/payload.bin
  |
  +-- secondary/payload_properties.txt
  |
  +-- ...

Such a package needs to be applied in a two-stage manner. During the
first stage, the updater applies the primary payload only. Upon
finishing, it reboots the device into the newly updated slot. It then
continues to install the secondary payload to the inactive slot, but
without switching the active slot at the end (needs the matching support
in update_engine, i.e. SWITCH_SLOT_ON_REBOOT flag).

Due to the special install procedure, the secondary payload will be
always generated as a full payload.

Bug: 35724498
Test: Generate full and incremental OTAs with --include_secondary. Check
      the generated OTAs.
Test: python -m unittest test_ota_from_target_files
Change-Id: I975e826bec492e86eb400f99de0c355a32420127

tools/releasetools/ota_from_target_files.py

@@ -92,6 +92,24 @@ Usage:  ota_from_target_files [flags] input_target_files output_ota_package
       first, so that any changes made to the system partition are done
       using the new recovery (new kernel, etc.).
 
+  --include_secondary
+      Additionally include the payload for secondary slot images (default:
+      False). Only meaningful when generating A/B OTAs.
+
+      By default, an A/B OTA package doesn't contain the images for the
+      secondary slot (e.g. system_other.img). Specifying this flag allows
+      generating a separate payload that will install secondary slot images.
+
+      Such a package needs to be applied in a two-stage manner, with a reboot
+      in-between. During the first stage, the updater applies the primary
+      payload only. Upon finishing, it reboots the device into the newly updated
+      slot. It then continues to install the secondary payload to the inactive
+      slot, but without switching the active slot at the end (needs the matching
+      support in update_engine, i.e. SWITCH_SLOT_ON_REBOOT flag).
+
+      Due to the special install procedure, the secondary payload will be always
+      generated as a full payload.
+
   --block
       Generate a block-based OTA for non-A/B device. We have deprecated the
       support for file-based OTA since O. Block-based OTA will be used by
@@ -159,6 +177,7 @@ OPTIONS.worker_threads = multiprocessing.cpu_count() // 2
 if OPTIONS.worker_threads == 0:
   OPTIONS.worker_threads = 1
 OPTIONS.two_step = False
+OPTIONS.include_secondary = False
 OPTIONS.no_signing = False
 OPTIONS.block_based = True
 OPTIONS.updater_binary = None
@@ -364,6 +383,8 @@ class Payload(object):
 
   PAYLOAD_BIN = 'payload.bin'
   PAYLOAD_PROPERTIES_TXT = 'payload_properties.txt'
+  SECONDARY_PAYLOAD_BIN = 'secondary/payload.bin'
+  SECONDARY_PAYLOAD_PROPERTIES_TXT = 'secondary/payload_properties.txt'
 
   def __init__(self):
     # The place where the output from the subprocess should go.
@@ -456,22 +477,31 @@ class Payload(object):
     self.payload_file = signed_payload_file
     self.payload_properties = properties_file
 
-  def WriteToZip(self, output_zip):
+  def WriteToZip(self, output_zip, secondary=False):
     """Writes the payload to the given zip.
 
     Args:
       output_zip: The output ZipFile instance.
+      secondary: Whether the payload should be packed as secondary payload
+          (default: False).
     """
     assert self.payload_file is not None
     assert self.payload_properties is not None
 
+    if secondary:
+      payload_arcname = Payload.SECONDARY_PAYLOAD_BIN
+      payload_properties_arcname = Payload.SECONDARY_PAYLOAD_PROPERTIES_TXT
+    else:
+      payload_arcname = Payload.PAYLOAD_BIN
+      payload_properties_arcname = Payload.PAYLOAD_PROPERTIES_TXT
+
     # Add the signed payload file and properties into the zip. In order to
     # support streaming, we pack them as ZIP_STORED. So these entries can be
     # read directly with the offset and length pairs.
-    common.ZipWrite(output_zip, self.payload_file, arcname=Payload.PAYLOAD_BIN,
+    common.ZipWrite(output_zip, self.payload_file, arcname=payload_arcname,
                     compress_type=zipfile.ZIP_STORED)
     common.ZipWrite(output_zip, self.payload_properties,
-                    arcname=Payload.PAYLOAD_PROPERTIES_TXT,
+                    arcname=payload_properties_arcname,
                     compress_type=zipfile.ZIP_STORED)
 
 
@@ -1162,6 +1192,47 @@ endif;
   WriteMetadata(metadata, output_zip)
 
 
+def GetTargetFilesZipForSecondaryImages(input_file):
+  """Returns a target-files.zip file for generating secondary payload.
+
+  Although the original target-files.zip already contains secondary slot
+  images (i.e. IMAGES/system_other.img), we need to rename the files to the
+  ones without _other suffix. Note that we cannot instead modify the names in
+  META/ab_partitions.txt, because there are no matching partitions on device.
+
+  For the partitions that don't have secondary images, the ones for primary
+  slot will be used. This is to ensure that we always have valid boot, vbmeta,
+  bootloader images in the inactive slot.
+
+  Args:
+    input_file: The input target-files.zip file.
+
+  Returns:
+    The filename of the target-files.zip for generating secondary payload.
+  """
+  target_file = common.MakeTempFile(prefix="targetfiles-", suffix=".zip")
+  target_zip = zipfile.ZipFile(target_file, 'w', allowZip64=True)
+
+  input_tmp, input_zip = common.UnzipTemp(input_file, UNZIP_PATTERN)
+  for info in input_zip.infolist():
+    unzipped_file = os.path.join(input_tmp, *info.filename.split('/'))
+    if info.filename == 'IMAGES/system_other.img':
+      common.ZipWrite(target_zip, unzipped_file, arcname='IMAGES/system.img')
+
+    # Primary images and friends need to be skipped explicitly.
+    elif info.filename in ('IMAGES/system.img',
+                           'IMAGES/system.map'):
+      pass
+
+    elif info.filename.startswith(('META/', 'IMAGES/')):
+      common.ZipWrite(target_zip, unzipped_file, arcname=info.filename)
+
+  common.ZipClose(input_zip)
+  common.ZipClose(target_zip)
+
+  return target_file
+
+
 def WriteABOTAPackageWithBrilloScript(target_file, output_file,
                                       source_file=None):
   """Generate an Android OTA package that has A/B update payload."""
@@ -1236,11 +1307,23 @@ def WriteABOTAPackageWithBrilloScript(target_file, output_file,
   payload.Generate(target_file, source_file)
 
   # Sign the payload.
-  payload.Sign(PayloadSigner())
+  payload_signer = PayloadSigner()
+  payload.Sign(payload_signer)
 
   # Write the payload into output zip.
   payload.WriteToZip(output_zip)
 
+  # Generate and include the secondary payload that installs secondary images
+  # (e.g. system_other.img).
+  if OPTIONS.include_secondary:
+    # We always include a full payload for the secondary slot, even when
+    # building an incremental OTA. See the comments for "--include_secondary".
+    secondary_target_file = GetTargetFilesZipForSecondaryImages(target_file)
+    secondary_payload = Payload()
+    secondary_payload.Generate(secondary_target_file)
+    secondary_payload.Sign(payload_signer)
+    secondary_payload.WriteToZip(output_zip, secondary=True)
+
   # If dm-verity is supported for the device, copy contents of care_map
   # into A/B OTA package.
   target_zip = zipfile.ZipFile(target_file, "r")
@@ -1339,6 +1422,8 @@ def main(argv):
                          "integers are allowed." % (a, o))
     elif o in ("-2", "--two_step"):
       OPTIONS.two_step = True
+    elif o == "--include_secondary":
+      OPTIONS.include_secondary = True
     elif o == "--no_signing":
       OPTIONS.no_signing = True
     elif o == "--verify":
@@ -1378,6 +1463,7 @@ def main(argv):
                                  "extra_script=",
                                  "worker_threads=",
                                  "two_step",
+                                 "include_secondary",
                                  "no_signing",
                                  "block",
                                  "binary=",

tools/releasetools/test_ota_from_target_files.py

@@ -23,10 +23,38 @@ import zipfile
 import common
 import test_utils
 from ota_from_target_files import (
-    _LoadOemDicts, BuildInfo, GetPackageMetadata, Payload, PayloadSigner,
+    _LoadOemDicts, BuildInfo, GetPackageMetadata,
+    GetTargetFilesZipForSecondaryImages, Payload, PayloadSigner,
     WriteFingerprintAssertion)
 
 
+def construct_target_files(secondary=False):
+  """Returns a target-files.zip file for generating OTA packages."""
+  target_files = common.MakeTempFile(prefix='target_files-', suffix='.zip')
+  with zipfile.ZipFile(target_files, 'w') as target_files_zip:
+    # META/update_engine_config.txt
+    target_files_zip.writestr(
+        'META/update_engine_config.txt',
+        "PAYLOAD_MAJOR_VERSION=2\nPAYLOAD_MINOR_VERSION=4\n")
+
+    # META/ab_partitions.txt
+    ab_partitions = ['boot', 'system', 'vendor']
+    target_files_zip.writestr(
+        'META/ab_partitions.txt',
+        '\n'.join(ab_partitions))
+
+    # Create dummy images for each of them.
+    for partition in ab_partitions:
+      target_files_zip.writestr('IMAGES/' + partition + '.img',
+                                os.urandom(len(partition)))
+
+    if secondary:
+      target_files_zip.writestr('IMAGES/system_other.img',
+                                os.urandom(len("system_other")))
+
+  return target_files
+
+
 class MockScriptWriter(object):
   """A class that mocks edify_generator.EdifyGenerator.
 
@@ -500,6 +528,21 @@ class OtaFromTargetFilesTest(unittest.TestCase):
         },
         metadata)
 
+  def test_GetTargetFilesZipForSecondaryImages(self):
+    input_file = construct_target_files(secondary=True)
+    target_file = GetTargetFilesZipForSecondaryImages(input_file)
+
+    with zipfile.ZipFile(target_file) as verify_zip:
+      namelist = verify_zip.namelist()
+
+    self.assertIn('META/ab_partitions.txt', namelist)
+    self.assertIn('IMAGES/boot.img', namelist)
+    self.assertIn('IMAGES/system.img', namelist)
+    self.assertIn('IMAGES/vendor.img', namelist)
+
+    self.assertNotIn('IMAGES/system_other.img', namelist)
+    self.assertNotIn('IMAGES/system.map', namelist)
+
 
 class PayloadSignerTest(unittest.TestCase):
 
@@ -598,36 +641,16 @@ class PayloadTest(unittest.TestCase):
     common.Cleanup()
 
   @staticmethod
-  def _construct_target_files():
+  def _create_payload_full(secondary=False):
-    target_files = common.MakeTempFile(prefix='target_files-', suffix='.zip')
+    target_file = construct_target_files(secondary)
-    with zipfile.ZipFile(target_files, 'w') as target_files_zip:
-      # META/update_engine_config.txt
-      target_files_zip.writestr(
-          'META/update_engine_config.txt',
-          "PAYLOAD_MAJOR_VERSION=2\nPAYLOAD_MINOR_VERSION=4\n")
-
-      # META/ab_partitions.txt
-      ab_partitions = ['boot', 'system', 'vendor']
-      target_files_zip.writestr(
-          'META/ab_partitions.txt',
-          '\n'.join(ab_partitions))
-
-      # Create dummy images for each of them.
-      for partition in ab_partitions:
-        target_files_zip.writestr('IMAGES/' + partition + '.img',
-                                  os.urandom(len(partition)))
-
-    return target_files
-
-  def _create_payload_full(self):
-    target_file = self._construct_target_files()
     payload = Payload()
     payload.Generate(target_file)
     return payload
 
-  def _create_payload_incremental(self):
+  @staticmethod
-    target_file = self._construct_target_files()
+  def _create_payload_incremental():
-    source_file = self._construct_target_files()
+    target_file = construct_target_files()
+    source_file = construct_target_files()
     payload = Payload()
     payload.Generate(target_file, source_file)
     return payload
@@ -641,8 +664,8 @@ class PayloadTest(unittest.TestCase):
     self.assertTrue(os.path.exists(payload.payload_file))
 
   def test_Generate_additionalArgs(self):
-    target_file = self._construct_target_files()
+    target_file = construct_target_files()
-    source_file = self._construct_target_files()
+    source_file = construct_target_files()
     payload = Payload()
     # This should work the same as calling payload.Generate(target_file,
     # source_file).
@@ -651,7 +674,7 @@ class PayloadTest(unittest.TestCase):
     self.assertTrue(os.path.exists(payload.payload_file))
 
   def test_Generate_invalidInput(self):
-    target_file = self._construct_target_files()
+    target_file = construct_target_files()
     common.ZipDelete(target_file, 'IMAGES/vendor.img')
     payload = Payload()
     self.assertRaises(AssertionError, payload.Generate, target_file)
@@ -732,3 +755,25 @@ class PayloadTest(unittest.TestCase):
     output_file = common.MakeTempFile(suffix='.zip')
     with zipfile.ZipFile(output_file, 'w') as output_zip:
       self.assertRaises(AssertionError, payload.WriteToZip, output_zip)
+
+  def test_WriteToZip_secondary(self):
+    payload = self._create_payload_full(secondary=True)
+    payload.Sign(PayloadSigner())
+
+    output_file = common.MakeTempFile(suffix='.zip')
+    with zipfile.ZipFile(output_file, 'w') as output_zip:
+      payload.WriteToZip(output_zip, secondary=True)
+
+    with zipfile.ZipFile(output_file) as verify_zip:
+      # First make sure we have the essential entries.
+      namelist = verify_zip.namelist()
+      self.assertIn(Payload.SECONDARY_PAYLOAD_BIN, namelist)
+      self.assertIn(Payload.SECONDARY_PAYLOAD_PROPERTIES_TXT, namelist)
+
+      # Then assert these entries are stored.
+      for entry_info in verify_zip.infolist():
+        if entry_info.filename not in (
+            Payload.SECONDARY_PAYLOAD_BIN,
+            Payload.SECONDARY_PAYLOAD_PROPERTIES_TXT):
+          continue
+        self.assertEqual(zipfile.ZIP_STORED, entry_info.compress_type)