Reduce memory usage for SourceBootObjectV4 images

4 files changed, 201 insertions, 91 deletions

diff --git a/src/idevicerestore.c b/src/idevicerestore.c
index 641eddb..308c6cb 100644
--- a/src/idevicerestore.c
+++ b/src/idevicerestore.c
@@ -2278,7 +2278,7 @@ int get_recoveryos_root_ticket_tss_response(struct idevicerestore_client_t* clie
 	}
 
 	/* add common tags from manifest */
-	/* Adds Ap,OSLongVersion, AppNonce, @ApImg4Ticket */
+	/* Adds Ap,OSLongVersion, ApNonce, @ApImg4Ticket */
 	if (tss_request_add_ap_img4_tags(request, parameters) < 0) {
 		error("ERROR: Unable to add AP IMG4 tags to TSS request\n");
 		plist_free(request);
diff --git a/src/ipsw.c b/src/ipsw.c
index d4ee5e5..5c88c0d 100644
--- a/src/ipsw.c
+++ b/src/ipsw.c
@@ -710,6 +710,147 @@ int ipsw_extract_to_memory(const char* ipsw, const char* infile, unsigned char**
 	return 0;
 }
 
+int ipsw_extract_send(const char* ipsw, const char* infile, int blocksize, ipsw_send_cb send_callback, void* ctx)
+{
+	unsigned char* buffer = NULL;
+	size_t done = 0;
+	size_t total_size = 0;
+
+	ipsw_archive* archive = ipsw_open(ipsw);
+	if (archive == NULL) {
+		error("ERROR: Invalid archive\n");
+		return -1;
+	}
+
+	if (archive->zip) {
+		int zindex = zip_name_locate(archive->zip, infile, 0);
+		if (zindex < 0) {
+			debug("NOTE: zip_name_locate: '%s' not found in archive.\n", infile);
+			ipsw_close(archive);
+			return -1;
+		}
+
+		struct zip_stat zstat;
+		zip_stat_init(&zstat);
+		if (zip_stat_index(archive->zip, zindex, 0, &zstat) != 0) {
+			error("ERROR: zip_stat_index: %s\n", infile);
+			ipsw_close(archive);
+			return -1;
+		}
+
+		struct zip_file* zfile = zip_fopen_index(archive->zip, zindex, 0);
+		if (zfile == NULL) {
+			error("ERROR: zip_fopen_index: %s\n", infile);
+			ipsw_close(archive);
+			return -1;
+		}
+
+		total_size = zstat.size;
+		buffer = (unsigned char*) malloc(blocksize);
+		if (buffer == NULL) {
+			error("ERROR: Out of memory\n");
+			zip_fclose(zfile);
+			ipsw_close(archive);
+			return -1;
+		}
+
+		while (done < total_size) {
+			size_t size = total_size-done;
+			if (size > blocksize) size = blocksize;
+			zip_int64_t zr = zip_fread(zfile, buffer, size);
+			if (zr < 0) {
+				error("ERROR: %s: zip_fread: %s\n", __func__, infile);
+				break;
+			} else if (zr == 0) {
+				// EOF
+				break;
+			}
+			if (send_callback(ctx, buffer, zr) < 0) {
+				error("ERROR: %s: send failed\n", __func__);
+				break;
+			}
+			done += zr;
+		}
+		zip_fclose(zfile);
+		free(buffer);
+	} else {
+		char *filepath = build_path(archive->path, infile);
+		struct stat fst;
+#ifdef WIN32
+		if (stat(filepath, &fst) != 0) {
+#else
+		if (lstat(filepath, &fst) != 0) {
+#endif
+			error("ERROR: %s: stat failed for %s: %s\n", __func__, filepath, strerror(errno));
+			free(filepath);
+			ipsw_close(archive);
+			return -1;
+		}
+		total_size = fst.st_size;
+		buffer = (unsigned char*)malloc(blocksize);
+		if (buffer == NULL) {
+			error("ERROR: Out of memory\n");
+			free(filepath);
+			ipsw_close(archive);
+			return -1;
+		}
+
+#ifndef WIN32
+		if (S_ISLNK(fst.st_mode)) {
+			ssize_t rl = readlink(filepath, (char*)buffer, (total_size > blocksize) ? blocksize : total_size);
+			if (rl < 0) {
+				error("ERROR: %s: readlink failed for %s: %s\n", __func__, filepath, strerror(errno));
+				free(filepath);
+				free(buffer);
+				ipsw_close(archive);
+				return -1;
+			}
+			send_callback(ctx, buffer, (size_t)rl);
+		} else {
+#endif
+			FILE *f = fopen(filepath, "rb");
+			if (!f) {
+				error("ERROR: %s: fopen failed for %s: %s\n", __func__, filepath, strerror(errno));
+				free(filepath);
+				free(buffer);
+				ipsw_close(archive);
+				return -2;
+			}
+
+			while (done < total_size) {
+				size_t size = total_size-done;
+				if (size > blocksize) size = blocksize;
+				size_t fr = fread(buffer, 1, size, f);
+				if (fr != size) {
+					error("ERROR: %s: fread failed for %s: %s\n", __func__, filepath, strerror(errno));
+					break;
+				}
+				if (send_callback(ctx, buffer, fr) < 0) {
+					error("ERROR: %s: send failed\n", __func__);
+					break;
+				}
+				done += fr;
+			}
+			fclose(f);
+#ifndef WIN32
+		}
+#endif
+		free(filepath);
+		free(buffer);
+	}
+	ipsw_close(archive);
+
+	if (done < total_size) {
+		error("ERROR: %s: Sending file data for %s failed (sent %" PRIu64 "/%" PRIu64 ")\n", __func__, infile, (uint64_t)done, (uint64_t)total_size);
+		return -1;
+	}
+
+	// send a NULL buffer to mark end of transfer
+	send_callback(ctx, NULL, 0);
+
+	return 0;
+}
+
 int ipsw_extract_build_manifest(const char* ipsw, plist_t* buildmanifest, int *tss_enabled)
 {
 	unsigned int size = 0;
diff --git a/src/ipsw.h b/src/ipsw.h
index 3b5da80..84ea7a9 100644
--- a/src/ipsw.h
+++ b/src/ipsw.h
@@ -35,6 +35,7 @@ extern "C" {
 int ipsw_print_info(const char* ipsw);
 
 typedef int (*ipsw_list_cb)(void *ctx, const char* ipsw, const char *name, struct stat *stat);
+typedef int (*ipsw_send_cb)(void *ctx, void *data, size_t size);
 
 int ipsw_is_directory(const char* ipsw);
 int ipsw_file_exists(const char* ipsw, const char* infile);
@@ -42,6 +43,7 @@ int ipsw_get_file_size(const char* ipsw, const char* infile, uint64_t* size);
 int ipsw_extract_to_file(const char* ipsw, const char* infile, const char* outfile);
 int ipsw_extract_to_file_with_progress(const char* ipsw, const char* infile, const char* outfile, int print_progress);
 int ipsw_extract_to_memory(const char* ipsw, const char* infile, unsigned char** pbuffer, unsigned int* psize);
+int ipsw_extract_send(const char* ipsw, const char* infile, int blocksize, ipsw_send_cb send_callback, void* ctx);
 int ipsw_extract_build_manifest(const char* ipsw, plist_t* buildmanifest, int *tss_enabled);
 int ipsw_extract_restore_plist(const char* ipsw, plist_t* restore_plist);
 int ipsw_list_contents(const char* ipsw, ipsw_list_cb cb, void *ctx);
diff --git a/src/restore.c b/src/restore.c
index 17c9b24..b086380 100644
--- a/src/restore.c
+++ b/src/restore.c
@@ -3326,18 +3326,18 @@ int extract_macos_variant(plist_t build_identity, char** output)
 	return 0;
 }
 
-int extract_global_manifest(struct idevicerestore_client_t* client, plist_t build_identity, char *variant, unsigned char** pbuffer, unsigned int* psize)
+static char* extract_global_manifest_path(plist_t build_identity, char *variant)
 {
 	plist_t build_info = plist_dict_get_item(build_identity, "Info");
 	if (!build_info) {
 		error("ERROR: build identity does not contain an 'Info' element\n");
-		return -1;
+		return NULL;
 	}
 
 	plist_t device_class_node = plist_dict_get_item(build_info, "DeviceClass");
 	if (!device_class_node) {
 		error("ERROR: build identity info does not contain a DeviceClass\n");
-		return -1;
+		return NULL;
 	}
 	char *device_class = NULL;
 	plist_get_string_val(device_class_node, &device_class);
@@ -3350,7 +3350,7 @@ int extract_global_manifest(struct idevicerestore_client_t* client, plist_t buil
 		ret = extract_macos_variant(build_identity, &macos_variant);
 		if (ret != 0) {
 			free(device_class);
-			return -1;
+			return NULL;
 		}
 	}
 
@@ -3361,7 +3361,17 @@ int extract_global_manifest(struct idevicerestore_client_t* client, plist_t buil
 	free(device_class);
 	free(macos_variant);
 
-	ret = ipsw_extract_to_memory(client->ipsw, ticket_path, pbuffer, psize);
+	return ticket_path;
+}
+
+int extract_global_manifest(struct idevicerestore_client_t* client, plist_t build_identity, char *variant, unsigned char** pbuffer, unsigned int* psize)
+{
+	char* ticket_path = extract_global_manifest_path(build_identity, variant);
+	if (!ticket_path) {
+		error("ERROR: failed to get global manifest path\n");
+		return -1;
+	}
+	int ret = ipsw_extract_to_memory(client->ipsw, ticket_path, pbuffer, psize);
 	if (ret != 0) {
 		free(ticket_path);
 		error("ERROR: failed to read global manifest\n");
@@ -3372,6 +3382,26 @@ int extract_global_manifest(struct idevicerestore_client_t* client, plist_t buil
 	return 0;
 }
 
+static int _restore_send_file_data(restored_client_t restore, void* data, size_t size)
+{
+	plist_t dict = plist_new_dict();
+	if (data != NULL) {
+		// Send a chunk of file data
+		plist_dict_set_item(dict, "FileData", plist_new_data((char*)data, size));
+	} else {
+		// Send FileDataDone to mark end of transfer
+		plist_dict_set_item(dict, "FileDataDone", plist_new_bool(1));
+	}
+	restored_error_t restore_error = restored_send(restore, dict);
+	if (restore_error != RESTORE_E_SUCCESS) {
+		plist_free(dict);
+		error("ERROR: %s: Failed to send data (%d)\n", __func__, restore_error);
+		return -1;
+	}
+	plist_free(dict);
+	return 0;
+}
+
 int restore_send_personalized_boot_object_v3(restored_client_t restore, struct idevicerestore_client_t* client, plist_t msg, plist_t build_identity)
 {
 	debug_plist(msg);
@@ -3395,9 +3425,8 @@ int restore_send_personalized_boot_object_v3(restored_client_t restore, struct i
 	plist_t blob = NULL;
 	plist_t dict = NULL;
 	restored_error_t restore_error = RESTORE_E_SUCCESS;
-	char *component_name = component;
 
-	info("About to send %s...\n", component_name);
+	info("About to send %s...\n", component);
 
 	if (strcmp(image_name, "__GlobalManifest__") == 0) {
 		int ret = extract_global_manifest(client, build_identity, NULL, &data, &size);
@@ -3446,7 +3475,7 @@ int restore_send_personalized_boot_object_v3(restored_client_t restore, struct i
 			return -1;
 		}
 
-		// Personalize IMG40
+		// Personalize IMG4
 		ret = personalize_component(component, component_data, component_size, client->tss, &data, &size);
 		free(component_data);
 		component_data = NULL;
@@ -3456,43 +3485,23 @@ int restore_send_personalized_boot_object_v3(restored_client_t restore, struct i
 		}
 	}
 
-	// Make plist
-	info("Sending %s now...\n", component_name);
+	info("Sending %s now...\n", component);
 
 	int64_t i = size;
 	while (i > 0) {
 		int blob_size = i > 8192 ? 8192 : i;
-
-		dict = plist_new_dict();
-		blob = plist_new_data((char *) (data + size - i), blob_size);
-		plist_dict_set_item(dict, "FileData", blob);
-
-		restore_error = restored_send(restore, dict);
-		if (restore_error != RESTORE_E_SUCCESS) {
-			error("ERROR: Unable to send component %s data\n", component_name);
+		if (_restore_send_file_data(restore, (data + size - i), blob_size) < 0) {
+			free(data);
+			error("ERROR: Unable to send component %s data\n", component);
 			return -1;
 		}
-
-		plist_free(dict);
-
 		i -= blob_size;
 	}
-	debug("\n");
-
-	// Send FileDataDone
-	dict = plist_new_dict();
-	plist_dict_set_item(dict, "FileDataDone", plist_new_bool(1));
-
-	restore_error = restored_send(restore, dict);
-	if (restore_error != RESTORE_E_SUCCESS) {
-		error("ERROR: Unable to send component %s data\n", component_name);
-		return -1;
-	}
-
-	plist_free(dict);
 	free(data);
 
-	info("Done sending %s\n", component_name);
+	_restore_send_file_data(restore, NULL, 0);
+
+	info("Done sending %s\n", component);
 	return 0;
 }
 
@@ -3521,9 +3530,8 @@ int restore_send_source_boot_object_v4(restored_client_t restore, struct idevice
 	plist_t blob = NULL;
 	plist_t dict = NULL;
 	restored_error_t restore_error = RESTORE_E_SUCCESS;
-	char *component_name = component;
 
-	info("About to send %s...\n", component_name);
+	info("About to send %s...\n", component);
 
 	if (strcmp(image_name, "__GlobalManifest__") == 0) {
 		char *variant = NULL;
@@ -3538,22 +3546,11 @@ int restore_send_source_boot_object_v4(restored_client_t restore, struct idevice
 			return -1;
 		}
 
-		int ret = extract_global_manifest(client, build_identity, variant, &data, &size);
-		if (ret != 0) {
-			return -1;
-		}
+		path = extract_global_manifest_path(build_identity, variant);
 	} else if (strcmp(image_name, "__RestoreVersion__") == 0) {
-		int ret = ipsw_extract_to_memory(client->ipsw, "RestoreVersion.plist", &data, &size);
-		if (ret != 0) {
-			error("ERROR: failed to read global manifest\n");
-			return -1;
-		}
+		path = strdup("RestoreVersion.plist");
 	} else if (strcmp(image_name, "__SystemVersion__") == 0) {
-		int ret = ipsw_extract_to_memory(client->ipsw, "SystemVersion.plist", &data, &size);
-		if (ret != 0) {
-			error("ERROR: failed to read global manifest\n");
-			return -1;
-		}
+		path = strdup("SystemVersion.plist");
 	} else {
 		// Get component path
 		if (client->tss) {
@@ -3568,53 +3565,23 @@ int restore_send_source_boot_object_v4(restored_client_t restore, struct idevice
 				return -1;
 			}
 		}
-
-		int ret = extract_component(client->ipsw, path, &data, &size);
-		free(path);
-		path = NULL;
-		if (ret < 0) {
-			error("ERROR: Unable to extract component %s\n", component);
-			return -1;
-		}
 	}
 
-	// Make plist
-	info("Sending %s now...\n", component_name);
-
-	int64_t i = size;
-	while (i > 0) {
-		int blob_size = i > 8192 ? 8192 : i;
-
-		dict = plist_new_dict();
-		blob = plist_new_data((char *) (data + size - i), blob_size);
-		plist_dict_set_item(dict, "FileData", blob);
-
-		restore_error = restored_send(restore, dict);
-		if (restore_error != RESTORE_E_SUCCESS) {
-			error("ERROR: Unable to send component %s data\n", component_name);
-			return -1;
-		}
-
-		plist_free(dict);
-
-		i -= blob_size;
+	if (!path) {
+		error("ERROR: Failed to get path for component %s\n", component);
+		return -1;
 	}
-	debug("\n");
 
-	// Send FileDataDone
-	dict = plist_new_dict();
-	plist_dict_set_item(dict, "FileDataDone", plist_new_bool(1));
+	info("Sending %s now...\n", component);
 
-	restore_error = restored_send(restore, dict);
-	if (restore_error != RESTORE_E_SUCCESS) {
-		error("ERROR: Unable to send component %s data\n", component_name);
+	if (ipsw_extract_send(client->ipsw, path, 8192, (ipsw_send_cb)_restore_send_file_data, restore) < 0) {
+		free(path);
+		error("ERROR: Failed to send component %s\n", component);
 		return -1;
 	}
+	free(path);
 
-	plist_free(dict);
-	free(data);
-
-	info("Done sending %s\n", component_name);
+	info("Done sending %s\n", component);
 	return 0;
 }