From cfa875c4d42aff5c445313274a84c56d7394b162 Mon Sep 17 00:00:00 2001
From: bigbiff <bigbiff@teamw.in>
Date: Sun, 20 Jun 2021 16:20:27 -0400
Subject: [PATCH] vab support: merges and unmapping super devices

- check for merges before formatting data
- add advanced option for unmapping super devices

Change-Id: I38d4d3bbdfa071969016c3e000c86a4d03c71e45
---
 Android.mk                        |  5 +-
 data.cpp                          |  4 ++
 gui/action.cpp                    | 18 ++++++++
 gui/objects.hpp                   |  2 +
 gui/theme/common/landscape.xml    | 11 +++++
 gui/theme/common/languages/en.xml |  8 ++++
 gui/theme/common/portrait.xml     | 11 +++++
 gui/theme/common/watch.xml        | 11 +++++
 partition.cpp                     | 24 +++++++++-
 partitionmanager.cpp              | 77 +++++++++++++++++++++++++++----
 partitions.hpp                    |  9 +++-
 twrpApex.cpp                      |  4 ++
 twrpApex.hpp                      |  1 +
 twrpinstall/twinstall.cpp         | 12 +++--
 14 files changed, 178 insertions(+), 19 deletions(-)

diff --git a/Android.mk b/Android.mk
index effab7fa..aea8b0ca 100755
--- a/Android.mk
+++ b/Android.mk
@@ -79,8 +79,8 @@ else
     LOCAL_CFLAGS += -DTW_EXCLUDE_APEX
 endif
 
-LOCAL_STATIC_LIBRARIES += libavb libtwrpinstall libminadbd_services libinit
-LOCAL_SHARED_LIBRARIES += libfs_mgr
+LOCAL_STATIC_LIBRARIES += libavb libtwrpinstall libminadbd_services libinit libsnapshot_nobinder
+LOCAL_SHARED_LIBRARIES += libfs_mgr libhardware android.hardware.boot@1.0 android.hardware.boot@1.1 libprotobuf-cpp-lite liblp libutils libhidlbase
 LOCAL_C_INCLUDES += \
     system/core/fs_mgr/libfs_avb/include/ \
     system/core/fs_mgr/include_fstab/ \
@@ -148,7 +148,6 @@ endif
 
 ifeq ($(AB_OTA_UPDATER),true)
     LOCAL_CFLAGS += -DAB_OTA_UPDATER=1
-    LOCAL_SHARED_LIBRARIES += libhardware android.hardware.boot@1.0 android.hardware.boot@1.1
     TWRP_REQUIRED_MODULES += libhardware android.hardware.boot@1.0-service android.hardware.boot@1.0-service.rc \
     android.hardware.boot@1.1-service android.hardware.boot@1.1-service.rc android.hardware.boot@1.1.xml
 endif
diff --git a/data.cpp b/data.cpp
index ffb9b420..3d81bf6f 100755
--- a/data.cpp
+++ b/data.cpp
@@ -786,7 +786,11 @@ void DataManager::SetDefaultValues()
 #else
 	mPersist.SetValue(TW_NO_SHA2, "1");
 #endif
+#ifdef AB_OTA_UPDATER
+	mPersist.SetValue(TW_UNMOUNT_SYSTEM, "0");
+#else
 	mPersist.SetValue(TW_UNMOUNT_SYSTEM, "1");
+#endif
 
 #ifdef TW_NO_SCREEN_TIMEOUT
 	mConst.SetValue("tw_screen_timeout_secs", "0");
diff --git a/gui/action.cpp b/gui/action.cpp
index 89432340..a623ccd5 100755
--- a/gui/action.cpp
+++ b/gui/action.cpp
@@ -207,6 +207,7 @@ GUIAction::GUIAction(xml_node<>* node)
 		ADD_ACTION(enableadb);
 		ADD_ACTION(enablefastboot);
 		ADD_ACTION(changeterminal);
