Spi memory lock

src/memory/spi_mem.c

@@ -14,6 +14,8 @@
 
 #include "spi_mem.h"
 #include "bitbox02_pins.h"
+#include "random.h"
+#include "screen.h"
 #include "util.h"
 #include <hal_delay.h>
 #include <spi_lite.h>
@@ -24,23 +26,33 @@
 #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_SHIFT 2
+#define SR_PROTECT_BITS (SPI_MEM_PROTECTED_BLOCKS << SR_PROTECT_BITS_SHIFT)
+#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)
 {
     gpio_set_pin_level(PIN_MEM_CS, 0);
 }
 
+// Drives the chip select pin high
 static void _spi_mem_cs_high(void)
 {
     gpio_set_pin_level(PIN_MEM_CS, 1);
 }
 
+// Reads the status register
 static uint8_t _spi_mem_read_sr(void)
 {
     uint8_t buffer[2] = {0};
@@ -51,19 +63,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;
@@ -72,23 +84,67 @@ 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();
+}
+
+// Verify if the given address is protected
+static bool _spi_mem_verify_address_protected(uint32_t address)
+{
+    uint8_t protected_blocks = (_spi_mem_read_sr() & SR_PROTECT_BITS_MASK) >> SR_PROTECT_BITS_SHIFT;
+    if (address < (protected_blocks * SPI_MEM_BLOCK_SIZE)) {
+        return true;
+    }
+    return false;
+}
+
+// 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();
+}
+
+bool spi_mem_full_erase(void)
+{
+    if (_spi_mem_read_sr() & SR_PROTECT_BITS_MASK) {
+        util_log("Cannot erase with protected area locked.");
+        return false;
+    }
+    _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();
+    return true;
 }
 
 bool spi_mem_sector_erase(uint32_t sector_addr)
@@ -97,15 +153,15 @@ bool spi_mem_sector_erase(uint32_t sector_addr)
         util_log("Invalid sector address %p", (void*)(uintptr_t)sector_addr);
         return false;
     }
