memory/spi-mem: add spi memory api calls to protect the memory

This adds api calls to lock/unlock and write a protected area of the
memory.

Author: Beerosagos

src/bootloader/startup.c

@@ -55,7 +55,7 @@ int main(void)
 #ifdef BOOTLOADER_DEVDEVICE
     qtouch_init();
 #endif
-    spi_mem_test();
+    spi_mem_protected_area_lock();
     bootloader_jump();
 
     // If did not jump to firmware code, begin USB processing

src/memory/spi_mem.c

@@ -27,27 +27,36 @@
 #define SECTOR_MASK 0xFFFFF000
 #define MEMORY_LIMIT (SPI_MEM_MEMORY_SIZE - 1)
 #define SR_WIP 0x01
+#define SR_PROTECT_BITS_MASK 0x3C
+#define SR_PROTECT_BITS 0xC // protect the first 2 blocks - 128KB
+#define CR1_TB_BIT_BOTTOM 0x8
+#define CR1_TB_BIT_MASK 0x8
 #define CMD_READ 0x03
 #define CMD_WREN 0x06
+#define CMD_WRSR 0x01
 #define CMD_SE 0x20
 #define CMD_PP 0x02
 #define CMD_RDSR 0x05
+#define CMD_RDCR 0x15
 #define CMD_CE 0x60
 
+// Drives the chip select pin low
 static void _spi_mem_cs_low(void)
 {
 #ifndef TESTING
     gpio_set_pin_level(PIN_MEM_CS, 0);
 #endif
 }
 
+// Drives the chip select pin high
 static void _spi_mem_cs_high(void)
 {
 #ifndef TESTING
     gpio_set_pin_level(PIN_MEM_CS, 1);
 #endif
 }
 
+// Reads the status register
 static uint8_t _spi_mem_read_sr(void)
 {
     uint8_t buffer[2] = {0};
@@ -58,19 +67,19 @@ static uint8_t _spi_mem_read_sr(void)
     return buffer[1];
 }
 
-static void _spi_mem_read(uint32_t address, size_t size, uint8_t* buffer)
+// Reads the configuration register
+static void _spi_mem_read_cr(uint8_t* data_out)
 {
-    buffer[0] = CMD_READ;
-    buffer[1] = (address >> 16) & 0xFF;
-    buffer[2] = (address >> 8) & 0xFF;
-    buffer[3] = address & 0xFF;
-    memset(&buffer[4], 0x00, size);
-
+    uint8_t buffer[3] = {0};
+    buffer[0] = CMD_RDCR;
     _spi_mem_cs_low();
-    SPI_MEM_exchange_block(buffer, size + 4);
+    SPI_MEM_exchange_block(buffer, 3);
     _spi_mem_cs_high();
+
+    memcpy(data_out, &buffer[1], 2);
 }
 
+// Waits until the WIP bits goes low
 static void _spi_mem_wait(void)
 {
     uint8_t status;
@@ -79,20 +88,50 @@ static void _spi_mem_wait(void)
     } while (status & SR_WIP);
 }
 
