6 files changed, 102 insertions, 83 deletions

diff --git a/recovery.cpp b/recovery.cpp
index bd176da5a..7e539ce3d 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -37,6 +37,7 @@
 #include <unistd.h>
 
 #include <algorithm>
+#include <functional>
 #include <memory>
 #include <string>
 #include <vector>
@@ -495,57 +496,6 @@ static bool erase_volume(const char* volume) {
   return (result == 0);
 }
 
-// Display a menu with the specified 'headers' and 'items'. Device specific HandleMenuKey() may
-// return a positive number beyond the given range. Caller sets 'menu_only' to true to ensure only
-// a menu item gets selected. 'initial_selection' controls the initial cursor location. Returns the
-// (non-negative) chosen item number, or -1 if timed out waiting for input.
-static int get_menu_selection(const char* const* headers, const char* const* items, bool menu_only,
-                              int initial_selection, Device* device) {
-  // Throw away keys pressed previously, so user doesn't accidentally trigger menu items.
-  ui->FlushKeys();
-
-  ui->StartMenu(headers, items, initial_selection);
-
-  int selected = initial_selection;
-  int chosen_item = -1;
-  while (chosen_item < 0) {
-    int key = ui->WaitKey();
-    if (key == -1) {  // WaitKey() timed out.
-      if (ui->WasTextEverVisible()) {
-        continue;
-      } else {
-        LOG(INFO) << "Timed out waiting for key input; rebooting.";
-        ui->EndMenu();
-        return -1;
-      }
-    }
-
-    bool visible = ui->IsTextVisible();
-    int action = device->HandleMenuKey(key, visible);
-
-    if (action < 0) {
-      switch (action) {
-        case Device::kHighlightUp:
-          selected = ui->SelectMenu(--selected);
-          break;
-        case Device::kHighlightDown:
-          selected = ui->SelectMenu(++selected);
-          break;
-        case Device::kInvokeItem:
-          chosen_item = selected;
-          break;
-        case Device::kNoAction:
-          break;
-      }
-    } else if (!menu_only) {
-      chosen_item = action;
-    }
-  }
-
-  ui->EndMenu();
-  return chosen_item;
-}
-
 // Returns the selected filename, or an empty string.
 static std::string browse_directory(const std::string& path, Device* device) {
   ensure_path_mounted(path.c_str());
@@ -588,7 +538,9 @@ static std::string browse_directory(const std::string& path, Device* device) {
 
   int chosen_item = 0;
   while (true) {
-    chosen_item = get_menu_selection(headers, entries, true, chosen_item, device);
+    chosen_item = ui->ShowMenu(
+        headers, entries, chosen_item, true,
+        std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
 
     const std::string& item = zips[chosen_item];
     if (chosen_item == 0) {
@@ -612,15 +564,17 @@ static std::string browse_directory(const std::string& path, Device* device) {
 }
 
 static bool yes_no(Device* device, const char* question1, const char* question2) {
-    const char* headers[] = { question1, question2, NULL };
-    const char* items[] = { " No", " Yes", NULL };
+  const char* headers[] = { question1, question2, NULL };
+  const char* items[] = { " No", " Yes", NULL };
 
-    int chosen_item = get_menu_selection(headers, items, true, 0, device);
-    return (chosen_item == 1);
+  int chosen_item = ui->ShowMenu(
+      headers, items, 0, true,
+      std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
+  return (chosen_item == 1);
 }
 
 static bool ask_to_wipe_data(Device* device) {
-    return yes_no(device, "Wipe all user data?", "  THIS CAN NOT BE UNDONE!");
+  return yes_no(device, "Wipe all user data?", "  THIS CAN NOT BE UNDONE!");
 }
 
 // Return true on success.
@@ -660,7 +614,9 @@ static bool prompt_and_wipe_data(Device* device) {
     NULL
   };
   for (;;) {
-    int chosen_item = get_menu_selection(headers, items, true, 0, device);
+    int chosen_item = ui->ShowMenu(
+        headers, items, 0, true,
+        std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
     if (chosen_item != 1) {
       return true;  // Just reboot, no wipe; not a failure, user asked for it
     }
@@ -859,7 +815,9 @@ static void choose_recovery_file(Device* device) {
 
   int chosen_item = 0;
   while (true) {
-    chosen_item = get_menu_selection(headers, menu_entries.data(), true, chosen_item, device);
+    chosen_item = ui->ShowMenu(
+        headers, menu_entries.data(), chosen_item, true,
+        std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
     if (entries[chosen_item] == "Back") break;
 
     ui->ShowFile(entries[chosen_item].c_str());
@@ -1005,7 +963,9 @@ static Device::BuiltinAction prompt_and_wait(Device* device, int status) {
     }
     ui->SetProgressType(RecoveryUI::EMPTY);
 
-    int chosen_item = get_menu_selection(nullptr, device->GetMenuItems(), false, 0, device);
+    int chosen_item = ui->ShowMenu(
+        nullptr, device->GetMenuItems(), 0, false,
+        std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
 
     // Device-specific code may take some action here. It may return one of the core actions
     // handled in the switch statement below.
diff --git a/screen_ui.cpp b/screen_ui.cpp
index 317e5529c..aaeb18c7f 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -1009,6 +1009,53 @@ void ScreenRecoveryUI::EndMenu() {
   pthread_mutex_unlock(&updateMutex);
 }
 
+int ScreenRecoveryUI::ShowMenu(const char* const* headers, const char* const* items,
+                               int initial_selection, bool menu_only,
+                               const std::function<int(int, bool)>& key_handler) {
+  // Throw away keys pressed previously, so user doesn't accidentally trigger menu items.
+  FlushKeys();
+
+  StartMenu(headers, items, initial_selection);
+
+  int selected = initial_selection;
+  int chosen_item = -1;
+  while (chosen_item < 0) {
+    int key = WaitKey();
+    if (key == -1) {  // WaitKey() timed out.
+      if (WasTextEverVisible()) {
+        continue;
+      } else {
+        LOG(INFO) << "Timed out waiting for key input; rebooting.";
+        EndMenu();
+        return -1;
+      }
+    }
+
+    bool visible = IsTextVisible();
+    int action = key_handler(key, visible);
+    if (action < 0) {
+      switch (action) {
+        case Device::kHighlightUp:
+          selected = SelectMenu(--selected);
+          break;
+        case Device::kHighlightDown:
+          selected = SelectMenu(++selected);
+          break;
+        case Device::kInvokeItem:
+          chosen_item = selected;
+          break;
+        case Device::kNoAction:
+          break;
+      }
+    } else if (!menu_only) {
+      chosen_item = action;
+    }
+  }
+
+  EndMenu();
+  return chosen_item;
+}
+
 bool ScreenRecoveryUI::IsTextVisible() {
   pthread_mutex_lock(&updateMutex);
   int visible = show_text;
diff --git a/screen_ui.h b/screen_ui.h
index c1222a576..837d346a2 100644
--- a/screen_ui.h
+++ b/screen_ui.h
@@ -20,6 +20,7 @@
 #include <pthread.h>
 #include <stdio.h>
 
+#include <functional>
 #include <memory>
 #include <string>
 #include <vector>
@@ -135,10 +136,8 @@ class ScreenRecoveryUI : public RecoveryUI {
   void ShowFile(const char* filename) override;
 
   // menu display
-  void StartMenu(const char* const* headers, const char* const* items,
-                 int initial_selection) override;
-  int SelectMenu(int sel) override;
-  void EndMenu() override;
+  int ShowMenu(const char* const* headers, const char* const* items, int initial_selection,
+               bool menu_only, const std::function<int(int, bool)>& key_handler) override;
 
   void KeyLongPress(int) override;
 
@@ -164,6 +163,18 @@ class ScreenRecoveryUI : public RecoveryUI {
 
   virtual bool InitTextParams();
 
+  // Displays some header text followed by a menu of items, which appears at the top of the screen
+  // (in place of any scrolling ui_print() output, if necessary).
+  virtual void StartMenu(const char* const* headers, const char* const* items,
+                         int initial_selection);
+
+  // Sets the menu highlight to the given index, wrapping if necessary. Returns the actual item
+  // selected.
+  virtual int SelectMenu(int sel);
+
+  // Ends menu mode, resetting the text overlay so that ui_print() statements will be displayed.
+  virtual void EndMenu();
+
   virtual void draw_background_locked();
   virtual void draw_foreground_locked();
   virtual void draw_screen_locked();
diff --git a/stub_ui.h b/stub_ui.h
index 1f6b29acb..3c36fcfb1 100644
--- a/stub_ui.h
+++ b/stub_ui.h
@@ -54,12 +54,11 @@ class StubRecoveryUI : public RecoveryUI {
   void ShowFile(const char* /* filename */) override {}
 
   // menu display
-  void StartMenu(const char* const* /* headers */, const char* const* /* items */,
-                 int /* initial_selection */) override {}
-  int SelectMenu(int sel) override {
-    return sel;
+  int ShowMenu(const char* const* /* headers */, const char* const* /* items */,
+               int initial_selection, bool /* menu_only */,
+               const std::function<int(int, bool)>& /* key_handler */) override {
+    return initial_selection;
   }
-  void EndMenu() override {}
 };
 
 #endif  // RECOVERY_STUB_UI_H
diff --git a/ui.h b/ui.h
index 4c54d6915..636c2ff70 100644
--- a/ui.h
+++ b/ui.h
@@ -21,6 +21,7 @@
 #include <pthread.h>
 #include <time.h>
 
+#include <functional>
 #include <string>
 
 // Abstract class for controlling the user interface during recovery.
@@ -128,17 +129,18 @@ class RecoveryUI {
 
   // --- menu display ---
 
-  // Display some header text followed by a menu of items, which appears at the top of the screen
-  // (in place of any scrolling ui_print() output, if necessary).
-  virtual void StartMenu(const char* const* headers, const char* const* items,
-                         int initial_selection) = 0;
-
-  // Sets the menu highlight to the given index, wrapping if necessary. Returns the actual item
-  // selected.
-  virtual int SelectMenu(int sel) = 0;
-
-  // Ends menu mode, resetting the text overlay so that ui_print() statements will be displayed.
-  virtual void EndMenu() = 0;
+  // Displays a menu with the given 'headers' and 'items'. The supplied 'key_handler' callback,
+  // which is typically bound to Device::HandleMenuKey(), should return the expected action for the
+  // given key code and menu visibility (e.g. to move the cursor or to select an item). Caller sets
+  // 'menu_only' to true to ensure only a menu item gets selected and returned. Otherwise if
+  // 'menu_only' is false, ShowMenu() will forward any non-negative value returned from the
+  // key_handler, which may be beyond the range of menu items. This could be used to trigger a
+  // device-specific action, even without that being listed in the menu. Caller needs to handle
+  // such a case accordingly (e.g. by calling Device::InvokeMenuItem() to process the action).
+  // Returns a non-negative value (the chosen item number or device-specific action code), or -1 if
+  // timed out waiting for input.
+  virtual int ShowMenu(const char* const* headers, const char* const* items, int initial_selection,
+                       bool menu_only, const std::function<int(int, bool)>& key_handler) = 0;
 
  protected:
   void EnqueueKey(int key_code);
diff --git a/wear_ui.h b/wear_ui.h
index 8b24cb73e..fcbbee289 100644
--- a/wear_ui.h
+++ b/wear_ui.h
@@ -25,9 +25,6 @@ class WearRecoveryUI : public ScreenRecoveryUI {
 
   void SetStage(int current, int max) override;
 
-  void StartMenu(const char* const* headers, const char* const* items,
-                 int initial_selection) override;
-
  protected:
   // progress bar vertical position, it's centered horizontally
   const int kProgressBarBaseline;
@@ -36,6 +33,9 @@ class WearRecoveryUI : public ScreenRecoveryUI {
   // Recovery, build id and etc) and the bottom lines that may otherwise go out of the screen.
   const int kMenuUnusableRows;
 
+  void StartMenu(const char* const* headers, const char* const* items,
+                 int initial_selection) override;
+
   int GetProgressBaseline() const override;
 
   void update_progress_locked() override;