Use condition variable instead of active waiting for device event handling

With some devices and USB hardware the reconnect of a device might actually
be faster than the check interval of the active waiting loop. With mutexes
and a condition variable we will not miss the moment of reconnect anymore,
even if it is really quick (like 7ms, right DanyL?)

8 files changed, 157 insertions, 18 deletions

diff --git a/src/common.h b/src/common.h
index 41a7bc6..87ceaaf 100644
--- a/src/common.h
+++ b/src/common.h
@@ -38,6 +38,7 @@ extern "C" {
 #include <libirecovery.h>
 
 #include "idevicerestore.h"
+#include "thread.h"
 
 #define MODE_UNKNOWN        -1
 #define MODE_WTF             0
@@ -106,6 +107,8 @@ struct idevicerestore_client_t {
 	void* progress_cb_data;
 	irecv_device_event_context_t irecv_e_ctx;
 	void* idevice_e_ctx;
+	mutex_t device_event_mutex;
+	cond_t device_event_cond;
 };
 
 extern struct idevicerestore_mode_t idevicerestore_modes[];
@@ -155,10 +158,6 @@ char* realpath(const char *filename, char *resolved_name);
 
 void get_user_input(char *buf, int maxlen, int secure);
 
-#define WAIT_INTERVAL 200000
-#define WAIT_MAX(x) (x * (1000000 / WAIT_INTERVAL))
-#define WAIT_FOR(cond, timeout) { int __repeat = WAIT_MAX(timeout); while (!(cond) && __repeat-- > 0) { __usleep(WAIT_INTERVAL); } }
-
 uint8_t _plist_dict_get_bool(plist_t dict, const char *key);
 uint64_t _plist_dict_get_uint(plist_t dict, const char *key);
 
diff --git a/src/dfu.c b/src/dfu.c
index 167b374..8ef828e 100644
--- a/src/dfu.c
+++ b/src/dfu.c
@@ -336,6 +336,8 @@ int dfu_enter_recovery(struct idevicerestore_client_t* client, plist_t build_ide
 		return 0;
 	}
 
+	mutex_lock(&client->device_event_mutex);
+
 	if (dfu_send_component(client, build_identity, "iBSS") < 0) {
 		error("ERROR: Unable to send iBSS to device\n");
 		irecv_close(client->dfu->client);
@@ -347,21 +349,24 @@ int dfu_enter_recovery(struct idevicerestore_client_t* client, plist_t build_ide
 	if (client->build_major > 8) {
 		/* reconnect */
 		debug("Waiting for device to disconnect...\n");
-		WAIT_FOR(client->mode == &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT), 10);
+		cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 10000);
 		if (client->mode != &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT)) {
+			mutex_unlock(&client->device_event_mutex);
 			if (!(client->flags & FLAG_QUIT)) {
 				error("ERROR: Device did not disconnect. Possibly invalid iBSS. Reset device and try again.\n");
 			}
 			return -1;
 		}
 		debug("Waiting for device to reconnect...\n");
-		WAIT_FOR(client->mode == &idevicerestore_modes[MODE_DFU] || client->mode == &idevicerestore_modes[MODE_RECOVERY] || (client->flags & FLAG_QUIT), 10);
+		cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 10000);
 		if ((client->mode != &idevicerestore_modes[MODE_DFU] && client->mode != &idevicerestore_modes[MODE_RECOVERY]) || (client->flags & FLAG_QUIT)) {
+			mutex_unlock(&client->device_event_mutex);
 			if (!(client->flags & FLAG_QUIT)) {
 				error("ERROR: Device did not reconnect in DFU or recovery mode. Possibly invalid iBSS. Reset device and try again.\n");
 			}
 			return -1;
 		}
+		mutex_unlock(&client->device_event_mutex);
 		dfu_client_new(client);
 
 		/* get nonce */
@@ -409,8 +414,11 @@ int dfu_enter_recovery(struct idevicerestore_client_t* client, plist_t build_ide
 			error("ERROR: set configuration failed\n");
 		}
 
