win32: Improve irecv_event_handler to make sure events get delivered properly

1 files changed, 43 insertions, 7 deletions

diff --git a/src/libirecovery.c b/src/libirecovery.c
index 16330d2..2491b66 100644
--- a/src/libirecovery.c
+++ b/src/libirecovery.c
@@ -2491,9 +2491,12 @@ static void *_irecv_event_handler(void* data)
 {
 	struct _irecv_event_handler_info* info = (struct _irecv_event_handler_info*)data;
 #ifdef WIN32
+	struct collection newDevices;
 	const GUID *guids[] = { &GUID_DEVINTERFACE_DFU, &GUID_DEVINTERFACE_IBOOT, NULL };
 	int running = 1;
 
+	collection_init(&newDevices);
+
 	mutex_lock(&(info->startup_mutex));
 	cond_signal(&(info->startup_cond));
 	mutex_unlock(&(info->startup_mutex));
@@ -2512,6 +2515,13 @@ static void *_irecv_event_handler(void* data)
 			usbDevices = SetupDiGetClassDevs(guids[k], NULL, NULL, DIGCF_PRESENT | DIGCF_DEVICEINTERFACE);
 			if (!usbDevices) {
 				debug("%s: ERROR: SetupDiGetClassDevs failed\n", __func__);
+				// cleanup/free newDevices
+				FOREACH(struct irecv_win_dev_ctx *win_ctx, &newDevices) {
+					free(win_ctx->details);
+					collection_remove(&newDevices, win_ctx);
+					free(win_ctx);
+				} ENDFOREACH
+				collection_free(&newDevices);
 				return NULL;
 			}
 
@@ -2564,11 +2574,27 @@ static void *_irecv_event_handler(void* data)
 					}
 				} ENDFOREACH
 
-				if (!found) {
-					struct irecv_win_dev_ctx win_ctx;
-					win_ctx.details = details;
-					win_ctx.location = location;
-					_irecv_handle_device_add(&win_ctx);
+				unsigned int pid = 0;
+				if (sscanf(details->DevicePath, "\\\\?\\usb#vid_05ac&pid_%04x", &pid)!= 1) {
+					debug("%s: ERROR: failed to parse PID! path: %s\n", __func__, details->DevicePath);
+					free(details);
+					continue;
+				}
+				// make sure the current device is actually in the right mode for the given driver interface
+				int skip = 0;
+				if ((guids[k] == &GUID_DEVINTERFACE_DFU && pid != IRECV_K_DFU_MODE && pid != IRECV_K_WTF_MODE)
+				    || (guids[k] == &GUID_DEVINTERFACE_IBOOT && (pid < IRECV_K_RECOVERY_MODE_1 || pid > IRECV_K_RECOVERY_MODE_4))
+				) {
+					skip = 1;
+				}
+
+				if (!found && !skip) {
+					// Add device to newDevices list, and deliver the notification later, when removed devices are first handled.
+					struct irecv_win_dev_ctx *win_ctx = (struct irecv_win_dev_ctx*)malloc(sizeof(struct irecv_win_dev_ctx));
+					win_ctx->details = details;
+					win_ctx->location = location;
+					collection_add(&newDevices, win_ctx);
+					details = NULL;
 				}
 				free(details);
 			}
@@ -2577,19 +2603,29 @@ static void *_irecv_event_handler(void* data)
 
 		FOREACH(struct irecv_usb_device_info *devinfo, &devices) {
 			if (!devinfo->alive) {
+				debug("%s: removed ecid: %016" PRIx64 ", location: %d\n",__func__, (uint64_t)devinfo->device_info.ecid, devinfo->location);
 				_irecv_handle_device_remove(devinfo);
 			}
 		} ENDFOREACH
 
+		// handle newly added devices and remove from local list
+		FOREACH(struct irecv_win_dev_ctx *win_ctx, &newDevices) {
+			debug("%s: found new: %s, location: %d\n", __func__, win_ctx->details->DevicePath, win_ctx->location);
+			_irecv_handle_device_add(win_ctx);
+			free(win_ctx->details);
+			collection_remove(&newDevices, win_ctx);
+			free(win_ctx);
+		} ENDFOREACH
+
+		Sleep(500);
 		mutex_lock(&listener_mutex);
 		if (collection_count(&listeners) == 0) {
 			running = 0;
 		}
 		mutex_unlock(&listener_mutex);
-
-		Sleep(500);
 	} while (running);
 
+	collection_free(&newDevices);
 #else /* !WIN32 */
 #ifdef HAVE_IOKIT
 	kern_return_t kr;