From 92d95d0b8f1870d7cc9eb809a993b04b0e733205 Mon Sep 17 00:00:00 2001
From: Nikias Bassen
Date: Fri, 27 Sep 2019 20:44:03 +0200
Subject: mobile_image_mounter: Better error handling (e.g. when device is
 locked)

---
 include/libimobiledevice/mobile_image_mounter.h |  1 +
 src/mobile_image_mounter.c                      | 77 +++++++++++++------------
 tools/ideviceimagemounter.c                     | 13 ++++-
 3 files changed, 53 insertions(+), 38 deletions(-)

diff --git a/include/libimobiledevice/mobile_image_mounter.h b/include/libimobiledevice/mobile_image_mounter.h
index 8d783c4..e185ad5 100644
--- a/include/libimobiledevice/mobile_image_mounter.h
+++ b/include/libimobiledevice/mobile_image_mounter.h
@@ -40,6 +40,7 @@ typedef enum {
 	MOBILE_IMAGE_MOUNTER_E_PLIST_ERROR    = -2,
 	MOBILE_IMAGE_MOUNTER_E_CONN_FAILED    = -3,
 	MOBILE_IMAGE_MOUNTER_E_COMMAND_FAILED = -4,
+	MOBILE_IMAGE_MOUNTER_E_DEVICE_LOCKED  = -5,
 	MOBILE_IMAGE_MOUNTER_E_UNKNOWN_ERROR  = -256
 } mobile_image_mounter_error_t;
 
diff --git a/src/mobile_image_mounter.c b/src/mobile_image_mounter.c
index c8c4c6f..bbebbb4 100644
--- a/src/mobile_image_mounter.c
+++ b/src/mobile_image_mounter.c
@@ -2,7 +2,7 @@
  * mobile_image_mounter.c
  * com.apple.mobile.mobile_image_mounter service implementation.
  *
- * Copyright (c) 2010 Nikias Bassen, All Rights Reserved.
+ * Copyright (c) 2010-2019 Nikias Bassen, All Rights Reserved.
  *
  * This library is free software; you can redistribute it and/or
  * modify it under the terms of the GNU Lesser General Public
@@ -141,6 +141,43 @@ leave_unlock:
 	return res;
 }
 
