10 files changed, 212 insertions, 22 deletions

diff --git a/Android.mk b/Android.mk
index 94b2cd666..467ce6187 100644
--- a/Android.mk
+++ b/Android.mk
@@ -177,6 +177,8 @@ endif
 
 ifeq ($(AB_OTA_UPDATER),true)
     LOCAL_CFLAGS += -DAB_OTA_UPDATER=1
+    LOCAL_SHARED_LIBRARIES += libhardware
+    LOCAL_ADDITIONAL_DEPENDENCIES += libhardware
 endif
 
 LOCAL_MODULE_PATH := $(TARGET_RECOVERY_ROOT_OUT)/sbin
@@ -346,7 +348,7 @@ else
     LOCAL_CFLAGS += -DTW_DEFAULT_LANGUAGE=en
 endif
 
-LOCAL_ADDITIONAL_DEPENDENCIES := \
+LOCAL_ADDITIONAL_DEPENDENCIES += \
     dump_image \
     erase_image \
     flash_image \
diff --git a/data.cpp b/data.cpp
index b034bf802..62709a59c 100644
--- a/data.cpp
+++ b/data.cpp
@@ -820,6 +820,13 @@ void DataManager::SetDefaultValues()
 
 	mData.SetValue("tw_has_adopted_storage", "0");
 
+#ifdef AB_OTA_UPDATER
+	LOGINFO("AB_OTA_UPDATER := true\n");
+	mConst.SetValue("tw_has_boot_slots", "1");
+#else
+	mConst.SetValue("tw_has_boot_slots", "0");
+#endif
+
 	pthread_mutex_unlock(&m_valuesLock);
 }
 
diff --git a/gui/action.cpp b/gui/action.cpp
index 223d75e25..8600186ee 100644
--- a/gui/action.cpp
+++ b/gui/action.cpp
@@ -228,6 +228,7 @@ GUIAction::GUIAction(xml_node<>* node)
 		ADD_ACTION(changefilesystem);
 		ADD_ACTION(flashimage);
 		ADD_ACTION(twcmd);
+		ADD_ACTION(setbootslot);
 	}
 
 	// First, get the action
@@ -1866,3 +1867,15 @@ int GUIAction::setlanguage(std::string arg __unused)
 	operation_end(op_status);
 	return 0;
 }
+
+int GUIAction::setbootslot(std::string arg)
+{
+	operation_start("Set Boot Slot");
+	if (!simulate)
+	{
+		PartitionManager.Set_Active_Slot(arg);
+	} else
+		simulate_progress_bar();
+	operation_end(0);
+	return 0;
+}
diff --git a/gui/objects.hpp b/gui/objects.hpp
index 0d969279a..8d4484aec 100644
--- a/gui/objects.hpp
+++ b/gui/objects.hpp
@@ -357,6 +357,7 @@ protected:
 	int mountsystemtoggle(std::string arg);
 	int setlanguage(std::string arg);
 	int twcmd(std::string arg);
+	int setbootslot(std::string arg);
 
 	int simulate;
 };
diff --git a/gui/theme/common/languages/en.xml b/gui/theme/common/languages/en.xml
index b82aecff4..ee772bfd9 100644
--- a/gui/theme/common/languages/en.xml
+++ b/gui/theme/common/languages/en.xml
@@ -209,6 +209,11 @@
 		<string name="enable_backup_comp_chk">Enable compression</string>
 		<string name="skip_md5_backup_chk">Skip MD5 generation during backup</string>
 		<string name="disable_backup_space_chk" version="2">Disable free space check before backup</string>
+		<string name="current_boot_slot">Current Slot: %tw_active_slot%</string>
+		<string name="boot_slot_a">Slot A</string>
+		<string name="boot_slot_b">Slot B</string>
+		<string name="changing_boot_slot">Changing Boot Slot</string>
+		<string name="changing_boot_slot_complete">Changing Boot Slot Complete</string>
 		<string name="refresh_sizes_btn">Refresh Sizes</string>
 		<string name="swipe_backup">Swipe to Backup</string>
 		<string name="append_date_btn">Append Date</string>
