Merge branch 'bugfix/fix_md5_check_with_stub' into 'master'

fix: Add correct MD5 calculation with stub enabled

Closes ESF-144

See merge request espressif/esp-serial-flasher!112

Author: DNedic

private_include/protocol.h

@@ -29,17 +29,18 @@ extern "C" {
 #define READ_DIRECTION  1
 #define WRITE_DIRECTION 0
 
-#define MD5_SIZE 32
+#define MD5_SIZE_ROM  32
+#define MD5_SIZE_STUB 16
 
 // Maximum block sized for RAM and Flash writes, respectively.
 #define ESP_RAM_BLOCK 0x1800
 
 #ifndef MAX
-#define MAX(a, b) ((a) > (b)) ? (a) : (b)
+#define MAX(a, b) (((a) > (b)) ? (a) : (b))
 #endif
 
 #ifndef MIN
-#define MIN(a, b) ((a) < (b)) ? (a) : (b)
+#define MIN(a, b) (((a) < (b)) ? (a) : (b))
 #endif
 
 #ifndef ROUNDUP
@@ -195,10 +196,17 @@ typedef struct __attribute__((packed))
 typedef struct __attribute__((packed))
 {
     common_response_t common;
-    uint8_t md5[MD5_SIZE];     // ROM only
+    uint8_t md5[MD5_SIZE_ROM];
     response_status_t status;
 } rom_md5_response_t;
 
+typedef struct __attribute__((packed))
+{
+    common_response_t common;
+    uint8_t md5[MD5_SIZE_STUB];
+    response_status_t status;
+} stub_md5_response_t;
+
 typedef struct __attribute__((packed))
 {
     command_common_t common;

private_include/protocol_prv.h

@@ -28,4 +28,4 @@ esp_loader_error_t send_cmd(const void *cmd_data, uint32_t size, uint32_t *reg_v
 esp_loader_error_t send_cmd_with_data(const void *cmd_data, size_t cmd_size,
                                       const void *data, size_t data_size);
 
-esp_loader_error_t send_cmd_md5(const void *cmd_data, size_t cmd_size, uint8_t md5_out[MD5_SIZE]);
+esp_loader_error_t send_cmd_md5(const void *cmd_data, size_t cmd_size, uint8_t *md5_out);

src/esp_loader.c

@@ -462,37 +462,40 @@ static void hexify(const uint8_t raw_md5[16], uint8_t hex_md5_out[32])
     }
 }
 
-
 esp_loader_error_t esp_loader_flash_verify(void)
 {
     if (s_target == ESP8266_CHIP && !esp_stub_get_running()) {
         return ESP_LOADER_ERROR_UNSUPPORTED_FUNC;
     }
 
-    uint8_t raw_md5[16] = {0};
-
-    /* Zero termination and new line character require 2 bytes */
-    uint8_t hex_md5[MD5_SIZE + 2] = {0};
-    uint8_t received_md5[MD5_SIZE + 2] = {0};
+    /* Zero termination require 1 byte */
+    uint8_t received_md5[MAX(MD5_SIZE_ROM, MD5_SIZE_STUB) + 1] = {0};
+    uint8_t calculated_md5[MAX(MD5_SIZE_ROM, MD5_SIZE_STUB) + 1] = {0};
 
+    uint8_t raw_md5[16] = {0};
     md5_final(raw_md5);
-    hexify(raw_md5, hex_md5);
 
     loader_port_start_timer(timeout_per_mb(s_image_size, MD5_TIMEOUT_PER_MB));
 
     RETURN_ON_ERROR( loader_md5_cmd(s_start_address, s_image_size, received_md5) );
 
-    bool md5_match = memcmp(hex_md5, received_md5, MD5_SIZE) == 0;
+    bool md5_match;
+    if (esp_stub_get_running()) {
+        md5_match = memcmp(raw_md5, received_md5, MD5_SIZE_STUB) == 0;
+        memcpy(calculated_md5, raw_md5, MD5_SIZE_STUB);
+    } else {
+        hexify(raw_md5, calculated_md5);
+        md5_match = memcmp(calculated_md5, received_md5, MD5_SIZE_ROM) == 0;
+    }
 
     if (!md5_match) {
-        hex_md5[MD5_SIZE] = '\n';
-        received_md5[MD5_SIZE] = '\n';
-
         loader_port_debug_print("Error: MD5 checksum does not match:\n");
         loader_port_debug_print("Expected:\n");
         loader_port_debug_print((char *)received_md5);
+        loader_port_debug_print("\n");
         loader_port_debug_print("Actual:\n");
-        loader_port_debug_print((char *)hex_md5);
+        loader_port_debug_print((char *)calculated_md5);
+        loader_port_debug_print("\n");
 
         return ESP_LOADER_ERROR_INVALID_MD5;
     }

src/protocol_uart.c

@@ -125,23 +125,26 @@ esp_loader_error_t send_cmd_with_data(const void *cmd_data, size_t cmd_size,
 }
 
 
-esp_loader_error_t send_cmd_md5(const void *cmd_data, size_t cmd_size, uint8_t md5_out[MD5_SIZE])
+esp_loader_error_t send_cmd_md5(const void *cmd_data, size_t cmd_size, uint8_t *md5_out)
 {
-    rom_md5_response_t response;
     command_t command = ((const command_common_t *)cmd_data)->command;
 
     RETURN_ON_ERROR( SLIP_send_delimiter() );
     RETURN_ON_ERROR( SLIP_send((const uint8_t *)cmd_data, cmd_size) );
     RETURN_ON_ERROR( SLIP_send_delimiter() );
 
-    RETURN_ON_ERROR( check_response(command, NULL, &response, sizeof(response)) );
-
-    memcpy(md5_out, response.md5, MD5_SIZE);
-
+    if (esp_stub_get_running()) {
+        stub_md5_response_t response;
+        RETURN_ON_ERROR( check_response(command, NULL, &response, sizeof(response)) );
+        memcpy(md5_out, response.md5, MD5_SIZE_STUB);
+    } else {
+        rom_md5_response_t response;
+        RETURN_ON_ERROR( check_response(command, NULL, &response, sizeof(response)) );
+        memcpy(md5_out, response.md5, MD5_SIZE_ROM);
+    }
     return ESP_LOADER_SUCCESS;
 }
 
-
 static esp_loader_error_t check_response(command_t cmd, uint32_t *reg_value, void *resp, uint32_t resp_size)
 {
     esp_loader_error_t err;