+		mutex_lock(&client->device_event_mutex);
+
 		/* send iBEC */
 		if (dfu_send_component(client, build_identity, "iBEC") < 0) {
+			mutex_unlock(&client->device_event_mutex);
 			error("ERROR: Unable to send iBEC to device\n");
 			irecv_close(client->dfu->client);
 			client->dfu->client = NULL;
@@ -419,6 +427,7 @@ int dfu_enter_recovery(struct idevicerestore_client_t* client, plist_t build_ide
 
 		if (client->mode == &idevicerestore_modes[MODE_RECOVERY]) {
 			if (irecv_send_command(client->dfu->client, "go") != IRECV_E_SUCCESS) {
+				mutex_unlock(&client->device_event_mutex);
 				error("ERROR: Unable to execute iBEC\n");
 				return -1;
 			}
@@ -428,21 +437,24 @@ int dfu_enter_recovery(struct idevicerestore_client_t* client, plist_t build_ide
 	}
 
 	debug("Waiting for device to disconnect...\n");
-	WAIT_FOR(client->mode == &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT), 10);
+	cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 10000);
 	if (client->mode != &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT)) {
+		mutex_unlock(&client->device_event_mutex);
 		if (!(client->flags & FLAG_QUIT)) {
 			error("ERROR: Device did not disconnect. Possibly invalid %s. Reset device and try again.\n", (client->build_major > 8) ? "iBEC" : "iBSS");
 		}
 		return -1;
 	}
 	debug("Waiting for device to reconnect in recovery mode...\n");
-	WAIT_FOR(client->mode == &idevicerestore_modes[MODE_RECOVERY] || (client->flags & FLAG_QUIT), 10);
+	cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 10000);
 	if (client->mode != &idevicerestore_modes[MODE_RECOVERY] || (client->flags & FLAG_QUIT)) {
+		mutex_unlock(&client->device_event_mutex);
 		if (!(client->flags & FLAG_QUIT)) {
 			error("ERROR: Device did not reconnect in recovery mode. Possibly invalid %s. Reset device and try again.\n", (client->build_major > 8) ? "iBEC" : "iBSS");
 		}
 		return -1;
 	}
+	mutex_unlock(&client->device_event_mutex);
 
 	if (recovery_client_new(client) < 0) {
 		error("ERROR: Unable to connect to recovery device\n");
diff --git a/src/idevicerestore.c b/src/idevicerestore.c
index c55cb2a..f96fc61 100644
--- a/src/idevicerestore.c
+++ b/src/idevicerestore.c
@@ -209,16 +209,25 @@ static void idevice_event_cb(const idevice_event_t *event, void *userdata)
 	struct idevicerestore_client_t *client = (struct idevicerestore_client_t*)userdata;
 	if (event->event == IDEVICE_DEVICE_ADD) {
 		if (normal_check_mode(client) == 0) {
+			mutex_lock(&client->device_event_mutex);
 			client->mode = &idevicerestore_modes[MODE_NORMAL];
 			debug("%s: device %016llx (udid: %s) connected in normal mode\n", __func__, client->ecid, client->udid);
+			cond_signal(&client->device_event_cond);
+			mutex_unlock(&client->device_event_mutex);
 		} else if (client->ecid && restore_check_mode(client) == 0) {
+			mutex_lock(&client->device_event_mutex);
 			client->mode = &idevicerestore_modes[MODE_RESTORE];
 			debug("%s: device %016llx (udid: %s) connected in restore mode\n", __func__, client->ecid, client->udid);
+			cond_signal(&client->device_event_cond);
+			mutex_unlock(&client->device_event_mutex);
 		}
 	} else if (event->event == IDEVICE_DEVICE_REMOVE) {
 		if (client->udid && !strcmp(event->udid, client->udid)) {
+			mutex_lock(&client->device_event_mutex);
 			client->mode = &idevicerestore_modes[MODE_UNKNOWN];
 			debug("%s: device %016llx (udid: %s) disconnected\n", __func__, client->ecid, client->udid);
+			cond_signal(&client->device_event_cond);
+			mutex_unlock(&client->device_event_mutex);
 		}
 	}
 }
@@ -231,6 +240,7 @@ static void irecv_event_cb(const irecv_device_event_t* event, void *userdata)
 			client->ecid = event->device_info->ecid;
 		}
 		if (client->ecid && event->device_info->ecid == client->ecid) {
+			mutex_lock(&client->device_event_mutex);
 			switch (event->mode) {
 				case IRECV_K_WTF_MODE:
 					client->mode = &idevicerestore_modes[MODE_WTF];
@@ -248,11 +258,16 @@ static void irecv_event_cb(const irecv_device_event_t* event, void *userdata)
 					client->mode = &idevicerestore_modes[MODE_UNKNOWN];
 			}
 			debug("%s: device %016llx (udid: %s) connected in %s mode\n", __func__, client->ecid, client->udid, client->mode->string);
+			cond_signal(&client->device_event_cond);
+			mutex_unlock(&client->device_event_mutex);
 		}
 	} else if (event->type == IRECV_DEVICE_REMOVE) {
 		if (client->ecid && event->device_info->ecid == client->ecid) {
+			mutex_lock(&client->device_event_mutex);
 			client->mode = &idevicerestore_modes[MODE_UNKNOWN];
 			debug("%s: device %016llx (udid: %s) disconnected\n", __func__, client->ecid, client->udid);
+			cond_signal(&client->device_event_cond);
+			mutex_unlock(&client->device_event_mutex);
 		}
 	}
 }
