1 files changed, 56 insertions, 28 deletions

diff --git a/recovery.cpp b/recovery.cpp
index 0e6e4976d..f9b3bfc0b 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -509,11 +509,14 @@ static Device::BuiltinAction prompt_and_wait(Device* device, int status) {
       case Device::NO_ACTION:
         break;
 
-      case Device::REBOOT:
-      case Device::SHUTDOWN:
-      case Device::REBOOT_BOOTLOADER:
       case Device::ENTER_FASTBOOT:
       case Device::ENTER_RECOVERY:
+      case Device::REBOOT:
+      case Device::REBOOT_BOOTLOADER:
+      case Device::REBOOT_FASTBOOT:
+      case Device::REBOOT_RECOVERY:
+      case Device::REBOOT_RESCUE:
+      case Device::SHUTDOWN:
         return chosen_action;
 
       case Device::WIPE_DATA:
@@ -537,24 +540,36 @@ static Device::BuiltinAction prompt_and_wait(Device* device, int status) {
         if (!ui->IsTextVisible()) return Device::NO_ACTION;
         break;
       }
+
       case Device::APPLY_ADB_SIDELOAD:
-      case Device::APPLY_SDCARD: {
+      case Device::APPLY_SDCARD:
+      case Device::ENTER_RESCUE: {
         save_current_log = true;
-        bool adb = (chosen_action == Device::APPLY_ADB_SIDELOAD);
-        if (adb) {
-          status = apply_from_adb(ui);
+
+        bool adb = true;
+        Device::BuiltinAction reboot_action;
+        if (chosen_action == Device::ENTER_RESCUE) {
+          // Switch to graphics screen.
+          ui->ShowText(false);
+          status = ApplyFromAdb(ui, true /* rescue_mode */, &reboot_action);
+        } else if (chosen_action == Device::APPLY_ADB_SIDELOAD) {
+          status = ApplyFromAdb(ui, false /* rescue_mode */, &reboot_action);
         } else {
+          adb = false;
           status = ApplyFromSdcard(device, ui);
         }
 
+        ui->Print("\nInstall from %s completed with status %d.\n", adb ? "ADB" : "SD card", status);
+        if (status == INSTALL_REBOOT) {
+          return reboot_action;
+        }
+
         if (status != INSTALL_SUCCESS) {
           ui->SetBackground(RecoveryUI::ERROR);
           ui->Print("Installation aborted.\n");
           copy_logs(save_current_log, has_cache, sehandle);
         } else if (!ui->IsTextVisible()) {
           return Device::NO_ACTION;  // reboot if logs aren't visible
-        } else {
-          ui->Print("\nInstall from %s complete.\n", adb ? "ADB" : "SD card");
         }
         break;
       }
@@ -715,6 +730,7 @@ Device::BuiltinAction start_recovery(Device* device, const std::vector<std::stri
     { "locale", required_argument, nullptr, 0 },
     { "prompt_and_wipe_data", no_argument, nullptr, 0 },
     { "reason", required_argument, nullptr, 0 },
+    { "rescue", no_argument, nullptr, 0 },
     { "retry_count", required_argument, nullptr, 0 },
     { "security", no_argument, nullptr, 0 },
     { "show_text", no_argument, nullptr, 't' },
@@ -737,6 +753,7 @@ Device::BuiltinAction start_recovery(Device* device, const std::vector<std::stri
   size_t wipe_package_size = 0;
   bool sideload = false;
   bool sideload_auto_reboot = false;
+  bool rescue = false;
   bool just_exit = false;
   bool shutdown_after = false;
   bool fsck_unshare_blocks = false;
@@ -770,6 +787,8 @@ Device::BuiltinAction start_recovery(Device* device, const std::vector<std::stri
           should_prompt_and_wipe_data = true;
         } else if (option == "reason") {
           reason = optarg;
+        } else if (option == "rescue") {
+          rescue = true;
         } else if (option == "retry_count") {
           android::base::ParseInt(optarg, &retry_count, 0);
         } else if (option == "security") {
@@ -833,6 +852,9 @@ Device::BuiltinAction start_recovery(Device* device, const std::vector<std::stri
   ui->Print("Supported API: %d\n", kRecoveryApiVersion);
 
   int status = INSTALL_SUCCESS;
+  // next_action indicates the next target to reboot into upon finishing the install. It could be
+  // overridden to a different reboot target per user request.
+  Device::BuiltinAction next_action = shutdown_after ? Device::SHUTDOWN : Device::REBOOT;
 
   if (update_package != nullptr) {
     // It's not entirely true that we will modify the flash. But we want
@@ -916,21 +938,24 @@ Device::BuiltinAction start_recovery(Device* device, const std::vector<std::stri
       status = INSTALL_ERROR;
     }
   } else if (sideload) {
-    // 'adb reboot sideload' acts the same as user presses key combinations
-    // to enter the sideload mode. When 'sideload-auto-reboot' is used, text
-    // display will NOT be turned on by default. And it will reboot after
-    // sideload finishes even if there are errors. Unless one turns on the
-    // text display during the installation. This is to enable automated
+    // 'adb reboot sideload' acts the same as user presses key combinations to enter the sideload
+    // mode. When 'sideload-auto-reboot' is used, text display will NOT be turned on by default. And
+    // it will reboot after sideload finishes even if there are errors. This is to enable automated
     // testing.
     save_current_log = true;
     if (!sideload_auto_reboot) {
       ui->ShowText(true);
     }
-    status = apply_from_adb(ui);
+    status = ApplyFromAdb(ui, false /* rescue_mode */, &next_action);
     ui->Print("\nInstall from ADB complete (status: %d).\n", status);
     if (sideload_auto_reboot) {
+      status = INSTALL_REBOOT;
       ui->Print("Rebooting automatically.\n");
     }
+  } else if (rescue) {
+    save_current_log = true;
+    status = ApplyFromAdb(ui, true /* rescue_mode */, &next_action);
+    ui->Print("\nInstall from ADB complete (status: %d).\n", status);
   } else if (fsck_unshare_blocks) {
     if (!do_fsck_unshare_blocks()) {
       status = INSTALL_ERROR;
@@ -953,23 +978,26 @@ Device::BuiltinAction start_recovery(Device* device, const std::vector<std::stri
     }
   }
 
-  Device::BuiltinAction after = shutdown_after ? Device::SHUTDOWN : Device::REBOOT;
-  // 1. If the recovery menu is visible, prompt and wait for commands.
-  // 2. If the state is INSTALL_NONE, wait for commands. (i.e. In user build, manually reboot into
-  //    recovery to sideload a package.)
-  // 3. sideload_auto_reboot is an option only available in user-debug build, reboot the device
-  //    without waiting.
-  // 4. In all other cases, reboot the device. Therefore, normal users will observe the device
-  //    reboot after it shows the "error" screen for 5s.
-  if ((status == INSTALL_NONE && !sideload_auto_reboot) || ui->IsTextVisible()) {
-    Device::BuiltinAction temp = prompt_and_wait(device, status);
-    if (temp != Device::NO_ACTION) {
-      after = temp;
+  // Determine the next action.
+  //  - If the state is INSTALL_REBOOT, device will reboot into the target as specified in
+  //    `next_action`.
+  //  - If the recovery menu is visible, prompt and wait for commands.
+  //  - If the state is INSTALL_NONE, wait for commands (e.g. in user build, one manually boots
+  //    into recovery to sideload a package or to wipe the device).
+  //  - In all other cases, reboot the device. Therefore, normal users will observe the device
+  //    rebooting a) immediately upon successful finish (INSTALL_SUCCESS); or b) an "error" screen
+  //    for 5s followed by an automatic reboot.
+  if (status != INSTALL_REBOOT) {
+    if (status == INSTALL_NONE || ui->IsTextVisible()) {
+      Device::BuiltinAction temp = prompt_and_wait(device, status);
+      if (temp != Device::NO_ACTION) {
+        next_action = temp;
+      }
     }
   }
 
   // Save logs and clean up before rebooting or shutting down.
   finish_recovery();
 
-  return after;
+  return next_action;
 }