1 files changed, 43 insertions, 0 deletions
diff --git a/recovery.cpp b/recovery.cpp
index 75430a2b..4c1d8275 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -45,6 +45,8 @@
 #include <cutils/properties.h> /* for property_list */
 #include <fs_mgr/roots.h>
 #include <volume_manager/VolumeManager.h>
+#include <hardware/boot_control.h>
+#include <hardware/hardware.h>
 #include <ziparchive/zip_archive.h>
 
 #include "bootloader_message/bootloader_message.h"
@@ -193,6 +195,43 @@ bool ask_to_continue_downgrade(Device* device) {
   }
 }
 
+std::string get_chosen_slot(Device* device) {
+  std::vector<std::string> headers{ "Choose which slot to boot into on next boot." };
+  std::vector<std::string> items{ "A", "B" };
+  size_t chosen_item = device->GetUI()->ShowMenu(
+      headers, items, 0, true,
+      std::bind(&Device::HandleMenuKey, device, std::placeholders::_1, std::placeholders::_2));
+  if (chosen_item < 0)
+    return "";
+  return items[chosen_item];
+}
+
+int set_slot(Device* device) {
+  std::string slot = get_chosen_slot(device);
+  if (slot == "")
+    return 0;
+  const hw_module_t *hw_module;
+  boot_control_module_t *module;
+  int ret;
+  ret = hw_get_module("bootctrl", &hw_module);
+  if (ret != 0) {
+    device->GetUI()->Print("Error getting bootctrl module.\n");
+  } else {
+    module = (boot_control_module_t*) hw_module;
+    module->init(module);
+    int slot_number = 0;
+    if (slot == "B")
+      slot_number = 1;
+    if (module->setActiveBootSlot(module, slot_number))
+      device->GetUI()->Print("Error changing bootloader boot slot to %s", slot.c_str());
+    else {
+      device->GetUI()->Print("Switched slot to %s.\n", slot.c_str());
+      device->GoHome();
+    }
+  }
+  return ret;
+}
+
 static bool ask_to_wipe_data(Device* device) {
   std::vector<std::string> headers{ "Format user data?", "This includes internal storage.", "THIS CANNOT BE UNDONE!" };
   std::vector<std::string> items{ " Cancel", " Format data" };
@@ -597,6 +636,10 @@ change_menu:
         ui->Print("Enabled ADB.\n");
         break;
 
+      case Device::SWAP_SLOT:
+        set_slot(device);
+        break;
+
       case Device::RUN_GRAPHICS_TEST:
         run_graphics_test(ui);
         break;