1 files changed, 25 insertions, 1 deletions
diff --git a/ui.cpp b/ui.cpp
index 01a4dd4e..94c17c24 100644
--- a/ui.cpp
+++ b/ui.cpp
@@ -55,6 +55,7 @@ RecoveryUI::RecoveryUI()
       brightness_normal_(50),
       brightness_dimmed_(25),
       touch_screen_allowed_(true),
+      volumes_changed_(false),
       kTouchLowThreshold(RECOVERY_UI_TOUCH_LOW_THRESHOLD),
       kTouchHighThreshold(RECOVERY_UI_TOUCH_HIGH_THRESHOLD),
       key_queue_len(0),
@@ -427,6 +428,9 @@ void RecoveryUI::ProcessKey(int key_code, int updown) {
 
       case RecoveryUI::REBOOT:
         if (reboot_enabled) {
+#ifndef VERIFIER_TEST
+          vdc->unmountAll();
+#endif
           reboot("reboot,");
           while (true) {
             pause();
@@ -471,6 +475,7 @@ void RecoveryUI::EnqueueKey(int key_code) {
 
 int RecoveryUI::WaitKey() {
   pthread_mutex_lock(&key_queue_mutex);
+  int timeouts = UI_WAIT_KEY_TIMEOUT_SEC;
 
   // Time out after UI_WAIT_KEY_TIMEOUT_SEC, unless a USB cable is
   // plugged in.
@@ -484,7 +489,20 @@ int RecoveryUI::WaitKey() {
 
     int rc = 0;
     while (key_queue_len == 0 && rc != ETIMEDOUT) {
-      rc = pthread_cond_timedwait(&key_queue_cond, &key_queue_mutex, &timeout);
+      struct timespec key_timeout;
+      gettimeofday(&now, nullptr);
+      key_timeout.tv_sec = now.tv_sec + 1;
+      key_timeout.tv_nsec = now.tv_usec * 1000;
+      rc = pthread_cond_timedwait(&key_queue_cond, &key_queue_mutex, &key_timeout);
+      if (rc == ETIMEDOUT) {
+        if (VolumesChanged()) {
+          pthread_mutex_unlock(&key_queue_mutex);
+          return KEY_REFRESH;
+        }
+        if (key_timeout.tv_sec <= timeout.tv_sec) {
+          rc = 0;
+        }
+      }
     }
 
     if (screensaver_state_ != ScreensaverState::DISABLED) {
@@ -651,3 +669,9 @@ void RecoveryUI::SetLocale(const std::string& new_locale) {
     }
   }
 }
+
+bool RecoveryUI::VolumesChanged() {
+    bool ret = volumes_changed_;
+    volumes_changed_ = false;
+    return ret;
+}