-void spi_mem_full_erase(void)
+// Set write enable bit
+static void _spi_mem_write_enable(void)
+{
+    uint8_t cmd = CMD_WREN;
+    _spi_mem_cs_low();
+    SPI_MEM_exchange_block(&cmd, 1);
+    _spi_mem_cs_high();
+}
+
+// Write the status and configuration registers
+static void _spi_mem_write_sr(uint8_t* data_in)
 {
-    uint8_t buffer[2];
+    _spi_mem_write_enable();
+    uint8_t buffer[4] = {0};
+    buffer[0] = CMD_WRSR;
+    memcpy(&buffer[1], data_in, 3);
+    _spi_mem_cs_low();
+    SPI_MEM_exchange_block(buffer, 4);
+    _spi_mem_cs_high();
+    _spi_mem_wait();
+}
+
+// Reads `size` bytes starting from `address` and writes the data into `buffer`
+static void _spi_mem_read(uint32_t address, size_t size, uint8_t* buffer)
+{
+    buffer[0] = CMD_READ;
+    buffer[1] = (address >> 16) & 0xFF;
+    buffer[2] = (address >> 8) & 0xFF;
+    buffer[3] = address & 0xFF;
+    memset(&buffer[4], 0x00, size);
 
-    // --- Enable Write ---
-    buffer[0] = CMD_WREN;
     _spi_mem_cs_low();
-    SPI_MEM_exchange_block(buffer, 1);
+    SPI_MEM_exchange_block(buffer, size + 4);
     _spi_mem_cs_high();
+}
+
+void spi_mem_full_erase(void)
+{
+    _spi_mem_write_enable();
 
     // --- Chip Erase ---
-    buffer[0] = CMD_CE;
+    uint8_t cmd = CMD_CE;
     _spi_mem_cs_low();
-    SPI_MEM_exchange_block(buffer, 1);
+    SPI_MEM_exchange_block(&cmd, 1);
     _spi_mem_cs_high();
 
     _spi_mem_wait();
@@ -105,14 +144,10 @@ bool spi_mem_sector_erase(uint32_t sector_addr)
         return false;
     }
 
-    uint8_t buffer[SPI_MEM_PAGE_SIZE + 4];
-    // --- Enable Write ---
-    buffer[0] = CMD_WREN;
-    _spi_mem_cs_low();
-    SPI_MEM_exchange_block(buffer, 1);
-    _spi_mem_cs_high();
+    _spi_mem_write_enable();
 
     // --- Sector Erase (write 4 bytes) ---
+    uint8_t buffer[SPI_MEM_PAGE_SIZE + 4];
     buffer[0] = CMD_SE;
     buffer[1] = (sector_addr >> 16) & 0xFF;
     buffer[2] = (sector_addr >> 8) & 0xFF;
@@ -163,21 +198,18 @@ uint8_t* spi_mem_read(uint32_t address, size_t size)
     return buffer;
 }
 
+// Writes SPI_MEM_PAGE_SIZE bytes from `input` at `page_addr`
 static bool _spi_mem_page_write(uint32_t page_addr, const uint8_t* input)
 {
     if (page_addr % SPI_MEM_PAGE_SIZE != 0) {
         util_log("Invalid page write address %p", (void*)(uintptr_t)page_addr);
         return false;
     }
 
-    uint8_t buffer[SPI_MEM_PAGE_SIZE + 4];
-    // --- Enable Write ---
-    buffer[0] = CMD_WREN;
-    _spi_mem_cs_low();
-    SPI_MEM_exchange_block(buffer, 1);
-    _spi_mem_cs_high();
+    _spi_mem_write_enable();
 
     // --- Page Program (write 4 bytes) ---
+    uint8_t buffer[SPI_MEM_PAGE_SIZE + 4];
     buffer[0] = CMD_PP;
     buffer[1] = (page_addr >> 16) & 0xFF;
     buffer[2] = (page_addr >> 8) & 0xFF;
@@ -258,6 +290,47 @@ int32_t spi_mem_smart_erase(void)
     return erased_sectors;
 }
 
