1 files changed, 29 insertions, 17 deletions

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 {