1 files changed, 21 insertions, 0 deletions
diff --git a/install/install.cpp b/install/install.cpp
index dab919b2..ed5200be 100644
--- a/install/install.cpp
+++ b/install/install.cpp
@@ -46,6 +46,7 @@
 #include <android-base/strings.h>
 #include <android-base/unique_fd.h>
 
+#include "bootloader_message/bootloader_message.h"
 #include "install/snapshot_utils.h"
 #include "install/spl_check.h"
 #include "install/wipe_data.h"
@@ -62,6 +63,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);
 
@@ -382,6 +384,21 @@ static InstallResult TryUpdateBinary(Package* package, bool* wipe_cache,
     LOG(WARNING) << "This is SPL downgrade";
   }
 
+  const auto reboot_to_recovery = [] {
+    if (std::string err; !clear_bootloader_message(&err)) {
+      LOG(ERROR) << "Failed to clear BCB message: " << err;
+    }
+    Reboot("recovery");
+  };
+
+  static bool ab_package_installed = false;
+  if (ab_package_installed) {
+    if (ask_to_ab_reboot(device)) {
+      reboot_to_recovery();
+    }
+    return INSTALL_ERROR;
+  }
+
   if (package_is_ab) {
     CHECK(package->GetType() == PackageType::kFile);
   }
@@ -567,7 +584,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 (!ui->IsSideloadAutoReboot() && ask_to_ab_reboot(device)) {
+      reboot_to_recovery();
+    }
   }
 
   return INSTALL_SUCCESS;