twrpRepacker: support vendor_boot and recovery_ab

Change-Id: Ic4d108eb04430550fe0c88f2849463128be5a418
diff --git a/Android.mk b/Android.mk
index 4882688..87f92fc 100644
--- a/Android.mk
+++ b/Android.mk
@@ -170,6 +170,15 @@
     LOCAL_CFLAGS += -DTW_SYSTEM_BUILD_PROP_ADDITIONAL_PATHS='"$(TW_SYSTEM_BUILD_PROP_ADDITIONAL_PATHS)"'
 endif
 
+ifneq ($(BOARD_BOOT_HEADER_VERSION),)
+    LOCAL_CFLAGS += -DBOARD_BOOT_HEADER_VERSION=$(BOARD_BOOT_HEADER_VERSION)
+endif
+
+ifeq ($(BOARD_MOVE_RECOVERY_RESOURCES_TO_VENDOR_BOOT),true)
+    TW_INCLUDE_REPACKTOOLS := true
+    LOCAL_CFLAGS += -DBOARD_MOVE_RECOVERY_RESOURCES_TO_VENDOR_BOOT
+endif
+
 ifeq ($(TW_USES_VENDOR_LIBS),true)
     LOCAL_CFLAGS += -DUSE_VENDOR_LIBS=1
 endif
@@ -215,6 +224,7 @@
     LOCAL_SHARED_LIBRARIES += libtwrpmtp-ffs
 endif
 ifeq ($(BOARD_USES_RECOVERY_AS_BOOT), true)
+    TW_INCLUDE_REPACKTOOLS := true
     LOCAL_CFLAGS += -DBOARD_USES_RECOVERY_AS_BOOT
 endif
 ifeq ($(BOARD_BUILD_SYSTEM_ROOT_IMAGE), true)
diff --git a/data.cpp b/data.cpp
index fa299c1..0078524 100644
--- a/data.cpp
+++ b/data.cpp
@@ -805,6 +805,16 @@
 #else
 	mConst.SetValue("tw_uses_initramfs", "0");
 #endif
+#if defined BOARD_USES_RECOVERY_AS_BOOT || defined BOARD_MOVE_RECOVERY_RESOURCES_TO_VENDOR_BOOT
+	mConst.SetValue("tw_include_install_recovery_ramdisk", "1");
+#else
+	mConst.SetValue("tw_include_install_recovery_ramdisk", "0");
+#endif
+#ifdef BOARD_MOVE_RECOVERY_RESOURCES_TO_VENDOR_BOOT
+	mConst.SetValue("tw_is_vendor_boot", "1");
+#else
+	mConst.SetValue("tw_is_vendor_boot", "0");
+#endif
 #ifdef TW_NO_SCREEN_TIMEOUT
 	mConst.SetValue("tw_screen_timeout_secs", "0");
 	mConst.SetValue("tw_no_screen_timeout", "1");
@@ -812,6 +822,15 @@
 	mPersist.SetValue("tw_screen_timeout_secs", "60");
 	mPersist.SetValue("tw_no_screen_timeout", "0");
 #endif
+#ifdef BOARD_BOOT_HEADER_VERSION
+	mConst.SetValue("tw_boot_header_version", BOARD_BOOT_HEADER_VERSION);
+#endif
+
+	if (GetIntValue("tw_is_vendor_boot") == 1 && GetIntValue("tw_boot_header_version") == 3)
+		mConst.SetValue("tw_is_vendor_boot_header_v3", "1");
+	else
+		mConst.SetValue("tw_is_vendor_boot_header_v3", "0");
+
 	mData.SetValue("tw_gui_done", "0");
 	mData.SetValue("tw_encrypt_backup", "0");
 	mData.SetValue("tw_sleep_total", "5");
diff --git a/gui/theme/common/landscape.xml b/gui/theme/common/landscape.xml
index f7e366b..f56ea4d 100755
--- a/gui/theme/common/landscape.xml
+++ b/gui/theme/common/landscape.xml
@@ -463,8 +463,8 @@
 				<listitem name="{@auto_reflashtwrp_chk=Automatically Reflash TWRP after flashing a ROM}">
 					<conditions>
 						<condition var1="tw_has_boot_slots" var2="1"/>
-						<condition var1="tw_has_repack_tools" var2="1"/>
 						<condition var1="tw_no_flash_current_twrp" var2="0"/>
+						<condition var1="tw_is_vendor_boot_header_v3" var2="0"/>
 					</conditions>
 					<data variable="tw_auto_reflashtwrp"/>
 				</listitem>
