Flash erase + write: complete SPI flash driver

agent/main.c

@@ -183,19 +183,28 @@ static void handle_crc32_cmd(const uint8_t *data, uint32_t len) {
     proto_send(RSP_CRC32, resp, 4);
 }
 
+/* Forward declaration */
+static void handle_flash_write(const uint8_t *data, uint32_t len);
+
 /*
- * CMD_WRITE: receive data from host and write to RAM (or flash later).
+ * CMD_WRITE: receive data and write to RAM or flash.
  *
- * Protocol:
- *   Host sends: CMD_WRITE [addr:4LE] [size:4LE] [expected_crc:4LE]
- *   Agent ACKs ready, then receives RSP_DATA packets.
- *   After all data, verifies CRC32. ACKs OK or CRC_ERROR.
+ * If addr < flash_size, routes to flash write (erase assumed done).
+ * Otherwise, writes to RAM.
  */
 static void handle_write(const uint8_t *data, uint32_t len) {
     if (len < 12) { proto_send_ack(ACK_CRC_ERROR); return; }
 
     uint32_t addr = read_le32(&data[0]);
     uint32_t size = read_le32(&data[4]);
+
+    /* Route flash addresses to flash write handler */
+    if (flash_readable && addr < flash_info.size &&
+        (addr + size) <= flash_info.size) {
+        handle_flash_write(data, len);
+        return;
+    }
+
     uint32_t expected_crc = read_le32(&data[8]);
 
     /* Validate: must be in writable RAM, reasonable size */
@@ -241,6 +250,131 @@ static void handle_write(const uint8_t *data, uint32_t len) {
     proto_send_ack(ACK_OK);
 }
 
+/*
+ * CMD_ERASE: erase flash sectors.
+ *   Host sends: CMD_ERASE [addr:4LE] [size:4LE]
+ *   Agent erases sectors covering the range.
+ *   Sends RSP_DATA with [sectors_done:2LE] after each sector for progress.
+ *   Final ACK_OK when all sectors erased.
+ */
+static void handle_erase(const uint8_t *data, uint32_t len) {
+    if (len < 8) { proto_send_ack(ACK_CRC_ERROR); return; }
+    if (!flash_readable) { proto_send_ack(ACK_FLASH_ERROR); return; }
+
+    uint32_t addr = read_le32(&data[0]);
+    uint32_t size = read_le32(&data[4]);
+    uint32_t sector_sz = flash_info.sector_size;
+
+    /* Validate: must be within flash, sector-aligned */
+    if (size == 0 || addr + size > flash_info.size ||
+        (addr % sector_sz) != 0 || (size % sector_sz) != 0) {
+        proto_send_ack(ACK_FLASH_ERROR);
+        return;
+    }
+
+    uint32_t num_sectors = size / sector_sz;
+
+    /* Send ACK to start erasing */
+    proto_send_ack(ACK_OK);
+
+    for (uint32_t i = 0; i < num_sectors; i++) {
+        flash_erase_sector(addr + i * sector_sz);
+
+        /* Progress: [sectors_done:2LE][debug:3B] */
+        uint8_t progress[5];
+        progress[0] = ((i + 1) >> 0) & 0xFF;
+        progress[1] = ((i + 1) >> 8) & 0xFF;
+        progress[2] = flash_unlock_debug[0];  /* SR before unlock */
+        progress[3] = flash_unlock_debug[1];  /* SR after write_enable */
+        progress[4] = flash_unlock_debug[2];  /* SR after unlock */
+        proto_send(RSP_DATA, progress, 5);
+    }
+
+    proto_send_ack(ACK_OK);
+}
+
+/*
+ * CMD_FLASH_WRITE: receive data and program it to flash.
+ *   Host sends: CMD_WRITE [flash_addr:4LE] [size:4LE] [expected_crc:4LE]
+ *   with flash_addr in flash range (0 to flash_size).
+ *
+ *   Agent receives data into RAM staging area, verifies CRC,
+ *   then programs flash page-by-page. Assumes sectors already erased.
+ *
+ *   Sends ACK_OK after all pages written + verified.
+ */
+static void handle_flash_write(const uint8_t *data, uint32_t len) {
+    if (len < 12) { proto_send_ack(ACK_CRC_ERROR); return; }
+    if (!flash_readable) { proto_send_ack(ACK_FLASH_ERROR); return; }
+
+    uint32_t flash_addr = read_le32(&data[0]);
+    uint32_t size = read_le32(&data[4]);
+    uint32_t expected_crc = read_le32(&data[8]);
+
+    if (size == 0 || size > flash_info.size ||
+        flash_addr + size > flash_info.size) {
+        proto_send_ack(ACK_FLASH_ERROR);
+        return;
+    }
+
+    /* Receive data into RAM staging area */
+    uint8_t *staging = (uint8_t *)(RAM_BASE + 0x200000);
+    proto_send_ack(ACK_OK);
+
+    uint32_t received = 0;
+    uint8_t pkt[MAX_PAYLOAD + 16];
+    while (received < size) {
+        uint32_t pkt_len = 0;
+        uint8_t cmd = proto_recv(pkt, &pkt_len, 10000);
+        if (cmd == RSP_DATA && pkt_len > 2) {
+            uint32_t chunk = pkt_len - 2;
+            for (uint32_t i = 0; i < chunk && received < size; i++)
+                staging[received++] = pkt[2 + i];
+        } else if (cmd == 0) {
+            uint8_t err[5];
+            err[0] = ACK_FLASH_ERROR;
+            write_le32(&err[1], received);
+            proto_send(RSP_ACK, err, 5);
+            return;
+        }
+    }
+
+    /* Verify received data CRC */
+    uint32_t actual_crc = crc32(0, staging, size);
+    if (actual_crc != expected_crc) {
+        uint8_t err[9];
+        err[0] = ACK_CRC_ERROR;
+        write_le32(&err[1], actual_crc);
+        write_le32(&err[5], received);
+        proto_send(RSP_ACK, err, 9);
+        return;
+    }
+
+    /* Program flash page-by-page */
+    uint32_t page_sz = flash_info.page_size;
+    uint32_t offset = 0;
+    while (offset < size) {
+        uint32_t chunk = size - offset;
+        if (chunk > page_sz) chunk = page_sz;
+        flash_write_page(flash_addr + offset, &staging[offset], chunk);
+        offset += chunk;
+    }
+
+    /* Verify written data by reading back from flash and comparing CRC */
+    const uint8_t *flash_ptr = (const uint8_t *)(FLASH_MEM + flash_addr);
+    uint32_t verify_crc = crc32(0, flash_ptr, size);
+    if (verify_crc != expected_crc) {
+        uint8_t err[9];
+        err[0] = ACK_CRC_ERROR;
+        write_le32(&err[1], verify_crc);
+        write_le32(&err[5], size);
+        proto_send(RSP_ACK, err, 9);
+        return;
+    }
+
+    proto_send_ack(ACK_OK);
+}
+
 /*
  * ARM32 position-independent trampoline (machine code).
  * Copies r2 bytes from r1 to r0, then branches to r0-r2 (original dst).
@@ -541,6 +675,9 @@ int main(void) {
             case CMD_WRITE:
                 handle_write(cmd_buf, data_len);
                 break;
+            case CMD_ERASE:
+                handle_erase(cmd_buf, data_len);
+                break;
             case CMD_CRC32:
                 handle_crc32_cmd(cmd_buf, data_len);
                 break;

agent/spi_flash.c

@@ -70,6 +70,48 @@ static uint32_t detect_size(uint8_t id2) {
     return 0x1000000; /* Default 16MB */
 }
 
