Build recovery-two-step.img for two-step OTAs.

In two-step OTAs, we write recovery image to /boot as the first step so
that we can reboot from there and install a new recovery image to
/recovery. However, bootloader will show "Your device is corrupt"
message when booting /boot with the recovery image. Because the recovery
image encodes the path of "/recovery" as part of the signature metadata,
which fails the verified boot.

This CL generates a special "recovery-two-step.img" in addition to the
regular recovery.img. This image encodes "/boot" when being signed,
which will be flashed to /boot at stage 1/3 in a two-step OTA.

Here are the desired changes:

- 'IMAGES/recovery-two-step.img' exists in target_files.zip for non-A/B
targets (e.g. bullhead). The image should not exist for targets that
don't have a recovery partition (e.g. A/B devices like sailfish).

- <device>-img.zip should not contain 'recovery-two-step.img'.

- Nothing should change when building non-two-step OTAs. For two-step
OTAs, 'recovery-two-step.img' should be included in the OTA package;
'updater-script' should flash this image to /boot at stage 1/3.

- When building a two-step OTA with an input TF.zip that doesn't have
  IMAGES/recovery-two-step.img, it should use the existing
  IMAGES/recovery.img instead.

Bug: 32986477
Test: Tested the steps above on bullhead and sailfish.
Change-Id: I34e6c599bcf2011d4cd5c926999418b3975d6d0f
(cherry picked from commit d42e97ebb45fdc5a30799a3f37e482948d318010)
diff --git a/tools/releasetools/ota_from_target_files.py b/tools/releasetools/ota_from_target_files.py
index d3d4974..d409d94 100755
--- a/tools/releasetools/ota_from_target_files.py
+++ b/tools/releasetools/ota_from_target_files.py
@@ -467,6 +467,39 @@
       script.AssertOemProperty(prop, oem_dict.get(prop))
 
 
