asr: implemented checksumming for filesystem upload

3 files changed, 130 insertions, 34 deletions

diff --git a/src/asr.c b/src/asr.c
index cfa8fb3..3d5aada 100644
--- a/src/asr.c
+++ b/src/asr.c
@@ -23,6 +23,7 @@
 #include <stdlib.h>
 #include <string.h>
 #include <libimobiledevice/libimobiledevice.h>
+#include <openssl/sha.h>
 
 #include "asr.h"
 #include "idevicerestore.h"
@@ -30,8 +31,9 @@
 #define ASR_PORT 12345
 #define ASR_BUFFER_SIZE 65536
 #define ASR_PAYLOAD_PACKET_SIZE 1450
+#define ASR_CHECKSUM_CHUNK_SIZE 131072
 
-int asr_open_with_timeout(idevice_t device, idevice_connection_t* asr) {
+int asr_open_with_timeout(idevice_t device, asr_client_t* asr) {
 	int i = 0;
 	int attempts = 10;
 	idevice_connection_t connection = NULL;
@@ -59,11 +61,45 @@ int asr_open_with_timeout(idevice_t device, idevice_connection_t* asr) {
 		debug("Retrying connection...\n");
 	}
 
-	*asr = connection;
+	asr_client_t asr_loc = (asr_client_t)malloc(sizeof(struct asr_client));
+	memset(asr_loc, '\0', sizeof(struct asr_client));
+	asr_loc->connection = connection;
+
+	/* receive Initiate command message */
+	plist_t data = NULL;
+	asr_loc->checksum_chunks = 0;
+	if (asr_receive(asr_loc, &data) < 0) {
+		error("ERROR: Unable to receive data from ASR\n");
+		asr_free(asr_loc);
+		plist_free(data);
+		return -1;
+	}
+	plist_t node;
+	node = plist_dict_get_item(data, "Command");
+	if (node && (plist_get_node_type(node) == PLIST_STRING)) {
+		char* strval = NULL;
+		plist_get_string_val(node, &strval);
+		if (strval && (strcmp(strval, "Initiate") != 0)) {
+			error("ERROR: unexpected ASR plist received:\n");
+			debug_plist(data);
+			plist_free(data);
+			asr_free(asr_loc);
+			return -1;
+		}
+	}
+
+	node = plist_dict_get_item(data, "Checksum Chunks");
+	if (node && (plist_get_node_type(node) == PLIST_BOOLEAN)) {
+		plist_get_bool_val(node, &(asr_loc->checksum_chunks));
+	}
+	plist_free(data);
+
+	*asr = asr_loc;
+
 	return 0;
 }
 
-int asr_receive(idevice_connection_t asr, plist_t* data) {
+int asr_receive(asr_client_t asr, plist_t* data) {
 	uint32_t size = 0;
 	char* buffer = NULL;
 	plist_t request = NULL;
@@ -78,7 +114,7 @@ int asr_receive(idevice_connection_t asr, plist_t* data) {
 	}
 	memset(buffer, '\0', ASR_BUFFER_SIZE);
 
-	device_error = idevice_connection_receive(asr, buffer, ASR_BUFFER_SIZE, &size);
+	device_error = idevice_connection_receive(asr->connection, buffer, ASR_BUFFER_SIZE, &size);
 	if (device_error != IDEVICE_E_SUCCESS) {
 		error("ERROR: Unable to receive data from ASR\n");
 		free(buffer);
@@ -95,7 +131,7 @@ int asr_receive(idevice_connection_t asr, plist_t* data) {
 	return 0;
 }
 
-int asr_send(idevice_connection_t asr, plist_t* data) {
+int asr_send(asr_client_t asr, plist_t* data) {
 	uint32_t size = 0;
 	char* buffer = NULL;
 
@@ -115,11 +151,11 @@ int asr_send(idevice_connection_t asr, plist_t* data) {
 	return 0;
 }
 
-int asr_send_buffer(idevice_connection_t asr, const char* data, uint32_t size) {
+int asr_send_buffer(asr_client_t asr, const char* data, uint32_t size) {
 	uint32_t bytes = 0;
 	idevice_error_t device_error = IDEVICE_E_SUCCESS;
 
-	device_error = idevice_connection_send(asr, data, size, &bytes);
+	device_error = idevice_connection_send(asr->connection, data, size, &bytes);
 	if (device_error != IDEVICE_E_SUCCESS || bytes != size) {
 		error("ERROR: Unable to send data to ASR\n");
 		return -1;
@@ -130,14 +166,18 @@ int asr_send_buffer(idevice_connection_t asr, const char* data, uint32_t size) {
 	return 0;
 }
 
-void asr_close(idevice_connection_t asr) {
+void asr_free(asr_client_t asr) {
 	if (asr != NULL) {
-		idevice_disconnect(asr);
+		if (asr->connection != NULL) {
+			idevice_disconnect(asr->connection);
+			asr->connection = NULL;
+		}
+		free(asr);
 		asr = NULL;
 	}
 }
 
-int asr_perform_validation(idevice_connection_t asr, const char* filesystem) {
+int asr_perform_validation(asr_client_t asr, const char* filesystem) {
 	FILE* file = NULL;
 	uint64_t length = 0;
 	char* command = NULL;
@@ -161,6 +201,9 @@ int asr_perform_validation(idevice_connection_t asr, const char* filesystem) {
 	plist_dict_insert_item(payload_info, "Size", plist_new_uint(length));
 
 	packet_info = plist_new_dict();
+	if (asr->checksum_chunks) {
+		plist_dict_insert_item(packet_info, "Checksum Chunk Size", plist_new_uint(ASR_CHECKSUM_CHUNK_SIZE));
+	}
 	plist_dict_insert_item(packet_info, "FEC Slice Stride", plist_new_uint(40));
 	plist_dict_insert_item(packet_info, "Packet Payload Size", plist_new_uint(ASR_PAYLOAD_PACKET_SIZE));
 	plist_dict_insert_item(packet_info, "Packets Per FEC", plist_new_uint(25));
@@ -216,7 +259,7 @@ int asr_perform_validation(idevice_connection_t asr, const char* filesystem) {
 	return 0;
 }
 
-int asr_handle_oob_data_request(idevice_connection_t asr, plist_t packet, FILE* file) {
+int asr_handle_oob_data_request(asr_client_t asr, plist_t packet, FILE* file) {
 	char* oob_data = NULL;
 	uint64_t oob_offset = 0;
 	uint64_t oob_length = 0;
@@ -262,7 +305,7 @@ int asr_handle_oob_data_request(idevice_connection_t asr, plist_t packet, FILE*
 	return 0;
 }
 
-int asr_send_payload(idevice_connection_t asr, const char* filesystem) {
+int asr_send_payload(asr_client_t asr, const char* filesystem) {
 	int i = 0;
 	char data[ASR_PAYLOAD_PACKET_SIZE];
 	FILE* file = NULL;
@@ -280,12 +323,32 @@ int asr_send_payload(idevice_connection_t asr, const char* filesystem) {
 	length = ftell(file);
 	fseek(file, 0, SEEK_SET);
 
-	for(i = length; i > 0; i -= ASR_PAYLOAD_PACKET_SIZE) {
-		int size = ASR_PAYLOAD_PACKET_SIZE;
+	int chunk = 0;
+	int add_checksum = 0;
+
+	SHA_CTX sha1;
+
+	if (asr->checksum_chunks) {
+		SHA1_Init(&sha1);
+	}
+
+	int size = 0;
+	for (i = length; i > 0; i -= size) {
+		size = ASR_PAYLOAD_PACKET_SIZE;
 		if (i < ASR_PAYLOAD_PACKET_SIZE) {
 			size = i;
 		}
 
+		if (add_checksum) {
+			add_checksum = 0;
+		}
+
+		if (asr->checksum_chunks && ((chunk + size) >= ASR_CHECKSUM_CHUNK_SIZE)) {
+			// reduce packet size to match checksum chunk size
+			size -= ((chunk + size) - ASR_CHECKSUM_CHUNK_SIZE);
+			add_checksum = 1;
+		}
+
 		if (fread(data, 1, size, file) != (unsigned int) size) {
 			error("Error reading filesystem\n");
 			fclose(file);
@@ -298,10 +361,45 @@ int asr_send_payload(idevice_connection_t asr, const char* filesystem) {
 			return -1;
 		}
 
+		if (asr->checksum_chunks) {
+			SHA1_Update(&sha1, data, size);
+			chunk += size;
+
+			if (add_checksum) {
+				// get sha1 of last chunk
+				SHA1_Final(data, &sha1);
+
+				// send checksum
+				if (asr_send_buffer(asr, data, 20) < 0) {
+					error("ERROR: Unable to send chunk checksum\n");
+					fclose(file);
+					return -1;
+				}
+
+				// reset SHA1 context
+				SHA1_Init(&sha1);
+
+				// reset chunk byte counter
+				chunk = 0;
+			}
+		}
+
 		bytes += size;
 		progress = ((double) bytes/ (double) length) * 100.0;
 		print_progress_bar(progress);
+	}
 
+	// if last chunk wasn't terminated with a checksum we do it here
+	if (asr->checksum_chunks && !add_checksum) {
+		// get sha1 of last chunk
+		SHA1_Final(data, &sha1);
+
+		// send checksum
+		if (asr_send_buffer(asr, data, 20) < 0) {
+			error("ERROR: Unable to send chunk checksum\n");
+			fclose(file);
+			return -1;
+		}
 	}
 
 	fclose(file);
diff --git a/src/asr.h b/src/asr.h
index 1423496..f4752b1 100644
--- a/src/asr.h
+++ b/src/asr.h
@@ -28,14 +28,20 @@ extern "C" {
 
 #include <libimobiledevice/libimobiledevice.h>
 
-int asr_open_with_timeout(idevice_t device, idevice_connection_t* asr);
-int asr_send(idevice_connection_t asr, plist_t* data);
-int asr_receive(idevice_connection_t asr, plist_t* data);
-int asr_send_buffer(idevice_connection_t asr, const char* data, uint32_t size);
-void asr_close(idevice_connection_t asr);
-int asr_perform_validation(idevice_connection_t asr, const char* filesystem);
-int asr_send_payload(idevice_connection_t asr, const char* filesystem);
-int asr_handle_oob_data_request(idevice_connection_t asr, plist_t packet, FILE* file);
+struct asr_client {
+	idevice_connection_t connection;
+	uint8_t checksum_chunks;
+};
+typedef struct asr_client *asr_client_t;
+
+int asr_open_with_timeout(idevice_t device, asr_client_t* asr);
+int asr_send(asr_client_t asr, plist_t* data);
+int asr_receive(asr_client_t asr, plist_t* data);
+int asr_send_buffer(asr_client_t asr, const char* data, uint32_t size);
+void asr_free(asr_client_t asr);
+int asr_perform_validation(asr_client_t asr, const char* filesystem);
+int asr_send_payload(asr_client_t asr, const char* filesystem);
+int asr_handle_oob_data_request(asr_client_t asr, plist_t packet, FILE* file);
 
 
 #ifdef __cplusplus
diff --git a/src/restore.c b/src/restore.c
index c1d5e25..25d51f6 100644
--- a/src/restore.c
+++ b/src/restore.c
@@ -640,7 +640,7 @@ int restore_send_filesystem(idevice_t device, const char* filesystem) {
 	int i = 0;
 	FILE* file = NULL;
 	plist_t data = NULL;
-	idevice_connection_t asr = NULL;
+	asr_client_t asr = NULL;
 	idevice_error_t device_error = IDEVICE_E_UNKNOWN_ERROR;
 
 	info("About to send filesystem...\n");
@@ -651,20 +651,12 @@ int restore_send_filesystem(idevice_t device, const char* filesystem) {
 	}
 	info("Connected to ASR\n");
 
-	/* receive Initiate command message */
-	if (asr_receive(asr, &data) < 0) {
-		error("ERROR: Unable to receive data from ASR\n");
-		asr_close(asr);
-		return -1;
-	}
-	plist_free(data);
-
 	// this step sends requested chunks of data from various offsets to asr so
 	// it can validate the filesystem before installing it
 	info("Validating the filesystem\n");
 	if (asr_perform_validation(asr, filesystem) < 0) {
 		error("ERROR: ASR was unable to validate the filesystem\n");
-		asr_close(asr);
+		asr_free(asr);
 		return -1;
 	}
 	info("Filesystem validated\n");
@@ -674,12 +666,12 @@ int restore_send_filesystem(idevice_t device, const char* filesystem) {
 	info("Sending filesystem now...\n");
 	if (asr_send_payload(asr, filesystem) < 0) {
 		error("ERROR: Unable to send payload to ASR\n");
-		asr_close(asr);
+		asr_free(asr);
 		return -1;
 	}
 	info("Done sending filesystem\n");
 
-	asr_close(asr);
+	asr_free(asr);
 	return 0;
 }