@@ -290,13 +305,16 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 	client->idevice_e_ctx = idevice_event_cb;
 
 	// check which mode the device is currently in so we know where to start
-	WAIT_FOR(client->mode != &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT), 10);
+	mutex_lock(&client->device_event_mutex);
+	cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 10000);
 	if (client->mode == &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT)) {
+		mutex_unlock(&client->device_event_mutex);
 		error("ERROR: Unable to discover device mode. Please make sure a device is attached.\n");
 		return -1;
 	}
 	idevicerestore_progress(client, RESTORE_STEP_DETECT, 0.1);
 	info("Found device in %s mode\n", client->mode->string);
+	mutex_unlock(&client->device_event_mutex);
 
 	if (client->mode->index == MODE_WTF) {
 		unsigned int cpid = 0;
@@ -362,6 +380,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 			}
 		}
 
+		mutex_lock(&client->device_event_mutex);
 		if (wtftmp) {
 			if (dfu_send_buffer(client, wtftmp, wtfsize) != 0) {
 				error("ERROR: Could not send WTF...\n");
@@ -371,7 +390,14 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 
 		free(wtftmp);
 
-		WAIT_FOR(client->mode == &idevicerestore_modes[MODE_DFU] || (client->flags & FLAG_QUIT), 10); /* TODO: verify if it actually goes from 0x1222 -> 0x1227 */
+		cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 10000);
+		if (client->mode != &idevicerestore_modes[MODE_DFU] || (client->flags & FLAG_QUIT)) {
+			mutex_unlock(&client->device_event_mutex);
+			/* TODO: verify if it actually goes from 0x1222 -> 0x1227 */
+			error("ERROR: Failed to put device into DFU from WTF mode\n");
+			return -1;
+		}
+		mutex_unlock(&client->device_event_mutex);
 	}
 
 	// discover the device type
@@ -533,12 +559,15 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 		}
 
 		// we need to refresh the current mode again
-		WAIT_FOR(client->mode != &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT), 60);
+		mutex_lock(&client->device_event_mutex);
+		cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 60000);
 		if (client->mode == &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT)) {
+			mutex_unlock(&client->device_event_mutex);
 			error("ERROR: Unable to discover device mode. Please make sure a device is attached.\n");
 			return -1;
 		}
 		info("Found device in %s mode\n", client->mode->string);
+		mutex_unlock(&client->device_event_mutex);
 	}
 
 	// verify if ipsw file exists
@@ -1152,8 +1181,11 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 			}
 		}
 
+		mutex_lock(&client->device_event_mutex);
+
 		/* now we load the iBEC */
 		if (recovery_send_ibec(client, build_identity) < 0) {
+			mutex_unlock(&client->device_event_mutex);
 			error("ERROR: Unable to send iBEC\n");
 			if (delete_fs && filesystem)
 				unlink(filesystem);
@@ -1162,8 +1194,10 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 		recovery_client_free(client);
 
 		debug("Waiting for device to disconnect...\n");
-		WAIT_FOR(client->mode == &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT), 10);
+		cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 10000);
 		if (client->mode != &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT)) {
+			mutex_unlock(&client->device_event_mutex);
+
 			if (!(client->flags & FLAG_QUIT)) {
 				error("ERROR: Device did not disconnect. Possibly invalid iBEC. Reset device and try again.\n");
 			}
@@ -1172,8 +1206,9 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 			return -2;
 		}
 		debug("Waiting for device to reconnect in recovery mode...\n");
-		WAIT_FOR(client->mode == &idevicerestore_modes[MODE_RECOVERY] || (client->flags & FLAG_QUIT), 10);
+		cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 10000);
 		if (client->mode != &idevicerestore_modes[MODE_RECOVERY] || (client->flags & FLAG_QUIT)) {
+			mutex_unlock(&client->device_event_mutex);
 			if (!(client->flags & FLAG_QUIT)) {
 				error("ERROR: Device did not reconnect in recovery mode. Possibly invalid iBEC. Reset device and try again.\n");
 			}
@@ -1181,6 +1216,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 				unlink(filesystem);
 			return -2;
 		}
+		mutex_unlock(&client->device_event_mutex);
 	}
 	idevicerestore_progress(client, RESTORE_STEP_PREPARE, 0.5);
 	if (client->flags & FLAG_QUIT) {
@@ -1259,14 +1295,17 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 	}
 	idevicerestore_progress(client, RESTORE_STEP_PREPARE, 0.9);
 
