2 files changed, 62 insertions, 30 deletions

diff --git a/partitionmanager.cpp b/partitionmanager.cpp
index f094734c4..220ed3a76 100644
--- a/partitionmanager.cpp
+++ b/partitionmanager.cpp
@@ -1510,21 +1510,43 @@ int TWPartitionManager::Fix_Permissions(void) {
 	return result;
 }
 
+TWPartition* TWPartitionManager::Find_Next_Storage(string Path, string Exclude) {
+	std::vector<TWPartition*>::iterator iter = Partitions.begin();
+
+	if (!Path.empty()) {
+		string Search_Path = TWFunc::Get_Root_Path(Path);
+		for (; iter != Partitions.end(); iter++) {
+			if ((*iter)->Mount_Point == Search_Path) {
+				iter++;
+				break;
+			}
+		}
+	}
+
+	for (; iter != Partitions.end(); iter++) {
+		if ((*iter)->Is_Storage && (*iter)->Is_Present && (*iter)->Mount_Point != Exclude) {
+			return (*iter);
+		}
+	}
+
+	return NULL;
+}
+
 int TWPartitionManager::Open_Lun_File(string Partition_Path, string Lun_File) {
 	TWPartition* Part = Find_Partition_By_Path(Partition_Path);
 
 	if (Part == NULL) {
-		LOGERR("Unable to locate volume information for USB storage mode.");
+		LOGERR("Unable to locate '%s' for USB storage mode.", Partition_Path.c_str());
 		return false;
 	}
-	if (!Part->UnMount(true))
+	LOGINFO("USB mount '%s', '%s' > '%s'\n", Partition_Path.c_str(), Part->Actual_Block_Device.c_str(), Lun_File.c_str());
+	if (!Part->UnMount(true) || !Part->Is_Present)
 		return false;
 
 	if (TWFunc::write_file(Lun_File, Part->Actual_Block_Device)) {
 		LOGERR("Unable to write to ums lunfile '%s': (%s)\n", Lun_File.c_str(), strerror(errno));
 		return false;
 	}
-	property_set("sys.storage.ums_enabled", "1");
 	return true;
 }
 
@@ -1534,37 +1556,48 @@ int TWPartitionManager::usb_storage_enable(void) {
 	string ext_path;
 	bool has_multiple_lun = false;
 
-	DataManager::GetValue(TW_HAS_DUAL_STORAGE, has_dual);
 	DataManager::GetValue(TW_HAS_DATA_MEDIA, has_data_media);
-	if (has_dual == 1 && has_data_media == 0) {
-		string Lun_File_str = CUSTOM_LUN_FILE;
-		size_t found = Lun_File_str.find("%");
-		if (found != string::npos) {
-			sprintf(lun_file, CUSTOM_LUN_FILE, 1);
-			if (TWFunc::Path_Exists(lun_file))
-				has_multiple_lun = true;
-		}
-		if (!has_multiple_lun) {
-			// Device doesn't have multiple lun files, mount current storage
-			sprintf(lun_file, CUSTOM_LUN_FILE, 0);
-			return Open_Lun_File(DataManager::GetCurrentStoragePath(), lun_file);
-		} else {
-			// Device has multiple lun files
-			sprintf(lun_file, CUSTOM_LUN_FILE, 0);
-			if (!Open_Lun_File(DataManager::GetSettingsStoragePath(), lun_file))
+	string Lun_File_str = CUSTOM_LUN_FILE;
+	size_t found = Lun_File_str.find("%");
+	if (found != string::npos) {
+		sprintf(lun_file, CUSTOM_LUN_FILE, 1);
+		if (TWFunc::Path_Exists(lun_file))
+			has_multiple_lun = true;
+	}
+	if (!has_multiple_lun) {
+		LOGINFO("Device doesn't have multiple lun files, mount current storage\n");
+		sprintf(lun_file, CUSTOM_LUN_FILE, 0);
+		if (TWFunc::Get_Root_Path(DataManager::GetCurrentStoragePath()) == "/data") {
+			TWPartition* Mount = Find_Next_Storage("", "/data");
+			if (Mount) {
+				if (!Open_Lun_File(Mount->Mount_Point, lun_file))
+					return false;
+			} else {
+				LOGERR("Unable to find storage partition to mount to USB\n");
 				return false;
-			DataManager::GetValue(TW_EXTERNAL_PATH, ext_path);
-			sprintf(lun_file, CUSTOM_LUN_FILE, 1);
-			return Open_Lun_File(ext_path, lun_file);
+			}
+		} else if (!Open_Lun_File(DataManager::GetCurrentStoragePath(), lun_file)) {
+			return false;
 		}
 	} else {
-		if (has_data_media == 0)
-			ext_path = DataManager::GetCurrentStoragePath();
-		else
-			DataManager::GetValue(TW_EXTERNAL_PATH, ext_path);
+		LOGINFO("Device has multiple lun files\n");
+		TWPartition* Mount1;
+		TWPartition* Mount2;
 		sprintf(lun_file, CUSTOM_LUN_FILE, 0);
-		return Open_Lun_File(ext_path, lun_file);
+		Mount1 = Find_Next_Storage("", "/data");
+		if (Mount1) {
+			if (!Open_Lun_File(Mount1->Mount_Point, lun_file))
+				return false;
+			Mount2 = Find_Next_Storage(Mount1->Mount_Point, "/data");
+			if (Mount2) {
+				Open_Lun_File(ext_path, lun_file);
+			}
+		} else {
+			LOGERR("Unable to find storage partition to mount to USB\n");
+			return false;
+		}
 	}
+	property_set("sys.storage.ums_enabled", "1");
 	return true;
 }
 
@@ -1576,8 +1609,6 @@ int TWPartitionManager::usb_storage_disable(void) {
 	for (index=0; index<2; index++) {
 		sprintf(lun_file, CUSTOM_LUN_FILE, index);
 		ret = TWFunc::write_file(lun_file, str);
-		Mount_All_Storage();
-		Update_System_Details();
 		if (ret < 0) {
 			break;
 		}
diff --git a/partitions.hpp b/partitions.hpp
index bec8903d8..650f2f60b 100644
--- a/partitions.hpp
+++ b/partitions.hpp
@@ -218,6 +218,7 @@ private:
 	bool Backup_Partition(TWPartition* Part, string Backup_Folder, bool generate_md5, unsigned long long* img_bytes_remaining, unsigned long long* file_bytes_remaining, unsigned long *img_time, unsigned long *file_time, unsigned long long *img_bytes, unsigned long long *file_bytes);
 	bool Restore_Partition(TWPartition* Part, string Restore_Name, int partition_count);
 	void Output_Partition(TWPartition* Part);
+	TWPartition* Find_Next_Storage(string Path, string Exclude);
 	int Open_Lun_File(string Partition_Path, string Lun_File);
 
 private: