6 files changed, 61 insertions, 66 deletions

diff --git a/mtp/MtpServer.cpp b/mtp/MtpServer.cpp
index 2c6d37648..9cd67324f 100755
--- a/mtp/MtpServer.cpp
+++ b/mtp/MtpServer.cpp
@@ -96,10 +96,9 @@ static const MtpEventCode kSupportedEventCodes[] = {
 	MTP_EVENT_OBJECT_PROP_CHANGED,
 };
 
-MtpServer::MtpServer(int fd, MtpDatabase* database, bool ptp,
+MtpServer::MtpServer(MtpDatabase* database, bool ptp,
 					int fileGroup, int filePerm, int directoryPerm)
-	:	mFD(fd),
-		mDatabase(database),
+	:	mDatabase(database),
 		mPtp(ptp),
 		mFileGroup(fileGroup),
 		mFilePermission(filePerm),
@@ -110,6 +109,7 @@ MtpServer::MtpServer(int fd, MtpDatabase* database, bool ptp,
 		mSendObjectFormat(0),
 		mSendObjectFileSize(0)
 {
+	mFD = -1;
 }
 
 MtpServer::~MtpServer() {
@@ -183,9 +183,11 @@ bool MtpServer::hasStorage(MtpStorageID id) {
 	return (getStorage(id) != NULL);
 }
 
-void MtpServer::run() {
-	int fd = mFD;
+void MtpServer::run(int fd) {
+	if (fd < 0)
+		return;
 
+	mFD = fd;
 	MTPI("MtpServer::run fd: %d\n", fd);
 
 	while (1) {
@@ -275,7 +277,7 @@ void MtpServer::run() {
 	mObjectEditList.clear();
 
 	if (mSessionOpen)
-		mDatabase->sessionEnded();
+		mDatabase->sessionEnded(); // This doesn't actually do anything but was carry over from AOSP
 	close(fd);
 	mFD = -1;
 }
diff --git a/mtp/MtpServer.h b/mtp/MtpServer.h
index 61f5ccfac..944331134 100755
--- a/mtp/MtpServer.h
+++ b/mtp/MtpServer.h
@@ -92,7 +92,7 @@ private:
     android::Vector<ObjectEdit*>  mObjectEditList;
 
 public:
-                        MtpServer(int fd, MtpDatabase* database, bool ptp,
+                        MtpServer(MtpDatabase* database, bool ptp,
                                     int fileGroup, int filePerm, int directoryPerm);
     virtual             ~MtpServer();
 
@@ -102,7 +102,7 @@ public:
     void                addStorage(MtpStorage* storage);
     void                removeStorage(MtpStorage* storage);
 
-    void                run();
+    void                run(int fd);
 
     void                sendObjectAdded(MtpObjectHandle handle);
     void                sendObjectRemoved(MtpObjectHandle handle);
diff --git a/mtp/mtp_MtpServer.cpp b/mtp/mtp_MtpServer.cpp
index b53d07a8f..8d6038c77 100755
--- a/mtp/mtp_MtpServer.cpp
+++ b/mtp/mtp_MtpServer.cpp
@@ -37,23 +37,6 @@
 
 void twmtp_MtpServer::start()
 {
-	if (setup() == 0) {
-		add_storage();
-		MTPD("Starting add / remove mtppipe monitor thread\n");
-		pthread_t thread;
-		ThreadPtr mtpptr = &twmtp_MtpServer::mtppipe_thread;
-		PThreadPtr p = *(PThreadPtr*)&mtpptr;
-		pthread_create(&thread, NULL, p, this);
-		server->run();
-	}
-}
-
-void twmtp_MtpServer::set_storages(storages* mtpstorages) {
-	stores = mtpstorages;
-}
-
-int twmtp_MtpServer::setup()
-{
 	usePtp =  false;
 	MyMtpDatabase* mtpdb = new MyMtpDatabase();
 	/* Sleep for a bit before we open the MTP USB device because some
@@ -64,27 +47,36 @@ int twmtp_MtpServer::setup()
 #ifdef USB_MTP_DEVICE
 #define STRINGIFY(x) #x
 #define EXPAND(x) STRINGIFY(x)
+	const char* mtp_device = EXPAND(USB_MTP_DEVICE);
 	MTPI("Using '%s' for MTP device.\n", EXPAND(USB_MTP_DEVICE));
-	int fd = open(EXPAND(USB_MTP_DEVICE), O_RDWR);
 #else
-	int fd = open("/dev/mtp_usb", O_RDWR);
+	const char* mtp_device = "/dev/mtp_usb";
 #endif
-	if (fd >= 0) {
-		MTPD("fd: %d\n", fd);
-		server = new MtpServer(fd, mtpdb, usePtp, 0, 0664, 0775);
-		refserver = server;
-		MTPI("created new mtpserver object\n");
-	} else {
+	int fd = open(mtp_device, O_RDWR);
+	if (fd < 0) {
 		MTPE("could not open MTP driver, errno: %d\n", errno);
-		return -1;
+		return;
+	}
+	MTPD("fd: %d\n", fd);
+	server = new MtpServer(mtpdb, usePtp, 0, 0664, 0775);
+	refserver = server;
+	MTPI("created new mtpserver object\n");
+	add_storage();
+	MTPD("Starting add / remove mtppipe monitor thread\n");
+	pthread_t thread;
+	ThreadPtr mtpptr = &twmtp_MtpServer::mtppipe_thread;
+	PThreadPtr p = *(PThreadPtr*)&mtpptr;
+	pthread_create(&thread, NULL, p, this);
+	// This loop restarts the MTP process if the device is unplugged and replugged in
+	while (true) {
+		server->run(fd);
+		fd = open(mtp_device, O_RDWR);
+		usleep(800000);
 	}
-	return 0;
 }
 
-void twmtp_MtpServer::run()
-{
-	MTPD("running in twmtp\n");
-	server->run();
+void twmtp_MtpServer::set_storages(storages* mtpstorages) {
+	stores = mtpstorages;
 }
 
 void twmtp_MtpServer::cleanup()
@@ -126,7 +118,7 @@ void twmtp_MtpServer::add_storage()
 	android::Mutex sMutex;
 	android::Mutex::Autolock autoLock(sMutex);
 
-	MTPI("adding internal storage\n");
+	MTPD("twmtp_MtpServer::add_storage count of storage devices: %i\n", stores->size());
 	for (unsigned int i = 0; i < stores->size(); ++i) {
 			std::string pathStr = stores->at(i)->mount;
 
diff --git a/mtp/mtp_MtpServer.hpp b/mtp/mtp_MtpServer.hpp
index 622ed6d6f..99f63d510 100755
--- a/mtp/mtp_MtpServer.hpp
+++ b/mtp/mtp_MtpServer.hpp
@@ -44,8 +44,6 @@ typedef std::vector<storage*> storages;
 class twmtp_MtpServer {
 	public:
 		void start();
-		int setup();
-		void run();
 		void cleanup();
 		void send_object_added(int handle);
 		void send_object_removed(int handle);
diff --git a/partitionmanager.cpp b/partitionmanager.cpp
index 8e75a9a06..25f9dc7ae 100644
--- a/partitionmanager.cpp
+++ b/partitionmanager.cpp
@@ -1945,7 +1945,6 @@ bool TWPartitionManager::Enable_MTP(void) {
 	}
 	//Launch MTP Responder
 	LOGINFO("Starting MTP\n");
-	int count = 0;
 
 	int mtppipe[2];
 
@@ -1968,36 +1967,23 @@ bool TWPartitionManager::Enable_MTP(void) {
 		TWFunc::write_file("/sys/class/android_usb/android0/idProduct", productstr);
 		property_set("sys.usb.config", "mtp,adb");
 	}
-	std::vector<TWPartition*>::iterator iter;
 	/* To enable MTP debug, use the twrp command line feature to
 	 * twrp set tw_mtp_debug 1
 	 */
 	twrpMtp *mtp = new twrpMtp(DataManager::GetIntValue("tw_mtp_debug"));
-	for (iter = Partitions.begin(); iter != Partitions.end(); iter++) {
-		if ((*iter)->Is_Storage && (*iter)->Is_Present && (*iter)->Mount(false)) {
-			printf("twrp addStorage %s, mtpstorageid: %u, maxFileSize: %lld\n", (*iter)->Storage_Path.c_str(), (*iter)->MTP_Storage_ID, (*iter)->Get_Max_FileSize());
-			mtp->addStorage((*iter)->Storage_Name, (*iter)->Storage_Path, (*iter)->MTP_Storage_ID, (*iter)->Get_Max_FileSize());
-			count++;
-		}
-	}
-	if (count) {
-		mtppid = mtp->forkserver(mtppipe);
-		if (mtppid) {
-			close(mtppipe[0]); // Host closes read side
-			mtp_write_fd = mtppipe[1];
-			DataManager::SetValue("tw_mtp_enabled", 1);
-			return true;
-		} else {
-			close(mtppipe[0]);
-			close(mtppipe[1]);
-			LOGERR("Failed to enable MTP\n");
-			return false;
-		}
+	mtppid = mtp->forkserver(mtppipe);
+	if (mtppid) {
+		close(mtppipe[0]); // Host closes read side
+		mtp_write_fd = mtppipe[1];
+		DataManager::SetValue("tw_mtp_enabled", 1);
+		Add_All_MTP_Storage();
+		return true;
 	} else {
 		close(mtppipe[0]);
 		close(mtppipe[1]);
+		LOGERR("Failed to enable MTP\n");
+		return false;
 	}
-	LOGERR("No valid storage partitions found for MTP.\n");
 #else
 	LOGERR("MTP support not included\n");
 #endif
@@ -2005,6 +1991,22 @@ bool TWPartitionManager::Enable_MTP(void) {
 	return false;
 }
 
+void TWPartitionManager::Add_All_MTP_Storage(void) {
+#ifdef TW_HAS_MTP
+	std::vector<TWPartition*>::iterator iter;
+
+	if (!mtppid)
+		return; // MTP is not enabled
+
+	for (iter = Partitions.begin(); iter != Partitions.end(); iter++) {
+		if ((*iter)->Is_Storage && (*iter)->Is_Present && (*iter)->Mount(false))
+			Add_Remove_MTP_Storage((*iter), MTP_MESSAGE_ADD_STORAGE);
+	}
+#else
+	return;
+#endif
+}
+
 bool TWPartitionManager::Disable_MTP(void) {
 	char old_value[PROPERTY_VALUE_MAX];
 	property_get("sys.usb.config", old_value, "error");
@@ -2082,7 +2084,7 @@ bool TWPartitionManager::Add_Remove_MTP_Storage(TWPartition* Part, int message_t
 			mtp_message.path = Part->Storage_Path.c_str();
 			mtp_message.display = Part->Storage_Name.c_str();
 			mtp_message.maxFileSize = Part->Get_Max_FileSize();
-			LOGINFO("sending message to add %i '%s'\n", Part->MTP_Storage_ID, mtp_message.path);
+			LOGINFO("sending message to add %i '%s' '%s'\n", mtp_message.storage_id, mtp_message.path, mtp_message.display);
 			if (write(mtp_write_fd, &mtp_message, sizeof(mtp_message)) <= 0) {
 				LOGINFO("error sending message to add storage %i\n", Part->MTP_Storage_ID);
 				return false;
diff --git a/partitions.hpp b/partitions.hpp
index 8379de9b3..3cbd8c3b5 100644
--- a/partitions.hpp
+++ b/partitions.hpp
@@ -219,6 +219,7 @@ public:
 	int Fstab_Processed();                                                    // Indicates if the fstab has been processed or not
 	void Output_Storage_Fstab();                                              // Creates a /cache/recovery/storage.fstab file with a list of all potential storage locations for app use
 	bool Enable_MTP();                                                        // Enables MTP
+	void Add_All_MTP_Storage();                                               // Adds all storage objects for MTP
 	bool Disable_MTP();                                                       // Disables MTP
 	bool Add_MTP_Storage(string Mount_Point);                                 // Adds or removes an MTP Storage partition
 	bool Add_MTP_Storage(unsigned int Storage_ID);                            // Adds or removes an MTP Storage partition