From a18ad835cb84268fa07118cf6abed01d1e93856a Mon Sep 17 00:00:00 2001
From: Nikias Bassen
Date: Thu, 26 Sep 2019 04:36:39 +0200
Subject: Make sure CTRL+C is working at specific stages of the process

---
 src/idevicerestore.c | 20 ++++++++++++++++----
 src/ipsw.c           | 18 ++++++++++++++++++
 src/ipsw.h           |  2 ++
 src/normal.c         |  4 ++--
 src/recovery.c       |  2 +-
 5 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/src/idevicerestore.c b/src/idevicerestore.c
index c7aeb81..edb977c 100644
--- a/src/idevicerestore.c
+++ b/src/idevicerestore.c
@@ -287,8 +287,8 @@ 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], 10);
-	if (client->mode == &idevicerestore_modes[MODE_UNKNOWN]) {
+	WAIT_FOR(client->mode != &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT), 10);
+	if (client->mode == &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT)) {
 		error("ERROR: Unable to discover device mode. Please make sure a device is attached.\n");
 		return -1;
 	}
@@ -452,6 +452,9 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 					plist_free(signed_fws);
 					return -1;
 				}
+				if (client->flags & FLAG_QUIT) {
+					return -1;
+				}
 				unsigned long selected = strtoul(input, NULL, 10);
 				if (selected == 0 || selected > count) {
 					printf("Invalid input value. Must be in range: 1..%d\n", count);
@@ -528,7 +531,7 @@ 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);
-		if (client->mode == &idevicerestore_modes[MODE_UNKNOWN]) {
+		if (client->mode == &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT)) {
 			error("ERROR: Unable to discover device mode. Please make sure a device is attached.\n");
 			return -1;
 		}
@@ -779,6 +782,9 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 					fflush(stdin);
 					input[0] = '\0';
 					get_user_input(input, 63, 0);
+					if (client->flags & FLAG_QUIT) {
+						return -1;
+					}
 					if (*input != '\0' && !strcmp(input, "YES")) {
 						break;
 					} else {
@@ -806,6 +812,9 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 			fflush(stdin);
 			input[0] = '\0';
 			get_user_input(input, 63, 0);
+			if (client->flags & FLAG_QUIT) {
+				return -1;
+			}
 			if (*input != '\0' && !strcmp(input, "YES")) {
 				break;
 			} else {
@@ -916,6 +925,8 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 			if (client->tss)
 				plist_free(client->tss);
 			plist_free(buildmanifest);
+			info("Removing %s\n", filesystem);
+			unlink(filesystem);
 			return -1;
 		}
 
@@ -1205,7 +1216,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 
 	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]) {
+	if (client->mode != &idevicerestore_modes[MODE_RESTORE] || (client->flags & FLAG_QUIT)) {
 		error("ERROR: Device failed to enter restore mode.\n");
 		if (delete_fs && filesystem)
 			unlink(filesystem);
@@ -1381,6 +1392,7 @@ static void handle_signal(int sig)
 {
 	if (idevicerestore_client) {
 		idevicerestore_client->flags |= FLAG_QUIT;
+		ipsw_cancel();
 	}
 }
 
diff --git a/src/ipsw.c b/src/ipsw.c
index 0dfd9c0..bd02bf9 100644
--- a/src/ipsw.c
+++ b/src/ipsw.c
@@ -48,6 +48,8 @@ typedef struct {
 	char *path;
 } ipsw_archive;
 
+static int cancel_flag = 0;
+
 ipsw_archive* ipsw_open(const char* ipsw);
 void ipsw_close(ipsw_archive* archive);
 
@@ -155,6 +157,8 @@ int ipsw_extract_to_file_with_progress(const char* ipsw, const char* infile, con
 		return -1;
 	}
 
+	cancel_flag = 0;
+
 	if (archive->zip) {
 		int zindex = zip_name_locate(archive->zip, infile, 0);
 		if (zindex < 0) {
@@ -192,6 +196,9 @@ int ipsw_extract_to_file_with_progress(const char* ipsw, const char* infile, con
 		int count, size = BUFSIZE;
 		double progress;
 		for(i = zstat.size; i > 0; i -= count) {
+			if (cancel_flag) {
+				break;
+			}
 			if (i < BUFSIZE)
 				size = i;
 			count = zip_fread(zfile, buffer, size);
@@ -260,6 +267,9 @@ int ipsw_extract_to_file_with_progress(const char* ipsw, const char* infile, con
 				uint64_t bytes = 0;
 				double progress;
 				while (!feof(fi)) {
+					if (cancel_flag) {
+						break;
+					}
 					ssize_t r = fread(buffer, 1, BUFSIZE, fi);
 					if (r < 0) {
 						error("ERROR: fread failed: %s\n", strerror(errno));
@@ -287,6 +297,9 @@ int ipsw_extract_to_file_with_progress(const char* ipsw, const char* infile, con
 		free(filepath);
 	}
 	ipsw_close(archive);
+	if (cancel_flag) {
+		ret = -2;
+	}
 	return ret;
 }
 
@@ -789,3 +802,8 @@ int ipsw_download_latest_fw(plist_t version_data, const char* product, const cha
 
 	return res;
 }
+
+void ipsw_cancel(void)
+{
+	cancel_flag++;
+}
diff --git a/src/ipsw.h b/src/ipsw.h
index 1d00d6f..a664f2b 100644
--- a/src/ipsw.h
+++ b/src/ipsw.h
@@ -46,6 +46,8 @@ int ipsw_download_fw(const char *fwurl, unsigned char* isha1, const char* todir,
 int ipsw_get_latest_fw(plist_t version_data, const char* product, char** fwurl, unsigned char* sha1buf);
 int ipsw_download_latest_fw(plist_t version_data, const char* product, const char* todir, char** ipswfile);
 
+void ipsw_cancel(void);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/normal.c b/src/normal.c
index 58d907c..14fa0f6 100644
--- a/src/normal.c
+++ b/src/normal.c
@@ -239,14 +239,14 @@ int normal_enter_recovery(struct idevicerestore_client_t* client)
 
 	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]) {
+	if (client->mode == &idevicerestore_modes[MODE_NORMAL] || (client->flags & FLAG_QUIT)) {
 		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]) {
+	if (client->mode != &idevicerestore_modes[MODE_RECOVERY] || (client->flags & FLAG_QUIT)) {
 		error("ERROR: Failed to enter recovery mode\n");
 		return -1;
 	}
diff --git a/src/recovery.c b/src/recovery.c
index c079d90..4a1e819 100644
--- a/src/recovery.c
+++ b/src/recovery.c
@@ -233,7 +233,7 @@ int recovery_enter_restore(struct idevicerestore_client_t* client, plist_t build
 
 	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]) {
+	if (client->mode == &idevicerestore_modes[MODE_RECOVERY] || (client->flags & FLAG_QUIT)) {
 		error("ERROR: Failed to place device in restore mode\n");
 		return -1;
 	}
-- 
cgit v1.1-32-gdbae