6 files changed, 163 insertions, 33 deletions

diff --git a/device.h b/device.h
index 9c433715b..cbecc4371 100644
--- a/device.h
+++ b/device.h
@@ -47,6 +47,7 @@ class Device {
     MOUNT_SYSTEM = 10,
     RUN_GRAPHICS_TEST = 11,
     RUN_LOCALE_TEST = 12,
+    KEY_INTERRUPTED = 13,
   };
 
   explicit Device(RecoveryUI* ui);
diff --git a/recovery.cpp b/recovery.cpp
index 9588cc130..cc30035df 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -326,6 +326,11 @@ static std::string browse_directory(const std::string& path, Device* device) {
         headers, entries, chosen_item, true,
         std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
 
+    // Return if WaitKey() was interrupted.
+    if (chosen_item == static_cast<size_t>(RecoveryUI::KeyError::INTERRUPTED)) {
+      return "";
+    }
+
     const std::string& item = entries[chosen_item];
     if (chosen_item == 0) {
       // Go up but continue browsing (if the caller is browse_directory).
@@ -401,6 +406,11 @@ static bool prompt_and_wipe_data(Device* device) {
     size_t chosen_item = ui->ShowMenu(
         headers, items, 0, true,
         std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
+
+    // If ShowMenu() returned RecoveryUI::KeyError::INTERRUPTED, WaitKey() was interrupted.
+    if (chosen_item == static_cast<size_t>(RecoveryUI::KeyError::INTERRUPTED)) {
+      return false;
+    }
     if (chosen_item != 1) {
       return true;  // Just reboot, no wipe; not a failure, user asked for it
     }
@@ -597,6 +607,11 @@ static void choose_recovery_file(Device* device) {
     chosen_item = ui->ShowMenu(
         headers, entries, chosen_item, true,
         std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
+
+    // Handle WaitKey() interrupt.
+    if (chosen_item == static_cast<size_t>(RecoveryUI::KeyError::INTERRUPTED)) {
+      break;
+    }
     if (entries[chosen_item] == "Back") break;
 
     ui->ShowFile(entries[chosen_item]);
@@ -745,12 +760,16 @@ static Device::BuiltinAction prompt_and_wait(Device* device, int status) {
     size_t chosen_item = ui->ShowMenu(
         {}, device->GetMenuItems(), 0, false,
         std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
-
+    // Handle Interrupt key
+    if (chosen_item == static_cast<size_t>(RecoveryUI::KeyError::INTERRUPTED)) {
+      return Device::KEY_INTERRUPTED;
+    }
     // Device-specific code may take some action here. It may return one of the core actions
     // handled in the switch statement below.
-    Device::BuiltinAction chosen_action = (chosen_item == static_cast<size_t>(-1))
-                                              ? Device::REBOOT
-                                              : device->InvokeMenuItem(chosen_item);
+    Device::BuiltinAction chosen_action =
+        (chosen_item == static_cast<size_t>(RecoveryUI::KeyError::TIMED_OUT))
+            ? Device::REBOOT
+            : device->InvokeMenuItem(chosen_item);
 
     bool should_wipe_cache = false;
     switch (chosen_action) {
@@ -831,6 +850,9 @@ static Device::BuiltinAction prompt_and_wait(Device* device, int status) {
           }
         }
         break;
+
+      case Device::KEY_INTERRUPTED:
+        return Device::KEY_INTERRUPTED;
     }
   }
 }
@@ -1092,6 +1114,7 @@ Device::BuiltinAction start_recovery(Device* device, const std::vector<std::stri
   title_lines.insert(std::begin(title_lines), "Android Recovery");
   ui->SetTitle(title_lines);
 
+  ui->ResetKeyInterruptStatus();
   device->StartRecovery();
 
   printf("Command:");
diff --git a/screen_ui.cpp b/screen_ui.cpp
index f9c4a06c1..c14f29d49 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -417,6 +417,7 @@ void ScreenRecoveryUI::CheckBackgroundTextImages() {
   FlushKeys();
   while (true) {
     int key = WaitKey();
+    if (key == static_cast<int>(KeyError::INTERRUPTED)) break;
     if (key == KEY_POWER || key == KEY_ENTER) {
       break;
     } else if (key == KEY_UP || key == KEY_VOLUMEUP) {
@@ -925,6 +926,7 @@ void ScreenRecoveryUI::ShowFile(FILE* fp) {
       while (show_prompt) {
         show_prompt = false;
         int key = WaitKey();
+        if (key == static_cast<int>(KeyError::INTERRUPTED)) return;
         if (key == KEY_POWER || key == KEY_ENTER) {
           return;
         } else if (key == KEY_UP || key == KEY_VOLUMEUP) {
@@ -1017,19 +1019,26 @@ size_t ScreenRecoveryUI::ShowMenu(const std::vector<std::string>& headers,
   // Throw away keys pressed previously, so user doesn't accidentally trigger menu items.
   FlushKeys();
 
+  // If there is a key interrupt in progress, return KeyError::INTERRUPTED without starting the
+  // menu.
+  if (IsKeyInterrupted()) return static_cast<size_t>(KeyError::INTERRUPTED);
+
   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 (key == static_cast<int>(KeyError::INTERRUPTED)) {  // WaitKey() was interrupted.
+      return static_cast<size_t>(KeyError::INTERRUPTED);
+    }
+    if (key == static_cast<int>(KeyError::TIMED_OUT)) {  // WaitKey() timed out.
       if (WasTextEverVisible()) {
         continue;
       } else {
         LOG(INFO) << "Timed out waiting for key input; rebooting.";
         EndMenu();
-        return static_cast<size_t>(-1);
+        return static_cast<size_t>(KeyError::TIMED_OUT);
       }
     }
 
diff --git a/tests/unit/screen_ui_test.cpp b/tests/unit/screen_ui_test.cpp
index 4c0a868f0..7d97a006b 100644
--- a/tests/unit/screen_ui_test.cpp
+++ b/tests/unit/screen_ui_test.cpp
@@ -264,6 +264,10 @@ int TestableScreenRecoveryUI::KeyHandler(int key, bool) const {
 }
 
 int TestableScreenRecoveryUI::WaitKey() {
+  if (IsKeyInterrupted()) {
+    return static_cast<int>(RecoveryUI::KeyError::INTERRUPTED);
+  }
+
   CHECK_LT(key_buffer_index_, key_buffer_.size());
   return static_cast<int>(key_buffer_[key_buffer_index_++]);
 }
@@ -391,7 +395,8 @@ TEST_F(ScreenRecoveryUITest, ShowMenu_TimedOut) {
   ui_->SetKeyBuffer({
       KeyCode::TIMEOUT,
   });
-  ASSERT_EQ(static_cast<size_t>(-1), ui_->ShowMenu(HEADERS, ITEMS, 3, true, nullptr));
+  ASSERT_EQ(static_cast<size_t>(RecoveryUI::KeyError::TIMED_OUT),
+            ui_->ShowMenu(HEADERS, ITEMS, 3, true, nullptr));
 }
 
 TEST_F(ScreenRecoveryUITest, ShowMenu_TimedOut_TextWasEverVisible) {
@@ -412,6 +417,38 @@ TEST_F(ScreenRecoveryUITest, ShowMenu_TimedOut_TextWasEverVisible) {
                                         std::placeholders::_1, std::placeholders::_2)));
 }
 
+TEST_F(ScreenRecoveryUITest, ShowMenuWithInterrupt) {
+  RETURN_IF_NO_GRAPHICS;
+
+  ASSERT_TRUE(ui_->Init(kTestLocale));
+  ui_->SetKeyBuffer({
+      KeyCode::UP,
+      KeyCode::DOWN,
+      KeyCode::UP,
+      KeyCode::DOWN,
+      KeyCode::ENTER,
+  });
+
+  ui_->InterruptKey();
+  ASSERT_EQ(static_cast<size_t>(RecoveryUI::KeyError::INTERRUPTED),
+            ui_->ShowMenu(HEADERS, ITEMS, 3, true,
+                          std::bind(&TestableScreenRecoveryUI::KeyHandler, ui_.get(),
+                                    std::placeholders::_1, std::placeholders::_2)));
+
+  ui_->SetKeyBuffer({
+      KeyCode::UP,
+      KeyCode::UP,
+      KeyCode::NO_OP,
+      KeyCode::NO_OP,
+      KeyCode::UP,
+      KeyCode::ENTER,
+  });
+  ASSERT_EQ(static_cast<size_t>(RecoveryUI::KeyError::INTERRUPTED),
+            ui_->ShowMenu(HEADERS, ITEMS, 0, true,
+                          std::bind(&TestableScreenRecoveryUI::KeyHandler, ui_.get(),
+                                    std::placeholders::_1, std::placeholders::_2)));
+}
+
 TEST_F(ScreenRecoveryUITest, LoadAnimation) {
   RETURN_IF_NO_GRAPHICS;
 
diff --git a/ui.cpp b/ui.cpp
index 6c91d01b8..a2c160f7a 100644
--- a/ui.cpp
+++ b/ui.cpp
@@ -58,6 +58,7 @@ RecoveryUI::RecoveryUI()
       touch_screen_allowed_(false),
       kTouchLowThreshold(RECOVERY_UI_TOUCH_LOW_THRESHOLD),
       kTouchHighThreshold(RECOVERY_UI_TOUCH_HIGH_THRESHOLD),
+      key_interrupted_(false),
       key_queue_len(0),
       key_last_down(-1),
       key_long_press(false),
@@ -404,34 +405,69 @@ void RecoveryUI::EnqueueKey(int key_code) {
   }
 }
 
+void RecoveryUI::SetScreensaverState(ScreensaverState state) {
+  switch (state) {
+    case ScreensaverState::NORMAL:
+      if (android::base::WriteStringToFile(std::to_string(brightness_normal_value_),
+                                           brightness_file_)) {
+        screensaver_state_ = ScreensaverState::NORMAL;
+        LOG(INFO) << "Brightness: " << brightness_normal_value_ << " (" << brightness_normal_
+                  << "%)";
+      } else {
+        LOG(ERROR) << "Unable to set brightness to normal";
+      }
+      break;
+    case ScreensaverState::DIMMED:
+      if (android::base::WriteStringToFile(std::to_string(brightness_dimmed_value_),
+                                           brightness_file_)) {
+        LOG(INFO) << "Brightness: " << brightness_dimmed_value_ << " (" << brightness_dimmed_
+                  << "%)";
+        screensaver_state_ = ScreensaverState::DIMMED;
+      } else {
+        LOG(ERROR) << "Unable to set brightness to dim";
+      }
+      break;
+    case ScreensaverState::OFF:
+      if (android::base::WriteStringToFile("0", brightness_file_)) {
+        LOG(INFO) << "Brightness: 0 (off)";
+        screensaver_state_ = ScreensaverState::OFF;
+      } else {
+        LOG(ERROR) << "Unable to set brightness to off";
+      }
+      break;
+    default:
+      LOG(ERROR) << "Invalid screensaver state";
+  }
+}
+
 int RecoveryUI::WaitKey() {
   std::unique_lock<std::mutex> lk(key_queue_mutex);
 
+  // Check for a saved key queue interruption.
+  if (key_interrupted_) {
+    SetScreensaverState(ScreensaverState::NORMAL);
+    return static_cast<int>(KeyError::INTERRUPTED);
+  }
+
   // Time out after UI_WAIT_KEY_TIMEOUT_SEC, unless a USB cable is
   // plugged in.
   do {
-    std::cv_status rc = std::cv_status::no_timeout;
-    while (key_queue_len == 0 && rc != std::cv_status::timeout) {
-      rc = key_queue_cond.wait_for(lk, std::chrono::seconds(UI_WAIT_KEY_TIMEOUT_SEC));
+    bool rc = key_queue_cond.wait_for(lk, std::chrono::seconds(UI_WAIT_KEY_TIMEOUT_SEC), [this] {
+      return this->key_queue_len != 0 || key_interrupted_;
+    });
+    if (key_interrupted_) {
+      SetScreensaverState(ScreensaverState::NORMAL);
+      return static_cast<int>(KeyError::INTERRUPTED);
     }
-
     if (screensaver_state_ != ScreensaverState::DISABLED) {
-      if (rc == std::cv_status::timeout) {
+      if (!rc) {
         // Lower the brightness level: NORMAL -> DIMMED; DIMMED -> OFF.
         if (screensaver_state_ == ScreensaverState::NORMAL) {
-          if (android::base::WriteStringToFile(std::to_string(brightness_dimmed_value_),
-                                               brightness_file_)) {
-            LOG(INFO) << "Brightness: " << brightness_dimmed_value_ << " (" << brightness_dimmed_
-                      << "%)";
-            screensaver_state_ = ScreensaverState::DIMMED;
-          }
+          SetScreensaverState(ScreensaverState::DIMMED);
         } else if (screensaver_state_ == ScreensaverState::DIMMED) {
-          if (android::base::WriteStringToFile("0", brightness_file_)) {
-            LOG(INFO) << "Brightness: 0 (off)";
-            screensaver_state_ = ScreensaverState::OFF;
-          }
+          SetScreensaverState(ScreensaverState::OFF);
         }
-      } else if (screensaver_state_ != ScreensaverState::NORMAL) {
+      } else {
         // Drop the first key if it's changing from OFF to NORMAL.
         if (screensaver_state_ == ScreensaverState::OFF) {
           if (key_queue_len > 0) {
@@ -440,17 +476,12 @@ int RecoveryUI::WaitKey() {
         }
 
         // Reset the brightness to normal.
-        if (android::base::WriteStringToFile(std::to_string(brightness_normal_value_),
-                                             brightness_file_)) {
-          screensaver_state_ = ScreensaverState::NORMAL;
-          LOG(INFO) << "Brightness: " << brightness_normal_value_ << " (" << brightness_normal_
-                    << "%)";
-        }
+        SetScreensaverState(ScreensaverState::NORMAL);
       }
     }
   } while (IsUsbConnected() && key_queue_len == 0);
 
-  int key = -1;
+  int key = static_cast<int>(KeyError::TIMED_OUT);
   if (key_queue_len > 0) {
     key = key_queue[0];
     memcpy(&key_queue[0], &key_queue[1], sizeof(int) * --key_queue_len);
@@ -458,6 +489,14 @@ int RecoveryUI::WaitKey() {
   return key;
 }
 
+void RecoveryUI::InterruptKey() {
+  {
+    std::lock_guard<std::mutex> lg(key_queue_mutex);
+    key_interrupted_ = true;
+  }
+  key_queue_cond.notify_one();
+}
+
 bool RecoveryUI::IsUsbConnected() {
   int fd = open("/sys/class/android_usb/android0/state", O_RDONLY);
   if (fd < 0) {
diff --git a/ui.h b/ui.h
index 32e28099e..a1e18cc10 100644
--- a/ui.h
+++ b/ui.h
@@ -51,6 +51,11 @@ class RecoveryUI {
     IGNORE
   };
 
+  enum class KeyError : int {
+    TIMED_OUT = -1,
+    INTERRUPTED = -2,
+  };
+
   RecoveryUI();
 
   virtual ~RecoveryUI();
@@ -99,9 +104,13 @@ class RecoveryUI {
 
   // --- key handling ---
 
-  // Waits for a key and return it. May return -1 after timeout.
+  // Waits for a key and return it. May return TIMED_OUT after timeout and
+  // KeyError::INTERRUPTED on a key interrupt.
   virtual int WaitKey();
 
+  // Wakes up the UI if it is waiting on key input, causing WaitKey to return KeyError::INTERRUPTED.
+  virtual void InterruptKey();
+
   virtual bool IsKeyPressed(int key);
   virtual bool IsLongPress();
 
@@ -147,11 +156,22 @@ class RecoveryUI {
   // 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
-  // static_cast<size_t>(-1) if timed out waiting for input.
+  // static_cast<size_t>(TIMED_OUT) if timed out waiting for input or
+  // static_cast<size_t>(ERR_KEY_INTERTUPT) if interrupted, such as by InterruptKey().
   virtual size_t ShowMenu(const std::vector<std::string>& headers,
                           const std::vector<std::string>& items, size_t initial_selection,
                           bool menu_only, const std::function<int(int, bool)>& key_handler) = 0;
 
+  // Resets the key interrupt status.
+  void ResetKeyInterruptStatus() {
+    key_interrupted_ = false;
+  }
+
+  // Returns the key interrupt status.
+  bool IsKeyInterrupted() const {
+    return key_interrupted_;
+  }
+
  protected:
   void EnqueueKey(int key_code);
 
@@ -187,10 +207,11 @@ class RecoveryUI {
   bool IsUsbConnected();
 
   bool InitScreensaver();
-
+  void SetScreensaverState(ScreensaverState state);
   // Key event input queue
   std::mutex key_queue_mutex;
   std::condition_variable key_queue_cond;
+  bool key_interrupted_;
   int key_queue[256], key_queue_len;
   char key_pressed[KEY_MAX + 1];  // under key_queue_mutex
   int key_last_down;              // under key_queue_mutex