From 84fd5ff2d62e5554c00896b9c9310a715f4fa1b0 Mon Sep 17 00:00:00 2001
From: Gabriele M <moto.falcon.git@gmail.com>
Date: Mon, 4 Dec 2017 22:10:06 +0100
Subject: [PATCH] Don't allow any operation if a reboot is needed

Change-Id: I175525071faa44f30b8f2bbf36da56a2e6f2dd1b
---
 .../updater/controller/ABUpdateInstaller.java | 24 +++++++++++--------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/org/lineageos/updater/controller/ABUpdateInstaller.java b/src/org/lineageos/updater/controller/ABUpdateInstaller.java
index 3eb69e12..0fcc31eb 100644
--- a/src/org/lineageos/updater/controller/ABUpdateInstaller.java
+++ b/src/org/lineageos/updater/controller/ABUpdateInstaller.java
@@ -42,6 +42,7 @@ class ABUpdateInstaller {
     private static final String TAG = "ABUpdateInstaller";
 
     private static final String PREF_INSTALLING_AB_ID = "installing_ab_id";
+    private static final String PREF_NEEDS_REBOOT = "needs_reboot";
 
     private final UpdaterController mUpdaterController;
     private final Context mContext;
@@ -57,7 +58,7 @@ class ABUpdateInstaller {
             Update update = mUpdaterController.getActualUpdate(mDownloadId);
             if (update == null) {
                 // We read the id from a preference, the update could no longer exist
-                installationDone();
+                installationDone(status == UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT);
                 return;
             }
 
@@ -71,7 +72,7 @@ class ABUpdateInstaller {
                 break;
 
                 case UpdateEngine.UpdateStatusConstants.REPORTING_ERROR_EVENT: {
-                    installationDone();
+                    installationDone(false);
                     update.setInstallProgress(0);
                     update.setStatus(UpdateStatus.INSTALLATION_FAILED);
                     mUpdaterController.notifyUpdateChange(mDownloadId);
@@ -79,7 +80,7 @@ class ABUpdateInstaller {
                 break;
 
                 case UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT: {
-                    installationDone();
+                    installationDone(true);
                     update.setInstallProgress(0);
                     update.setStatus(UpdateStatus.INSTALLED);
                     mUpdaterController.notifyUpdateChange(mDownloadId);
@@ -97,7 +98,7 @@ class ABUpdateInstaller {
                     if (mReconnecting) {
                         // The service was restarted because we thought we were installing an
                         // update, but we aren't, so clear everything.
-                        installationDone();
+                        installationDone(false);
                         mReconnecting = false;
                     }
                 }
@@ -111,13 +112,15 @@ class ABUpdateInstaller {
     };
 
     static synchronized boolean isInstallingUpdate(Context context) {
-        return PreferenceManager.getDefaultSharedPreferences(context)
-                .getString(ABUpdateInstaller.PREF_INSTALLING_AB_ID, null) != null;
+        SharedPreferences pref = PreferenceManager.getDefaultSharedPreferences(context);
+        return pref.getString(ABUpdateInstaller.PREF_INSTALLING_AB_ID, null) != null ||
+                pref.getBoolean(ABUpdateInstaller.PREF_NEEDS_REBOOT, false);
     }
 
     static synchronized boolean isInstallingUpdate(Context context, String downloadId) {
-        return downloadId.equals(PreferenceManager.getDefaultSharedPreferences(context)
-                .getString(ABUpdateInstaller.PREF_INSTALLING_AB_ID, null));
+        SharedPreferences pref = PreferenceManager.getDefaultSharedPreferences(context);
+        return downloadId.equals(pref.getString(ABUpdateInstaller.PREF_INSTALLING_AB_ID, null)) ||
+                pref.getBoolean(ABUpdateInstaller.PREF_NEEDS_REBOOT, false);
     }
 
     ABUpdateInstaller(Context context, UpdaterController updaterController) {
@@ -200,8 +203,9 @@ class ABUpdateInstaller {
         return mUpdateEngine.bind(mUpdateEngineCallback);
     }
 
-    private void installationDone() {
+    private void installationDone(boolean needsReboot) {
         PreferenceManager.getDefaultSharedPreferences(mContext).edit()
+                .putBoolean(PREF_NEEDS_REBOOT, needsReboot)
                 .remove(PREF_INSTALLING_AB_ID)
                 .apply();
     }
@@ -218,7 +222,7 @@ class ABUpdateInstaller {
         }
 
         mUpdateEngine.cancel();
-        installationDone();
+        installationDone(false);
 
         mUpdaterController.getActualUpdate(mDownloadId)
                 .setStatus(UpdateStatus.INSTALLATION_CANCELLED);