recovery: Ask to reboot before installing anything after A/B update

Change-Id: Id3e44ddbe11119fcc66d379216f8f5e683979a9f

2 files changed, 19 insertions, 0 deletions

diff --git a/install/install.cpp b/install/install.cpp
index 61bab17b..ff256a0a 100644
--- a/install/install.cpp
+++ b/install/install.cpp
@@ -62,6 +62,7 @@
 
 using namespace std::chrono_literals;
 
+bool ask_to_ab_reboot(Device* device);
 bool ask_to_continue_unverified(Device* device);
 bool ask_to_continue_downgrade(Device* device);
 
@@ -383,6 +384,14 @@ static InstallResult TryUpdateBinary(Package* package, bool* wipe_cache,
     return INSTALL_ERROR;
   }
 
+  static bool ab_package_installed = false;
+  if (ab_package_installed) {
+    if (ask_to_ab_reboot(device)) {
+      Reboot("userrequested,recovery,ui");
+    }
+    return INSTALL_ERROR;
+  }
+
   if (package_is_ab) {
     CHECK(package->GetType() == PackageType::kFile);
   }
@@ -568,7 +577,11 @@ static InstallResult TryUpdateBinary(Package* package, bool* wipe_cache,
     LOG(FATAL) << "Invalid status code " << status;
   }
   if (package_is_ab) {
+    ab_package_installed = true;
     PerformPowerwashIfRequired(zip, device);
+    if (ask_to_ab_reboot(device)) {
+      Reboot("userrequested,recovery,ui");
+    }
   }
 
   return INSTALL_SUCCESS;
diff --git a/recovery.cpp b/recovery.cpp
index 75430a2b..36cb78b5 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -175,6 +175,12 @@ static bool yes_no(Device* device, const char* question1, const char* question2)
   return (chosen_item == 1);
 }
 
+bool ask_to_ab_reboot(Device* device) {
+  device->GetUI()->SetProgressType(RecoveryUI::EMPTY);
+  return yes_no(device, "To install additional packages, you need to reboot recovery first",
+                "Do you want to reboot to recovery now?");
+}
+
 bool ask_to_continue_unverified(Device* device) {
   if (get_build_type() == "user") {
     return false;