4 files changed, 63 insertions, 20 deletions

diff --git a/install/adb_install.cpp b/install/adb_install.cpp
index d79f6f4b0..9dfe0407f 100644
--- a/install/adb_install.cpp
+++ b/install/adb_install.cpp
@@ -121,19 +121,20 @@ static auto AdbInstallPackageHandler(RecoveryUI* ui, int* result) {
 
 static auto AdbRebootHandler(MinadbdCommand command, int* result,
                              Device::BuiltinAction* reboot_action) {
+  // Use Device::REBOOT_{FASTBOOT,RECOVERY,RESCUE}, instead of the ones with ENTER_. This allows
+  // rebooting back into fastboot/recovery/rescue mode through bootloader, which may use a newly
+  // installed bootloader/recovery image.
   switch (command) {
     case MinadbdCommand::kRebootBootloader:
       *reboot_action = Device::REBOOT_BOOTLOADER;
       break;
     case MinadbdCommand::kRebootFastboot:
-      *reboot_action = Device::ENTER_FASTBOOT;
+      *reboot_action = Device::REBOOT_FASTBOOT;
       break;
     case MinadbdCommand::kRebootRecovery:
-      *reboot_action = Device::ENTER_RECOVERY;
+      *reboot_action = Device::REBOOT_RECOVERY;
       break;
     case MinadbdCommand::kRebootRescue:
-      // Use Device::REBOOT_RESCUE instead of Device::ENTER_RESCUE. This allows rebooting back into
-      // rescue mode (potentially using a newly installed recovery image).
       *reboot_action = Device::REBOOT_RESCUE;
       break;
     case MinadbdCommand::kRebootAndroid:
@@ -180,7 +181,7 @@ static bool HandleMessageFromMinadbd(int socket_fd,
 
 // TODO(xunchang) add a wrapper function and kill the minadbd service there.
 static void ListenAndExecuteMinadbdCommands(
-    pid_t minadbd_pid, android::base::unique_fd&& socket_fd,
+    RecoveryUI* ui, pid_t minadbd_pid, android::base::unique_fd&& socket_fd,
     const std::map<MinadbdCommand, CommandFunction>& command_map) {
   android::base::unique_fd epoll_fd(epoll_create1(O_CLOEXEC));
   if (epoll_fd == -1) {
@@ -203,6 +204,10 @@ static void ListenAndExecuteMinadbdCommands(
   // Set the timeout to be 300s when waiting for minadbd commands.
   constexpr int TIMEOUT_MILLIS = 300 * 1000;
   while (true) {
+    // Reset the progress bar and the background image before each command.
+    ui->SetProgressType(RecoveryUI::EMPTY);
+    ui->SetBackground(RecoveryUI::NO_COMMAND);
+
     // Poll for the status change of the socket_fd, and handle the message if the fd is ready to
     // read.
     int event_count =
@@ -266,7 +271,8 @@ static void ListenAndExecuteMinadbdCommands(
 //                               b11. exit the listening loop
 //
 static void CreateMinadbdServiceAndExecuteCommands(
-    const std::map<MinadbdCommand, CommandFunction>& command_map, bool rescue_mode) {
+    RecoveryUI* ui, const std::map<MinadbdCommand, CommandFunction>& command_map,
+    bool rescue_mode) {
   signal(SIGPIPE, SIG_IGN);
 
   android::base::unique_fd recovery_socket;
@@ -305,8 +311,8 @@ static void CreateMinadbdServiceAndExecuteCommands(
     return;
   }
 
-  std::thread listener_thread(ListenAndExecuteMinadbdCommands, child, std::move(recovery_socket),
-                              std::ref(command_map));
+  std::thread listener_thread(ListenAndExecuteMinadbdCommands, ui, child,
+                              std::move(recovery_socket), std::ref(command_map));
   if (listener_thread.joinable()) {
     listener_thread.join();
   }
@@ -357,7 +363,7 @@ int ApplyFromAdb(RecoveryUI* ui, bool rescue_mode, Device::BuiltinAction* reboot
       std::bind(&AdbRebootHandler, MinadbdCommand::kRebootRescue, &install_result, reboot_action) },
   };
 
-  CreateMinadbdServiceAndExecuteCommands(command_map, rescue_mode);
+  CreateMinadbdServiceAndExecuteCommands(ui, command_map, rescue_mode);
 
   // Clean up before switching to the older state, for example setting the state
   // to none sets sys/class/android_usb/android0/enable to 0.
diff --git a/recovery.cpp b/recovery.cpp
index 5bd9b1728..f9b3bfc0b 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -509,12 +509,14 @@ static Device::BuiltinAction prompt_and_wait(Device* device, int status) {
       case Device::NO_ACTION:
         break;
 
+      case Device::ENTER_FASTBOOT:
+      case Device::ENTER_RECOVERY:
       case Device::REBOOT:
-      case Device::SHUTDOWN:
       case Device::REBOOT_BOOTLOADER:
+      case Device::REBOOT_FASTBOOT:
+      case Device::REBOOT_RECOVERY:
       case Device::REBOOT_RESCUE:
-      case Device::ENTER_FASTBOOT:
-      case Device::ENTER_RECOVERY:
+      case Device::SHUTDOWN:
         return chosen_action;
 
       case Device::WIPE_DATA:
@@ -728,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' },
@@ -750,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;
@@ -783,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") {
@@ -946,6 +952,10 @@ Device::BuiltinAction start_recovery(Device* device, const std::vector<std::stri
       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;
diff --git a/recovery_main.cpp b/recovery_main.cpp
index 18abff765..de8ac1f42 100644
--- a/recovery_main.cpp
+++ b/recovery_main.cpp
@@ -155,9 +155,11 @@ static std::vector<std::string> get_args(const int argc, char** const argv) {
   }
 
   // Finally, if no arguments were specified, check whether we should boot
-  // into fastboot.
+  // into fastboot or rescue mode.
   if (args.size() == 1 && boot_command == "boot-fastboot") {
     args.emplace_back("--fastboot");
+  } else if (args.size() == 1 && boot_command == "boot-rescue") {
+    args.emplace_back("--rescue");
   }
 
   return args;
@@ -470,6 +472,7 @@ int main(int argc, char** argv) {
     switch (ret) {
       case Device::SHUTDOWN:
         ui->Print("Shutting down...\n");
+        // TODO: Move all the reboots to reboot(), which should conditionally set quiescent flag.
         android::base::SetProperty(ANDROID_RB_PROPERTY, "shutdown,");
         break;
 
@@ -478,11 +481,32 @@ int main(int argc, char** argv) {
         android::base::SetProperty(ANDROID_RB_PROPERTY, "reboot,bootloader");
         break;
 
-      case Device::REBOOT_RESCUE:
-        ui->Print("Rebooting to rescue...\n");
-        android::base::SetProperty(ANDROID_RB_PROPERTY, "reboot,rescue");
+      case Device::REBOOT_FASTBOOT:
+        ui->Print("Rebooting to recovery/fastboot...\n");
+        android::base::SetProperty(ANDROID_RB_PROPERTY, "reboot,fastboot");
         break;
 
+      case Device::REBOOT_RECOVERY:
+        ui->Print("Rebooting to recovery...\n");
+        reboot("reboot,recovery");
+        break;
+
+      case Device::REBOOT_RESCUE: {
+        // Not using `reboot("reboot,rescue")`, as it requires matching support in kernel and/or
+        // bootloader.
+        bootloader_message boot = {};
+        strlcpy(boot.command, "boot-rescue", sizeof(boot.command));
+        std::string err;
+        if (!write_bootloader_message(boot, &err)) {
+          LOG(ERROR) << "Failed to write bootloader message: " << err;
+          // Stay under recovery on failure.
+          continue;
+        }
+        ui->Print("Rebooting to recovery/rescue...\n");
+        reboot("reboot,recovery");
+        break;
+      }
+
       case Device::ENTER_FASTBOOT:
         if (logical_partitions_mapped()) {
           ui->Print("Partitions may be mounted - rebooting to enter fastboot.");
diff --git a/recovery_ui/include/recovery_ui/device.h b/recovery_ui/include/recovery_ui/device.h
index 09b5d1f4d..7c76cdb0a 100644
--- a/recovery_ui/include/recovery_ui/device.h
+++ b/recovery_ui/include/recovery_ui/device.h
@@ -33,6 +33,10 @@ class Device {
   static constexpr const int kHighlightDown = -3;
   static constexpr const int kInvokeItem = -4;
 
+  // ENTER vs REBOOT: The latter will trigger a reboot that goes through bootloader, which allows
+  // using a new bootloader / recovery image if applicable. For example, REBOOT_RESCUE goes from
+  // rescue -> bootloader -> rescue, whereas ENTER_RESCUE switches from recovery -> rescue
+  // directly.
   enum BuiltinAction {
     NO_ACTION = 0,
     REBOOT = 1,
@@ -50,11 +54,10 @@ class Device {
     KEY_INTERRUPTED = 13,
     ENTER_FASTBOOT = 14,
     ENTER_RECOVERY = 15,
-    // ENTER vs REBOOT: The latter will trigger a reboot that uses `rescue` as the reboot target.
-    // So it goes from rescue -> bootloader -> rescue, whereas ENTER_RESCUE switches from recovery
-    // -> rescue directly.
     ENTER_RESCUE = 16,
-    REBOOT_RESCUE = 17,
+    REBOOT_FASTBOOT = 17,
+    REBOOT_RECOVERY = 18,
+    REBOOT_RESCUE = 19,
   };
 
   explicit Device(RecoveryUI* ui);