+	mutex_lock(&client->device_event_mutex);
 	info("Waiting for device to enter restore mode...\n");
-	WAIT_FOR(client->mode == &idevicerestore_modes[MODE_RESTORE] || (client->flags & FLAG_QUIT), 180);
+	cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 180000);
 	if (client->mode != &idevicerestore_modes[MODE_RESTORE] || (client->flags & FLAG_QUIT)) {
+		mutex_unlock(&client->device_event_mutex);
 		error("ERROR: Device failed to enter restore mode.\n");
 		if (delete_fs && filesystem)
 			unlink(filesystem);
 		return -1;
 	}
+	mutex_unlock(&client->device_event_mutex);
 
 	// device is finally in restore mode, let's do this
 	if (client->mode->index == MODE_RESTORE) {
@@ -1321,6 +1360,8 @@ struct idevicerestore_client_t* idevicerestore_client_new(void)
 	}
 	memset(client, '\0', sizeof(struct idevicerestore_client_t));
 	client->mode = &idevicerestore_modes[MODE_UNKNOWN];
+	mutex_init(&client->device_event_mutex);
+	cond_init(&client->device_event_cond);
 	return client;
 }
 
@@ -1336,6 +1377,9 @@ void idevicerestore_client_free(struct idevicerestore_client_t* client)
 	if (client->idevice_e_ctx) {
 		idevice_event_unsubscribe();
 	}
+	cond_destroy(&client->device_event_cond);
+	mutex_destroy(&client->device_event_mutex);
+
 	if (client->tss_url) {
 		free(client->tss_url);
 	}
diff --git a/src/normal.c b/src/normal.c
index 89ec462..4104f33 100644
--- a/src/normal.c
+++ b/src/normal.c
@@ -33,6 +33,7 @@
 #include "common.h"
 #include "normal.h"
 #include "recovery.h"
+#include "thread.h"
 
 static int normal_idevice_new(struct idevicerestore_client_t* client, idevice_t* device)
 {
@@ -237,19 +238,23 @@ int normal_enter_recovery(struct idevicerestore_client_t* client)
 	lockdown = NULL;
 	device = NULL;
 
+	mutex_lock(&client->device_event_mutex);
 	debug("DEBUG: Waiting for device to disconnect...\n");
-	WAIT_FOR(client->mode != &idevicerestore_modes[MODE_NORMAL] || (client->flags & FLAG_QUIT), 60);
+	cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 60000);
 	if (client->mode == &idevicerestore_modes[MODE_NORMAL] || (client->flags & FLAG_QUIT)) {
+		mutex_unlock(&client->device_event_mutex);
 		error("ERROR: Failed to place device in recovery mode\n");
 		return -1;
 	}
 
 	debug("DEBUG: Waiting for device to connect in recovery mode...\n");
-	WAIT_FOR(client->mode == &idevicerestore_modes[MODE_RECOVERY] || (client->flags & FLAG_QUIT), 60);
+	cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 60000);
 	if (client->mode != &idevicerestore_modes[MODE_RECOVERY] || (client->flags & FLAG_QUIT)) {
+		mutex_unlock(&client->device_event_mutex);
 		error("ERROR: Failed to enter recovery mode\n");
 		return -1;
 	}
+	mutex_unlock(&client->device_event_mutex);
 
 	if (recovery_client_new(client) < 0) {
 		error("ERROR: Unable to enter recovery mode\n");
diff --git a/src/recovery.c b/src/recovery.c
index 4a1e819..88eeab5 100644
--- a/src/recovery.c
+++ b/src/recovery.c
@@ -226,17 +226,21 @@ int recovery_enter_restore(struct idevicerestore_client_t* client, plist_t build
 		return -1;
 	}
 
+	mutex_lock(&client->device_event_mutex);
 	if (recovery_send_kernelcache(client, build_identity) < 0) {
+		mutex_unlock(&client->device_event_mutex);
 		error("ERROR: Unable to send KernelCache\n");
 		return -1;
 	}
 
 	debug("DEBUG: Waiting for device to disconnect...\n");
-	WAIT_FOR(client->mode != &idevicerestore_modes[MODE_RECOVERY] || (client->flags & FLAG_QUIT), 30);
+	cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 30000);
 	if (client->mode == &idevicerestore_modes[MODE_RECOVERY] || (client->flags & FLAG_QUIT)) {
+		mutex_unlock(&client->device_event_mutex);
 		error("ERROR: Failed to place device in restore mode\n");
 		return -1;
 	}