@@ -3298,8 +3298,8 @@
 				<listitem name="{@auto_reflashtwrp_chk=Automatically Reflash TWRP after flashing a ROM}">
 					<conditions>
 						<condition var1="tw_has_boot_slots" var2="1"/>
-						<condition var1="tw_has_repack_tools" var2="1"/>
 						<condition var1="tw_no_flash_current_twrp" var2="0"/>
+						<condition var1="tw_is_vendor_boot_header_v3" var2="0"/>
 					</conditions>
 					<data variable="tw_auto_reflashtwrp"/>
 				</listitem>
@@ -3785,6 +3785,7 @@
 					<conditions>
 						<condition var1="tw_has_boot_slots" var2="1"/>
 						<condition var1="tw_has_repack_tools" var2="1"/>
+						<condition var1="tw_include_install_recovery_ramdisk" var2="1"/>
 					</conditions>
 					<actions>
 						<action function="set">tw_repack_kernel=0</action>
@@ -3796,6 +3797,7 @@
 						<condition var1="tw_has_boot_slots" var2="1"/>
 						<condition var1="tw_has_repack_tools" var2="1"/>
 						<condition var1="tw_no_flash_current_twrp" var2="0"/>
+						<condition var1="tw_is_vendor_boot_header_v3" var2="0"/>
 					</conditions>
 					<actions>
 						<action function="set">tw_repack_kernel=0</action>
diff --git a/gui/theme/common/portrait.xml b/gui/theme/common/portrait.xml
index 99eeef6..8b5324b 100755
--- a/gui/theme/common/portrait.xml
+++ b/gui/theme/common/portrait.xml
@@ -415,8 +415,8 @@
 				<listitem name="{@auto_reflashtwrp_chk=Automatically Reflash TWRP after flashing a ROM}">
 					<conditions>
 						<condition var1="tw_has_boot_slots" var2="1"/>
-						<condition var1="tw_has_repack_tools" var2="1"/>
 						<condition var1="tw_no_flash_current_twrp" var2="0"/>
+						<condition var1="tw_is_vendor_boot_header_v3" var2="0"/>
 					</conditions>
 					<data variable="tw_auto_reflashtwrp"/>
 				</listitem>
@@ -3473,8 +3473,8 @@
 				<listitem name="{@auto_reflashtwrp_chk=Automatically Reflash TWRP after flashing a ROM}">
 					<conditions>
 						<condition var1="tw_has_boot_slots" var2="1"/>
-						<condition var1="tw_has_repack_tools" var2="1"/>
 						<condition var1="tw_no_flash_current_twrp" var2="0"/>
+						<condition var1="tw_is_vendor_boot_header_v3" var2="0"/>
 					</conditions>
 					<data variable="tw_auto_reflashtwrp"/>
 				</listitem>
@@ -3908,6 +3908,7 @@
 					<conditions>
 						<condition var1="tw_has_boot_slots" var2="1"/>
 						<condition var1="tw_has_repack_tools" var2="1"/>
+						<condition var1="tw_include_install_recovery_ramdisk" var2="1"/>
 					</conditions>
 					<actions>
 						<action function="set">tw_repack_kernel=0</action>
@@ -3919,6 +3920,7 @@
 						<condition var1="tw_has_boot_slots" var2="1"/>
 						<condition var1="tw_has_repack_tools" var2="1"/>
 						<condition var1="tw_no_flash_current_twrp" var2="0"/>
+						<condition var1="tw_is_vendor_boot_header_v3" var2="0"/>
 					</conditions>
 					<actions>
 						<action function="set">tw_repack_kernel=0</action>
diff --git a/gui/theme/common/watch.xml b/gui/theme/common/watch.xml
index e6d0880..9e3a232 100755
--- a/gui/theme/common/watch.xml
+++ b/gui/theme/common/watch.xml
@@ -578,8 +578,8 @@
 				<listitem name="{@auto_reflashtwrp_chk=Automatically Reflash TWRP after flashing a ROM}">
 					<conditions>
 						<condition var1="tw_has_boot_slots" var2="1"/>
-						<condition var1="tw_has_repack_tools" var2="1"/>
 						<condition var1="tw_no_flash_current_twrp" var2="0"/>
+						<condition var1="tw_is_vendor_boot_header_v3" var2="0"/>
 					</conditions>
 					<data variable="tw_auto_reflashtwrp"/>
 				</listitem>
