Use global event handlers for normale/restore and recovery/dfu mode devices

7 files changed, 200 insertions, 168 deletions

diff --git a/configure.ac b/configure.ac
index e216535..1d37315 100644
--- a/configure.ac
+++ b/configure.ac
@@ -16,7 +16,7 @@ AC_CONFIG_MACRO_DIR([m4])
 
 # Minimum package versions
 LIBIRECOVERY_VERSION=0.2.0
-LIBIMOBILEDEVICE_VERSION=1.1.6
+LIBIMOBILEDEVICE_VERSION=1.2.1
 LIBPLIST_VERSION=1.12
 LIBZIP_VERSION=0.8
 LIBCURL_VERSION=7.0
diff --git a/src/common.h b/src/common.h
index 2c8d07f..41a7bc6 100644
--- a/src/common.h
+++ b/src/common.h
@@ -32,6 +32,8 @@ extern "C" {
 #include <config.h>
 #endif
 
+#include <unistd.h>
+
 #include <plist/plist.h>
 #include <libirecovery.h>
 
@@ -90,7 +92,6 @@ struct idevicerestore_client_t {
 	char* ipsw;
 	const char* filesystem;
 	struct dfu_client_t* dfu;
-	struct normal_client_t* normal;
 	struct restore_client_t* restore;
 	struct recovery_client_t* recovery;
 	irecv_device_t device;
@@ -103,6 +104,8 @@ struct idevicerestore_client_t {
 	char* cache_dir;
 	idevicerestore_progress_cb_t progress_cb;
 	void* progress_cb_data;
+	irecv_device_event_context_t irecv_e_ctx;
+	void* idevice_e_ctx;
 };
 
 extern struct idevicerestore_mode_t idevicerestore_modes[];
@@ -128,10 +131,12 @@ char *generate_guid(void);
 #ifndef sleep
 #define sleep(x) Sleep(x*1000)
 #endif
+#define __usleep(x) Sleep(x/1000)
 #else
 #include <sys/stat.h>
 #define __mkdir(path, mode) mkdir(path, mode)
 #define FMT_qu "%qu"
+#define __usleep(x) usleep(x)
 #endif
 
 int mkdir_with_parents(const char *dir, int mode);
@@ -150,6 +155,10 @@ 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/idevicerestore.c b/src/idevicerestore.c
index 34e8edb..c7aeb81 100644
--- a/src/idevicerestore.c
+++ b/src/idevicerestore.c
@@ -34,6 +34,7 @@
 #include <plist/plist.h>
 #include <zlib.h>
 #include <libgen.h>
+#include <signal.h>
 
 #include <curl/curl.h>
 
@@ -203,6 +204,56 @@ static int compare_versions(const char *s_ver1, const char *s_ver2)
 	return (get_version_num(s_ver1) & 0xFFFF00) - (get_version_num(s_ver2) & 0xFFFF00);
 }
 
+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) {
+			client->mode = &idevicerestore_modes[MODE_NORMAL];
+			debug("%s: device %016llx (udid: %s) connected in normal mode\n", __func__, client->ecid, client->udid);
+		} else if (client->ecid && restore_check_mode(client) == 0) {
+			client->mode = &idevicerestore_modes[MODE_RESTORE];
+			debug("%s: device %016llx (udid: %s) connected in restore mode\n", __func__, client->ecid, client->udid);
+		}
+	} else if (event->event == IDEVICE_DEVICE_REMOVE) {
+		if (client->udid && !strcmp(event->udid, client->udid)) {
+			client->mode = &idevicerestore_modes[MODE_UNKNOWN];
+			debug("%s: device %016llx (udid: %s) disconnected\n", __func__, client->ecid, client->udid);
+		}
+	}
+}
+
+static void irecv_event_cb(const irecv_device_event_t* event, void *userdata)
+{
+	struct idevicerestore_client_t *client = (struct idevicerestore_client_t*)userdata;
+	if (event->type == IRECV_DEVICE_ADD) {
+		if (client->ecid && event->device_info->ecid == client->ecid) {
+			switch (event->mode) {
+				case IRECV_K_WTF_MODE:
+					client->mode = &idevicerestore_modes[MODE_WTF];
+					break;
+				case IRECV_K_DFU_MODE:
+					client->mode = &idevicerestore_modes[MODE_DFU];
+					break;
+				case IRECV_K_RECOVERY_MODE_1:
+				case IRECV_K_RECOVERY_MODE_2:
+				case IRECV_K_RECOVERY_MODE_3:
+				case IRECV_K_RECOVERY_MODE_4:
+					client->mode = &idevicerestore_modes[MODE_RECOVERY];
+					break;
+				default:
+					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);
+		}
+	} else if (event->type == IRECV_DEVICE_REMOVE) {
+		if (client->ecid && event->device_info->ecid == client->ecid) {
+			client->mode = &idevicerestore_modes[MODE_UNKNOWN];
+			debug("%s: device %016llx (udid: %s) disconnected\n", __func__, client->ecid, client->udid);
+		}
+	}
+}
+
 int idevicerestore_start(struct idevicerestore_client_t* client)
 {
 	int tss_enabled = 0;
@@ -230,8 +281,14 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 
 	idevicerestore_progress(client, RESTORE_STEP_DETECT, 0.0);
 
+	irecv_device_event_subscribe(&client->irecv_e_ctx, irecv_event_cb, client);
+
+	idevice_event_subscribe(idevice_event_cb, client);
+	client->idevice_e_ctx = idevice_event_cb;
+
 	// check which mode the device is currently in so we know where to start
-	if (check_mode(client) < 0) {
+	WAIT_FOR(client->mode != &idevicerestore_modes[MODE_UNKNOWN], 10);
+	if (client->mode == &idevicerestore_modes[MODE_UNKNOWN]) {
 		error("ERROR: Unable to discover device mode. Please make sure a device is attached.\n");
 		return -1;
 	}
@@ -309,10 +366,9 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 		}
 		dfu_client_free(client);
 
-		sleep(1);
-
 		free(wtftmp);
-		client->mode = &idevicerestore_modes[MODE_DFU];
+
+		WAIT_FOR(client->mode == &idevicerestore_modes[MODE_DFU] || (client->flags & FLAG_QUIT), 10); /* TODO: verify if it actually goes from 0x1222 -> 0x1227 */
 	}
 
 	// discover the device type