+	mutex_unlock(&client->device_event_mutex);
 
 	return 0;
 }
diff --git a/src/restore.c b/src/restore.c
index 0ea8a45..fea57fb 100644
--- a/src/restore.c
+++ b/src/restore.c
@@ -306,15 +306,19 @@ int restore_reboot(struct idevicerestore_client_t* client)
 		}
 	}
 
+	mutex_lock(&client->device_event_mutex);
+
 	info("Rebooting restore mode device...\n");
 	restored_reboot(client->restore->client);
 
 	restored_client_free(client->restore->client);
 
-	WAIT_FOR(client->mode != &idevicerestore_modes[MODE_RESTORE] || (client->flags & FLAG_QUIT), 30);
+	cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 30000);
 	if (client->mode == &idevicerestore_modes[MODE_RESTORE]) {
+		mutex_unlock(&client->device_event_mutex);
 		return -1;
 	}
+	mutex_unlock(&client->device_event_mutex);
 
 	return 0;
 }
diff --git a/src/thread.c b/src/thread.c
index b84ee66..8ccdd7f 100644
--- a/src/thread.c
+++ b/src/thread.c
@@ -115,3 +115,63 @@ void thread_once(thread_once_t *once_control, void (*init_routine)(void))
 	pthread_once(once_control, init_routine);
 #endif
 }
+
+void cond_init(cond_t* cond)
+{
+#ifdef WIN32
+	cond->sem = CreateSemaphore(NULL, 0, 32767, NULL);
+#else
+	pthread_cond_init(cond, NULL);
+#endif
+}
+
+void cond_destroy(cond_t* cond)
+{
+#ifdef WIN32
+	CloseHandle(cond->sem);
+#else
+	pthread_cond_destroy(cond);
+#endif
+}
+
+int cond_signal(cond_t* cond)
+{
+#ifdef WIN32
+	int result = 0;
+	if (!ReleaseSemaphore(cond->sem, 1, NULL)) {
+		result = -1;
+	}
+	return result;
+#else
+	return pthread_cond_signal(cond);
+#endif
+}
+
+int cond_wait(cond_t* cond, mutex_t* mutex)
+{
+#ifdef WIN32
+	mutex_unlock(mutex);
+	WaitForSingleObject(cond->sem, INFINITE);
+#else
+	return pthread_cond_wait(cond, mutex);
+#endif
+}
+
+int cond_wait_timeout(cond_t* cond, mutex_t* mutex, unsigned int timeout_ms)
+{
+#ifdef WIN32
+	mutex_unlock(mutex);
+	WaitForSingleObject(cond->sem, timeout_ms);
+#else
+	struct timespec ts;
+	struct timeval now;
+	gettimeofday(&now, NULL);
+
+	ts.tv_sec = now.tv_sec + timeout_ms / 1000;
+	ts.tv_nsec = now.tv_usec * 1000 + 1000 * 1000 * (timeout_ms % 1000);
+	ts.tv_sec += ts.tv_nsec / (1000 * 1000 * 1000);
+	ts.tv_nsec %= (1000 * 1000 * 1000);
+
+	return pthread_cond_timedwait(cond, mutex, &ts);
+#endif
+}
diff --git a/src/thread.h b/src/thread.h
index feccc94..69c6632 100644
--- a/src/thread.h
+++ b/src/thread.h
@@ -26,6 +26,9 @@
 #include <windows.h>
 typedef HANDLE thread_t;
 typedef CRITICAL_SECTION mutex_t;
+typedef struct {
+	HANDLE sem;
+} cond_t;
 typedef volatile struct {
 	LONG lock;
 	int state;
@@ -35,8 +38,10 @@ typedef volatile struct {
 #else
 #include <pthread.h>
 #include <signal.h>
+#include <sys/time.h>
 typedef pthread_t thread_t;
 typedef pthread_mutex_t mutex_t;
+typedef pthread_cond_t cond_t;
 typedef pthread_once_t thread_once_t;
 #define THREAD_ONCE_INIT PTHREAD_ONCE_INIT
 #define THREAD_ID pthread_self()
@@ -56,4 +61,10 @@ void mutex_unlock(mutex_t* mutex);
 
 void thread_once(thread_once_t *once_control, void (*init_routine)(void));
 
+void cond_init(cond_t* cond);
+void cond_destroy(cond_t* cond);
+int cond_signal(cond_t* cond);
+int cond_wait(cond_t* cond, mutex_t* mutex);
+int cond_wait_timeout(cond_t* cond, mutex_t* mutex, unsigned int timeout_ms);
+
 #endif