Added flash size detection.

include/esp_loader.h

@@ -31,8 +31,10 @@ typedef enum
     ESP_LOADER_SUCCESS,                /*!< Success */
     ESP_LOADER_ERROR_FAIL,             /*!< Unspecified error */
     ESP_LOADER_ERROR_TIMEOUT,          /*!< Timeout elapsed */
+    ESP_LOADER_ERROR_IMAGE_SIZE,       /*!< Image size to flash is larger than flash size */
     ESP_LOADER_ERROR_INVALID_MD5,      /*!< Computed and receied MD5 does not match */
     ESP_LOADER_ERROR_INVALID_TARGET,   /*!< Connected target is invalid */
+    ESP_LOADER_ERROR_UNSUPPORTED_CHIP, /*!< Attached chip is not supported */
     ESP_LOADER_ERROR_INVALID_RESPONSE  /*!< Internal error */
 } esp_loader_error_t;
 

private_include/serial_comm.h

@@ -52,6 +52,7 @@ esp_loader_error_t loader_change_baudrate_cmd(uint32_t baudrate);
 
 esp_loader_error_t loader_md5_cmd(uint32_t address, uint32_t size, uint8_t *md5_out);
 
+esp_loader_error_t loader_spi_parameters(uint32_t total_size);
 
 #ifdef __cplusplus
 }

private_include/serial_comm_prv.h

@@ -181,6 +181,18 @@ typedef struct __attribute__((packed))
     response_status_t status;
 } rom_md5_response_t;
 
+typedef struct __attribute__((packed))
+{
+    command_common_t common;
+    uint32_t id;
+    uint32_t total_size;
+    uint32_t block_size;
+    uint32_t sector_size;
+    uint32_t page_size;
+    uint32_t status_mask;
+} write_spi_command_t;  
+
+
 #ifdef __cplusplus
 }
 #endif

src/esp_loader.c

@@ -20,11 +20,16 @@
 #include "esp_loader.h"
 #include "md5_hash.h"
 #include <string.h>
+#include <assert.h>
 
 #ifndef MAX
 #define MAX(a, b) ((a) > (b)) ? (a) : (b)
 #endif
 
+#ifndef MIN
+#define MIN(a, b) ((a) < (b)) ? (a) : (b)
+#endif
+
 typedef struct {
     uint32_t reg_1;
     uint32_t reg_2;
@@ -51,9 +56,53 @@ static const uint32_t DEFAULT_FLASH_TIMEOUT = 3000;       // timeout for most fl
 static const uint32_t ERASE_REGION_TIMEOUT_PER_MB = 3000; // timeout (per megabyte) for erasing a region
 static const uint8_t  PADDING_PATTERN = 0xFF;
 
-
 static uint32_t s_flash_write_size = 0;
 
+#define MEGABYTE  1024 * 1024
+
+size_t size_id_to_flash_size[] = {
+    MEGABYTE / 4,  // 256KB,
+    MEGABYTE / 2,  // 512KB,
+    1 * MEGABYTE,  // 1MB,
+    2 * MEGABYTE,  // 2MB,
+    4 * MEGABYTE,  // 4MB,
+    8 * MEGABYTE,  // 8MB,
+    16 * MEGABYTE  // 16MB
+};
+
+typedef enum {
+    SPI_FLASH_READ_ID = 0x9F
+} spi_flash_cmd_t;
+
+#if defined TARGET_ESP8266
+#define SPI_REG_BASE    0x60000200
+#define SPI_CMD_REG     SPI_REG_BASE + 0x00
+#define SPI_USR_REG     SPI_REG_BASE + 0x1c
+#define SPI_USR1_REG    SPI_REG_BASE + 0x20
+#define SPI_USR2_REG    SPI_REG_BASE + 0x24
+#define SPI_W0_REG      SPI_REG_BASE + 0x40
+#define SPI_MOSI_DLEN_REG 0
+#define SPI_MISO_DLEN_REG 0
+#elif defined TARGET_ESP32_S2
+#define SPI_REG_BASE    0x3f402000
+#define SPI_CMD_REG     SPI_REG_BASE + 0x00
+#define SPI_USR_REG     SPI_REG_BASE + 0x18
+#define SPI_USR1_REG    SPI_REG_BASE + 0x1c
+#define SPI_USR2_REG    SPI_REG_BASE + 0x20
+#define SPI_W0_REG      SPI_REG_BASE + 0x58
+#define SPI_MOSI_DLEN_REG SPI_REG_BASE + 0x24
+#define SPI_MISO_DLEN_REG SPI_REG_BASE + 0x28
+#elif defined TARGET_ESP32
+#define SPI_REG_BASE    0x60002000
+#define SPI_CMD_REG     SPI_REG_BASE + 0x00
+#define SPI_USR_REG     SPI_REG_BASE + 0x1c
+#define SPI_USR1_REG    SPI_REG_BASE + 0x20
+#define SPI_USR2_REG    SPI_REG_BASE + 0x24
+#define SPI_W0_REG      SPI_REG_BASE + 0x80
+#define SPI_MOSI_DLEN_REG SPI_REG_BASE + 0x28
+#define SPI_MISO_DLEN_REG SPI_REG_BASE + 0x2c
+#endif
+
 #if MD5_ENABLED
 
 static const uint32_t MD5_TIMEOUT_PER_MB = 800;
@@ -138,17 +187,142 @@ esp_loader_error_t esp_loader_connect(esp_loader_connect_args_t *connect_args)
     return err;
 }
 
