Merge "Change the naming policy of system_$(VER)"
diff --git a/tools/releasetools/ota_from_target_files.py b/tools/releasetools/ota_from_target_files.py
index 12b01c4..cd497b2 100755
--- a/tools/releasetools/ota_from_target_files.py
+++ b/tools/releasetools/ota_from_target_files.py
@@ -92,6 +92,24 @@
       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 @@
 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 @@
 
   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 @@
     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 @@
   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 @@
   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 @@
                          "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 @@
                                  "extra_script=",
                                  "worker_threads=",
                                  "two_step",
+                                 "include_secondary",
                                  "no_signing",
                                  "block",
                                  "binary=",
diff --git a/tools/releasetools/test_ota_from_target_files.py b/tools/releasetools/test_ota_from_target_files.py
index 849ca1d..6edf80c 100644
--- a/tools/releasetools/test_ota_from_target_files.py
+++ b/tools/releasetools/test_ota_from_target_files.py
@@ -23,10 +23,38 @@
 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 @@
         },
         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 @@
     common.Cleanup()
 
   @staticmethod
-  def _construct_target_files():
-    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)))
-
-    return target_files
-
-  def _create_payload_full(self):
-    target_file = self._construct_target_files()
+  def _create_payload_full(secondary=False):
+    target_file = construct_target_files(secondary)
     payload = Payload()
     payload.Generate(target_file)
     return payload
 
-  def _create_payload_incremental(self):
-    target_file = self._construct_target_files()
-    source_file = self._construct_target_files()
+  @staticmethod
+  def _create_payload_incremental():
+    target_file = construct_target_files()
+    source_file = construct_target_files()
     payload = Payload()
     payload.Generate(target_file, source_file)
     return payload
@@ -641,8 +664,8 @@
     self.assertTrue(os.path.exists(payload.payload_file))
 
   def test_Generate_additionalArgs(self):
-    target_file = self._construct_target_files()
-    source_file = self._construct_target_files()
+    target_file = 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 @@
     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 @@
     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)