+// Writes the `protection` bits into the status register and sets the
+// Top/Bottom bit to bottom in the configuration register.
+static void _spi_mem_set_protection(uint8_t protection)
+{
+    uint8_t reg[3];
+    reg[0] = _spi_mem_read_sr();
+    _spi_mem_read_cr(&reg[1]);
+
+    // clean and update status register with protection bits
+    reg[0] &= ~SR_PROTECT_BITS_MASK;
+    reg[0] |= protection & SR_PROTECT_BITS_MASK;
+
+    // set the top/bottom protection bit.
+    // This is an OTP bit,so the write will have an effect
+    // only the first time.
+    reg[1] = reg[1] | CR1_TB_BIT_BOTTOM;
+
+    _spi_mem_write_sr(reg);
+}
+
+void spi_mem_protected_area_lock(void)
+{
+    _spi_mem_set_protection(SR_PROTECT_BITS);
+}
+
+void spi_mem_protected_area_unlock(void)
+{
+    _spi_mem_set_protection(0x0);
+}
+
+bool spi_mem_protected_area_write(uint32_t address, const uint8_t* input, size_t size)
+{
+    bool result;
+    uint8_t protection = _spi_mem_read_sr() & SR_PROTECT_BITS_MASK;
+    _spi_mem_set_protection(0x0);
+    result = spi_mem_write(address, input, size);
+    _spi_mem_set_protection(protection);
+
+    return result;
+}
+
 // This is an utility function, useful to test the code, but not to be merged.
 void spi_mem_test(void)
 {
@@ -268,6 +341,7 @@ void spi_mem_test(void)
     // --- Setup test buffers ---
     uint8_t write_data[SPI_MEM_PAGE_SIZE];
     uint8_t read_data[SPI_MEM_PAGE_SIZE];
+    uint32_t sector_addr = 0x00010000; // block 01
 
     for (size_t i = 0; i < SPI_MEM_PAGE_SIZE; i++) {
         write_data[i] = (uint8_t)i;
@@ -324,7 +398,6 @@ void spi_mem_test(void)
     }
 
     // === Test 3: Full sector write and read ===
-    uint32_t sector_addr = 0x00020000;
     uint8_t sector_write_data[SPI_MEM_SECTOR_SIZE];
     uint8_t* sector_read_data;
 
@@ -628,6 +701,148 @@ void spi_mem_test(void)
         util_log("Test 10.6: Zero-size read correctly rejected");
     }
 