@@ -3846,8 +3846,8 @@
 				<listitem name="{@auto_reflashtwrp_chk=Automatically Reflash TWRP after flashing a ROM}">
 					<conditions>
 						<condition var1="tw_has_boot_slots" var2="1"/>
-						<condition var1="tw_has_repack_tools" var2="1"/>
 						<condition var1="tw_no_flash_current_twrp" var2="0"/>
+						<condition var1="tw_is_vendor_boot_header_v3" var2="0"/>
 					</conditions>
 					<data variable="tw_auto_reflashtwrp"/>
 				</listitem>
@@ -4395,6 +4395,7 @@
 					<conditions>
 						<condition var1="tw_has_boot_slots" var2="1"/>
 						<condition var1="tw_has_repack_tools" var2="1"/>
+						<condition var1="tw_include_install_recovery_ramdisk" var2="1"/>
 					</conditions>
 					<actions>
 						<action function="set">tw_repack_kernel=0</action>
@@ -4406,6 +4407,7 @@
 						<condition var1="tw_has_boot_slots" var2="1"/>
 						<condition var1="tw_has_repack_tools" var2="1"/>
 						<condition var1="tw_no_flash_current_twrp" var2="0"/>
+						<condition var1="tw_is_vendor_boot_header_v3" var2="0"/>
 					</conditions>
 					<actions>
 						<action function="set">tw_repack_kernel=0</action>
diff --git a/partitionmanager.cpp b/partitionmanager.cpp
index f25552b..c878e1b 100755
--- a/partitionmanager.cpp
+++ b/partitionmanager.cpp
@@ -2688,8 +2688,13 @@
 				Partition_List->push_back(part);
 			}
 		}
