3 files changed, 44 insertions, 18 deletions

diff --git a/default_device.cpp b/default_device.cpp
index 1f181318e..a25f05f8e 100644
--- a/default_device.cpp
+++ b/default_device.cpp
@@ -29,6 +29,8 @@ static const char* ITEMS[] =  {"reboot system now",
                                "apply update from ADB",
                                "wipe data/factory reset",
                                "wipe cache partition",
+                               "reboot to bootloader",
+                               "power down",
                                NULL };
 
 class DefaultDevice : public Device {
@@ -65,6 +67,8 @@ class DefaultDevice : public Device {
           case 1: return APPLY_ADB_SIDELOAD;
           case 2: return WIPE_DATA;
           case 3: return WIPE_CACHE;
+          case 4: return REBOOT_BOOTLOADER;
+          case 5: return SHUTDOWN;
           default: return NO_ACTION;
         }
     }
diff --git a/device.h b/device.h
index 583de75ef..df71377c1 100644
--- a/device.h
+++ b/device.h
@@ -66,7 +66,8 @@ class Device {
     virtual int HandleMenuKey(int key, int visible) = 0;
 
     enum BuiltinAction { NO_ACTION, REBOOT, APPLY_EXT, APPLY_CACHE,
-                         APPLY_ADB_SIDELOAD, WIPE_DATA, WIPE_CACHE };
+                         APPLY_ADB_SIDELOAD, WIPE_DATA, WIPE_CACHE,
+                         REBOOT_BOOTLOADER, SHUTDOWN };
 
     // Perform a recovery action selected from the menu.
     // 'menu_position' will be the item number of the selected menu
diff --git a/recovery.cpp b/recovery.cpp
index 8d3731540..09af5f8b1 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -771,7 +771,10 @@ wipe_data(int confirm, Device* device) {
     ui->Print("Data wipe complete.\n");
 }
 
-static void
+// Return REBOOT, SHUTDOWN, or REBOOT_BOOTLOADER.  Returning NO_ACTION
+// means to take the default, which is to reboot or shutdown depending
+// on if the --shutdown_after flag was passed to recovery.
+static Device::BuiltinAction
 prompt_and_wait(Device* device, int status) {
     const char* const* headers = prepend_title(device->GetMenuHeaders());
 
@@ -795,23 +798,28 @@ prompt_and_wait(Device* device, int status) {
         // device-specific code may take some action here.  It may
         // return one of the core actions handled in the switch
         // statement below.
-        chosen_item = device->InvokeMenuItem(chosen_item);
+        Device::BuiltinAction chosen_action = device->InvokeMenuItem(chosen_item);
 
         int wipe_cache;
-        switch (chosen_item) {
+        switch (chosen_action) {
+            case Device::NO_ACTION:
+                break;
+
             case Device::REBOOT:
-                return;
+            case Device::SHUTDOWN:
+            case Device::REBOOT_BOOTLOADER:
+                return chosen_action;
 
             case Device::WIPE_DATA:
                 wipe_data(ui->IsTextVisible(), device);
-                if (!ui->IsTextVisible()) return;
+                if (!ui->IsTextVisible()) return Device::NO_ACTION;
                 break;
 
             case Device::WIPE_CACHE:
                 ui->Print("\n-- Wiping cache...\n");
                 erase_volume("/cache");
                 ui->Print("Cache wipe complete.\n");
-                if (!ui->IsTextVisible()) return;
+                if (!ui->IsTextVisible()) return Device::NO_ACTION;
                 break;
 
             case Device::APPLY_EXT:
@@ -829,7 +837,7 @@ prompt_and_wait(Device* device, int status) {
                         ui->SetBackground(RecoveryUI::ERROR);
                         ui->Print("Installation aborted.\n");
                     } else if (!ui->IsTextVisible()) {
-                        return;  // reboot if logs aren't visible
+                        return Device::NO_ACTION;  // reboot if logs aren't visible
                     } else {
                         ui->Print("\nInstall from sdcard complete.\n");
                     }
@@ -852,7 +860,7 @@ prompt_and_wait(Device* device, int status) {
                         ui->SetBackground(RecoveryUI::ERROR);
                         ui->Print("Installation aborted.\n");
                     } else if (!ui->IsTextVisible()) {
-                        return;  // reboot if logs aren't visible
+                        return Device::NO_ACTION;  // reboot if logs aren't visible
                     } else {
                         ui->Print("\nInstall from cache complete.\n");
                     }
@@ -867,7 +875,7 @@ prompt_and_wait(Device* device, int status) {
                         ui->Print("Installation aborted.\n");
                         copy_logs();
                     } else if (!ui->IsTextVisible()) {
-                        return;  // reboot if logs aren't visible
+                        return Device::NO_ACTION;  // reboot if logs aren't visible
                     } else {
                         ui->Print("\nInstall from ADB complete.\n");
                     }
@@ -1074,18 +1082,31 @@ main(int argc, char **argv) {
         copy_logs();
         ui->SetBackground(RecoveryUI::ERROR);
     }
+    Device::BuiltinAction after = shutdown_after ? Device::SHUTDOWN : Device::REBOOT;
     if (status != INSTALL_SUCCESS || ui->IsTextVisible()) {
-        prompt_and_wait(device, status);
+        Device::BuiltinAction temp = prompt_and_wait(device, status);
+        if (temp != Device::NO_ACTION) after = temp;
     }
 
-    // Otherwise, get ready to boot the main system...
+    // Save logs and clean up before rebooting or shutting down.
     finish_recovery(send_intent);
-    if (shutdown_after) {
-        ui->Print("Shutting down...\n");
-        property_set(ANDROID_RB_PROPERTY, "shutdown,");
-    } else {
-        ui->Print("Rebooting...\n");
-        property_set(ANDROID_RB_PROPERTY, "reboot,");
+
+    switch (after) {
+        case Device::SHUTDOWN:
+            ui->Print("Shutting down...\n");
+            property_set(ANDROID_RB_PROPERTY, "shutdown,");
+            break;
+
+        case Device::REBOOT_BOOTLOADER:
+            ui->Print("Rebooting to bootloader...\n");
+            property_set(ANDROID_RB_PROPERTY, "reboot,bootloader");
+            break;
+
+        default:
+            ui->Print("Rebooting...\n");
+            property_set(ANDROID_RB_PROPERTY, "reboot,");
+            break;
     }
+    sleep(5); // should reboot before this finishes
     return EXIT_SUCCESS;
 }