@@ -332,6 +388,11 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 	if (client->flags & FLAG_PWN) {
 		recovery_client_free(client);
 
+		if (client->mode->index != MODE_DFU) {
+			error("ERROR: Device needs to be in DFU mode for this option.\n");
+			return -1;
+		}
+
 		info("connecting to DFU\n");
 		if (dfu_client_new(client) < 0) {
 			return -1;
@@ -466,7 +527,8 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 		}
 
 		// we need to refresh the current mode again
-		if (check_mode(client) < 0) {
+		WAIT_FOR(client->mode != &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT), 60);
+		if (client->mode == &idevicerestore_modes[MODE_UNKNOWN]) {
 			error("ERROR: Unable to discover device mode. Please make sure a device is attached.\n");
 			return -1;
 		}
@@ -1141,6 +1203,15 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 	}
 	idevicerestore_progress(client, RESTORE_STEP_PREPARE, 0.9);
 
+	info("Waiting for device to enter restore mode...\n");
+	WAIT_FOR(client->mode == &idevicerestore_modes[MODE_RESTORE] || (client->flags & FLAG_QUIT), 180);
+	if (client->mode != &idevicerestore_modes[MODE_RESTORE]) {
+		error("ERROR: Device failed to enter restore mode.\n");
+		if (delete_fs && filesystem)
+			unlink(filesystem);
+		return -1;
+	}
+
 	// device is finally in restore mode, let's do this
 	if (client->mode->index == MODE_RESTORE) {
 		info("About to restore device... \n");
@@ -1193,6 +1264,7 @@ struct idevicerestore_client_t* idevicerestore_client_new(void)
 		return NULL;
 	}
 	memset(client, '\0', sizeof(struct idevicerestore_client_t));
