4 files changed, 47 insertions, 4 deletions

diff --git a/install.cpp b/install.cpp
index 0bd7945c4..0f059602a 100644
--- a/install.cpp
+++ b/install.cpp
@@ -160,6 +160,11 @@ try_update_binary(const char *path, ZipArchive *zip, int* wipe_cache) {
             *wipe_cache = 1;
         } else if (strcmp(command, "clear_display") == 0) {
             ui->SetBackground(RecoveryUI::NONE);
+        } else if (strcmp(command, "enable_reboot") == 0) {
+            // packages can explicitly request that they want the user
+            // to be able to reboot during installation (useful for
+            // debugging packages that don't exit).
+            ui->SetEnableReboot(true);
         } else {
             LOGE("unknown command [%s]\n", command);
         }
@@ -236,7 +241,9 @@ really_install_package(const char *path, int* wipe_cache)
     /* Verify and install the contents of the package.
      */
     ui->Print("Installing update...\n");
+    ui->SetEnableReboot(false);
     int result = try_update_binary(path, &zip, wipe_cache);
+    ui->SetEnableReboot(true);
 
     sysReleaseMap(&map);
 
diff --git a/ui.cpp b/ui.cpp
index 091012f29..98442539b 100644
--- a/ui.cpp
+++ b/ui.cpp
@@ -49,6 +49,7 @@ RecoveryUI::RecoveryUI() :
     key_last_down(-1),
     key_long_press(false),
     key_down_count(0),
+    enable_reboot(true),
     consecutive_power_keys(0),
     consecutive_alternate_keys(0),
     last_key(-1) {
@@ -116,6 +117,7 @@ int RecoveryUI::input_callback(int fd, uint32_t epevents, void* data)
 void RecoveryUI::process_key(int key_code, int updown) {
     bool register_key = false;
     bool long_press = false;
+    bool reboot_enabled;
 
     pthread_mutex_lock(&key_queue_mutex);
     key_pressed[key_code] = updown;
@@ -137,6 +139,7 @@ void RecoveryUI::process_key(int key_code, int updown) {
         }
         key_last_down = -1;
     }
+    reboot_enabled = enable_reboot;
     pthread_mutex_unlock(&key_queue_mutex);
 
     if (register_key) {
@@ -150,7 +153,9 @@ void RecoveryUI::process_key(int key_code, int updown) {
             break;
 
           case RecoveryUI::REBOOT:
-            android_reboot(ANDROID_RB_RESTART, 0, 0);
+            if (reboot_enabled) {
+                android_reboot(ANDROID_RB_RESTART, 0, 0);
+            }
             break;
 
           case RecoveryUI::ENQUEUE:
@@ -281,9 +286,15 @@ RecoveryUI::KeyAction RecoveryUI::CheckKey(int key) {
     }
 
     if (key == KEY_POWER) {
-        ++consecutive_power_keys;
-        if (consecutive_power_keys >= 7) {
-            return REBOOT;
+        pthread_mutex_lock(&key_queue_mutex);
+        bool reboot_enabled = enable_reboot;
+        pthread_mutex_unlock(&key_queue_mutex);
+
+        if (reboot_enabled) {
+            ++consecutive_power_keys;
+            if (consecutive_power_keys >= 7) {
+                return REBOOT;
+            }
         }
     } else {
         consecutive_power_keys = 0;
@@ -311,3 +322,9 @@ void RecoveryUI::NextCheckKeyIsLong(bool is_long_press) {
 
 void RecoveryUI::KeyLongPress(int key) {
 }
+
+void RecoveryUI::SetEnableReboot(bool enabled) {
+    pthread_mutex_lock(&key_queue_mutex);
+    enable_reboot = enabled;
+    pthread_mutex_unlock(&key_queue_mutex);
+}
diff --git a/ui.h b/ui.h
index faa0acd46..31a8a7fb1 100644
--- a/ui.h
+++ b/ui.h
@@ -93,6 +93,13 @@ class RecoveryUI {
     // be called with "true".
     virtual void KeyLongPress(int key);
 
+    // Normally in recovery there's a key sequence that triggers
+    // immediate reboot of the device, regardless of what recovery is
+    // doing (with the default CheckKey implementation, it's pressing
+    // the power button 7 times in row).  Call this to enable or
+    // disable that feature.  It is enabled by default.
+    virtual void SetEnableReboot(bool enabled);
+
     // --- menu display ---
 
     // Display some header text followed by a menu of items, which appears
@@ -121,6 +128,7 @@ private:
     int key_last_down;                 // under key_queue_mutex
     bool key_long_press;               // under key_queue_mutex
     int key_down_count;                // under key_queue_mutex
+    bool enable_reboot;                // under key_queue_mutex
     int rel_sum;
 
     int consecutive_power_keys;
diff --git a/updater/install.c b/updater/install.c
index 8defc7720..e1071c90c 100644
--- a/updater/install.c
+++ b/updater/install.c
@@ -1640,6 +1640,15 @@ Value* WipeBlockDeviceFn(const char* name, State* state, int argc, Expr* argv[])
     return StringValue(strdup(success ? "t" : ""));
 }
 
+Value* EnableRebootFn(const char* name, State* state, int argc, Expr* argv[]) {
+    if (argc != 0) {
+        return ErrorAbort(state, "%s() expects no args, got %d", name, argc);
+    }
+    UpdaterInfo* ui = (UpdaterInfo*)(state->cookie);
+    fprintf(ui->cmd_pipe, "enable_reboot\n");
+    return StringValue(strdup("t"));
+}
+
 void RegisterInstallFunctions() {
     RegisterFunction("mount", MountFn);
     RegisterFunction("is_mounted", IsMountedFn);
@@ -1689,4 +1698,6 @@ void RegisterInstallFunctions() {
     RegisterFunction("reboot_now", RebootNowFn);
     RegisterFunction("get_stage", GetStageFn);
     RegisterFunction("set_stage", SetStageFn);
+
+    RegisterFunction("enable_reboot", EnableRebootFn);
 }