+/* Forward declarations */
+static void fmc_wait_ready(void);
+static void spi_wait_wip(void);
+
+/* Mode switching: normal mode for register commands, boot mode for reads */
+/* I/O pad configuration base for SPI flash pins */
+#define IO_BASE  0x100C0000
+#define io_reg(off) (*(volatile uint32_t *)(IO_BASE + (off)))
+
+static void fmc_enter_normal(void) {
+    /* Full FMC init for register-mode operations (matching U-Boot) */
+
+    /* Configure SPI flash I/O pads (SPL may have left them in boot-mode config) */
+    io_reg(0x14) = 0x401;  /* sfc_clk */
+    io_reg(0x18) = 0x461;  /* sfc_hold_io0 */
+    io_reg(0x1C) = 0x461;  /* sfc_miso_io1 */
+    io_reg(0x20) = 0x461;  /* sfc_wp_io2 */
+    io_reg(0x24) = 0x461;  /* sfc_mosi_io3 */
+    io_reg(0x28) = 0x461;  /* sfc_csn */
+
+    fmc_reg(FMC_CFG) = 0x1821;  /* OP_MODE_NORMAL | bootrom defaults */
+    fmc_reg(FMC_SPI_TIMING_CFG) = SPI_TIMING_VAL;
+    fmc_reg(FMC_INT_CLR) = 0xFF;
+    fmc_reg(FMC_GLOBAL_CFG) = 0;  /* Disable all write protection */
+}
+
+static void fmc_enter_boot(void) {
+    /* FMC_CFG bit 0 can't be cleared by writing — once in normal mode,
+     * the only way back to boot mode is a soft reset via CRG register.
+     * This resets the FMC controller state, restoring boot mode. */
+    fmc_wait_ready();
+    spi_wait_wip();
+
+    /* Soft reset: set bit 0, then clear it */
+    uint32_t crg = REG_FMC_CRG;
+    REG_FMC_CRG = crg | FMC_SOFT_RESET;
+    /* Brief delay for reset to take effect */
+    for (volatile int i = 0; i < 1000; i++) {}
+    REG_FMC_CRG = crg & ~FMC_SOFT_RESET;
+    for (volatile int i = 0; i < 1000; i++) {}
+}
+
 static void fmc_wait_ready(void) {
     volatile uint32_t timeout = 400000;
     while ((fmc_reg(FMC_OP) & FMC_OP_REG_OP_START) && timeout > 0)
@@ -78,21 +120,69 @@ static void fmc_wait_ready(void) {
 
 static void spi_wait_wip(void) {
     for (int i = 0; i < 10000000; i++) {
-        fmc_reg(FMC_CMD) = SPI_CMD_READ_STATUS;
-        fmc_reg(FMC_OP_CFG) = OP_CFG_OEN_EN | OP_CFG_CS(0);
-        fmc_reg(FMC_OP) = FMC_OP_CMD1_EN | FMC_OP_READ_STATUS | FMC_OP_REG_OP_START;
-        fmc_wait_ready();
-        if (!(fmc_reg(FMC_STATUS) & SPI_STATUS_WIP)) return;
+        uint8_t sr = flash_read_status();
+        if (!(sr & SPI_STATUS_WIP)) return;
     }
 }
 
 static void spi_write_enable(void) {
+    /* Ensure hardware WP is deasserted before write enable.
+     * FMC_GLOBAL_CFG bit 2 = WP_ENABLE, bit 6 = WP_LEVEL.
+     * Clear both to fully disable write protection. */
+    fmc_reg(FMC_GLOBAL_CFG) = 0;
+
     fmc_reg(FMC_CMD) = SPI_CMD_WRITE_ENABLE;
     fmc_reg(FMC_OP_CFG) = OP_CFG_OEN_EN | OP_CFG_CS(0);
     fmc_reg(FMC_OP) = FMC_OP_CMD1_EN | FMC_OP_REG_OP_START;
     fmc_wait_ready();
 }
 
+/*
+ * Clear block protection bits in the flash status register.
+ * HiSilicon SPL locks all sectors (BP0-BP3 = 0x38 in status register).
+ * Must be called before any erase/write operation.
+ */
+#define SPI_CMD_WRITE_STATUS  0x01
+#define SPI_STATUS_BP_MASK    0x7C  /* BP0-BP4: bits 2-6 */
+
+static void flash_unlock(void) {
+    /* Read current status */
+    uint8_t sr_before = flash_read_status();
+
+    if (!(sr_before & SPI_STATUS_BP_MASK)) return;  /* Already unlocked */
+
+    /* Disable write protect in FMC_GLOBAL_CFG (bit 2 = WP_ENABLE).
+     * U-Boot does this before any status register write. */
+    uint32_t gcfg = fmc_reg(FMC_GLOBAL_CFG);
+    fmc_reg(FMC_GLOBAL_CFG) = gcfg & ~(1 << 2);
+
+    /* Write enable required before status register write */
+    spi_write_enable();
+
+    /* Verify WEL is set */
+    uint8_t sr_wel = flash_read_status();
+    (void)sr_wel;  /* Available for debugging */
+
+    /* Write status register = 0x00 (clear all BP bits) */
+    volatile uint8_t *fmc_buf = (volatile uint8_t *)(FLASH_MEM);
+    fmc_buf[0] = 0x00;
+
+    fmc_reg(FMC_CMD) = SPI_CMD_WRITE_STATUS;
+    fmc_reg(FMC_DATA_NUM) = 1;
+    fmc_reg(FMC_OP_CFG) = OP_CFG_OEN_EN | OP_CFG_CS(0);
+    fmc_reg(FMC_OP) = FMC_OP_CMD1_EN | FMC_OP_WRITE_DATA | FMC_OP_REG_OP_START;
+    fmc_wait_ready();
+    spi_wait_wip();
+
+    /* Verify status cleared */
+    uint8_t sr_after = flash_read_status();
+
+    /* Store results for diagnostic (accessible via flash_info) */
+    flash_unlock_debug[0] = sr_before;
+    flash_unlock_debug[1] = sr_wel;
+    flash_unlock_debug[2] = sr_after;
+}
+
 void flash_read_id(uint8_t id[3]) {
     fmc_reg(FMC_CMD) = SPI_CMD_READ_ID;
     fmc_reg(FMC_OP_CFG) = OP_CFG_OEN_EN | OP_CFG_CS(0);
@@ -117,28 +207,16 @@ int flash_init(flash_info_t *info) {
     /* Verify FMC IP version */
     if (fmc_reg(FMC_VERSION) != 0x100) return -1;
 
-    /* Switch to normal mode.
-     * Keep SPI NOR selected. Preserve page/ECC config from bootrom.
-     * FMC_CFG typical bootrom value: 0x1820. We only set bit 0 (normal). */
-    fmc_reg(FMC_CFG) = 0x1821;  /* 0x1820 (bootrom default) | OP_MODE_NORMAL */
-
     /* Set SPI timing parameters */
     fmc_reg(FMC_SPI_TIMING_CFG) = SPI_TIMING_VAL;
 
     /* Clear pending interrupts */
     fmc_reg(FMC_INT_CLR) = 0xFF;
 
     /* Read JEDEC ID (requires normal mode) */
+    fmc_enter_normal();
     flash_read_id(info->jedec_id);
-
-    /* Switch back to boot mode for transparent memory-mapped reads.
-     * In boot mode, FLASH_MEM window reads directly from flash.
-     * In normal mode, it's the FMC I/O buffer. */
-    {
-        uint32_t c = fmc_reg(FMC_CFG);
-        c &= ~FMC_CFG_OP_MODE_NORMAL;
-        fmc_reg(FMC_CFG) = c;
-    }
+    fmc_enter_boot();
 
     info->size = detect_size(info->jedec_id[2]);
     info->sector_size = 0x10000;  /* 64KB */
@@ -153,7 +231,28 @@ void flash_read(uint32_t addr, uint8_t *buf, uint32_t len) {
         buf[i] = flash[addr + i];
 }
 
+uint8_t flash_unlock_debug[3];
+
+/* Read SPI flash status register (must be in normal mode).
+ * Uses READ_DATA path (reads into I/O buffer) instead of READ_STATUS
+ * path (FMC_STATUS register) for more reliable results. */
+uint8_t flash_read_status(void) {
+    fmc_reg(FMC_CMD) = SPI_CMD_READ_STATUS;
+    fmc_reg(FMC_OP_CFG) = OP_CFG_OEN_EN | OP_CFG_CS(0);
+    fmc_reg(FMC_DATA_NUM) = 1;
+    fmc_reg(FMC_OP) = FMC_OP_CMD1_EN | FMC_OP_READ_DATA | FMC_OP_REG_OP_START;
+    fmc_wait_ready();
+    return (uint8_t)(*(volatile uint32_t *)(FLASH_MEM) & 0xFF);
+}
+
 int flash_erase_sector(uint32_t addr) {
+    fmc_enter_normal();
+
+    /* Disable hardware write protect — must be done every time because
+     * fmc_enter_boot (soft reset) restores the GLOBAL_CFG register. */
+    fmc_reg(FMC_GLOBAL_CFG) = fmc_reg(FMC_GLOBAL_CFG) & ~(1 << 2);
+
+    flash_unlock();
     spi_write_enable();
 
     fmc_reg(FMC_CMD) = SPI_CMD_SECTOR_ERASE;
@@ -162,16 +261,25 @@ int flash_erase_sector(uint32_t addr) {
     fmc_reg(FMC_OP) = FMC_OP_CMD1_EN | FMC_OP_ADDR_EN | FMC_OP_REG_OP_START;
     fmc_wait_ready();
 
+    /* Erase takes up to 150ms for 64KB sector — poll WIP */
     spi_wait_wip();
+
+    /* Soft-reset FMC to return to boot mode for memory-mapped reads */
+    fmc_enter_boot();
+
+    /* Re-verify status — soft reset may re-load flash status from hardware */
     return 0;
 }
 
 int flash_write_page(uint32_t addr, const uint8_t *data, uint32_t len) {
     if (len > 256) len = 256;
 
+    fmc_enter_normal();
+    fmc_reg(FMC_GLOBAL_CFG) = fmc_reg(FMC_GLOBAL_CFG) & ~(1 << 2);
+    flash_unlock();
     spi_write_enable();
 
-    /* Copy data to FMC I/O buffer (memory window) */
+    /* Copy data to FMC I/O buffer (memory window in normal mode) */
     volatile uint8_t *fmc_buf = (volatile uint8_t *)(FLASH_MEM);
     for (uint32_t i = 0; i < len; i++)
         fmc_buf[i] = data[i];
@@ -184,6 +292,7 @@ int flash_write_page(uint32_t addr, const uint8_t *data, uint32_t len) {
     fmc_wait_ready();
 
     spi_wait_wip();
+    fmc_enter_boot();
     return 0;
 }