@@ -675,5 +680,6 @@
 		<string name="partition_not_found">path: {1} not found in partititon list</string>
 		<string name="copy_kernel_log">Copied kernel log to {1}</string>
 		<string name="include_kernel_log">Include Kernel Log</string>
+		<string name="unable_set_boot_slot">Error changing bootloader boot slot to {1}</string>
 	</resources>
 </language>
diff --git a/gui/theme/common/portrait.xml b/gui/theme/common/portrait.xml
index 11ae8740e..f0ba791fb 100644
--- a/gui/theme/common/portrait.xml
+++ b/gui/theme/common/portrait.xml
@@ -1699,6 +1699,44 @@
 				<data variable="tw_disable_free_space"/>
 			</checkbox>
 
+			<text style="text_m">
+				<condition var1="tw_has_boot_slots" var2="1"/>
+				<placement x="%center_x%" y="%row17_y%" placement="5"/>
+				<text>{@current_boot_slot=Current Slot: %tw_active_slot%}</text>
+			</text>
+
+			<button style="main_button">
+				<condition var1="tw_has_boot_slots" var2="1"/>
+				<placement x="%indent%" y="%row19_y%"/>
+				<text>{@boot_slot_a=Slot A}</text>
+				<actions>
+					<action function="set">tw_back=backup_options</action>
+					<action function="set">tw_action=setbootslot</action>
+					<action function="set">tw_action_param=A</action>
+					<action function="set">tw_has_action2=0</action>
+					<action function="set">tw_action_text1={@changing_boot_slot=Changing Boot Slot}</action>
+					<action function="set">tw_action_text2=</action>
+					<action function="set">tw_complete_text1={@changing_boot_slot_complete=Change Boot Slot Complete}</action>
+					<action function="page">action_page</action>
+				</actions>
+			</button>
+
+			<button style="main_button">
+				<condition var1="tw_has_boot_slots" var2="1"/>
+				<placement x="%center_x%" y="%row19_y%"/>
+				<text>{@boot_slot_b=Slot B}</text>
+				<actions>
+					<action function="set">tw_back=backup_options</action>
+					<action function="set">tw_action=setbootslot</action>
+					<action function="set">tw_action_param=B</action>
+					<action function="set">tw_has_action2=0</action>
+					<action function="set">tw_action_text1={@changing_boot_slot=Changing Boot Slot}</action>
+					<action function="set">tw_action_text2=</action>
+					<action function="set">tw_complete_text1={@changing_boot_slot_complete=Change Boot Slot Complete}</action>
+					<action function="page">action_page</action>
+				</actions>
+			</button>
+
 			<action>
 				<touch key="home"/>
 				<action function="page">main</action>
@@ -2779,6 +2817,44 @@
 				</actions>
 			</button>
 
+			<text style="text_m">
+				<condition var1="tw_has_boot_slots" var2="1"/>
+				<placement x="%center_x%" y="%row17_y%" placement="5"/>
+				<text>{@current_boot_slot=Current Slot: %tw_active_slot%}</text>
+			</text>
+
+			<button style="main_button">
+				<condition var1="tw_has_boot_slots" var2="1"/>
+				<placement x="%indent%" y="%row19_y%"/>
+				<text>{@boot_slot_a=Slot A}</text>
+				<actions>
+					<action function="set">tw_back=reboot</action>
+					<action function="set">tw_action=setbootslot</action>
+					<action function="set">tw_action_param=A</action>
+					<action function="set">tw_has_action2=0</action>
+					<action function="set">tw_action_text1={@changing_boot_slot=Changing Boot Slot}</action>
+					<action function="set">tw_action_text2=</action>
+					<action function="set">tw_complete_text1={@changing_boot_slot_complete=Change Boot Slot Complete}</action>
+					<action function="page">action_page</action>
+				</actions>
+			</button>
+
+			<button style="main_button">
+				<condition var1="tw_has_boot_slots" var2="1"/>
+				<placement x="%center_x%" y="%row19_y%"/>
+				<text>{@boot_slot_b=Slot B}</text>
+				<actions>
+					<action function="set">tw_back=reboot</action>
+					<action function="set">tw_action=setbootslot</action>
+					<action function="set">tw_action_param=B</action>
+					<action function="set">tw_has_action2=0</action>
+					<action function="set">tw_action_text1={@changing_boot_slot=Changing Boot Slot}</action>
+					<action function="set">tw_action_text2=</action>
+					<action function="set">tw_complete_text1={@changing_boot_slot_complete=Change Boot Slot Complete}</action>
+					<action function="page">action_page</action>
+				</actions>
+			</button>
+
 			<action>
 				<touch key="home"/>
 				<action function="page">main</action>