+#ifndef TARGET_ESP8266
+
+static esp_loader_error_t spi_set_data_lengths(size_t mosi_bits, size_t miso_bits)
+{
+    if (mosi_bits > 0) {
+        RETURN_ON_ERROR( esp_loader_write_register(SPI_MOSI_DLEN_REG, mosi_bits - 1) );
+    }
+    if (miso_bits > 0) {
+        RETURN_ON_ERROR( esp_loader_write_register(SPI_MISO_DLEN_REG, miso_bits - 1) );
+    }
+
+    return ESP_LOADER_SUCCESS;
+}
+
+#else
+
+static esp_loader_error_t spi_set_data_lengths(size_t mosi_bits, size_t miso_bits)
+{
+    uint32_t mosi_bitlen_shift = 17;
+    uint32_t miso_bitlen_shift = 8;
+    uint32_t mosi_mask = (mosi_bits == 0) ? 0 : mosi_bits - 1;
+    uint32_t miso_mask = (miso_bits == 0) ? 0 : miso_bits - 1;
+    uint32_t usr_reg = (miso_mask << miso_bitlen_shift) | (mosi_mask << mosi_bitlen_shift);
+
+    return esp_loader_write_register(SPI_USR1_REG, usr_reg);
+}
+
+#endif
+
+static esp_loader_error_t spi_flash_command(spi_flash_cmd_t cmd, void *data_tx, size_t tx_size, void *data_rx, size_t rx_size)
+{
+    assert(rx_size <= 32); // Reading more than 32 bits back from a SPI flash operation is unsupported
+    assert(tx_size <= 64); // Writing more than 64 bytes of data with one SPI command is unsupported
+
+    uint32_t SPI_USR_CMD  = (1 << 31);
+    uint32_t SPI_USR_MISO = (1 << 28);
+    uint32_t SPI_USR_MOSI = (1 << 27);
+    uint32_t SPI_CMD_USR  = (1 << 18);
+    uint32_t CMD_LEN_SHIFT = 28;
+
+    // Save SPI configuration
+    uint32_t old_spi_usr;
+    uint32_t old_spi_usr2;
+    RETURN_ON_ERROR( esp_loader_read_register(SPI_USR_REG, &old_spi_usr) );
+    RETURN_ON_ERROR( esp_loader_read_register(SPI_USR2_REG, &old_spi_usr2) );
+
+    RETURN_ON_ERROR( spi_set_data_lengths(tx_size, rx_size) );
+
+    uint32_t usr_reg_2 = (7 << CMD_LEN_SHIFT) | cmd;
+    uint32_t usr_reg = SPI_USR_CMD;
+    if (rx_size > 0) {
+        usr_reg |= SPI_USR_MISO;
+    }
+    if (tx_size > 0) {
+        usr_reg |= SPI_USR_MOSI;
+    }
+
+    RETURN_ON_ERROR( esp_loader_write_register(SPI_USR_REG, usr_reg) );
+    RETURN_ON_ERROR( esp_loader_write_register(SPI_USR2_REG, usr_reg_2 ) );
+
+    if (tx_size == 0) {
+        // clear data register before we read it
+        RETURN_ON_ERROR( esp_loader_write_register(SPI_W0_REG, 0) );
+    } else {
+        uint32_t *data = (uint32_t *)data_tx;
+        uint32_t words_to_write = MIN((tx_size + 31) / 8 * 4, 1);
+        uint32_t data_reg_addr = SPI_W0_REG;
+
+        while (words_to_write--) {
+            uint32_t word = *data++;
+            RETURN_ON_ERROR( esp_loader_write_register(data_reg_addr, word) );
+            data_reg_addr += 4;
+        }
+    }
+
+    RETURN_ON_ERROR( esp_loader_write_register(SPI_CMD_REG, SPI_CMD_USR) );
+
+    uint32_t trials = 10;
+    while (trials--) {
+        uint32_t cmd_reg;
+        RETURN_ON_ERROR( esp_loader_read_register(SPI_CMD_REG, &cmd_reg) );
+        if ((cmd_reg & SPI_CMD_USR) == 0) {
+            break;
+        }
+    }
+
+    if (trials == 0) {
+        return ESP_LOADER_ERROR_TIMEOUT;
+    }
+
+    RETURN_ON_ERROR( esp_loader_read_register(SPI_W0_REG, data_rx) );
+
+    // Restore SPI configuration
+    RETURN_ON_ERROR( esp_loader_write_register(SPI_USR_REG, old_spi_usr) );
+    RETURN_ON_ERROR( esp_loader_write_register(SPI_USR2_REG, old_spi_usr2) );
+
+    return ESP_LOADER_SUCCESS;
+}
+
+static esp_loader_error_t detect_flash_size(size_t *flash_size)
+{
+    uint32_t flash_id = 0;
+
+    RETURN_ON_ERROR( spi_flash_command(SPI_FLASH_READ_ID, NULL, 0, &flash_id, 24) );
+    uint32_t size_id = flash_id >> 16;
+
+    if (size_id < 0x12 && size_id > 0x18) {
+        return ESP_LOADER_ERROR_UNSUPPORTED_CHIP;
+    }
+
+    *flash_size = size_id_to_flash_size[size_id - 0x12];
+
+    return ESP_LOADER_SUCCESS;
+}
+
 
 esp_loader_error_t esp_loader_flash_start(uint32_t offset, uint32_t image_size, uint32_t block_size)
 {
     uint32_t blocks_to_write = (image_size + block_size - 1) / block_size;
     uint32_t erase_size = block_size * blocks_to_write;
     s_flash_write_size = block_size;
+    size_t flash_size;
+
+    if (detect_flash_size(&flash_size) == ESP_LOADER_SUCCESS) {
+        if (image_size > flash_size) {
+            return ESP_LOADER_ERROR_IMAGE_SIZE;
+        }
+        loader_port_start_timer(DEFAULT_TIMEOUT);
+        RETURN_ON_ERROR( loader_spi_parameters(flash_size) );
+    } else {
+        loader_port_debug_print("Flash size detection failed, falling back to default");
+    }
 
     init_md5(offset, image_size);
 
     loader_port_start_timer(timeout_per_mb(erase_size, ERASE_REGION_TIMEOUT_PER_MB));
-
     return loader_flash_begin_cmd(offset, erase_size, block_size, blocks_to_write);
 }
 