+    // === Test 11: Lock/Unlock Protection Behavior ===
+    util_log("Test 11: Lock/Unlock Protection Verification");
+
+    // Backup original SR and CR
+    uint8_t sr_before = _spi_mem_read_sr();
+    uint8_t cr_before[2];
+    _spi_mem_read_cr(cr_before);
+
+    // unlock memory if it was already locked
+    if ((sr_before & SR_PROTECT_BITS_MASK) != 0) {
+        _spi_mem_set_protection(0);
+        sr_before = _spi_mem_read_sr();
+    }
+
+    // 11.1 - Lock memory
+    spi_mem_protected_area_lock();
+    uint8_t sr_locked = _spi_mem_read_sr();
+    uint8_t cr_locked[2];
+    _spi_mem_read_cr(cr_locked);
+
+    // Verify that the protection bits were set
+    if ((sr_locked & SR_PROTECT_BITS_MASK) != SR_PROTECT_BITS) {
+        util_log("Test 11.1: Lock bits not set correctly");
+        success = false;
+    }
+
+    // Verify that the remaining part of the SR register didn't change
+    if ((sr_locked & ~SR_PROTECT_BITS_MASK) != (sr_before & ~SR_PROTECT_BITS_MASK)) {
+        util_log("Test 11.1: Status register corrupted");
+        success = false;
+    }
+
+    // Verify that the TB bit was set
+    if ((cr_locked[0] & CR1_TB_BIT_MASK) != CR1_TB_BIT_BOTTOM) {
+        util_log("Test 11.1: Lock bits not set correctly");
+        success = false;
+    }
+
+    // Verify that the remaining part of the CR register didn't change
+    bool cr1_corrupted = (cr_before[0] & ~CR1_TB_BIT_MASK) != (cr_locked[0] & ~CR1_TB_BIT_MASK);
+    if (cr1_corrupted || (cr_locked[1] != cr_before[1])) {
+        util_log("Test 11.1: CR register changed unexpectedly during lock");
+        success = false;
+    }
+
+    // 11.2 - Verify that writing locked memory doesn't change the data
+    uint8_t* read_page = spi_mem_read(sector_addr, SPI_MEM_PAGE_SIZE);
+    if (!read_page) {
+        success = false;
+    }
+    if (!spi_mem_write(sector_addr, write_data, SPI_MEM_PAGE_SIZE)) {
+        success = false;
+    }
+    uint8_t* read_page2 = spi_mem_read(sector_addr, SPI_MEM_PAGE_SIZE);
+    if (!read_page2) {
+        success = false;
+    }
+    if (memcmp(read_page, read_page2, SPI_MEM_PAGE_SIZE) != 0) {
+        util_log("Test 11.2: Write to protected area incorrectly allowed");
+        success = false;
+    }
+    free(read_page2);
+
+    // 11.3 - Verify that we can still write unprotected sectors
+    uint32_t unprotected_sector = 0x001D0000; // block 29
+    uint8_t* read_unprotected_page = spi_mem_read(unprotected_sector, SPI_MEM_PAGE_SIZE);
+    if (!read_page) {
+        success = false;
+    }
+    if (!spi_mem_write(unprotected_sector, write_data, SPI_MEM_PAGE_SIZE)) {
+        success = false;
+    }
+    read_page2 = spi_mem_read(unprotected_sector, SPI_MEM_PAGE_SIZE);
+    if (!read_page2) {
+        success = false;
+    }
+    if (memcmp(read_unprotected_page, read_page2, SPI_MEM_PAGE_SIZE) == 0) {
+        util_log("Test 11.3: Write to unprotected area failed");
+        success = false;
+    }
+    free(read_page2);
+    free(read_unprotected_page);
+
+    // 11.4 - Verify that writing locked memory with write_locked changes the data
+    if (!spi_mem_protected_area_write(sector_addr, write_data, SPI_MEM_PAGE_SIZE)) {
+        success = false;
+    }
+    read_page2 = spi_mem_read(sector_addr, SPI_MEM_PAGE_SIZE);
+    if (!read_page2) {
+        success = false;
+    }
+    if (memcmp(read_page, read_page2, SPI_MEM_PAGE_SIZE) == 0) {
+        util_log("Test 11.4: Write to protected area failed");
+        success = false;
+    }
+    free(read_page2);
+
+    // 11.5 - Verify that protected area is still locked after write
+    sr_locked = _spi_mem_read_sr();
+    if ((sr_locked & SR_PROTECT_BITS_MASK) != SR_PROTECT_BITS) {
+        util_log("Test 11.5: Lock bits not set correctly");
+        success = false;
+    }
+    // 11.6 - Unlock the memory and write back the initial data, with normal write
+    spi_mem_protected_area_unlock();
+    if (!spi_mem_write(sector_addr, read_page, SPI_MEM_PAGE_SIZE)) {
+        success = false;
+    }
+    read_page2 = spi_mem_read(sector_addr, SPI_MEM_PAGE_SIZE);
+    if (!read_page2) {
+        success = false;
+    }
+    if (memcmp(read_page, read_page2, SPI_MEM_PAGE_SIZE) != 0) {
+        util_log("Test 11.6: Write to unprotected area failed");
+        success = false;
+    }
+    free(read_page2);
+    free(read_page);
+
+    // 11.7 - Verify SR are restored or unchanged in unexpected ways
+    uint8_t sr_final = _spi_mem_read_sr();
+    uint8_t cr_final[2];
+    _spi_mem_read_cr(cr_final);
+
+    if (sr_final != sr_before) {
+        util_log("Test 11.7: Status register corrupted during unlock");
+        success = false;
+    }
+
+    // Verify that the remaining part of the CR register didn't change
+    cr1_corrupted = (cr_before[0] & ~CR1_TB_BIT_MASK) != (cr_final[0] & ~CR1_TB_BIT_MASK);
+    if (cr1_corrupted || (cr_locked[1] != cr_before[1])) {
+        util_log("Test 11.7: CR register changed unexpectedly during unlock");
+        success = false;
+    }
+
+    // Finally clean up the memory
+    if (!spi_mem_smart_erase()) {
+        util_log("Final smart erased failed");
+        success = false;
+    }
+
     util_log("==== spi_mem_test %s ====", success ? "PASSED ✅" : "FAILED ❌");
     screen_sprintf_debug(10000, "Test result: %s", success ? "OK" : "FAILED");
 }