+    if (_spi_mem_verify_address_protected(sector_addr)) {
+        util_log("Sector address %p protected", (void*)(uintptr_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;
@@ -156,21 +212,23 @@ 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();
+    if (_spi_mem_verify_address_protected(page_addr)) {
+        util_log("Page address %p protected", (void*)(uintptr_t)page_addr);
+        return false;
+    }
+
+    _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;
@@ -194,6 +252,11 @@ bool spi_mem_write(uint32_t address, const uint8_t* input, size_t size)
         return false;
     }
 
+    if (_spi_mem_verify_address_protected(address)) {
+        util_log("Address %p protected", (void*)(uintptr_t)address);
+        return false;
+    }
+
     uint32_t initial_sector_addr = address & SECTOR_MASK;
     uint32_t final_sector_addr = ((address + size - 1) & SECTOR_MASK) + SPI_MEM_SECTOR_SIZE;
     uint16_t sectors = (final_sector_addr - initial_sector_addr) / SPI_MEM_SECTOR_SIZE;
@@ -250,3 +313,52 @@ 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)
+{
+    // Additional assert to simplify debug.
+    ASSERT(_spi_mem_verify_address_protected(address + size));
+    if (!_spi_mem_verify_address_protected(address + size)) {
+        util_log(
+            "Write address %p and size %i outside protected area",
+            (void*)(uintptr_t)address,
+            (int)size);
+        return false;
+    }
+    uint8_t protection = _spi_mem_read_sr() & SR_PROTECT_BITS_MASK;
+    _spi_mem_set_protection(0x0);
+    bool result = spi_mem_write(address, input, size);
+    _spi_mem_set_protection(protection);
+
+    return result;
+}

src/memory/spi_mem.h

@@ -22,15 +22,18 @@
 
 #define SPI_MEM_PAGE_SIZE 0x100 // 256
 #define SPI_MEM_SECTOR_SIZE 0x1000 // 4k
-#define SPI_MEM_BLOCK_SIZE 0x8000 // 32k
+#define SPI_MEM_BLOCK_SIZE 0x10000 // 64k
 #define SPI_MEM_MEMORY_SIZE 0x200000 // 2M
+#define SPI_MEM_PROTECTED_BLOCKS 2 // First 2 blocks - 128KB
 
 /**
  * @brief Erase the entire flash memory chip.
  *
  * This operation is blocking and will take 30-60 seconds.
+ *
+ * @return true on success, false on error.
  */
-void spi_mem_full_erase(void);
+bool spi_mem_full_erase(void);
 
 /**
  * @brief Erase a single SECTOR_SIZE sector from flash memory.
@@ -88,4 +91,37 @@ USE_RESULT bool spi_mem_write(uint32_t address, const uint8_t* input, size_t siz
  */
 USE_RESULT int32_t spi_mem_smart_erase(void);
 
+/**
+ * @brief Enables write protection for the first SPI_MEM_PROTECTED_BLOCKS of the memory.
+ *
+ * Sets the protection bits in the status register to lock the configured protected
+ * memory region, preventing accidental writes or erases.
+ *
+ */
+void spi_mem_protected_area_lock(void);
+
+/**
+ * @brief Disables write protection for the protected memory region.
+ *
+ * Clears the protection bits in the status register, allowing writes and erases
+ * to proceed in the previously locked memory area.
+ *
+ */
+void spi_mem_protected_area_unlock(void);
+
+/**
+ * @brief Temporarily unlocks and writes to a protected flash memory region.
+ *
+ * This function reads the current protection configuration, disables protection,
+ * writes the specified data, and restores the previous protection settings.
+ *
+ * @param[in] address  Start address to write to.
+ * @param[in] input    Pointer to the data to write.
+ * @param[in] size     Number of bytes to write.
+ * @return true if the write operation succeeds, false otherwise.
+ */
+USE_RESULT bool spi_mem_protected_area_write(uint32_t address, const uint8_t* input, size_t size);
+
+void spi_mem_test(void);
+
 #endif // _SPI_MEM_H

src/platform/platform_init.c

@@ -13,6 +13,8 @@
 // limitations under the License.
 
 #include "platform_init.h"
+#include "memory/memory_shared.h"
+#include "memory/spi_mem.h"
 #include <driver_init.h>
 #include <ui/oled/oled.h>
 #if !defined(BOOTLOADER)
@@ -38,4 +40,7 @@ void platform_init(void)
 #if !defined(BOOTLOADER)
     sd_mmc_start();
 #endif
+    if (memory_get_platform() == MEMORY_PLATFORM_BITBOX02_PLUS) {
+        spi_mem_protected_area_lock();
+    }
 }

test/simulator/framework/mock_spi_mem.c

@@ -18,9 +18,10 @@
 
 __extension__ static uint8_t _memory[] = {[0 ... SPI_MEM_MEMORY_SIZE] = 0xFF};
 
-void spi_mem_full_erase(void)
+bool spi_mem_full_erase(void)
 {
     memset(_memory, 0xFF, sizeof(_memory));
+    return true;
 }
 
 bool spi_mem_write(uint32_t address, const uint8_t* input, size_t size)
@@ -38,3 +39,12 @@ uint8_t* spi_mem_read(uint32_t address, size_t size)
     memcpy(result, &_memory[address], size);
     return result;
 }
+
+void spi_mem_protected_area_lock(void) {}
+
+void spi_mem_protected_area_unlock(void) {}
+
+bool spi_mem_protected_area_write(uint32_t address, const uint8_t* input, size_t size)
+{
+    return spi_mem_write(address, input, size);
+}

test/unit-test/framework/mock_spi_mem.c

@@ -18,9 +18,10 @@
 
 __extension__ static uint8_t _memory[] = {[0 ... SPI_MEM_MEMORY_SIZE] = 0xFF};
 
-void spi_mem_full_erase(void)
+bool spi_mem_full_erase(void)
 {
     memset(_memory, 0xFF, sizeof(_memory));
+    return true;
 }
 
 bool spi_mem_write(uint32_t address, const uint8_t* input, size_t size)
@@ -38,3 +39,12 @@ uint8_t* spi_mem_read(uint32_t address, size_t size)
     memcpy(result, &_memory[address], size);
     return result;
 }
+
+void spi_mem_protected_area_lock(void) {}
+
+void spi_mem_protected_area_unlock(void) {}
+
+bool spi_mem_protected_area_write(uint32_t address, const uint8_t* input, size_t size)
+{
+    return spi_mem_write(address, input, size);
+}