-		if (DataManager::GetIntValue("tw_has_repack_tools") != 0 && DataManager::GetIntValue("tw_has_boot_slots") != 0) {
-			TWPartition* boot = Find_Partition_By_Path("/boot");
+		if (DataManager::GetIntValue("tw_has_repack_tools") != 0 && DataManager::GetIntValue("tw_has_boot_slots") != 0 && DataManager::GetIntValue("tw_include_install_recovery_ramdisk") != 0) {
+			std::string dest_partition = "/boot";
+			#ifdef BOARD_MOVE_RECOVERY_RESOURCES_TO_VENDOR_BOOT
+				dest_partition = "/vendor_boot";
+			#endif
+
+			TWPartition* boot = Find_Partition_By_Path(dest_partition);
 			if (boot) {
 				// Allow flashing kernels and ramdisks
 				struct PartitionList repack_ramdisk;
@@ -2697,6 +2702,7 @@
 				repack_ramdisk.Mount_Point = "/repack_ramdisk";
 				repack_ramdisk.selected = 0;
 				Partition_List->push_back(repack_ramdisk);
+				LOGINFO("Install Recovery Ramdisk: target partition=%s\n", dest_partition.c_str());
 				/*struct PartitionList repack_kernel; For now let's leave repacking kernels under advanced only
 				repack_kernel.Display_Name = gui_lookup("install_kernel", "Install Kernel");
 				repack_kernel.Mount_Point = "/repack_kernel";
diff --git a/twrpRepacker.cpp b/twrpRepacker.cpp
index 1da19e2..687295c 100755
--- a/twrpRepacker.cpp
+++ b/twrpRepacker.cpp
@@ -110,20 +110,54 @@
 	return ramdisk_format;
 }
 
-bool twrpRepacker::Repack_Image_And_Flash(const std::string& Target_Image, const struct Repack_Options_struct& Repack_Options) {
-	bool recompress = false;
+static bool is_AB_for_repacker() {
+	std::string slot = android::base::GetProperty("ro.boot.slot_suffix", "");
+	if (slot.empty())
+		slot = android::base::GetProperty("ro.boot.slot", "");
+	return !slot.empty();
+}
 
+bool twrpRepacker::Repack_Image_And_Flash(const std::string& Target_Image, const struct Repack_Options_struct& Repack_Options) {
 	if (!TWFunc::Path_Exists("/system/bin/magiskboot")) {
 		LOGERR("Image repacking tool not present in this TWRP build!");
 		return false;
 	}
+
+	bool recompress = false;
+	bool is_vendor_boot = false;
+	bool is_vendor_boot_v4 = false;
+	std::string dest_partition = "/boot";
+	std::string ramdisk_cpio = "ramdisk.cpio";
+
+	#ifdef BOARD_MOVE_RECOVERY_RESOURCES_TO_VENDOR_BOOT
+		dest_partition = "/vendor_boot";
+		is_vendor_boot = true;
+		if (DataManager::GetIntValue("tw_boot_header_version") == 4) {
+			is_vendor_boot_v4 = true;
+			ramdisk_cpio = "vendor_ramdisk_recovery.cpio";
+			LOGINFO("Vendor_boot with v4 header\n");
+		} else {
+			LOGINFO("Vendor_boot with v3 header\n");
+		}
+	#else
+		// we shouldn't reach here, because of the code in twrpRepacker::Flash_Current_Twrp(); but if we do, then handle it
+		if (PartitionManager.Find_Partition_By_Path("/recovery") && is_AB_for_repacker()) {
+			dest_partition = "/recovery";
+		}
+	#endif
+
+	if (is_vendor_boot || is_vendor_boot_v4) {
+		// placeholder for any specific vendor_boot stuff;
+		// in the meantime, stop the compiler's complaints about unused variables
+	}
+
 	DataManager::SetProgress(0);
 	PartitionManager.Update_System_Details();
-	TWPartition* part = PartitionManager.Find_Partition_By_Path("/boot");
+	TWPartition* part = PartitionManager.Find_Partition_By_Path(dest_partition);
 	if (part)
 		gui_msg(Msg("unpacking_image=Unpacking {1}...")(part->Get_Display_Name()));
 	else {
-		gui_msg(Msg(msg::kError, "unable_to_locate=Unable to locate {1}.")("/boot"));
+		gui_msg(Msg(msg::kError, "unable_to_locate=Unable to locate {1}.")(dest_partition.c_str()));
 		return false;
 	}
 	if (!Backup_Image_For_Repack(part, REPACK_ORIG_DIR, Repack_Options.Backup_First, gui_lookup("repack", "Repack")))
@@ -144,23 +178,23 @@
 	std::string path = REPACK_NEW_DIR;
 	if (Repack_Options.Type == REPLACE_KERNEL) {
 		// When we replace the kernel, what we really do is copy the boot partition ramdisk into the new image's folder
-		if (TWFunc::copy_file(REPACK_ORIG_DIR "ramdisk.cpio", REPACK_NEW_DIR "ramdisk.cpio", 0644)) {
+		if (TWFunc::copy_file(REPACK_ORIG_DIR + ramdisk_cpio, REPACK_NEW_DIR + ramdisk_cpio, 0644)) {
 			LOGERR("Failed to copy ramdisk\n");
 			return false;
 		}
 	} else if (Repack_Options.Type == REPLACE_RAMDISK_UNPACKED) {
-			if (TWFunc::copy_file(Target_Image, REPACK_ORIG_DIR "ramdisk.cpio", 0644)) {
+			if (TWFunc::copy_file(Target_Image, REPACK_ORIG_DIR + ramdisk_cpio, 0644)) {
 				LOGERR("Failed to copy ramdisk\n");
 				return false;
 			}
-			if (TWFunc::copy_file(Target_Image, REPACK_NEW_DIR "ramdisk.cpio", 0644)) {
+			if (TWFunc::copy_file(Target_Image, REPACK_NEW_DIR + ramdisk_cpio, 0644)) {
 				LOGERR("Failed to copy ramdisk\n");
 				return false;
 			}
 		path = REPACK_ORIG_DIR;
 	} else if (Repack_Options.Type == REPLACE_RAMDISK) {
 		// Repack the ramdisk
-		if (TWFunc::copy_file(REPACK_NEW_DIR "ramdisk.cpio", REPACK_ORIG_DIR "ramdisk.cpio", 0644)) {
+		if (TWFunc::copy_file(REPACK_NEW_DIR + ramdisk_cpio, REPACK_ORIG_DIR + ramdisk_cpio, 0644)) {
 			LOGERR("Failed to copy ramdisk\n");
 			return false;
 		}
@@ -181,7 +215,7 @@
 	command += path + "boot.img";
 
 	std::string orig_compressed_image(REPACK_ORIG_DIR);
-	orig_compressed_image += "ramdisk.cpio";
+	orig_compressed_image += ramdisk_cpio;
 	std::string copy_compressed_image(REPACK_ORIG_DIR);
 	copy_compressed_image += "ramdisk-1.cpio";
 
@@ -201,7 +235,7 @@
 
 	DataManager::SetProgress(.75);
 	std::string file = "new-boot.img";
-	DataManager::SetValue("tw_flash_partition", "/boot;");
+	DataManager::SetValue("tw_flash_partition", dest_partition + ";");
 	if (!PartitionManager.Flash_Image(path, file)) {
 		LOGINFO("Error flashing new image\n");
 		return false;
@@ -219,7 +253,7 @@
 			DataManager::SetProgress(.25);
 			if (!Backup_Image_For_Repack(part, REPACK_ORIG_DIR, Repack_Options.Backup_First, gui_lookup("repack", "Repack")))
 				return false;
-			if (TWFunc::copy_file(REPACK_NEW_DIR "ramdisk.cpio", REPACK_ORIG_DIR "ramdisk.cpio", 0644)) {
+			if (TWFunc::copy_file(REPACK_NEW_DIR + ramdisk_cpio, REPACK_ORIG_DIR + ramdisk_cpio, 0644)) {
 				LOGERR("Failed to copy ramdisk\n");
 				return false;
 			}
@@ -246,7 +280,7 @@
 			}
 			DataManager::SetProgress(.75);
 			std::string file = "new-boot.img";
-			DataManager::SetValue("tw_flash_partition", "/boot;");
+			DataManager::SetValue("tw_flash_partition", dest_partition + ";");
 			if (!PartitionManager.Flash_Image(path, file)) {
 				LOGINFO("Error flashing new image\n");
 				return false;
@@ -256,21 +290,56 @@
 		}
 	}
 	TWFunc::removeDir(REPACK_NEW_DIR, false);
-	gui_msg(Msg(msg::kWarning, "repack_overwrite_warning=If device was previously rooted, then root has been overwritten and will need to be reinstalled."));
+	if (dest_partition == "/boot")
+		gui_msg(Msg(msg::kWarning, "repack_overwrite_warning=If device was previously rooted, then root has been overwritten and will need to be reinstalled."));
 	string Current_Slot = PartitionManager.Get_Active_Slot_Display();
-		if (Current_Slot == "A")
-			PartitionManager.Override_Active_Slot("B");
-		else
-			PartitionManager.Override_Active_Slot("A");
+	if (Current_Slot == "A")
+		PartitionManager.Override_Active_Slot("B");
+	else
+		PartitionManager.Override_Active_Slot("A");
 	return true;
 }
 
 bool twrpRepacker::Flash_Current_Twrp() {
-if (!TWFunc::Path_Exists("/ramdisk-files.txt")) {
+	// A/B with dedicated recovery partition
+	std::string slot = android::base::GetProperty("ro.boot.slot_suffix", "");
+	if (slot.empty())
+		slot = android::base::GetProperty("ro.boot.slot", "");
+	if (!slot.empty() && PartitionManager.Find_Partition_By_Path("/recovery")) {
+		std::string root,src, dest;
+		std::string dest_partition = "/recovery";
+		root = "/dev/block/bootdevice/by-name" + dest_partition;
+		if (slot == "_a" || slot == "a") {
+			src = root + "_a";
+			dest= root + "_b";
+		}
+		else {
+			src = root + "_b";
+			dest= root + "_a";
+		}
+		PartitionManager.Unlock_Block_Partitions();
+
+		// only copy the relevant active slot to the inactive slot, on the basis that the recovery currently running
+		// in the active slot can simply be copied over to the inactive slot, so that both have the same recovery image
+		std::string command = "dd bs=1048576 if=" + src + " of=" + dest;
+		LOGINFO("Command=%s\n", command.c_str());
+
+		if (TWFunc::Exec_Cmd(command) != 0) {
+			LOGERR("Failed to flash the %s image\n", dest_partition.c_str());
+			return false;
+		}
+		else {
+			gui_print("Finished flashing the %s image\n", dest_partition.c_str());
+			return true;
+		}
+		// if we reach here, something is awry - bale out
+		return false;
+	}
+
+	if (!TWFunc::Path_Exists("/ramdisk-files.txt")) {
 			LOGERR("can not find ramdisk-files.txt");
 			return false;
 		}
-		PartitionManager.Unlock_Block_Partitions();
 		Repack_Options_struct Repack_Options;
 		Repack_Options.Disable_Verity = false;
 		Repack_Options.Disable_Force_Encrypt = false;
@@ -288,6 +357,6 @@
 		}
 		if (!Repack_Image_And_Flash("/tmp/currentramdisk.cpio.gz", Repack_Options))
 			return false;
-       else
-       return true;
+		else
+			return true;
 }