3 files changed, 26 insertions, 1 deletions

diff --git a/device.cpp b/device.cpp
index fd1a987..c77ebf6 100644
--- a/device.cpp
+++ b/device.cpp
@@ -29,6 +29,8 @@ static const char* MENU_ITEMS[] = {
     NULL
 };
 
+extern int ui_root_menu;
+
 const char* const* Device::GetMenuItems() {
   return MENU_ITEMS;
 }
@@ -54,18 +56,31 @@ int Device::HandleMenuKey(int key, int visible) {
   }
 
   switch (key) {
+    case KEY_RIGHTSHIFT:
     case KEY_DOWN:
     case KEY_VOLUMEDOWN:
+    case KEY_MENU:
       return kHighlightDown;
 
+    case KEY_LEFTSHIFT:
     case KEY_UP:
     case KEY_VOLUMEUP:
+    case KEY_SEARCH:
       return kHighlightUp;
 
     case KEY_ENTER:
     case KEY_POWER:
+    case BTN_MOUSE:
+    case KEY_HOME:
+    case KEY_HOMEPAGE:
+    case KEY_SEND:
       return kInvokeItem;
 
+    case KEY_BACKSPACE:
+    case KEY_BACK:
+      if (!ui_root_menu)
+        return kGoBack;
+
     default:
       // If you have all of the above buttons, any other buttons
       // are ignored. Otherwise, any button cycles the highlight.
diff --git a/device.h b/device.h
index 393b507..1faad5f 100644
--- a/device.h
+++ b/device.h
@@ -90,6 +90,7 @@ class Device {
     static const int kHighlightUp = -2;
     static const int kHighlightDown = -3;
     static const int kInvokeItem = -4;
+    static const int kGoBack = -5;
 
     // Called before and after we do a wipe data/factory reset operation,
     // either via a reboot from the main system with the --wipe_data flag,
diff --git a/recovery.cpp b/recovery.cpp
index 5288dab..b7641aa 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -600,7 +600,7 @@ get_menu_selection(const char* const * headers, const char* const * items,
     int selected = initial_selection;
     int chosen_item = -1;
 
-    while (chosen_item < 0) {
+    while (chosen_item < 0 && chosen_item != Device::kGoBack) {
         int key = ui->WaitKey();
         int visible = ui->IsTextVisible();
 
@@ -612,6 +612,8 @@ get_menu_selection(const char* const * headers, const char* const * items,
                 ui->EndMenu();
                 return 0; // XXX fixme
             }
+        } else if (key == -2) { // we are returning from ui_cancel_wait_key(): trigger a GO_BACK
+            return Device::kGoBack;
         }
 
         int action = device->HandleMenuKey(key, visible);
@@ -629,6 +631,9 @@ get_menu_selection(const char* const * headers, const char* const * items,
                     break;
                 case Device::kNoAction:
                     break;
+                case Device::kGoBack:
+                    chosen_item = Device::kGoBack;
+                    break;
             }
         } else if (!menu_only) {
             chosen_item = action;
@@ -864,6 +869,8 @@ static int apply_from_sdcard(Device* device, bool* wipe_cache) {
     return status;
 }
 
+int ui_root_menu = 0;
+
 // 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.
@@ -871,6 +878,7 @@ static Device::BuiltinAction
 prompt_and_wait(Device* device, int status) {
     for (;;) {
         finish_recovery(NULL);
+        ui_root_menu = 1;
         switch (status) {
             case INSTALL_SUCCESS:
             case INSTALL_NONE:
@@ -885,6 +893,7 @@ prompt_and_wait(Device* device, int status) {
         ui->SetProgressType(RecoveryUI::EMPTY);
 
         int chosen_item = get_menu_selection(nullptr, device->GetMenuItems(), 0, 0, device);
+        ui_root_menu = 0;
 
         // device-specific code may take some action here.  It may
         // return one of the core actions handled in the switch