diff --git a/partition.cpp b/partition.cpp
index 78b11ecab..00b23b79f 100644
--- a/partition.cpp
+++ b/partition.cpp
@@ -139,6 +139,7 @@ enum TW_FSTAB_FLAGS {
 	TWFLAG_USERMRF,
 	TWFLAG_WIPEDURINGFACTORYRESET,
 	TWFLAG_WIPEINGUI,
+	TWFLAG_SLOTSELECT,
 };
 
 /* Flags without a trailing '=' are considered dual format flags and can be
@@ -172,6 +173,7 @@ const struct flag_list tw_flags[] = {
 	{ "usermrf",                TWFLAG_USERMRF },
 	{ "wipeduringfactoryreset", TWFLAG_WIPEDURINGFACTORYRESET },
 	{ "wipeingui",              TWFLAG_WIPEINGUI },
+	{ "slotselect",             TWFLAG_SLOTSELECT },
 	{ 0,                        0 },
 };
 
@@ -231,6 +233,7 @@ TWPartition::TWPartition() {
 	Mount_Read_Only = false;
 	Is_Adopted_Storage = false;
 	Adopted_GUID = "";
+	SlotSelect = false;
 }
 
 TWPartition::~TWPartition(void) {
@@ -411,7 +414,7 @@ bool TWPartition::Process_Fstab_Line(const char *fstab_line, bool Display_Error)
 #endif
 	} else if (Is_Image(Fstab_File_System)) {
 		Find_Actual_Block_Device();
-		Setup_Image(Display_Error);
+		Setup_Image();
 		if (Mount_Point == "/boot") {
 			Display_Name = "Boot";
 			Backup_Display_Name = Display_Name;
@@ -498,7 +501,7 @@ void TWPartition::Setup_Data_Partition(bool Display_Error) {
 				Is_Decrypted = false;
 				Can_Be_Mounted = false;
 				Current_File_System = "emmc";
-				Setup_Image(Display_Error);
+				Setup_Image();
 				DataManager::SetValue(TW_IS_ENCRYPTED, 1);
 				DataManager::SetValue(TW_CRYPTO_PWTYPE, cryptfs_get_password_type());
 				DataManager::SetValue(TW_CRYPTO_PASSWORD, "");
@@ -665,6 +668,9 @@ void TWPartition::Apply_TW_Flag(const unsigned flag, const char* str, const bool
 			if (Wipe_Available_in_GUI)
 				Can_Be_Wiped = true;
 			break;
+		case TWFLAG_SLOTSELECT:
+			SlotSelect = true;
+			break;
 		default:
 			// Should not get here
 			LOGINFO("Flag identified for processing, but later unmatched: %i\n", flag);
@@ -808,7 +814,7 @@ void TWPartition::Setup_File_System(bool Display_Error) {
 	Backup_Method = BM_FILES;
 }
 
-void TWPartition::Setup_Image(bool Display_Error) {
+void TWPartition::Setup_Image() {
 	Display_Name = Mount_Point.substr(1, Mount_Point.size() - 1);
 	Backup_Name = Display_Name;
 	if (Current_File_System == "emmc")
@@ -817,15 +823,6 @@ void TWPartition::Setup_Image(bool Display_Error) {
 		Backup_Method = BM_FLASH_UTILS;
 	else
 		LOGINFO("Unhandled file system '%s' on image '%s'\n", Current_File_System.c_str(), Display_Name.c_str());
-	if (Find_Partition_Size()) {
-		Used = Size;
-		Backup_Size = Size;
-	} else {
-		if (Display_Error)
-			LOGERR("Unable to find partition size for '%s'\n", Mount_Point.c_str());
-		else
-			LOGINFO("Unable to find partition size for '%s'\n", Mount_Point.c_str());
-	}
 }
 
 void TWPartition::Setup_AndSec(void) {
@@ -1727,7 +1724,7 @@ bool TWPartition::Wipe_Encryption() {
 	Has_Data_Media = false;
 	Decrypted_Block_Device = "";
 #ifdef TW_INCLUDE_CRYPTO
-	if (Is_Decrypted) {
+	if (Is_Decrypted && !Decrypted_Block_Device.empty()) {
 		if (!UnMount(true))
 			return false;
 		if (delete_crypto_blk_dev((char*)("userdata")) != 0) {
@@ -2474,8 +2471,16 @@ bool TWPartition::Restore_Image(PartitionSettings *part_settings) {
 bool TWPartition::Update_Size(bool Display_Error) {
 	bool ret = false, Was_Already_Mounted = false;
 
-	if (!Can_Be_Mounted && !Is_Encrypted)
+	Find_Actual_Block_Device();
+
+	if (!Can_Be_Mounted && !Is_Encrypted) {
+		if (TWFunc::Path_Exists(Actual_Block_Device) && Find_Partition_Size()) {
+			Used = Size;
+			Backup_Size = Size;
+			return true;
+		}
 		return false;
+	}
 
 	Was_Already_Mounted = Is_Mounted();
 	if (Removable || Is_Encrypted) {
@@ -2523,15 +2528,22 @@ bool TWPartition::Update_Size(bool Display_Error) {
 void TWPartition::Find_Actual_Block_Device(void) {
 	if (Is_Decrypted && !Decrypted_Block_Device.empty()) {
 		Actual_Block_Device = Decrypted_Block_Device;
-		if (TWFunc::Path_Exists(Decrypted_Block_Device))
+		if (TWFunc::Path_Exists(Decrypted_Block_Device)) {
 			Is_Present = true;
+			return;
+		}
+	} else if (SlotSelect && TWFunc::Path_Exists(Primary_Block_Device + PartitionManager.Get_Active_Slot_Suffix())) {
+		Actual_Block_Device = Primary_Block_Device + PartitionManager.Get_Active_Slot_Suffix();
+		unlink(Primary_Block_Device.c_str());
+		symlink(Actual_Block_Device.c_str(), Primary_Block_Device.c_str()); // we create a non-slot symlink pointing to the currently selected slot which may assist zips with installing
+		Is_Present = true;
+		return;
 	} else if (TWFunc::Path_Exists(Primary_Block_Device)) {
 		Is_Present = true;
 		Actual_Block_Device = Primary_Block_Device;
 		return;
 	}
-	if (Is_Decrypted) {
-	} else if (!Alternate_Block_Device.empty() && TWFunc::Path_Exists(Alternate_Block_Device)) {
+	if (!Alternate_Block_Device.empty() && TWFunc::Path_Exists(Alternate_Block_Device)) {
 		Actual_Block_Device = Alternate_Block_Device;
 		Is_Present = true;
 	} else {
diff --git a/partitionmanager.cpp b/partitionmanager.cpp
index 079d476b3..9e23ccb94 100644
--- a/partitionmanager.cpp
+++ b/partitionmanager.cpp
@@ -64,12 +64,26 @@ extern "C" {
 	#include "gui/pages.hpp"
 #endif
 
+#ifdef AB_OTA_UPDATER
+#include <hardware/hardware.h>
+#include <hardware/boot_control.h>
+#endif
+
 extern bool datamedia;
 
 TWPartitionManager::TWPartitionManager(void) {
 	mtp_was_enabled = false;
 	mtp_write_fd = -1;
 	stop_backup.set_value(0);
+#ifdef AB_OTA_UPDATER
+	char slot_suffix[PROPERTY_VALUE_MAX];
+	property_get("ro.boot.slot_suffix", slot_suffix, "_a");
+	Active_Slot_Display = "";
+	if (strcmp(slot_suffix, "_a") == 0)
+		Set_Active_Slot("A");
+	else
+		Set_Active_Slot("B");
+#endif
 }
 
 int TWPartitionManager::Process_Fstab(string Fstab_Filename, bool Display_Error) {
@@ -182,6 +196,9 @@ int TWPartitionManager::Process_Fstab(string Fstab_Filename, bool Display_Error)
 #endif
 	Update_System_Details();
 	UnMount_Main_Partitions();
+#ifdef AB_OTA_UPDATER
+	DataManager::SetValue("tw_active_slot", Get_Active_Slot_Display());
+#endif
 	return true;
 }
 
@@ -290,6 +307,8 @@ void TWPartitionManager::Output_Partition(TWPartition* Part) {
 		printf("Can_Flash_Img ");
 	if (Part->Is_Adopted_Storage)
 		printf("Is_Adopted_Storage ");
+	if (Part->SlotSelect)
+		printf("SlotSelect ");
 	printf("\n");
 	if (!Part->SubPartition_Of.empty())
 		printf("   SubPartition_Of: %s\n", Part->SubPartition_Of.c_str());
@@ -1376,8 +1395,8 @@ void TWPartitionManager::Update_System_Details(void) {
 
 	gui_msg("update_part_details=Updating partition details...");
 	for (iter = Partitions.begin(); iter != Partitions.end(); iter++) {
+		(*iter)->Update_Size(true);
 		if ((*iter)->Can_Be_Mounted) {
-			(*iter)->Update_Size(true);
 			if ((*iter)->Mount_Point == "/system") {
 				int backup_display_size = (int)((*iter)->Backup_Size / 1048576LLU);
 				DataManager::SetValue(TW_BACKUP_SYSTEM_SIZE, backup_display_size);
@@ -2473,3 +2492,46 @@ void TWPartitionManager::Remove_Partition_By_Path(string Path) {
 		}
 	}
 }
+
+void TWPartitionManager::Set_Active_Slot(const string& Slot) {
+	if (Slot != "A" && Slot != "B") {
+		LOGERR("Set_Active_Slot invalid slot '%s'\n", Slot.c_str());
+		return;
+	}
+	if (Active_Slot_Display == Slot)
+		return;
+	LOGINFO("Setting active slot %s\n", Slot.c_str());
+#ifdef AB_OTA_UPDATER
+	if (!Active_Slot_Display.empty()) {
+		const hw_module_t *hw_module;
+		boot_control_module_t *module;
+		int ret;
+		ret = hw_get_module("bootctrl", &hw_module);
+		if (ret != 0) {
+			LOGERR("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))
+				gui_msg(Msg(msg::kError, "unable_set_boot_slot=Error changing bootloader boot slot to {1}")(Slot));
+		}
+		DataManager::SetValue("tw_active_slot", Slot); // Doing this outside of this if block may result in a seg fault because the DataManager may not be ready yet
+	}
+#else
+	LOGERR("Boot slot feature not present\n");
+#endif
+	Active_Slot_Display = Slot;
+	if (Fstab_Processed())
+		Update_System_Details();
+}
+string TWPartitionManager::Get_Active_Slot_Suffix() {
+	if (Active_Slot_Display == "A")
+		return "_a";
+	return "_b";
+}
+string TWPartitionManager::Get_Active_Slot_Display() {
+	return Active_Slot_Display;
+}
diff --git a/partitions.hpp b/partitions.hpp
index 3f2e40a5d..a6af5b1a4 100644
--- a/partitions.hpp
+++ b/partitions.hpp
@@ -133,7 +133,7 @@ private:
 	bool Is_File_System(string File_System);                                  // Checks to see if the file system given is considered a file system
 	bool Is_Image(string File_System);                                        // Checks to see if the file system given is considered an image
 	void Setup_File_System(bool Display_Error);                               // Sets defaults for a file system partition
-	void Setup_Image(bool Display_Error);                                     // Sets defaults for an image partition
+	void Setup_Image();                                                       // Sets defaults for an image partition
 	void Setup_AndSec(void);                                                  // Sets up .android_secure settings
 	void Find_Real_Block_Device(string& Block_Device, bool Display_Error);    // Checks the block device given and follows symlinks until it gets to the real block device
 	unsigned long long IOCTL_Get_Block_Size();                                // Finds the partition size using ioctl
@@ -215,8 +215,9 @@ private:
 	bool Can_Flash_Img;                                                       // Indicates if this partition can have images flashed to it via the GUI
 	bool Mount_Read_Only;                                                     // Only mount this partition as read-only
 	bool Is_Adopted_Storage;                                                  // Indicates that this partition is for adopted storage (android_expand)
-	TWExclude backup_exclusions;
-	TWExclude wipe_exclusions;
+	bool SlotSelect;                                                          // Partition has A/B slots
+	TWExclude backup_exclusions;                                              // Exclusions for file based backups
+	TWExclude wipe_exclusions;                                                // Exclusions for file based wipes (data/media devices only)
 
 friend class TWPartitionManager;
 friend class DataManager;
@@ -289,6 +290,9 @@ public:
 	bool Flash_Image(string& path, string& filename);                         // Flashes an image to a selected partition from the partition list
 	bool Restore_Partition(struct PartitionSettings *part_settings);          // Restore the partitions based on type
 	TWAtomicInt stop_backup;
+	void Set_Active_Slot(const string& Slot);                                 // Sets the active slot to A or B
+	string Get_Active_Slot_Suffix();                                          // Returns active slot _a or _b
+	string Get_Active_Slot_Display();                                         // Returns active slot A or B for display purposes
 
 private:
 	void Setup_Settings_Storage_Partition(TWPartition* Part);                 // Sets up settings storage
@@ -308,6 +312,7 @@ private:
 
 private:
 	std::vector<TWPartition*> Partitions;                                     // Vector list of all partitions
+	string Active_Slot_Display;                                               // Current Active Slot (A or B) for display purposes
 };
 
 extern TWPartitionManager PartitionManager;
diff --git a/prebuilt/Android.mk b/prebuilt/Android.mk
index a81781bf7..390c3f1f0 100644
--- a/prebuilt/Android.mk
+++ b/prebuilt/Android.mk
@@ -177,6 +177,12 @@ ifeq ($(TW_INCLUDE_CRYPTO), true)
         RELINK_SOURCE_FILES += $(TARGET_OUT_VENDOR_SHARED_LIBRARIES)/libcryptfs_hw.so
     endif
 endif
+ifeq ($(AB_OTA_UPDATER), true)
+    RELINK_SOURCE_FILES += $(TARGET_OUT_EXECUTABLES)/bootctl
+    ifneq ($(TW_INCLUDE_CRYPTO), true)
+        RELINK_SOURCE_FILES += $(TARGET_OUT_SHARED_LIBRARIES)/libhardware.so
+    endif
+endif
 ifeq ($(TARGET_USERIMAGES_USE_EXT4), true)
     RELINK_SOURCE_FILES += $(TARGET_OUT_EXECUTABLES)/make_ext4fs
 endif