+def _WriteRecoveryImageToBoot(script, output_zip):
+  """Find and write recovery image to /boot in two-step OTA.
+
+  In two-step OTAs, we write recovery image to /boot as the first step so that
+  we can reboot to there and install a new recovery image to /recovery.
+  A special "recovery-two-step.img" will be preferred, which encodes the correct
+  path of "/boot". Otherwise the device may show "device is corrupt" message
+  when booting into /boot.
+
+  Fall back to using the regular recovery.img if the two-step recovery image
+  doesn't exist. Note that rebuilding the special image at this point may be
+  infeasible, because we don't have the desired boot signer and keys when
+  calling ota_from_target_files.py.
+  """
+
+  recovery_two_step_img_name = "recovery-two-step.img"
+  recovery_two_step_img_path = os.path.join(
+      OPTIONS.input_tmp, "IMAGES", recovery_two_step_img_name)
+  if os.path.exists(recovery_two_step_img_path):
+    recovery_two_step_img = common.GetBootableImage(
+        recovery_two_step_img_name, recovery_two_step_img_name,
+        OPTIONS.input_tmp, "RECOVERY")
+    common.ZipWriteStr(
+        output_zip, recovery_two_step_img_name, recovery_two_step_img.data)
+    print "two-step package: using %s in stage 1/3" % (
+        recovery_two_step_img_name,)
+    script.WriteRawImage("/boot", recovery_two_step_img_name)
+  else:
+    print "two-step package: using recovery.img in stage 1/3"
+    # The "recovery.img" entry has been written into package earlier.
+    script.WriteRawImage("/boot", "recovery.img")
+
+
 def HasRecoveryPatch(target_files_zip):
   namelist = [name for name in target_files_zip.namelist()]
   return ("SYSTEM/recovery-from-boot.p" in namelist or
@@ -616,6 +649,9 @@
     script.AppendExtra("""
 if get_stage("%(bcb_dev)s") == "2/3" then
 """ % bcb_dev)
+
+    # Stage 2/3: Write recovery image to /recovery (currently running /boot).
+    script.Comment("Stage 2/3")
     script.WriteRawImage("/recovery", "recovery.img")
     script.AppendExtra("""
 set_stage("%(bcb_dev)s", "3/3");
@@ -623,6 +659,9 @@
 else if get_stage("%(bcb_dev)s") == "3/3" then
 """ % bcb_dev)
 
+    # Stage 3/3: Make changes.
+    script.Comment("Stage 3/3")
+
   # Dump fingerprints
   script.Print("Target: %s" % CalculateFingerprint(
       oem_props, oem_dict, OPTIONS.info_dict))
@@ -722,7 +761,11 @@
 set_stage("%(bcb_dev)s", "");
 """ % bcb_dev)
     script.AppendExtra("else\n")
-    script.WriteRawImage("/boot", "recovery.img")
+
+    # Stage 1/3: Nothing to verify for full OTA. Write recovery image to /boot.
+    script.Comment("Stage 1/3")
+    _WriteRecoveryImageToBoot(script, output_zip)
+
     script.AppendExtra("""
 set_stage("%(bcb_dev)s", "2/3");
 reboot_now("%(bcb_dev)s", "");
@@ -945,6 +988,9 @@
     script.AppendExtra("""
 if get_stage("%(bcb_dev)s") == "2/3" then
 """ % bcb_dev)
+
+    # Stage 2/3: Write recovery image to /recovery (currently running /boot).
+    script.Comment("Stage 2/3")
     script.AppendExtra("sleep(20);\n")
     script.WriteRawImage("/recovery", "recovery.img")
     script.AppendExtra("""
@@ -953,6 +999,9 @@
 else if get_stage("%(bcb_dev)s") != "3/3" then
 """ % bcb_dev)
 
+    # Stage 1/3: (a) Verify the current system.
+    script.Comment("Stage 1/3")
+
   # Dump fingerprints
   script.Print("Source: %s" % CalculateFingerprint(
       oem_props, oem_dict, OPTIONS.source_info_dict))
@@ -1016,13 +1065,18 @@
   device_specific.IncrementalOTA_VerifyEnd()
 
   if OPTIONS.two_step:
-    script.WriteRawImage("/boot", "recovery.img")
+    # Stage 1/3: (b) Write recovery image to /boot.
+    _WriteRecoveryImageToBoot(script, output_zip)
+
     script.AppendExtra("""
 set_stage("%(bcb_dev)s", "2/3");
 reboot_now("%(bcb_dev)s", "");
 else
 """ % bcb_dev)
 
+    # Stage 3/3: Make changes.
+    script.Comment("Stage 3/3")
+
   # Verify the existing partitions.
   system_diff.WriteVerifyScript(script, touched_blocks_only=True)
   if vendor_diff:
@@ -1616,6 +1670,9 @@
     script.AppendExtra("""
 if get_stage("%(bcb_dev)s") == "2/3" then
 """ % bcb_dev)
+
+    # Stage 2/3: Write recovery image to /recovery (currently running /boot).
+    script.Comment("Stage 2/3")
     script.AppendExtra("sleep(20);\n")
     script.WriteRawImage("/recovery", "recovery.img")
     script.AppendExtra("""
@@ -1624,6 +1681,9 @@
 else if get_stage("%(bcb_dev)s") != "3/3" then
 """ % bcb_dev)
 
+    # Stage 1/3: (a) Verify the current system.
+    script.Comment("Stage 1/3")
+
   # Dump fingerprints
   script.Print("Source: %s" % (source_fp,))
   script.Print("Target: %s" % (target_fp,))
@@ -1668,13 +1728,18 @@
   device_specific.IncrementalOTA_VerifyEnd()
 
   if OPTIONS.two_step:
-    script.WriteRawImage("/boot", "recovery.img")
+    # Stage 1/3: (b) Write recovery image to /boot.
+    _WriteRecoveryImageToBoot(script, output_zip)
+
     script.AppendExtra("""
 set_stage("%(bcb_dev)s", "2/3");
 reboot_now("%(bcb_dev)s", "");
 else
 """ % bcb_dev)
 
+    # Stage 3/3: Make changes.
+    script.Comment("Stage 3/3")
+
   script.Comment("---- start making changes here ----")
 
   device_specific.IncrementalOTA_InstallBegin()