Prevent possible crash if device mode can't be detected

1 files changed, 47 insertions, 8 deletions

diff --git a/src/idevicerestore.c b/src/idevicerestore.c
index ed1f53f..5263c76 100644
--- a/src/idevicerestore.c
+++ b/src/idevicerestore.c
@@ -198,7 +198,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 	load_version_data(client);
 
 	// check which mode the device is currently in so we know where to start
-	if (check_mode(client) < 0 || client->mode->index == MODE_UNKNOWN) {
+	if (check_mode(client) < 0) {
 		error("ERROR: Unable to discover device mode. Please make sure a device is attached.\n");
 		return -1;
 	}
@@ -337,7 +337,10 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 		}
 
 		// we need to refresh the current mode again
-		check_mode(client);
+		if (check_mode(client) < 0) {
+			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);
 	}
 
@@ -1171,14 +1174,23 @@ int check_mode(struct idevicerestore_client_t* client) {
 		mode = MODE_RESTORE;
 	}
 
-	client->mode = &idevicerestore_modes[mode];
+	if (mode == MODE_UNKNOWN) {
+		client->mode = NULL;
+	} else {
+		client->mode = &idevicerestore_modes[mode];
+	}
 	return mode;
 }
 
 const char* check_hardware_model(struct idevicerestore_client_t* client) {
 	const char* hw_model = NULL;
+	int mode = MODE_UNKNOWN;
+
+	if (client->mode) {
+		mode = client->mode->index;
+	}
 
-	switch (client->mode->index) {
+	switch (mode) {
 	case MODE_RESTORE:
 		hw_model = restore_check_hardware_model(client);
 		break;
@@ -1205,8 +1217,13 @@ const char* check_hardware_model(struct idevicerestore_client_t* client) {
 int is_image4_supported(struct idevicerestore_client_t* client)
 {
 	int res = 0;
+	int mode = MODE_UNKNOWN;
+
+	if (client->mode) {
+		mode = client->mode->index;
+	}
 
-	switch (client->mode->index) {
+	switch (mode) {
 	case MODE_NORMAL:
 		res = normal_is_image4_supported(client);
 		break;
@@ -1224,7 +1241,13 @@ int is_image4_supported(struct idevicerestore_client_t* client)
 }
 
 int get_ecid(struct idevicerestore_client_t* client, uint64_t* ecid) {
-	switch (client->mode->index) {
+	int mode = MODE_UNKNOWN;
+
+	if (client->mode) {
+		mode = client->mode->index;
+	}
+
+	switch (mode) {
 	case MODE_NORMAL:
 		if (normal_get_ecid(client, ecid) < 0) {
 			*ecid = 0;
@@ -1238,6 +1261,7 @@ int get_ecid(struct idevicerestore_client_t* client, uint64_t* ecid) {
 			return -1;
 		}
 		break;
+
 	case MODE_RECOVERY:
 		if (recovery_get_ecid(client, ecid) < 0) {
 			*ecid = 0;
@@ -1247,6 +1271,7 @@ int get_ecid(struct idevicerestore_client_t* client, uint64_t* ecid) {
 
 	default:
 		error("ERROR: Device is in an invalid state\n");
+		*ecid = 0;
 		return -1;
 	}
 
@@ -1254,12 +1279,18 @@ int get_ecid(struct idevicerestore_client_t* client, uint64_t* ecid) {
 }
 
 int get_ap_nonce(struct idevicerestore_client_t* client, unsigned char** nonce, int* nonce_size) {
+	int mode = MODE_UNKNOWN;
+
 	*nonce = NULL;
 	*nonce_size = 0;
 
 	info("Getting ApNonce ");
 
-	switch (client->mode->index) {
+	if (client->mode) {
+		mode = client->mode->index;
+	}
+
+	switch (mode) {
 	case MODE_NORMAL:
 		info("in normal mode... ");
 		if (normal_get_ap_nonce(client, nonce, nonce_size) < 0) {
@@ -1283,6 +1314,7 @@ int get_ap_nonce(struct idevicerestore_client_t* client, unsigned char** nonce,
 		break;
 
 	default:
+		info("failed\n");
 		error("ERROR: Device is in an invalid state\n");
 		return -1;
 	}
@@ -1297,12 +1329,18 @@ int get_ap_nonce(struct idevicerestore_client_t* client, unsigned char** nonce,
 }
 
 int get_sep_nonce(struct idevicerestore_client_t* client, unsigned char** nonce, int* nonce_size) {
+	int mode = MODE_UNKNOWN;
+
 	*nonce = NULL;
 	*nonce_size = 0;
 
 	info("Getting SepNonce ");
 
-	switch (client->mode->index) {
+	if (client->mode) {
+		mode = client->mode->index;
+	}
+
+	switch (mode) {
 	case MODE_NORMAL:
 		info("in normal mode... ");
 		if (normal_get_sep_nonce(client, nonce, nonce_size) < 0) {
@@ -1326,6 +1364,7 @@ int get_sep_nonce(struct idevicerestore_client_t* client, unsigned char** nonce,
 		break;
 
 	default:
+		info("failed\n");
 		error("ERROR: Device is in an invalid state\n");
 		return -1;
 	}