+	client->mode = &idevicerestore_modes[MODE_UNKNOWN];
 	return client;
 }
 
@@ -1202,6 +1274,12 @@ void idevicerestore_client_free(struct idevicerestore_client_t* client)
 		return;
 	}
 
+	if (client->irecv_e_ctx) {
+		irecv_device_event_unsubscribe(client->irecv_e_ctx);
+	}
+	if (client->idevice_e_ctx) {
+		idevice_event_unsubscribe();
+	}
 	if (client->tss_url) {
 		free(client->tss_url);
 	}
@@ -1297,6 +1375,15 @@ void idevicerestore_set_progress_callback(struct idevicerestore_client_t* client
 }
 
 #ifndef IDEVICERESTORE_NOMAIN
+static struct idevicerestore_client_t* idevicerestore_client = NULL;
+
+static void handle_signal(int sig)
+{
+	if (idevicerestore_client) {
+		idevicerestore_client->flags |= FLAG_QUIT;
+	}
+}
+
 int main(int argc, char* argv[]) {
 	int opt = 0;
 	int optindex = 0;
@@ -1309,6 +1396,19 @@ int main(int argc, char* argv[]) {
 		return -1;
 	}
 
+	idevicerestore_client = client;
+
+	struct sigaction sa;
+	memset(&sa, 0, sizeof(struct sigaction));
+	sa.sa_handler = handle_signal;
+	sigaction(SIGINT, &sa, NULL);
+	sigaction(SIGTERM, &sa, NULL);
+#ifndef WIN32
+	sigaction(SIGQUIT, &sa, NULL);
+	sa.sa_handler = SIG_IGN;
+	sigaction(SIGPIPE, &sa, NULL);
+#endif
+
 	if (!isatty(fileno(stdin)) || !isatty(fileno(stdout))) {
 		client->flags &= ~FLAG_INTERACTIVE;
 	} else {
diff --git a/src/normal.c b/src/normal.c
index 4dd5d5c..0f64108 100644
--- a/src/normal.c
+++ b/src/normal.c
@@ -34,69 +34,40 @@
 #include "normal.h"
 #include "recovery.h"
 
-static int normal_device_connected = 0;
-
-void normal_device_callback(const idevice_event_t* event, void* userdata) {
-	struct idevicerestore_client_t* client = (struct idevicerestore_client_t*) userdata;
-	if (event->event == IDEVICE_DEVICE_ADD) {
-		normal_device_connected = 1;
-
-	} else if (event->event == IDEVICE_DEVICE_REMOVE) {
-		normal_device_connected = 0;
-		client->flags &= FLAG_QUIT;
-	}
-}
-
-int normal_client_new(struct idevicerestore_client_t* client) {
-	struct normal_client_t* normal = (struct normal_client_t*) malloc(sizeof(struct normal_client_t));
-	if (normal == NULL) {
-		error("ERROR: Out of memory\n");
-		return -1;
-	}
-
-	if (normal_open_with_timeout(client) < 0) {
-		normal_client_free(client);
-		return -1;
-	}
-
-	client->normal = normal;
-	return 0;
-}
-
-void normal_client_free(struct idevicerestore_client_t* client) {
-	struct normal_client_t* normal = NULL;
-	if (client) {
-		normal = client->normal;
-		if(normal) {
-			if(normal->client) {
-				lockdownd_client_free(normal->client);
-				normal->client = NULL;
-			}
-			if(normal->device) {
-				idevice_free(normal->device);
-				normal->device = NULL;
-			}
-		}
-		free(normal);
-		client->normal = NULL;
-	}
-}
-
 static int normal_idevice_new(struct idevicerestore_client_t* client, idevice_t* device)
 {
 	int num_devices = 0;
 	char **devices = NULL;
 	idevice_t dev = NULL;
 	idevice_error_t device_error;
+	lockdownd_client_t lockdown = NULL;
 
 	*device = NULL;
 
 	if (client->udid) {
 		device_error = idevice_new(&dev, client->udid);
 		if (device_error != IDEVICE_E_SUCCESS) {
-			error("ERROR: %s: can't open device with UDID %s\n", __func__, client->udid);
+			debug("%s: can't open device with UDID %s\n", __func__, client->udid);
+			return -1;
+		}
+
+		if (lockdownd_client_new(dev, &lockdown, "idevicerestore") != LOCKDOWN_E_SUCCESS) {
+			error("ERROR: %s: can't connect to lockdownd on device with UDID %s\n", __func__, client->udid);
+			return -1;
+
+		}
+		char* type = NULL;
+		if (lockdownd_query_type(lockdown, &type) != LOCKDOWN_E_SUCCESS) {
+			return -1;
+		}
+		if (strcmp(type, "com.apple.mobile.lockdown") != 0) {
+			free(type);
 			return -1;
 		}
+		free(type);
+		lockdownd_client_free(lockdown);
+		lockdown = NULL;
+
 		*device = dev;
 		return 0;
 	}
@@ -105,7 +76,6 @@ static int normal_idevice_new(struct idevicerestore_client_t* client, idevice_t*
 	if (num_devices == 0) {
 		return -1;
 	}
-	lockdownd_client_t lockdown = NULL;
 	int j;
 	for (j = 0; j < num_devices; j++) {
 		if (lockdown != NULL) {
@@ -118,7 +88,7 @@ static int normal_idevice_new(struct idevicerestore_client_t* client, idevice_t*
 		}
 		device_error = idevice_new(&dev, devices[j]);
 		if (device_error != IDEVICE_E_SUCCESS) {
-			error("ERROR: %s: can't open device with UDID %s\n", __func__, devices[j]);
+			debug("%s: can't open device with UDID %s\n", __func__, devices[j]);
 			continue;
 		}
 
@@ -179,47 +149,8 @@ int normal_check_mode(struct idevicerestore_client_t* client) {
 	return 0;
 }
 
-int normal_open_with_timeout(struct idevicerestore_client_t* client) {
-	int i = 0;
-	int attempts = 10;
-	idevice_t device = NULL;
-
-	// no context exists so bail
-	if(client == NULL) {
-		return -1;
-	}
-
-	normal_device_connected = 0;
-
-	// create our normal client if it doesn't yet exist
-	if(client->normal == NULL) {
-		client->normal = (struct normal_client_t*) malloc(sizeof(struct normal_client_t));
-		if(client->normal == NULL) {
-			error("ERROR: Out of memory\n");
-			return -1;
-		}
-	}
-
-	for (i = 1; i <= attempts; i++) {
-		normal_idevice_new(client, &device);
-		if (device) {
-			normal_device_connected = 1;
-			break;
-		}
-
-		if (i == attempts) {
-			error("ERROR: Unable to connect to device in normal mode\n");
-			return -1;
-		}
-		sleep(2);
-	}
-
-	client->normal->device = device;
-
-	return 0;
-}
-
-irecv_device_t normal_get_irecv_device(struct idevicerestore_client_t* client) {
+irecv_device_t normal_get_irecv_device(struct idevicerestore_client_t* client)
+{
 	idevice_t device = NULL;
 	lockdownd_client_t lockdown = NULL;
 	lockdownd_error_t lockdown_error = LOCKDOWN_E_SUCCESS;
@@ -267,7 +198,8 @@ irecv_device_t normal_get_irecv_device(struct idevicerestore_client_t* client) {
 	return irecv_device;
 }
 
-int normal_enter_recovery(struct idevicerestore_client_t* client) {
+int normal_enter_recovery(struct idevicerestore_client_t* client)
+{
 	idevice_t device = NULL;
 	lockdownd_client_t lockdown = NULL;
 	idevice_error_t device_error = IDEVICE_E_SUCCESS;
@@ -299,12 +231,25 @@ int normal_enter_recovery(struct idevicerestore_client_t* client) {
 	lockdown = NULL;
 	device = NULL;
 
+	debug("DEBUG: Waiting for device to disconnect...\n");
+	WAIT_FOR(client->mode != &idevicerestore_modes[MODE_NORMAL] || (client->flags & FLAG_QUIT), 30);
+	if (client->mode == &idevicerestore_modes[MODE_NORMAL]) {
+		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), 30);
+	if (client->mode != &idevicerestore_modes[MODE_RECOVERY]) {
+		error("ERROR: Failed to enter recovery mode\n");
+		return -1;
+	}
+
 	if (recovery_client_new(client) < 0) {
 		error("ERROR: Unable to enter recovery mode\n");
 		return -1;
 	}
 
-	client->mode = &idevicerestore_modes[MODE_RECOVERY];
 	return 0;
 }
 
diff --git a/src/normal.h b/src/normal.h
index bafafcc..f438f9b 100644
--- a/src/normal.h
+++ b/src/normal.h
@@ -32,19 +32,8 @@ extern "C" {
 #include <libimobiledevice/lockdown.h>
 #include <libimobiledevice/libimobiledevice.h>
 
-struct normal_client_t {
-	idevice_t device;
-	lockdownd_client_t client;
-	const char* ipsw;
-	plist_t tss;
-};
-
-
 int normal_check_mode(struct idevicerestore_client_t* client);
 irecv_device_t normal_get_irecv_device(struct idevicerestore_client_t* client);
-int normal_client_new(struct idevicerestore_client_t* client);
-void normal_client_free(struct idevicerestore_client_t* client);
-int normal_open_with_timeout(struct idevicerestore_client_t* client);
 int normal_enter_recovery(struct idevicerestore_client_t* client);
 int normal_get_ecid(struct idevicerestore_client_t* client, uint64_t* ecid);
 int normal_is_image4_supported(struct idevicerestore_client_t* client);
diff --git a/src/recovery.c b/src/recovery.c
index 436261f..c079d90 100644
--- a/src/recovery.c
+++ b/src/recovery.c
@@ -231,7 +231,13 @@ int recovery_enter_restore(struct idevicerestore_client_t* client, plist_t build
 		return -1;
 	}
 
-	client->mode = &idevicerestore_modes[MODE_RESTORE];
+	debug("DEBUG: Waiting for device to disconnect...\n");
+	WAIT_FOR(client->mode != &idevicerestore_modes[MODE_RECOVERY] || (client->flags & FLAG_QUIT), 30);
+	if (client->mode == &idevicerestore_modes[MODE_RECOVERY]) {
+		error("ERROR: Failed to place device in restore mode\n");
+		return -1;
+	}
+
 	return 0;
 }
 
diff --git a/src/restore.c b/src/restore.c
index 557ffc0..1cc8b9a 100644
--- a/src/restore.c
+++ b/src/restore.c
@@ -103,7 +103,8 @@ static int restore_finished = 0;
 
 static int restore_device_connected = 0;
 
-int restore_client_new(struct idevicerestore_client_t* client) {
+int restore_client_new(struct idevicerestore_client_t* client)
+{
 	struct restore_client_t* restore = (struct restore_client_t*) malloc(sizeof(struct restore_client_t));
 	if (restore == NULL) {
 		error("ERROR: Out of memory\n");
@@ -119,7 +120,8 @@ int restore_client_new(struct idevicerestore_client_t* client) {
 	return 0;
 }
 
-void restore_client_free(struct idevicerestore_client_t* client) {
+void restore_client_free(struct idevicerestore_client_t* client)
+{
 	if (client && client->restore) {
 		if(client->restore->client) {
 			restored_client_free(client->restore->client);
@@ -162,12 +164,12 @@ static int restore_idevice_new(struct idevicerestore_client_t* client, idevice_t
 		}
 		device_error = idevice_new(&dev, devices[j]);
 		if (device_error != IDEVICE_E_SUCCESS) {
-			error("ERROR: %s: can't open device with UDID %s", __func__, devices[j]);
+			debug("%s: can't open device with UDID %s\n", __func__, devices[j]);
 			continue;
 		}
 
 		if (restored_client_new(dev, &restore, "idevicerestore") != RESTORE_E_SUCCESS) {
-			error("ERROR: %s: can't connect to restored on device with UDID %s\n", __func__, devices[j]);
+			debug("%s: can't connect to restored on device with UDID %s\n", __func__, devices[j]);
 			continue;
 
 		}
@@ -222,7 +224,8 @@ static int restore_idevice_new(struct idevicerestore_client_t* client, idevice_t
 	return 0;
 }
 
-int restore_check_mode(struct idevicerestore_client_t* client) {
+int restore_check_mode(struct idevicerestore_client_t* client)
+{
 	idevice_t device = NULL;
 
 	restore_idevice_new(client, &device);
@@ -234,7 +237,8 @@ int restore_check_mode(struct idevicerestore_client_t* client) {
 	return 0;
 }
 
-irecv_device_t restore_get_irecv_device(struct idevicerestore_client_t* client) {
+irecv_device_t restore_get_irecv_device(struct idevicerestore_client_t* client)
+{
 	char* model = NULL;
 	plist_t node = NULL;
 	idevice_t device = NULL;
@@ -290,18 +294,8 @@ irecv_device_t restore_get_irecv_device(struct idevicerestore_client_t* client)
 	return irecv_device;
 }
 
-void restore_device_callback(const idevice_event_t* event, void* userdata) {
-	struct idevicerestore_client_t* client = (struct idevicerestore_client_t*) userdata;
-	if (event->event == IDEVICE_DEVICE_ADD) {
-		restore_device_connected = 1;
-		client->udid = strdup(event->udid);
-	} else if (event->event == IDEVICE_DEVICE_REMOVE) {
-		restore_device_connected = 0;
-		client->flags |= FLAG_QUIT;
-	}
-}
-
-int restore_reboot(struct idevicerestore_client_t* client) {
+int restore_reboot(struct idevicerestore_client_t* client)
+{
 	if(client->restore == NULL) {
 		if (restore_open_with_timeout(client) < 0) {
 			error("ERROR: Unable to open device in restore mode\n");
@@ -314,8 +308,10 @@ int restore_reboot(struct idevicerestore_client_t* client) {
 
 	restored_client_free(client->restore->client);
 
-	// FIXME: wait for device disconnect here
-	sleep(10);
+	WAIT_FOR(client->mode != &idevicerestore_modes[MODE_RESTORE] || (client->flags & FLAG_QUIT), 30);
+	if (client->mode == &idevicerestore_modes[MODE_RESTORE]) {
+		return -1;
+	}
 
 	return 0;
 }
@@ -384,18 +380,8 @@ static int restore_is_current_device(struct idevicerestore_client_t* client, con
 	return (strcasecmp(this_srnm, client->srnm) == 0);
 }
 
-static void restore_device_event_cb(const idevice_event_t *event, void *user_data)
+int restore_open_with_timeout(struct idevicerestore_client_t* client)
 {
-	if (event->event == IDEVICE_DEVICE_ADD) {
-		struct idevicerestore_client_t* client = (struct idevicerestore_client_t*)user_data;
-		if (!restore_device_connected && restore_is_current_device(client, event->udid)) {
-			restore_device_connected = 1;
-			client->udid = strdup(event->udid);
-		}
-	}
-}
-
-int restore_open_with_timeout(struct idevicerestore_client_t* client) {
 	int i = 0;
 	int attempts = 180;
 	char *type = NULL;
@@ -427,22 +413,9 @@ int restore_open_with_timeout(struct idevicerestore_client_t* client) {
 
 	restore_device_connected = 0;
 
-	info("Waiting for device...\n");
-	idevice_event_subscribe(restore_device_event_cb, client);
-	i = 0;
-	while (i++ < attempts) {
-		debug("Attempt %d to connect to restore mode device...\n", i);
-		if (restore_device_connected) {
-			info("Device %s is now connected in restore mode...\n", client->udid);
-			break;
-		}
-		sleep(1);
-	}
-	idevice_event_unsubscribe();
-
-	if (!restore_device_connected) {
+	if (!restore_is_current_device(client, client->udid)) {
 		error("ERROR: Unable to connect to device in restore mode\n");
-		return (i == attempts ? -2:-1);
+		return -1;
 	}
 
 	info("Connecting now...\n");
@@ -597,7 +570,8 @@ const char* restore_progress_string(unsigned int operation)
 
 static int lastop = 0;
 
-int restore_handle_previous_restore_log_msg(restored_client_t client, plist_t msg) {
+int restore_handle_previous_restore_log_msg(restored_client_t client, plist_t msg)
+{
 	plist_t node = NULL;
 	char* restorelog = NULL;
 
@@ -614,7 +588,8 @@ int restore_handle_previous_restore_log_msg(restored_client_t client, plist_t ms
 	return 0;
 }
 
-int restore_handle_progress_msg(struct idevicerestore_client_t* client, plist_t msg) {
+int restore_handle_progress_msg(struct idevicerestore_client_t* client, plist_t msg)
+{
 	plist_t node = NULL;
 	uint64_t progress = 0;
 	uint64_t operation = 0;
@@ -674,7 +649,8 @@ int restore_handle_progress_msg(struct idevicerestore_client_t* client, plist_t
 	return 0;
 }
 
-int restore_handle_status_msg(restored_client_t client, plist_t msg) {
+int restore_handle_status_msg(restored_client_t client, plist_t msg)
+{
 	int result = 0;
 	uint64_t value = 0;
 	char* log = NULL;
@@ -786,7 +762,8 @@ static void restore_asr_progress_cb(double progress, void* userdata)
 	}
 }
 
-int restore_send_filesystem(struct idevicerestore_client_t* client, idevice_t device, const char* filesystem) {
+int restore_send_filesystem(struct idevicerestore_client_t* client, idevice_t device, const char* filesystem)
+{
 	asr_client_t asr = NULL;
 
 	info("About to send filesystem...\n");
@@ -930,7 +907,8 @@ int restore_send_component(restored_client_t restore, struct idevicerestore_clie
 	return 0;
 }
 
-int restore_send_nor(restored_client_t restore, struct idevicerestore_client_t* client, plist_t build_identity) {
+int restore_send_nor(restored_client_t restore, struct idevicerestore_client_t* client, plist_t build_identity)
+{
 	char* llb_path = NULL;
 	char* llb_filename = NULL;
 	char* sep_path = NULL;
@@ -2596,7 +2574,8 @@ int restore_handle_data_request_msg(struct idevicerestore_client_t* client, idev
 	return 0;
 }
 
-int restore_device(struct idevicerestore_client_t* client, plist_t build_identity, const char* filesystem) {
+int restore_device(struct idevicerestore_client_t* client, plist_t build_identity, const char* filesystem)
+{
 	int err = 0;
 	char* type = NULL;
 	plist_t node = NULL;
@@ -2802,7 +2781,7 @@ int restore_device(struct idevicerestore_client_t* client, plist_t build_identit
 
 	// this is the restore process loop, it reads each message in from
 	// restored and passes that data on to it's specific handler
-	while ((client->flags & FLAG_QUIT) == 0) {
+	while (!(client->flags & FLAG_QUIT)) {
 		// finally, if any of these message handlers returned -1 then we encountered
 		// an unrecoverable error, so we need to bail.
 		if (err < 0) {
@@ -2811,10 +2790,14 @@ int restore_device(struct idevicerestore_client_t* client, plist_t build_identit
 		}
 
 		restore_error = restored_receive(restore, &message);
-		if (restore_error != RESTORE_E_SUCCESS) {
-			debug("No data to read\n");
+		if (restore_error == RESTORE_E_RECEIVE_TIMEOUT) {
+			debug("No data to read (timeout)\n");
 			message = NULL;
 			continue;
+		} else if (restore_error != RESTORE_E_SUCCESS) {
+			error("ERROR: Could not read data (%d). Aborting.\n", restore_error);
+			err = -11;
+			break;
 		}
 
 		// discover what kind of message has been received