7 files changed, 96 insertions, 48 deletions

diff --git a/minui/events.cpp b/minui/events.cpp
index 2c41eb8a7..daa10c6cf 100644
--- a/minui/events.cpp
+++ b/minui/events.cpp
@@ -201,7 +201,7 @@ int ev_sync_key_state(ev_set_key_callback set_key_cb, void* data) {
     return 0;
 }
 
-void ev_iterate_available_keys(ev_key_callback cb, void* data) {
+void ev_iterate_available_keys(std::function<void(int)> f) {
     unsigned long ev_bits[BITS_TO_LONGS(EV_MAX)];
     unsigned long key_bits[BITS_TO_LONGS(KEY_MAX)];
 
@@ -224,7 +224,7 @@ void ev_iterate_available_keys(ev_key_callback cb, void* data) {
 
         for (int key_code = 0; key_code <= KEY_MAX; ++key_code) {
             if (test_bit(key_code, key_bits)) {
-                cb(key_code, data);
+                f(key_code);
             }
         }
     }
diff --git a/minui/minui.h b/minui/minui.h
index 8f2ff1139..82abb8a63 100644
--- a/minui/minui.h
+++ b/minui/minui.h
@@ -68,13 +68,11 @@ struct input_event;
 
 typedef int (*ev_callback)(int fd, uint32_t epevents, void* data);
 typedef int (*ev_set_key_callback)(int code, int value, void* data);
-typedef void (*ev_key_callback)(int code, void* data);
 
 int ev_init(ev_callback input_cb, void* data);
 void ev_exit(void);
 int ev_add_fd(int fd, ev_callback cb, void* data);
 int ev_sync_key_state(ev_set_key_callback set_key_cb, void* data);
-void ev_iterate_available_keys(ev_key_callback cb, void* data);
 
 // 'timeout' has the same semantics as poll(2).
 //    0 : don't block
@@ -130,4 +128,11 @@ void res_free_surface(gr_surface surface);
 }
 #endif
 
+#ifdef __cplusplus
+
+#include <functional>
+void ev_iterate_available_keys(std::function<void(int)> f);
+
+#endif
+
 #endif
diff --git a/recovery.cpp b/recovery.cpp
index cdbb598e3..43a42eab8 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -545,12 +545,10 @@ get_menu_selection(const char* const * headers, const char* const * items,
         if (action < 0) {
             switch (action) {
                 case Device::kHighlightUp:
-                    --selected;
-                    selected = ui->SelectMenu(selected);
+                    selected = ui->SelectMenu(--selected);
                     break;
                 case Device::kHighlightDown:
-                    ++selected;
-                    selected = ui->SelectMenu(selected);
+                    selected = ui->SelectMenu(++selected);
                     break;
                 case Device::kInvokeItem:
                     chosen_item = selected;
diff --git a/screen_ui.cpp b/screen_ui.cpp
index 6d8df68b3..b62417f5f 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -185,6 +185,9 @@ void ScreenRecoveryUI::SetColor(UIElement e) {
         case MENU_SEL_BG:
             gr_color(0, 106, 157, 255);
             break;
+        case MENU_SEL_BG_ACTIVE:
+            gr_color(0, 156, 100, 255);
+            break;
         case MENU_SEL_FG:
             gr_color(255, 255, 255, 255);
             break;
@@ -220,7 +223,7 @@ void ScreenRecoveryUI::draw_screen_locked() {
 
                 if (i == menu_top + menu_sel) {
                     // draw the highlight bar
-                    SetColor(MENU_SEL_BG);
+                    SetColor(IsLongPress() ? MENU_SEL_BG_ACTIVE : MENU_SEL_BG);
                     gr_fill(0, y-2, gr_fb_width(), y+char_height+2);
                     // white text of selected item
                     SetColor(MENU_SEL_FG);
@@ -638,3 +641,9 @@ void ScreenRecoveryUI::Redraw() {
     update_screen_locked();
     pthread_mutex_unlock(&updateMutex);
 }
+
+void ScreenRecoveryUI::KeyLongPress(int) {
+    // Redraw so that if we're in the menu, the highlight
+    // will change color to indicate a successful long press.
+    Redraw();
+}
diff --git a/screen_ui.h b/screen_ui.h
index 41ff4af11..ea1a95be1 100644
--- a/screen_ui.h
+++ b/screen_ui.h
@@ -56,10 +56,12 @@ class ScreenRecoveryUI : public RecoveryUI {
     int SelectMenu(int sel);
     void EndMenu();
 
+    void KeyLongPress(int);
+
     void Redraw();
 
-    enum UIElement { HEADER, MENU, MENU_SEL_BG, MENU_SEL_FG, LOG, TEXT_FILL };
-    virtual void SetColor(UIElement e);
+    enum UIElement { HEADER, MENU, MENU_SEL_BG, MENU_SEL_BG_ACTIVE, MENU_SEL_FG, LOG, TEXT_FILL };
+    void SetColor(UIElement e);
 
   private:
     Icon currentIcon;
diff --git a/ui.cpp b/ui.cpp
index 871624520..064890ea4 100644
--- a/ui.cpp
+++ b/ui.cpp
@@ -51,26 +51,40 @@ RecoveryUI::RecoveryUI()
           key_down_count(0),
           enable_reboot(true),
           consecutive_power_keys(0),
-          last_key(-1) {
+          last_key(-1),
+          has_power_key(false),
+          has_up_key(false),
+          has_down_key(false) {
     pthread_mutex_init(&key_queue_mutex, NULL);
     pthread_cond_init(&key_queue_cond, NULL);
     self = this;
     memset(key_pressed, 0, sizeof(key_pressed));
 }
 
+void RecoveryUI::OnKeyDetected(int key_code) {
+    if (key_code == KEY_POWER) {
+        has_power_key = true;
+    } else if (key_code == KEY_DOWN || key_code == KEY_VOLUMEDOWN) {
+        has_down_key = true;
+    } else if (key_code == KEY_UP || key_code == KEY_VOLUMEUP) {
+        has_up_key = true;
+    }
+}
+
 void RecoveryUI::Init() {
     ev_init(input_callback, NULL);
+
+    using namespace std::placeholders;
+    ev_iterate_available_keys(std::bind(&RecoveryUI::OnKeyDetected, this, _1));
+
     pthread_create(&input_t, NULL, input_thread, NULL);
 }
 
-
 int RecoveryUI::input_callback(int fd, uint32_t epevents, void* data) {
     struct input_event ev;
-    int ret;
-
-    ret = ev_get_input(fd, epevents, &ev);
-    if (ret)
+    if (ev_get_input(fd, epevents, &ev) == -1) {
         return -1;
+    }
 
     if (ev.type == EV_SYN) {
         return 0;
@@ -95,8 +109,9 @@ int RecoveryUI::input_callback(int fd, uint32_t epevents, void* data) {
         self->rel_sum = 0;
     }
 
-    if (ev.type == EV_KEY && ev.code <= KEY_MAX)
+    if (ev.type == EV_KEY && ev.code <= KEY_MAX) {
         self->process_key(ev.code, ev.value);
+    }
 
     return 0;
 }
@@ -142,8 +157,7 @@ void RecoveryUI::process_key(int key_code, int updown) {
     pthread_mutex_unlock(&key_queue_mutex);
 
     if (register_key) {
-        NextCheckKeyIsLong(long_press);
-        switch (CheckKey(key_code)) {
+        switch (CheckKey(key_code, long_press)) {
           case RecoveryUI::IGNORE:
             break;
 
@@ -257,23 +271,44 @@ bool RecoveryUI::IsKeyPressed(int key) {
     return pressed;
 }
 
+bool RecoveryUI::IsLongPress() {
+    pthread_mutex_lock(&key_queue_mutex);
+    bool result = key_long_press;
+    pthread_mutex_unlock(&key_queue_mutex);
+    return result;
+}
+
 void RecoveryUI::FlushKeys() {
     pthread_mutex_lock(&key_queue_mutex);
     key_queue_len = 0;
     pthread_mutex_unlock(&key_queue_mutex);
 }
 
-// The default CheckKey implementation assumes the device has power,
-// volume up, and volume down keys.
-//
-// - Hold power and press vol-up to toggle display.
-// - Press power seven times in a row to reboot.
-// - Alternate vol-up and vol-down seven times to mount /system.
-RecoveryUI::KeyAction RecoveryUI::CheckKey(int key) {
-    if ((IsKeyPressed(KEY_POWER) && key == KEY_VOLUMEUP) || key == KEY_HOME) {
-        return TOGGLE;
+RecoveryUI::KeyAction RecoveryUI::CheckKey(int key, bool is_long_press) {
+    pthread_mutex_lock(&key_queue_mutex);
+    key_long_press = false;
+    pthread_mutex_unlock(&key_queue_mutex);
+
+    // If we have power and volume up keys, that chord is the signal to toggle the text display.
+    if (has_power_key && has_up_key) {
+        if (key == KEY_VOLUMEUP && IsKeyPressed(KEY_POWER)) {
+            return TOGGLE;
+        }
+    } else {
+        // Otherwise long press of any button toggles to the text display,
+        // and there's no way to toggle back (but that's pretty useless anyway).
+        if (is_long_press && !IsTextVisible()) {
+            return TOGGLE;
+        }
+
+        // Also, for button-limited devices, a long press is translated to KEY_ENTER.
+        if (is_long_press && IsTextVisible()) {
+            EnqueueKey(KEY_ENTER);
+            return IGNORE;
+        }
     }
 
+    // Press power seven times in a row to reboot.
     if (key == KEY_POWER) {
         pthread_mutex_lock(&key_queue_mutex);
         bool reboot_enabled = enable_reboot;
@@ -290,14 +325,10 @@ RecoveryUI::KeyAction RecoveryUI::CheckKey(int key) {
     }
 
     last_key = key;
-
-    return ENQUEUE;
-}
-
-void RecoveryUI::NextCheckKeyIsLong(bool is_long_press) {
+    return IsTextVisible() ? ENQUEUE : IGNORE;
 }
 
-void RecoveryUI::KeyLongPress(int key) {
+void RecoveryUI::KeyLongPress(int) {
 }
 
 void RecoveryUI::SetEnableReboot(bool enabled) {
diff --git a/ui.h b/ui.h
index bc728b05a..0d5ab557d 100644
--- a/ui.h
+++ b/ui.h
@@ -69,30 +69,27 @@ class RecoveryUI {
 
     // --- key handling ---
 
-    // Wait for keypress and return it.  May return -1 after timeout.
+    // Wait for a key and return it.  May return -1 after timeout.
     virtual int WaitKey();
 
     virtual bool IsKeyPressed(int key);
+    virtual bool IsLongPress();
 
     // Erase any queued-up keys.
     virtual void FlushKeys();
 
-    // Called on each keypress, even while operations are in progress.
+    // Called on each key press, even while operations are in progress.
     // Return value indicates whether an immediate operation should be
     // triggered (toggling the display, rebooting the device), or if
     // the key should be enqueued for use by the main thread.
     enum KeyAction { ENQUEUE, TOGGLE, REBOOT, IGNORE };
-    virtual KeyAction CheckKey(int key);
-
-    // Called immediately before each call to CheckKey(), tell you if
-    // the key was long-pressed.
-    virtual void NextCheckKeyIsLong(bool is_long_press);
+    virtual KeyAction CheckKey(int key, bool is_long_press);
 
     // Called when a key is held down long enough to have been a
     // long-press (but before the key is released).  This means that
     // if the key is eventually registered (released without any other
-    // keys being pressed in the meantime), NextCheckKeyIsLong() will
-    // be called with "true".
+    // keys being pressed in the meantime), CheckKey will be called with
+    // 'is_long_press' true.
     virtual void KeyLongPress(int key);
 
     // Normally in recovery there's a key sequence that triggers
@@ -110,8 +107,8 @@ class RecoveryUI {
     virtual void StartMenu(const char* const * headers, const char* const * items,
                            int initial_selection) = 0;
 
-    // Set the menu highlight to the given index, and return it (capped to
-    // the range [0..numitems).
+    // Set the menu highlight to the given index, wrapping if necessary.
+    // Returns the actual item selected.
     virtual int SelectMenu(int sel) = 0;
 
     // End menu mode, resetting the text overlay so that ui_print()
@@ -136,14 +133,20 @@ private:
     int consecutive_power_keys;
     int last_key;
 
-    typedef struct {
+    bool has_power_key;
+    bool has_up_key;
+    bool has_down_key;
+
+    struct key_timer_t {
         RecoveryUI* ui;
         int key_code;
         int count;
-    } key_timer_t;
+    };
 
     pthread_t input_t;
 
+    void OnKeyDetected(int key_code);
+
     static void* input_thread(void* cookie);
     static int input_callback(int fd, uint32_t epevents, void* data);
     void process_key(int key_code, int updown);