+		ADD_ACTION(unmapsuperdevices);
 
 		// remember actions that run in the caller thread
 		for (mapFunc::const_iterator it = mf.begin(); it != mf.end(); ++it)
@@ -2266,6 +2267,23 @@ int GUIAction::changeterminal(std::string arg) {
 		gui_changePage("terminalcommand");
 	return 0;
 }
+
+int GUIAction::unmapsuperdevices(std::string arg __unused) {
+	int op_status = 1;
+
+	operation_start("Remove Super Devices");
+	if (simulate) {
+		simulate_progress_bar();
+	} else {
+		if (PartitionManager.Unmap_Super_Devices()) {
+			op_status = 0;
+		}
+	}
+
+	operation_end(op_status);
+	return 0;
+}
+
 #ifndef TW_EXCLUDE_NANO
 int GUIAction::editfile(std::string arg) {
 	if (term != NULL) {
diff --git a/gui/objects.hpp b/gui/objects.hpp
index 3948fb28..ec550bb3 100755
--- a/gui/objects.hpp
+++ b/gui/objects.hpp
@@ -328,6 +328,8 @@ protected:
 	int screenshot(std::string arg);
 	int setbrightness(std::string arg);
 	int checkforapp(std::string arg);
+	int unmapsuperdevices(std::string arg);
+	int removedynamicgroups(std:: string arg);
 
 	// (originally) threaded actions
 	int fileexists(std::string arg);
diff --git a/gui/theme/common/landscape.xml b/gui/theme/common/landscape.xml
index 3ac0d285..bf8912a3 100755
--- a/gui/theme/common/landscape.xml
+++ b/gui/theme/common/landscape.xml
@@ -3764,6 +3764,17 @@
 						<action function="page">confirm_action</action>
 					</actions>
 				</listitem>
+				<listitem name="{@unmap_super_devices=Unmap Super Devices}">
+					<actions>
+						<action function="set">tw_back=advanced</action>
+						<action function="set">tw_action=unmapsuperdevices</action>
+						<action function="set">tw_text1={@unmap_super_devices_confirm=Unmap all Super Devices?}</action>
+						<action function="set">tw_action_text1={@unmapping_super_devices=Unmapping Super Devices...}</action>
+						<action function="set">tw_complete_text1={@unmap_super_devices_complete=Unmapped all Super Devices}</action>
+						<action function="set">tw_slider_text={@swipe_to_confirm=Swipe to Confirm}</action>
+						<action function="page">confirm_action</action>
+					</actions>
+				</listitem>
 			</listbox>
 
 			<action>
diff --git a/gui/theme/common/languages/en.xml b/gui/theme/common/languages/en.xml
index 811f40c4..1dca4f61 100755
--- a/gui/theme/common/languages/en.xml
+++ b/gui/theme/common/languages/en.xml
@@ -760,5 +760,13 @@
 		<string name="change_twrp_folder_on_process">Changing TWRP folder</string>
 		<string name="change_twrp_folder_after_process">TWRP folder changed to</string>
 		<string name="tw_folder_exists">A folder with that name already exists!</string>
+		<string name="unmap_super_devices">Unmap Super Devices</strings>
+		<string name="unmap_super_devices_confirm">Unmap all Super Devices?</string>
+		<string name="unmapping_super_devices">Unmapping Super Devices...</string>
+		<string name="unmap_super_devices_complete">Unmapped all Super Devices</string>
+		<string name="remove_dynamic_groups_confirm">Unmap all Super Devices?</string>
+		<string name="remove_dynamic_groups">Unmapping Super Devices...</string>
+		<string name="remove_dynamic_groups_complete">Unmapped all Super Devices</string>
+		<string name="mount_vab_partitions">Devices on super may not mount until rebooting recovery.</string>
 	</resources>
 </language>
diff --git a/gui/theme/common/portrait.xml b/gui/theme/common/portrait.xml
index 78677cbf..9e0fec88 100755
--- a/gui/theme/common/portrait.xml
+++ b/gui/theme/common/portrait.xml
@@ -3886,6 +3886,17 @@
 						<action function="page">confirm_action</action>
 					</actions>
 				</listitem>
+				<listitem name="{@unmap_super_devices=Unmap Super Devices}">
+					<actions>
+						<action function="set">tw_back=advanced</action>
+						<action function="set">tw_action=unmapsuperdevices</action>
+						<action function="set">tw_text1={@unmap_super_devices_confirm=Unmap all Super Devices?}</action>
+						<action function="set">tw_action_text1={@unmapping_super_devices=Unmapping Super Devices...}</action>
+						<action function="set">tw_complete_text1={@unmap_super_devices_complete=Unmapped all Super Devices}</action>
+						<action function="set">tw_slider_text={@swipe_to_confirm=Swipe to Confirm}</action>
+						<action function="page">confirm_action</action>
+					</actions>
+				</listitem>
 			</listbox>
 
 			<action>
diff --git a/gui/theme/common/watch.xml b/gui/theme/common/watch.xml
index 165f329c..1b705470 100755
--- a/gui/theme/common/watch.xml
+++ b/gui/theme/common/watch.xml
@@ -4365,6 +4365,17 @@
 						<action function="page">confirm_action</action>
 					</actions>
 				</listitem>
+				<listitem name="{@unmap_super_devices=Unmap Super Devices}">
+					<actions>
+						<action function="set">tw_back=advanced</action>
+						<action function="set">tw_action=unmapsuperdevices</action>
+						<action function="set">tw_text1={@unmap_super_devices_confirm=Unmap all Super Devices?}</action>
+						<action function="set">tw_action_text1={@unmapping_super_devices=Unmapping Super Devices...}</action>
+						<action function="set">tw_complete_text1={@unmap_super_devices_complete=Unmapped all Super Devices}</action>
+						<action function="set">tw_slider_text={@swipe_to_confirm=Swipe to Confirm}</action>
+						<action function="page">confirm_action</action>
+					</actions>
+				</listitem>
 			</listbox>
 
 			<button>
diff --git a/partition.cpp b/partition.cpp
index ba4c9dd0..343ccb27 100755
--- a/partition.cpp
+++ b/partition.cpp
@@ -32,6 +32,8 @@
 #include <libgen.h>
 #include <zlib.h>
 #include <sstream>
+#include <android-base/properties.h>
+#include <libsnapshot/snapshot.h>
 
 #include "cutils/properties.h"
 #include "libblkid/include/blkid.h"
@@ -1714,7 +1716,6 @@ bool TWPartition::Wipe(string New_File_System) {
 		recreate_media = false;
 	} else {
 		DataManager::GetValue(TW_RM_RF_VAR, check);
-
 		if (check || Use_Rm_Rf)
 			wiped = Wipe_RMRF();
 		else if (New_File_System == "ext4")
@@ -3473,3 +3474,24 @@ std::string TWPartition::Get_Display_Name() {
 bool TWPartition::Is_SlotSelect() {
 	return SlotSelect;
 }
+
+bool TWPartition::Check_Pending_Merges() {
+	auto sm = android::snapshot::SnapshotManager::NewForFirstStageMount();
+	if (!sm) {
+		LOGERR("Unable to call snapshot manager\n");
+		return false;
+	}
+
+	auto callback = [&]() -> void {
+		double progress;
+		sm->GetUpdateState(&progress);
+		LOGINFO("waiting for merge to complete: %.2f\n", progress);
+	};
+
+	LOGINFO("checking for merges\n");
+	if (!sm->HandleImminentDataWipe(callback)) {
+		LOGERR("Unable to check merge status\n");
+		return false;
+	}
+	return true;
+}
\ No newline at end of file
diff --git a/partitionmanager.cpp b/partitionmanager.cpp
index a41fb0c8..86a56a02 100755
--- a/partitionmanager.cpp
+++ b/partitionmanager.cpp
@@ -36,10 +36,10 @@
 #include <sys/wait.h>
 #include <linux/fs.h>
 #include <sys/mount.h>
-
-
 #include <sys/poll.h>
 #include <sys/socket.h>
+#include <sys/stat.h>
+#include <sys/types.h>
 #include <linux/types.h>
 #include <linux/netlink.h>
 #include <android-base/chrono_utils.h>
@@ -52,8 +52,12 @@
 #include <fs_mgr.h>
 #include <fs_mgr_dm_linear.h>
 #include <fs_mgr_overlayfs.h>
+#include <fs_mgr/roots.h>
 #include <libgsi/libgsi.h>
 #include <liblp/liblp.h>
+#include <libgsi/libgsi.h>
+#include <liblp/builder.h>
+#include <libsnapshot/snapshot.h>
 
 #include "variables.h"
 #include "twcommon.h"
@@ -109,8 +113,10 @@ extern "C" {
 #include <hardware/boot_control.h>
 #endif
 
+using android::fs_mgr::DestroyLogicalPartition;
 using android::fs_mgr::Fstab;
 using android::fs_mgr::FstabEntry;
+using android::fs_mgr::MetadataBuilder;
 
 extern bool datamedia;
 std::vector<users_struct> Users_List;
@@ -279,8 +285,9 @@ int TWPartitionManager::Process_Fstab(string Fstab_Filename, bool Display_Error)
 		else
 			(*iter)->Has_Android_Secure = false;
 
-		if (Is_Super_Partition(TWFunc::Remove_Beginning_Slash((*iter)->Get_Mount_Point()).c_str()))
+		if (Is_Super_Partition(TWFunc::Remove_Beginning_Slash((*iter)->Get_Mount_Point()).c_str())) {
 			Prepare_Super_Volume((*iter));
+		}
 	}
 
 	//Setup Apex before decryption
@@ -553,6 +560,8 @@ void TWPartitionManager::Output_Partition(TWPartition* Part) {
 		printf("SlotSelect ");
 	if (Part->Mount_Read_Only)
 		printf("Mount_Read_Only ");
+	if (Part->Is_Super)
+		printf("Is_Super ");
 	printf("\n");
 	if (!Part->SubPartition_Of.empty())
 		printf("   SubPartition_Of: %s\n", Part->SubPartition_Of.c_str());
@@ -1570,8 +1579,21 @@ int TWPartitionManager::Wipe_Android_Secure(void) {
 
 int TWPartitionManager::Format_Data(void) {
 	TWPartition* dat = Find_Partition_By_Path("/data");
+	TWPartition* metadata = Find_Partition_By_Path("/metadata");
+	if (metadata != NULL)
+		metadata->UnMount(false);
 
 	if (dat != NULL) {
+		if (android::base::GetBoolProperty("ro.virtual_ab.enabled", true)) {
+#ifndef TW_EXCLUDE_APEX
+			twrpApex apex;
+			apex.Unmount();
+#endif
+			if (metadata != NULL)
+				metadata->Mount(true);
+			if (!dat->Check_Pending_Merges())
+				return false;
+		}
 		return dat->Wipe_Encryption();
 	} else {
 		gui_msg(Msg(msg::kError, "unable_to_locate=Unable to locate {1}.")("/data"));
@@ -3307,14 +3329,16 @@ bool TWPartitionManager::Prepare_Empty_Folder(const std::string& Folder) {
 	return TWFunc::Recursive_Mkdir(Folder);
 }
 
+std::string TWPartitionManager::Get_Bare_Partition_Name(std::string Mount_Point) {
+	if (Mount_Point == "/system_root")
+		return "system";
+	else
+		return TWFunc::Remove_Beginning_Slash(Mount_Point);
+}
+
 bool TWPartitionManager::Prepare_Super_Volume(TWPartition* twrpPart) {
     Fstab fstab;
-	std::string bare_partition_name;
-
-	if (twrpPart->Get_Mount_Point() == "/system_root")
-		bare_partition_name = "system";
-	else
-		bare_partition_name = TWFunc::Remove_Beginning_Slash(twrpPart->Get_Mount_Point());
+	std::string bare_partition_name = Get_Bare_Partition_Name(twrpPart->Get_Mount_Point());
 
 	LOGINFO("Trying to prepare %s from super partition\n", bare_partition_name.c_str());
 
@@ -3487,3 +3511,38 @@ void TWPartitionManager::Unlock_Block_Partitions() {
 		closedir(d);
 	}
 }
+
+bool TWPartitionManager::Unmap_Super_Devices() {
+	bool destroyed = false;
+#ifndef TW_EXCLUDE_APEX
+	twrpApex apex;
+	apex.Unmount();
+#endif
+	for (auto iter = Partitions.begin(); iter != Partitions.end();) {
+		LOGINFO("Checking partition: %s\n", (*iter)->Get_Mount_Point().c_str());
+		if ((*iter)->Is_Super) {
+			TWPartition *part = *iter;
+			std::string bare_partition_name = Get_Bare_Partition_Name((*iter)->Get_Mount_Point());
+			std::string blk_device_partition = bare_partition_name + PartitionManager.Get_Active_Slot_Suffix();
+			(*iter)->UnMount(false);
+			LOGINFO("removing dynamic partition: %s\n", blk_device_partition.c_str());
+			destroyed = DestroyLogicalPartition(blk_device_partition);
+			std::string cow_partition = blk_device_partition + "-cow";
+			std::string cow_partition_path = "/dev/block/mapper/" + cow_partition;
+			struct stat st;
+			if (lstat(cow_partition_path.c_str(), &st) == 0) {
+				LOGINFO("removing cow partition: %s\n", cow_partition.c_str());
+				destroyed = DestroyLogicalPartition(cow_partition);
+			}
+			rmdir((*iter)->Mount_Point.c_str());
+			iter = Partitions.erase(iter);
+			delete part;
+			if (!destroyed) {
+				return false;
+			}
+		} else {
+			++iter;
+		}
+	}
+	return true;
+}
diff --git a/partitions.hpp b/partitions.hpp
index 0a3fa0a1..3d15669f 100755
--- a/partitions.hpp
+++ b/partitions.hpp
@@ -176,7 +176,8 @@ public:
 protected:
 	bool Has_Data_Media;                                                      // Indicates presence of /data/media, may affect wiping and backup methods
 	void Setup_Data_Media();                                                  // Sets up a partition as a /data/media emulated storage partition
-	void Set_Block_Device(std::string block_device);						  // Allow super partition setup to change block device
+	void Set_Block_Device(std::string block_device);                          // Allow super partition setup to change block device
+	bool Check_Pending_Merges();                                              // Check and run pending merges on data for VAB devices
 
 private:
 	bool Process_Fstab_Line(const char *fstab_line, bool Display_Error, std::map<string, Flags_Map> *twrp_flags); // Processes a fstab line
@@ -391,6 +392,7 @@ public:
 	void read_uevent();                                                       // Reads uevent data into a buffer
 	void close_uevent();                                                      // Closes the uevent netlink socket
 	void Add_Partition(TWPartition* Part);                                    // Adds a new partition to the Partitions vector
+	std::string Get_Bare_Partition_Name(std::string Mount_Point);
     bool Prepare_Super_Volume(TWPartition* twrpPart);					  	  // Prepare logical super partition volume for mounting
 	std::string Get_Super_Partition();										  // Get Super Partition block device path
 	void Setup_Super_Devices();												  // Setup logical dm devices on super partition
@@ -400,6 +402,9 @@ public:
 	std::vector<users_struct>* Get_Users_List();                              // Returns pointer to list of users
 	int Set_FDE_Encrypt_Status();                                             // Sets encryption state for FDE devices (ro.crypto.state and ro.crypto.type)
 	void Unlock_Block_Partitions();                                           // Unlock all block devices after update_engine runs
+	bool Unmap_Cow_Devices();
+	bool Remove_Dynamic_Groups();                                             // Remove dynmic groups from super partition
+	bool Unmap_Super_Devices();                                               // Unmap super devices in TWRP
 
 private:
 	void Setup_Settings_Storage_Partition(TWPartition* Part);                 // Sets up settings storage
@@ -421,7 +426,7 @@ private:
 	std::string original_ramdisk_format;                                      // Ramdisk format of boot partition
 	std::string repacked_ramdisk_format;                                      // Ramdisk format of boot image to repack from
 	void Mark_User_Decrypted(int userID);                                     // Marks given user ID in Users_List as decrypted
-	void Check_Users_Decryption_Status();                                      // Checks to see if all users are decrypted
+	void Check_Users_Decryption_Status();                                     // Checks to see if all users are decrypted
 
 private:
 	std::vector<TWPartition*> Partitions;                                     // Vector list of all partitions
diff --git a/twrpApex.cpp b/twrpApex.cpp
index 5531445d..ca6cff42 100755
--- a/twrpApex.cpp
+++ b/twrpApex.cpp
@@ -185,3 +185,7 @@ bool twrpApex::loadApexImage(std::string fileToMount, size_t loop_device_number)
 
 	return true;
 }
+
+bool twrpApex::Unmount() {
+	return (umount2(APEX_BASE, MNT_DETACH) == 0);
+}
diff --git a/twrpApex.hpp b/twrpApex.hpp
index 1c85d06e..d03df106 100755
--- a/twrpApex.hpp
+++ b/twrpApex.hpp
@@ -29,6 +29,7 @@
 class twrpApex {
 public:
 	bool loadApexImages();
+	bool Unmount();
 
 private:
 	std::string unzipImage(std::string file);
diff --git a/twrpinstall/twinstall.cpp b/twrpinstall/twinstall.cpp
index 62bbcbb1..139e983a 100755
--- a/twrpinstall/twinstall.cpp
+++ b/twrpinstall/twinstall.cpp
@@ -334,17 +334,21 @@ int TWinstall_zip(const char* path, int* wipe_cache, bool check_for_digest) {
 			// We need this so backuptool can do its magic
 			bool system_mount_state = PartitionManager.Is_Mounted_By_Path(PartitionManager.Get_Android_Root_Path());
 			bool vendor_mount_state = PartitionManager.Is_Mounted_By_Path("/vendor");
-			PartitionManager.Mount_By_Path(PartitionManager.Get_Android_Root_Path(), true);
-			PartitionManager.Mount_By_Path("/vendor", true);
+			PartitionManager.Mount_By_Path(PartitionManager.Get_Android_Root_Path(), false);
+			PartitionManager.Mount_By_Path("/vendor", false);
 			TWFunc::copy_file("/system/bin/sh", "/tmp/sh", 0755);
 			mount("/tmp/sh", "/system/bin/sh", "auto", MS_BIND, NULL);
 			ret_val = Run_Update_Binary(path, wipe_cache, AB_OTA_ZIP_TYPE);
 			umount("/system/bin/sh");
 			unlink("/tmp/sh");
 			if (!vendor_mount_state)
-				PartitionManager.UnMount_By_Path("/vendor", true);
+				PartitionManager.UnMount_By_Path("/vendor", false);
 			if (!system_mount_state)
-				PartitionManager.UnMount_By_Path(PartitionManager.Get_Android_Root_Path(), true);
+				PartitionManager.UnMount_By_Path(PartitionManager.Get_Android_Root_Path(), false);
+			if (android::base::GetBoolProperty("ro.virtual_ab.enabled", true)) {
+				PartitionManager.Prepare_All_Super_Volumes();
+				gui_warn("mount_vab_partitions=Devices on super may not mount until rebooting recovery.");
+			}
 			gui_warn("flash_ab_reboot=To flash additional zips, please reboot recovery to switch to the updated slot.");
 		} else {
 			std::string binary_name("ui.xml");