src/serial_comm.c

@@ -400,6 +400,26 @@ esp_loader_error_t loader_md5_cmd(uint32_t address, uint32_t size, uint8_t *md5_
     return send_cmd_md5(&md5_cmd, sizeof(md5_cmd), md5_out);
 }
 
+esp_loader_error_t loader_spi_parameters(uint32_t total_size)
+{
+    write_spi_command_t spi_cmd = {
+        .common = {
+            .direction = WRITE_DIRECTION,
+            .command = SPI_SET_PARAMS,
+            .size = 24,
+            .checksum = 0
+        },
+        .id = 0,
+        .total_size = total_size,
+        .block_size = 64 * 1024,
+        .sector_size = 4 * 1024,
+        .page_size = 0x100,
+        .status_mask = 0xFFFF,
+    };
+
+    return send_cmd(&spi_cmd, sizeof(spi_cmd), NULL);
+}
+
 __attribute__ ((weak)) void loader_port_debug_print(const char *str)
 {
 

test/serial_io_mock.cpp

@@ -129,6 +129,11 @@ int8_t *write_buffer_data()
     return write_buffer.data();
 }
 
+size_t write_buffer_size()
+{
+    return write_buffer.size();
+}
+
 void set_read_buffer(const void *data, size_t size)
 {
     SLIP_encode((const int8_t *)data, size, read_buffer);
@@ -137,7 +142,7 @@ void set_read_buffer(const void *data, size_t size)
 void print_array(int8_t *data, uint32_t size)
 {
     for (uint32_t i = 0; i < size; i++) {
-        printf("%0x, ", (uint8_t)data[i]);
+        printf("%02x, ", (uint8_t)data[i]);
     }
     printf("\n");
 }

test/serial_io_mock.h

@@ -21,6 +21,7 @@
 void clear_buffers();
 
 void write_buffer_print();
+size_t write_buffer_size();
 int8_t* write_buffer_data();
 
 void set_read_buffer(const void *data, size_t size);

test/test.cpp

@@ -73,6 +73,7 @@ inline void queue_response(expected_response &response, size_t size = sizeof(exp
     set_read_buffer(&response, size);
 }
 
+expected_response set_params_response(SPI_SET_PARAMS);
 expected_response flash_begin_response(FLASH_BEGIN);
 expected_response flash_data_response(FLASH_DATA);
 expected_response flash_end_response(FLASH_END);
@@ -83,6 +84,24 @@ expected_response sync_response(SYNC);
 
 
 struct __attribute__((packed)) flash_start_frame {
+    
+    uint8_t delimiter_1 = 0xc0;
+    write_spi_command_t write_spi_command = {
+        .common = {
+            .direction = WRITE_DIRECTION,
+            .command = SPI_SET_PARAMS,
+            .size = 24,
+            .checksum = 0
+        },
+        .id = 0,
+        .total_size = 0x00400000, // Assume 4MB flash size
+        .block_size = 64 * 1024,
+        .sector_size = 4 * 1024,
+        .page_size = 0x100,
+        .status_mask = 0xFFFF,
+    };
+    uint8_t delimiter_2 = 0xc0;
+
     uint8_t delimiter_3 = 0xc0;
     begin_command_t begin_cmd  = {
         .common = {
@@ -188,6 +207,28 @@ template<size_t PAYLOAD_SIZE>
 uint32_t flash_write_frame<PAYLOAD_SIZE>::seq_num = 0;
 
 
+size_t queue_responses_to_ignore()
+{
+    auto flash_id_response = read_reg_response;
+    flash_id_response.data.common.value = 0x16 << 16; // emulate 4MB flash
+
+    queue_response(read_reg_response);  // Save SPI_USR_REG 
+    queue_response(read_reg_response);  // Save SPI_USR2_REG
+    queue_response(write_reg_response); // SPI_MISO_DLEN_REG
+    queue_response(write_reg_response); // Set new SPI_USR_REG
+    queue_response(write_reg_response); // Set new SPI_USR2_REG
+    queue_response(write_reg_response); // Zero out SPI_W0_REG
+    queue_response(write_reg_response); // SPI_CMD_REG Start transaction
+    queue_response(read_reg_response);  // SPI_CMD_REG Transaction done
+    queue_response(flash_id_response);  // SPI_W0_REG Read data
+    queue_response(write_reg_response); // Restore SPI_USR_REG
+    queue_response(write_reg_response); // Restore SPI_USR2_REG
+ 
+    // Delimiters are added manually to every packet (+2).
+    size_t bytes_to_ignore = (sizeof(write_reg_command_t) + 2) * 7 + (sizeof(read_reg_command_t) + 2) * 4 ;
+
+    return bytes_to_ignore;
+}
 
 TEST_CASE( "Large payload that does not fit BLOCK_SIZE is split into \
             multiple data frames. Last data frame is padded with 0xFF" )
@@ -206,11 +247,14 @@ TEST_CASE( "Large payload that does not fit BLOCK_SIZE is split into \
 
     // Check flash start operation 
     clear_buffers();
+    size_t bytes_to_ignore = queue_responses_to_ignore();
+
+    queue_response(set_params_response);
     queue_response(flash_begin_response);
 
     REQUIRE ( esp_loader_flash_start(0, sizeof(data) * 3, BLOCK_SIZE) == ESP_LOADER_SUCCESS );
-    
-    REQUIRE( memcmp(write_buffer_data(), &expected_start, sizeof(expected_start)) == 0 );
+    // Ignore read/write register commands in this test.
+    REQUIRE( memcmp(write_buffer_data() + bytes_to_ignore, &expected_start, sizeof(expected_start)) == 0 );
 
 
     // Check flash write operation 
@@ -224,8 +268,6 @@ TEST_CASE( "Large payload that does not fit BLOCK_SIZE is split into \
     REQUIRE( esp_loader_flash_write(data, 200) == ESP_LOADER_SUCCESS );
 
     REQUIRE( memcmp(write_buffer_data(), &expected_data, sizeof(expected_data)) == 0 );
-
-    // queue_response(flash_end_response);
 }