+static mobile_image_mounter_error_t process_result(plist_t result, const char *expected_status)
+{
+	mobile_image_mounter_error_t res = MOBILE_IMAGE_MOUNTER_E_COMMAND_FAILED;
+	char* strval = NULL;
+	plist_t node;
+
+	node = plist_dict_get_item(result, "Error");
+	if (node && plist_get_node_type(node) == PLIST_STRING) {
+		plist_get_string_val(node, &strval);
+	}
+	if (strval) {
+		if (!strcmp(strval, "DeviceLocked")) {
+			debug_info("Device is locked, can't mount");
+			res = MOBILE_IMAGE_MOUNTER_E_DEVICE_LOCKED;
+		} else {
+			debug_info("Unhandled error '%s' received", strval);
+		}
+		free(strval);
+		return res;
+	}
+
+	node = plist_dict_get_item(result, "Status");
+	if (node && plist_get_node_type(node) == PLIST_STRING) {
+		plist_get_string_val(node, &strval);
+	}
+	if (!strval) {
+		debug_info("Error: Unexpected response received!");
+	} else if (strcmp(strval, expected_status) == 0) {
+		res = MOBILE_IMAGE_MOUNTER_E_SUCCESS;
+	} else {
+		debug_info("Error: didn't get %s but %s", expected_status, strval);
+	}
+	free(strval);
+
+	return res;
+}
+
 LIBIMOBILEDEVICE_API mobile_image_mounter_error_t mobile_image_mounter_upload_image(mobile_image_mounter_client_t client, const char *image_type, size_t image_size, const char *signature, uint16_t signature_size, mobile_image_mounter_upload_cb_t upload_cb, void* userdata)
 {
 	if (!client || !image_type || (image_size == 0) || !upload_cb) {
@@ -169,23 +206,10 @@ LIBIMOBILEDEVICE_API mobile_image_mounter_error_t mobile_image_mounter_upload_im
 		debug_info("Error receiving response from device!");
 		goto leave_unlock;
 	}
-	res = MOBILE_IMAGE_MOUNTER_E_COMMAND_FAILED;
-
-	char* strval = NULL;
-	plist_t node = plist_dict_get_item(result, "Status");
-	if (node && plist_get_node_type(node) == PLIST_STRING) {
-		plist_get_string_val(node, &strval);
-	}
-	if (!strval) {
-		debug_info("Error: Unexpected response received!");
-		goto leave_unlock;
-	}
-	if (strcmp(strval, "ReceiveBytesAck") != 0) {
-		debug_info("Error: didn't get ReceiveBytesAck but %s", strval);
-		free(strval);
+	res = process_result(result, "ReceiveBytesAck");
+	if (res != MOBILE_IMAGE_MOUNTER_E_SUCCESS) {
 		goto leave_unlock;
 	}
-	free(strval);
 
 	size_t tx = 0;
 	size_t bufsize = 65536;
@@ -223,26 +247,7 @@ LIBIMOBILEDEVICE_API mobile_image_mounter_error_t mobile_image_mounter_upload_im
 		debug_info("Error receiving response from device!");
 		goto leave_unlock;
 	}
-	res = MOBILE_IMAGE_MOUNTER_E_COMMAND_FAILED;
-
-	strval = NULL;
-	node = plist_dict_get_item(result, "Status");
-	if (node && plist_get_node_type(node) == PLIST_STRING) {
-		plist_get_string_val(node, &strval);
-	}
-	if (!strval) {
-		debug_info("Error: Unexpected response received!");
-		goto leave_unlock;
-	}
-	if (strcmp(strval, "Complete") != 0) {
-		debug_info("Error: didn't get Complete but %s", strval);
-		free(strval);
-		goto leave_unlock;
-	} else {
-		res = MOBILE_IMAGE_MOUNTER_E_SUCCESS;
-	}
-	free(strval);
-
+	res = process_result(result, "Complete");
 
 leave_unlock:
 	mobile_image_mounter_unlock(client);
diff --git a/tools/ideviceimagemounter.c b/tools/ideviceimagemounter.c
index 7101c7e..93cab09 100644
--- a/tools/ideviceimagemounter.c
+++ b/tools/ideviceimagemounter.c
@@ -248,7 +248,7 @@ int main(int argc, char **argv)
 	lockdownd_client_free(lckd);
 	lckd = NULL;
 
-	mobile_image_mounter_error_t err;
+	mobile_image_mounter_error_t err = MOBILE_IMAGE_MOUNTER_E_UNKNOWN_ERROR;
 	plist_t result = NULL;
 
 	if (list_mode) {
@@ -344,7 +344,7 @@ int main(int argc, char **argv)
 						uint32_t written, total = 0;
 						while (total < amount) {
 							written = 0;
-							if (afc_file_write(afc, af, buf, amount, &written) !=
+							if (afc_file_write(afc, af, buf + total, amount - total, &written) !=
 								AFC_E_SUCCESS) {
 								fprintf(stderr, "AFC Write error!\n");
 								break;
@@ -368,6 +368,14 @@ int main(int argc, char **argv)
 
 		fclose(f);
 
+		if (err != MOBILE_IMAGE_MOUNTER_E_SUCCESS) {
+			if (err == MOBILE_IMAGE_MOUNTER_E_DEVICE_LOCKED) {
+				printf("ERROR: Device is locked, can't mount. Unlock device and try again.\n");
+			} else {
+				printf("ERROR: Unknown error occurred, can't mount.\n");
+			}
+			goto error_out;
+		}
 		printf("done.\n");
 
 		printf("Mounting...\n");
@@ -435,6 +443,7 @@ int main(int argc, char **argv)
 		plist_free(result);
 	}
 
+error_out:
 	/* perform hangup command */
 	mobile_image_mounter_hangup(mim);
 	/* free client */
-- 
cgit v1.1-32-gdbae