spi_flash_write_encrypted: Allow 16-byte aligned block writes

As each 32 byte write has two identical 16 byte AES blocks, it's
possible to write them separately.

components/spi_flash/flash_ops.c

@@ -90,7 +90,7 @@ size_t spi_flash_get_chip_size()
     return g_rom_flashchip.chip_size;
 }
 
-SpiFlashOpResult IRAM_ATTR spi_flash_unlock()
+static SpiFlashOpResult IRAM_ATTR spi_flash_unlock()
 {
     static bool unlocked = false;
     if (!unlocked) {
@@ -260,30 +260,58 @@ out:
 
 esp_err_t IRAM_ATTR spi_flash_write_encrypted(size_t dest_addr, const void *src, size_t size)
 {
-    if ((dest_addr % 32) != 0) {
+    const uint8_t *ssrc = (const uint8_t *)src;
+    if ((dest_addr % 16) != 0) {
         return ESP_ERR_INVALID_ARG;
     }
-    if ((size % 32) != 0) {
+    if ((size % 16) != 0) {
         return ESP_ERR_INVALID_SIZE;
     }
-    if ((uint32_t) src < 0x3ff00000) {
+
-        // if source address is in DROM, we won't be able to read it
-        // from within SPIWrite
-        // TODO: consider buffering source data using heap and writing it anyway?
-        return ESP_ERR_INVALID_ARG;
-    }
     COUNTER_START();
     spi_flash_disable_interrupts_caches_and_other_cpu();
     SpiFlashOpResult rc;
     rc = spi_flash_unlock();
+    spi_flash_enable_interrupts_caches_and_other_cpu();
+
     if (rc == SPI_FLASH_RESULT_OK) {
         /* SPI_Encrypt_Write encrypts data in RAM as it writes,
            so copy to a temporary buffer - 32 bytes at a time.
+
+           Each call to SPI_Encrypt_Write takes a 32 byte "row" of
+           data to encrypt, and each row is two 16 byte AES blocks
+           that share a key (as derived from flash address).
         */
-        uint32_t encrypt_buf[32/sizeof(uint32_t)];
+        uint8_t encrypt_buf[32] __attribute__((aligned(4)));
-        for (size_t i = 0; i < size; i += 32) {
+        uint32_t row_size;
-            memcpy(encrypt_buf, ((const uint8_t *)src) + i, 32);
+        for (size_t i = 0; i < size; i += row_size) {
-            rc = SPI_Encrypt_Write((uint32_t) dest_addr + i, encrypt_buf, 32);
+            uint32_t row_addr = dest_addr + i;
+            if (i == 0 && (row_addr % 32) != 0) {
+                /* writing to second block of a 32 byte row */
+                row_size = 16;
+                row_addr -= 16;
+                /* copy to second block in buffer */
+                memcpy(encrypt_buf + 16, ssrc + i, 16);
+                /* decrypt the first block from flash, will reencrypt to same bytes */
+                spi_flash_read_encrypted(row_addr, encrypt_buf, 16);
+            }
+            else if (size - i == 16) {
+                /* 16 bytes left, is first block of a 32 byte row */
+                row_size = 16;
+                /* copy to first block in buffer */
+                memcpy(encrypt_buf, ssrc + i, 16);
+                /* decrypt the second block from flash, will reencrypt to same bytes */
+                spi_flash_read_encrypted(row_addr + 16, encrypt_buf + 16, 16);
+            }
+            else {
+                /* Writing a full 32 byte row (2 blocks) */
+                row_size = 32;
+                memcpy(encrypt_buf, ssrc + i, 32);
+            }
+
+            spi_flash_disable_interrupts_caches_and_other_cpu();
+            rc = SPI_Encrypt_Write(row_addr, (uint32_t *)encrypt_buf, 32);
+            spi_flash_enable_interrupts_caches_and_other_cpu();
             if (rc != SPI_FLASH_RESULT_OK) {
                 break;
             }

components/spi_flash/include/esp_spi_flash.h

@@ -92,14 +92,16 @@ esp_err_t spi_flash_write(size_t dest_addr, const void *src, size_t size);
  *
  * @note Flash encryption must be enabled for this function to work.
  *
- * @note Address in flash, dest, has to be 32-byte aligned.
+ * @note Destination flash address and length must be 16-byte
+ * aligned. Due to hardware limitations, this function is more
+ * efficient if both these arguments are 32-byte aligned. This is
+ * because the encryption engine natively deals with 32-byte rows of
+ * two AES blocks. Writing half a row (16 bytes) requires reading out
+ * the other 16 bytes and re-encrypting them back to the same value.
  *
- * @note If source address is in DROM, this function will return
+ * @param  dest_addr destination address in Flash. Must be a multiple of 16 bytes.
- *       ESP_ERR_INVALID_ARG.
- *
- * @param  dest_addr destination address in Flash. Must be a multiple of 32 bytes.
  * @param  src       pointer to the source buffer.
- * @param  size      length of data, in bytes. Must be a multiple of 32 bytes.
+ * @param  size      length of data, in bytes. Must be a multiple of 16 bytes.
  *
  * @return esp_err_t
  */

components/spi_flash/test/test_flash_encryption.c

@@ -0,0 +1,91 @@
+#include <stdio.h>
+#include <freertos/FreeRTOS.h>
+#include <freertos/task.h>
+#include <freertos/semphr.h>
+
+#include <unity.h>
+#include <esp_spi_flash.h>
+#include <esp_attr.h>
+#include <esp_flash_encrypt.h>
+
+#include "test_config.h"
+
+static void test_encrypted_write(size_t offset, const uint8_t *data, size_t length);
+static void verify_erased_flash(size_t offset, size_t length);
+
+TEST_CASE("test 16 byte encrypted writes", "[spi_flash]")
+{
+    if (!esp_flash_encryption_enabled()) {
+        TEST_IGNORE_MESSAGE("flash encryption disabled, skipping spi_flash_write_encrypted() tests");
+    }
+
+    TEST_ASSERT_EQUAL_HEX(ESP_OK,
+                      spi_flash_erase_sector(TEST_REGION_START / SPI_FLASH_SEC_SIZE));
+
+    uint8_t fortyeight_bytes[0x30]; // 0, 1, 2, 3, 4... 47
+    for(int i = 0; i < sizeof(fortyeight_bytes); i++) {
+        fortyeight_bytes[i] = i;
+    }
+
+    /* Verify unaligned start or length fails */
+    TEST_ASSERT_EQUAL_HEX(ESP_ERR_INVALID_ARG,
+                      spi_flash_write_encrypted(TEST_REGION_START+1, fortyeight_bytes, 32));
+
+    TEST_ASSERT_EQUAL_HEX(ESP_ERR_INVALID_SIZE,
+                      spi_flash_write_encrypted(TEST_REGION_START, fortyeight_bytes, 15));
+
+    /* ensure nothing happened to the flash yet */
+    verify_erased_flash(TEST_REGION_START, 0x20);
+
+    /* Write 32 byte block, this is the "normal" encrypted write */
+    test_encrypted_write(TEST_REGION_START, fortyeight_bytes, 0x20);
+    verify_erased_flash(TEST_REGION_START + 0x20, 0x20);
+
+    /* Slip in an unaligned spi_flash_read_encrypted() test */
+    uint8_t buf[0x10];
+    spi_flash_read_encrypted(TEST_REGION_START+0x10, buf, 0x10);
+    TEST_ASSERT_EQUAL_HEX8_ARRAY(fortyeight_bytes+0x10, buf, 16);
+
+    /* Write 16 bytes unaligned */
+    test_encrypted_write(TEST_REGION_START + 0x30, fortyeight_bytes, 0x10);
+    /* the 16 byte regions before and after the 16 bytes we just wrote should still be 0xFF */
+    verify_erased_flash(TEST_REGION_START + 0x20, 0x10);
+    verify_erased_flash(TEST_REGION_START + 0x40, 0x10);
+
+    /* Write 48 bytes starting at a 32-byte aligned offset */
+    test_encrypted_write(TEST_REGION_START + 0x40, fortyeight_bytes, 0x30);
+    /* 16 bytes after this write should still be 0xFF -unencrypted- */
+    verify_erased_flash(TEST_REGION_START + 0x70, 0x10);
+
+    /* Write 48 bytes starting at a 16-byte aligned offset */
+    test_encrypted_write(TEST_REGION_START + 0x90, fortyeight_bytes, 0x30);
+    /* 16 bytes after this write should still be 0xFF -unencrypted- */
+    verify_erased_flash(TEST_REGION_START + 0x120, 0x10);
+}
+
+static void test_encrypted_write(size_t offset, const uint8_t *data, size_t length)
+{
+    uint8_t readback[length];
+    printf("encrypt %d bytes at 0x%x\n", length, offset);
+    TEST_ASSERT_EQUAL_HEX(ESP_OK,
+                          spi_flash_write_encrypted(offset, data, length));
+
+    TEST_ASSERT_EQUAL_HEX(ESP_OK,
+                          spi_flash_read_encrypted(offset, readback, length));
+
+    TEST_ASSERT_EQUAL_HEX8_ARRAY(data, readback, length);
+}
+
+static void verify_erased_flash(size_t offset, size_t length)
+{
+    uint8_t readback[length];
+    printf("verify erased 0x%x - 0x%x\n", offset, offset + length);
+    TEST_ASSERT_EQUAL_HEX(ESP_OK,
+                          spi_flash_read(offset, readback, length));
+    for (int i = 0; i < length; i++) {
+        char message[32];
+        sprintf(message, "unerased flash @ 0x%08x", offset + i);
+        TEST_ASSERT_EQUAL_HEX_MESSAGE(0xFF, readback[i], message);
+    }
+}
+