Handle connection failures in debugserver_client_receive_response; also simplify the code

1 files changed, 85 insertions, 92 deletions

diff --git a/src/debugserver.c b/src/debugserver.c
index 1c20c25..233af4c 100644
--- a/src/debugserver.c
+++ b/src/debugserver.c
@@ -375,130 +375,126 @@ LIBIMOBILEDEVICE_API debugserver_error_t debugserver_client_set_receive_params(d
 	return DEBUGSERVER_E_SUCCESS;
 }
 
-static int debugserver_client_receive_internal_check(debugserver_client_t client, char* received_char)
+static debugserver_error_t debugserver_client_receive_internal_char(debugserver_client_t client, char* received_char)
 {
 	debugserver_error_t res = DEBUGSERVER_E_SUCCESS;
-	int did_receive_char = 0;
-	char buffer = 0;
 	uint32_t bytes = 0;
 
 	/* we loop here as we expect an answer */
-	res = debugserver_client_receive(client, &buffer, sizeof(char), &bytes);
-	if (res == DEBUGSERVER_E_SUCCESS && received_char[0] != 0) {
-		if (memcmp(&buffer, received_char, sizeof(char)) == 0) {
-			did_receive_char = 1;
-		}
-	} else {
-		did_receive_char = 0;
+	res = debugserver_client_receive(client, received_char, sizeof(char), &bytes);
+	if (res != DEBUGSERVER_E_SUCCESS) {
+		return res;
 	}
-
-	if (!did_receive_char) {
-		memcpy(received_char, &buffer, sizeof(char));
+	if (bytes != 1) {
+		debug_info("received %d bytes when asking for %d!", bytes, sizeof(char));
+		return DEBUGSERVER_E_UNKNOWN_ERROR;
 	}
-
-	return did_receive_char;
+	return res;
 }
 
 LIBIMOBILEDEVICE_API debugserver_error_t debugserver_client_receive_response(debugserver_client_t client, char** response, size_t* response_size)
 {
 	debugserver_error_t res = DEBUGSERVER_E_SUCCESS;
 
-	int should_receive = 1;
+	char data = '\0';
 	int skip_prefix = 0;
-	char* command_prefix = strdup("$");
 
-	char* buffer = NULL;
+	char* buffer = malloc(1024);
 	uint32_t buffer_size = 0;
-	uint32_t buffer_capacity = 0;
+	uint32_t buffer_capacity = 1024;
 
 	if (response)
 		*response = NULL;
 
 	if (!client->noack_mode) {
-		char ack[2] = {'+', '\0'};
-		debug_info("attempting to receive ACK %c", *ack);
-		should_receive = debugserver_client_receive_internal_check(client, ack);
-		debug_info("received char: %c", *ack);
-		if (strncmp(ack, command_prefix, sizeof(char)) == 0) {
-			should_receive = 1;
+		debug_info("attempting to receive ACK (+)");
+		res = debugserver_client_receive_internal_char(client, &data);
+		if (res != DEBUGSERVER_E_SUCCESS) {
+			goto cleanup;
+		}
+		if (data == '+') {
+			debug_info("received ACK (+)");
+		} else if (data == '$') {
+			debug_info("received prefix ($)");
+			buffer[0] = '$';
+			buffer_size = 1;
 			skip_prefix = 1;
-			buffer = malloc(1024);
-			buffer_capacity = 1024;
-			strcpy(buffer, command_prefix);
-			buffer_size += sizeof(char);
-			debug_info("received ACK");
+		} else {
+			debug_info("unrecognized response when looking for ACK: %c", data);
+			goto cleanup;
 		}
 	}
 
-	debug_info("should_receive: %d, skip_prefix: %d", should_receive, skip_prefix);
-
-	if (should_receive && !skip_prefix) {
-		debug_info("attempting to receive prefix");
-		should_receive = debugserver_client_receive_internal_check(client, command_prefix);
-		debug_info("received command_prefix: %c", *command_prefix);
-		if (should_receive) {
-			if (buffer) {
-				strcpy(buffer, command_prefix);
-			} else {
-				buffer = malloc(1024);
-				buffer_capacity = 1024;
-				strcpy(buffer, command_prefix);
-				buffer_size += sizeof(char);
-			}
+	debug_info("skip_prefix: %d", skip_prefix);
+
+	if (!skip_prefix) {
+		debug_info("attempting to receive prefix ($)");
+		res = debugserver_client_receive_internal_char(client, &data);
+		if (res != DEBUGSERVER_E_SUCCESS) {
+			goto cleanup;
+		}
+		if (data == '$') {
+			debug_info("received prefix ($)");
+			buffer[0] = '$';
+			buffer_size = 1;
+		} else {
+			debug_info("unrecognized response when looking for prefix: %c", data);
+			goto cleanup;
 		}
 	}
 
-	debug_info("buffer: %*s, should_receive: %d, skip_prefix: %d", buffer_size, buffer, should_receive, skip_prefix);
+	debug_info("buffer: %*s", buffer_size, buffer);
 
-	if (should_receive) {
-		uint32_t checksum_length = DEBUGSERVER_CHECKSUM_HASH_LENGTH;
-		int receiving_checksum_response = 0;
-		debug_info("attempting to read up response until checksum");
+	uint32_t checksum_length = DEBUGSERVER_CHECKSUM_HASH_LENGTH;
+	int receiving_checksum_response = 0;
+	debug_info("attempting to read up response until checksum");
 
-		while ((checksum_length > 0)) {
-			char data[2] = {'#', '\0'};
-			if (debugserver_client_receive_internal_check(client, data)) {
-				receiving_checksum_response = 1;
-			}
-			if (receiving_checksum_response) {
-				checksum_length--;
-			}
-			if (buffer_size + 1 >= buffer_capacity) {
-				char* newbuffer = realloc(buffer, buffer_capacity+1024);
-				if (!newbuffer) {
-					return DEBUGSERVER_E_UNKNOWN_ERROR;
-				}
-				buffer = newbuffer;
-				buffer[buffer_capacity] = '\0';
-				buffer_capacity += 1024;
-			}
-			strcat(buffer, data);
-			buffer_size += sizeof(char);
+	while ((checksum_length > 0)) {
+		res = debugserver_client_receive_internal_char(client, &data);
+		if (res != DEBUGSERVER_E_SUCCESS) {
+			goto cleanup;
 		}
-		debug_info("validating response checksum...");
-		if (client->noack_mode || debugserver_response_is_checksum_valid(buffer, buffer_size)) {
-			if (response) {
-				/* assemble response string */
-				uint32_t resp_size = sizeof(char) * (buffer_size - DEBUGSERVER_CHECKSUM_HASH_LENGTH - 1);
-				*response = (char*)malloc(resp_size + 1);
-				memcpy(*response, buffer + 1, resp_size);
-				(*response)[resp_size] = '\0';
-				if (response_size) *response_size = resp_size;
-			}
-			if (!client->noack_mode) {
-				/* confirm valid command */
-				debugserver_client_send_ack(client);
-			}
-		} else {
-			/* response was invalid */
-			res = DEBUGSERVER_E_RESPONSE_ERROR;
-			if (!client->noack_mode) {
-				/* report invalid command */
-				debugserver_client_send_noack(client);
+		if (data == '#') {
+			receiving_checksum_response = 1;
+		}
+		if (receiving_checksum_response) {
+			checksum_length--;
+		}
+		if (buffer_size + 1 >= buffer_capacity) {
+			char* newbuffer = realloc(buffer, buffer_capacity+1024);
+			if (!newbuffer) {
+				return DEBUGSERVER_E_UNKNOWN_ERROR;
 			}
+			buffer = newbuffer;
+			buffer_capacity += 1024;
+		}
+		buffer[buffer_size] = data;
+		buffer_size += sizeof(char);
+	}
+	debug_info("validating response checksum...");
+	if (client->noack_mode || debugserver_response_is_checksum_valid(buffer, buffer_size)) {
+		if (response) {
+			/* assemble response string */
+			uint32_t resp_size = sizeof(char) * (buffer_size - DEBUGSERVER_CHECKSUM_HASH_LENGTH - 1);
+			*response = (char*)malloc(resp_size + 1);
+			memcpy(*response, buffer + 1, resp_size);
+			(*response)[resp_size] = '\0';
+			if (response_size) *response_size = resp_size;
+		}
+		if (!client->noack_mode) {
+			/* confirm valid command */
+			debugserver_client_send_ack(client);
+		}
+	} else {
+		/* response was invalid */
+		res = DEBUGSERVER_E_RESPONSE_ERROR;
+		if (!client->noack_mode) {
+			/* report invalid command */
+			debugserver_client_send_noack(client);
 		}
 	}
 
+cleanup:
 	if (response) {
 		debug_info("response: %s", *response);
 	}
@@ -506,9 +502,6 @@ LIBIMOBILEDEVICE_API debugserver_error_t debugserver_client_receive_response(deb
 	if (buffer)
 		free(buffer);
 
-	if (command_prefix)
-		free(command_prefix);
-
 	return res;
 }