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;
 }