spi_flash: refactoring flash encryption into new api

components/hal/CMakeLists.txt

@@ -25,6 +25,7 @@ if(NOT BOOTLOADER_BUILD)
                 "uart_hal_iram.c"
                 "spi_flash_hal.c"
                 "spi_flash_hal_iram.c"
+                "spi_flash_encrypt_hal_iram.c"
                 "soc_hal.c"
                 "interrupt_controller_hal.c"
                 "sha_hal.c"

components/hal/esp32/include/hal/spi_flash_encrypted_ll.h

@@ -0,0 +1,119 @@
+// Copyright 2015-2021 Espressif Systems (Shanghai) PTE LTD
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/*******************************************************************************
+ * NOTICE
+ * The ll is not public api, don't use in application code.
+ * See readme.md in hal/include/hal/readme.md
+ ******************************************************************************/
+
+// The Lowlevel layer for SPI Flash Encryption.
+
+#include "soc/dport_reg.h"
+#include "soc/flash_encryption_reg.h"
+#include <stdbool.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * Enable encryption in flash
+ */
+static inline void spi_flash_encrypt_ll_enable(void)
+{
+    DPORT_REG_SET_BIT(DPORT_SLAVE_SPI_CONFIG_REG, DPORT_SLAVE_SPI_MASK_PRO | DPORT_SPI_ENCRYPT_ENABLE);
+}
+
+/**
+ * Disable encryption in flash
+ */
+static inline void spi_flash_encrypt_ll_disable(void)
+{
+    DPORT_REG_CLR_BIT(DPORT_SLAVE_SPI_CONFIG_REG, DPORT_SLAVE_SPI_MASK_PRO | DPORT_SPI_ENCRYPT_ENABLE);
+}
+
+/**
+ * Copy the flash address to physical address
+ *
+ * @param flash_addr The flash address.
+ *
+ * @note the address must be 8-byte aligned
+ */
+static inline void spi_flash_encrypt_ll_address_save(uint32_t flash_addr)
+{
+    REG_WRITE(FLASH_ENCRYPTION_ADDRESS_REG, flash_addr);
+}
+
+/**
+ * Wait for flash encryption operation completeness.
+ */
+static inline void spi_flash_encrypt_ll_calculate_wait_idle(void)
+{
+    while(!(REG_READ(FLASH_ENCRYPTION_DONE_REG) & BIT(0))) {
+    }
+}
+
+/**
+ * Start encryption on data buffer.
+ */
+static inline void spi_flash_encrypt_ll_calculate_start(void)
+{
+    REG_WRITE(FLASH_ENCRYPTION_START_REG, BIT(0));
+}
+
+/**
+ * Save the plaintext for encryption
+ *
+ * @param address address of the partition to be written.
+ * @param buffer Buffer to store the input data.
+ * @param size Buffer size.
+ */
+static inline void spi_flash_encrypt_ll_plaintext_save(uint32_t address, const uint32_t* buffer, uint32_t size)
+{
+    for (int i = 0; i < 8; i++) {
+        REG_WRITE(FLASH_ENCRYPTION_BUFFER_REG + (i << 2), buffer[i]);
+    }
+}
+
+/**
+ * Finish the flash encryption and make encrypted result accessible to SPI.
+ */
+static inline void spi_flash_encrypt_ll_done(void)
+{
+    // Do nothing on ESP32
+}
+
+/**
+ * Set to destroy encrypted result
+ */
+static inline void spi_flash_encrypt_ll_destroy(void)
+{
+    // Do nothing on ESP32
+}
+
+/**
+ * Check if is qualified to encrypt the buffer
+ *
+ * @param address the address of written flash partition.
+ * @param length Buffer size.
+ */
+static inline bool spi_flash_encrypt_ll_check(uint32_t address, uint32_t length)
+{
+    return ((address % 16) == 0) ? true : false;
+}
+
+#ifdef __cplusplus
+}
+#endif

components/hal/esp32c3/include/hal/spi_flash_encrypted_ll.h

@@ -0,0 +1,157 @@
+// Copyright 2015-2021 Espressif Systems (Shanghai) PTE LTD
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/*******************************************************************************
+ * NOTICE
+ * The ll is not public api, don't use in application code.
+ * See readme.md in hal/include/hal/readme.md
+ ******************************************************************************/
+
+// The Lowlevel layer for SPI Flash Encryption.
+
+#include "soc/system_reg.h"
+#include "soc/hwcrypto_reg.h"
+#include "soc/soc.h"
+#include "string.h"
+#include "assert.h"
+#include <stdbool.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/// Choose type of chip you want to encrypt manully
+typedef enum
+{
+    FLASH_ENCRYPTION_MANU = 0, ///!< Manually encrypt the flash chip.
+    PSRAM_ENCRYPTION_MANU = 1  ///!< Manually encrypt the psram chip.
+} flash_encrypt_ll_type_t;
+
+/**
+ * Enable the flash encryption function under spi boot mode and download boot mode.
+ */
+static inline void spi_flash_encrypt_ll_enable(void)
+{
+    REG_SET_BIT(SYSTEM_EXTERNAL_DEVICE_ENCRYPT_DECRYPT_CONTROL_REG,
+                SYSTEM_ENABLE_DOWNLOAD_MANUAL_ENCRYPT |
+                SYSTEM_ENABLE_SPI_MANUAL_ENCRYPT);
+}
+
+/*
+ * Disable the flash encryption mode.
+ */
+static inline void spi_flash_encrypt_ll_disable(void)
+{
+    REG_CLR_BIT(SYSTEM_EXTERNAL_DEVICE_ENCRYPT_DECRYPT_CONTROL_REG,
+                SYSTEM_ENABLE_SPI_MANUAL_ENCRYPT);
+}
+
+/**
+ * Choose type of chip you want to encrypt manully
+ *
+ * @param type The type of chip to be encrypted
+ *
+ * @note The hardware currently support flash encryption.
+ */
+static inline void spi_flash_encrypt_ll_type(flash_encrypt_ll_type_t type)
+{
+    // Our hardware only support flash encryption
+    assert(type == FLASH_ENCRYPTION_MANU);
+    REG_WRITE(AES_XTS_DESTINATION_REG, type);
+}
+
+/**
+ * Configure the data size of a single encryption.
+ *
+ * @param block_size Size of the desired block.
+ */
+static inline void spi_flash_encrypt_ll_buffer_length(uint32_t size)
+{
+    // Desired block should not be larger than the block size.
+    REG_WRITE(AES_XTS_SIZE_REG, size >> 5);
+}
+
+/**
+ * Save 32-bit piece of plaintext.
+ *
+ * @param address the address of written flash partition.
+ * @param buffer Buffer to store the input data.
+ * @param size Buffer size.
+ *
+ */
+static inline void spi_flash_encrypt_ll_plaintext_save(uint32_t address, const uint32_t* buffer, uint32_t size)
+{
+    uint32_t plaintext_offs = (address % 64);
+    memcpy((void *)(AES_XTS_PLAIN_BASE + plaintext_offs), buffer, size);
+}
+
+/**
+ * Copy the flash address to XTS_AES physical address
+ *
+ * @param flash_addr flash address to write.
+ */
+static inline void spi_flash_encrypt_ll_address_save(uint32_t flash_addr)
+{
+    REG_WRITE(AES_XTS_PHYSICAL_ADDR_REG, flash_addr);
+}
+
+/**
+ * Start flash encryption
+ */
+static inline void spi_flash_encrypt_ll_calculate_start(void)
+{
+    REG_WRITE(AES_XTS_TRIGGER_REG, 1);
+}
+
+/**
+ * Wait for flash encryption termination
+ */
+static inline void spi_flash_encrypt_ll_calculate_wait_idle(void)
+{
+    while(REG_READ(AES_XTS_STATE_REG) == 0x1) {
+    }
+}
+
+/**
+ * Finish the flash encryption and make encrypted result accessible to SPI.
+ */
+static inline void spi_flash_encrypt_ll_done(void)
+{
+    REG_WRITE(AES_XTS_RELEASE_REG, 1);
+    while(REG_READ(AES_XTS_STATE_REG) != 0x3) {
+    }
+}
+
+/**
+ * Set to destroy encrypted result
+ */
+static inline void spi_flash_encrypt_ll_destroy(void)
+{
+    REG_WRITE(AES_XTS_DESTROY_REG, 1);
+}
+
+/**
+ * Check if is qualified to encrypt the buffer
+ *
+ * @param address the address of written flash partition.
+ * @param length Buffer size.
+ */
+static inline bool spi_flash_encrypt_ll_check(uint32_t address, uint32_t length)
+{
+    return ((address % length) == 0) ? true : false;
+}
+
+#ifdef __cplusplus
+}
+#endif

components/hal/esp32s2/include/hal/spi_flash_encrypted_ll.h

@@ -0,0 +1,166 @@
+// Copyright 2015-2021 Espressif Systems (Shanghai) PTE LTD
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/*******************************************************************************
+ * NOTICE
+ * The ll is not public api, don't use in application code.
+ * See readme.md in hal/include/hal/readme.md
+ ******************************************************************************/
+
+// The Lowlevel layer for SPI Flash Encryption.
+
+#include "soc/system_reg.h"
+#include "soc/hwcrypto_reg.h"
+#include "soc/soc.h"
+#include "string.h"
+#include "assert.h"
+#include <stdbool.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/// Choose type of chip you want to encrypt manully
+typedef enum
+{
+    FLASH_ENCRYPTION_MANU = 0, ///!< Manually encrypt the flash chip.
+    PSRAM_ENCRYPTION_MANU = 1  ///!< Manually encrypt the psram chip.
+} flash_encrypt_ll_type_t;
+
+/**
+ * Enable the flash encryption function under spi boot mode and download boot mode.
+ */
+static inline void spi_flash_encrypt_ll_enable(void)
+{
+    REG_SET_BIT(DPORT_EXTERNAL_DEVICE_ENCRYPT_DECRYPT_CONTROL_REG,
+                DPORT_ENABLE_DOWNLOAD_MANUAL_ENCRYPT |
+                DPORT_ENABLE_SPI_MANUAL_ENCRYPT);
+}
+
+/**
+ * Enable the AES accelerator.
+ * Also clear reset on digital signature unit, otherwise AES is held in resetop.
+ */
+static inline void spi_flash_encrypt_ll_aes_accelerator_enable(void)
+{
+    REG_SET_BIT(DPORT_CPU_PERIP_CLK_EN1_REG, DPORT_CRYPTO_AES_CLK_EN);
+    REG_CLR_BIT(DPORT_CPU_PERIP_RST_EN1_REG, DPORT_CRYPTO_AES_RST | DPORT_CRYPTO_DS_RST);
+}
+
+/*
+ * Disable the flash encryption mode.
+ */
+static inline void spi_flash_encrypt_ll_disable(void)
+{
+    REG_CLR_BIT(DPORT_EXTERNAL_DEVICE_ENCRYPT_DECRYPT_CONTROL_REG,
+                DPORT_ENABLE_SPI_MANUAL_ENCRYPT);
+}
+
+/**
+ * Choose type of chip you want to encrypt manully
+ *
+ * @param type The type of chip to be encrypted
+ *
+ * @note The hardware currently support flash encryption.
+ */
+static inline void spi_flash_encrypt_ll_type(flash_encrypt_ll_type_t type)
+{
+    // Our hardware only support flash encryption
+    assert(type == FLASH_ENCRYPTION_MANU);
+    REG_WRITE(AES_XTS_DESTINATION_REG, type);
+}
+
+/**
+ * Configure the data size of a single encryption.
+ *
+ * @param block_size Size of the desired block.
+ */
+static inline void spi_flash_encrypt_ll_buffer_length(uint32_t size)
+{
+    // Desired block should not be larger than the block size.
+    REG_WRITE(AES_XTS_SIZE_REG, size >> 5);
+}
+
+/**
+ * Save 32-bit piece of plaintext.
+ *
+ * @param address the address of written flash partition.
+ * @param buffer Buffer to store the input data.
+ * @param size Buffer size.
+ */
+static inline void spi_flash_encrypt_ll_plaintext_save(uint32_t address, const uint32_t* buffer, uint32_t size)
+{
+    uint32_t plaintext_offs = (address % 64);
+    memcpy((void *)(AES_XTS_PLAIN_BASE + plaintext_offs), buffer, size);
+}
+
+/**
+ * Copy the flash address to XTS_AES physical address
+ *
+ * @param flash_addr flash address to write.
+ */
+static inline void spi_flash_encrypt_ll_address_save(uint32_t flash_addr)
+{
+    REG_WRITE(AES_XTS_PHYSICAL_ADDR_REG, flash_addr);
+}
+
+/**
+ * Start flash encryption
+ */
+static inline void spi_flash_encrypt_ll_calculate_start(void)
+{
+    REG_WRITE(AES_XTS_TRIGGER_REG, 1);
+}
+
+/**
+ * Wait for flash encryption termination
+ */
+static inline void spi_flash_encrypt_ll_calculate_wait_idle(void)
+{
+    while(REG_READ(AES_XTS_STATE_REG) == 0x1) {
+    }
+}
+
+/**
+ * Finish the flash encryption and make encrypted result accessible to SPI.
+ */
+static inline void spi_flash_encrypt_ll_done(void)
+{
+    REG_WRITE(AES_XTS_RELEASE_REG, 1);
+    while(REG_READ(AES_XTS_STATE_REG) != 0x3) {
+    }
+}
+
+/**
+ * Set to destroy encrypted result
+ */
+static inline void spi_flash_encrypt_ll_destroy(void)
+{
+    REG_WRITE(AES_XTS_DESTROY_REG, 1);
+}
+
+/**
+ * Check if is qualified to encrypt the buffer
+ *
+ * @param address the address of written flash partition.
+ * @param length Buffer size.
+ */
+static inline bool spi_flash_encrypt_ll_check(uint32_t address, uint32_t length)
+{
+    return ((address % length) == 0) ? true : false;
+}
+
+#ifdef __cplusplus
+}
+#endif

components/hal/esp32s3/include/hal/spi_flash_encrypted_ll.h

@@ -0,0 +1,157 @@
+// Copyright 2015-2021 Espressif Systems (Shanghai) PTE LTD
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/*******************************************************************************
+ * NOTICE
+ * The ll is not public api, don't use in application code.
+ * See readme.md in hal/include/hal/readme.md
+ ******************************************************************************/
+
+// The Lowlevel layer for SPI Flash Encryption.
+
+#include "soc/system_reg.h"
+#include "soc/hwcrypto_reg.h"
+#include "soc/soc.h"
+#include "string.h"
+#include "assert.h"
+#include <stdbool.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/// Choose type of chip you want to encrypt manully
+typedef enum
+{
+    FLASH_ENCRYPTION_MANU = 0, ///!< Manually encrypt the flash chip.
+    PSRAM_ENCRYPTION_MANU = 1  ///!< Manually encrypt the psram chip.
+} flash_encrypt_ll_type_t;
+
+/**
+ * Enable the flash encryption function under spi boot mode and download boot mode.
+ */
+static inline void spi_flash_encrypt_ll_enable(void)
+{
+    REG_SET_BIT(SYSTEM_EXTERNAL_DEVICE_ENCRYPT_DECRYPT_CONTROL_REG,
+                SYSTEM_ENABLE_DOWNLOAD_MANUAL_ENCRYPT |
+                SYSTEM_ENABLE_SPI_MANUAL_ENCRYPT);
+}
+
+/*
+ * Disable the flash encryption mode.
+ */
+static inline void spi_flash_encrypt_ll_disable(void)
+{
+    REG_CLR_BIT(SYSTEM_EXTERNAL_DEVICE_ENCRYPT_DECRYPT_CONTROL_REG,
+                SYSTEM_ENABLE_SPI_MANUAL_ENCRYPT);
+}
+
+/**
+ * Choose type of chip you want to encrypt manully
+ *
+ * @param type The type of chip to be encrypted
+ *
+ * @note The hardware currently support flash encryption.
+ */
+static inline void spi_flash_encrypt_ll_type(flash_encrypt_ll_type_t type)
+{
+    // Our hardware only support flash encryption
+    assert(type == FLASH_ENCRYPTION_MANU);
+    REG_WRITE(AES_XTS_DESTINATION_REG, type);
+}
+
+/**
+ * Configure the data size of a single encryption.
+ *
+ * @param block_size Size of the desired block.
+ */
+static inline void spi_flash_encrypt_ll_buffer_length(uint32_t size)
+{
+    // Desired block should not be larger than the block size.
+    REG_WRITE(AES_XTS_SIZE_REG, size >> 5);
+}
+
+/**
+ * Save 32-bit piece of plaintext.
+ *
+ * @param address the address of written flash partition.
+ * @param buffer Buffer to store the input data.
+ * @param size Buffer size.
+ *
+ */
+static inline void spi_flash_encrypt_ll_plaintext_save(uint32_t address, const uint32_t* buffer, uint32_t size)
+{
+    uint32_t plaintext_offs = (address % 64);
+    memcpy((void *)(AES_XTS_PLAIN_BASE + plaintext_offs), buffer, size);
+}
+
+/**
+ * Copy the flash address to XTS_AES physical address
+ *
+ * @param flash_addr flash address to write.
+ */
+static inline void spi_flash_encrypt_ll_address_save(uint32_t flash_addr)
+{
+    REG_WRITE(AES_XTS_PHYSICAL_ADDR_REG, flash_addr);
+}
+
+/**
+ * Start flash encryption
+ */
+static inline void spi_flash_encrypt_ll_calculate_start(void)
+{
+    REG_WRITE(AES_XTS_TRIGGER_REG, 1);
+}
+
+/**
+ * Wait for flash encryption termination
+ */
+static inline void spi_flash_encrypt_ll_calculate_wait_idle(void)
+{
+    while(REG_READ(AES_XTS_STATE_REG) == 0x1) {
+    }
+}
+
+/**
+ * Finish the flash encryption and make encrypted result accessible to SPI.
+ */
+static inline void spi_flash_encrypt_ll_done(void)
+{
+    REG_WRITE(AES_XTS_RELEASE_REG, 1);
+    while(REG_READ(AES_XTS_STATE_REG) != 0x3) {
+    }
+}
+
+/**
+ * Set to destroy encrypted result
+ */
+static inline void spi_flash_encrypt_ll_destroy(void)
+{
+    REG_WRITE(AES_XTS_DESTROY_REG, 1);
+}
+
+/**
+ * Check if is qualified to encrypt the buffer
+ *
+ * @param address the address of written flash partition.
+ * @param length Buffer size.
+ */
+static inline bool spi_flash_encrypt_ll_check(uint32_t address, uint32_t length)
+{
+    return ((address % length) == 0) ? true : false;
+}
+
+#ifdef __cplusplus
+}
+#endif

components/hal/include/hal/spi_flash_encrypt_hal.h

@@ -0,0 +1,62 @@
+// Copyright 2021 Espressif Systems (Shanghai) PTE LTD
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/*******************************************************************************
+ * NOTICE
+ * The HAL is not public api, don't use in application code.
+ * See readme.md in hal/include/hal/readme.md
+ ******************************************************************************/
+
+// The HAL layer for SPI Flash Encryption
+
+#include "hal/spi_flash_encrypted_ll.h"
+
+/**
+ * @brief Enable the flash encryption
+ */
+void spi_flash_encryption_hal_enable(void);
+
+/**
+ * @brief Disable the flash encryption
+ */
+void spi_flash_encryption_hal_disable(void);
+
+/**
+ * Prepare flash encryption before operation.
+ *
+ * @param address The destination address in flash for the write operation.
+ * @param buffer Data for programming
+ * @param size Size to program.
+ *
+ * @note address and buffer must be 8-word aligned.
+ */
+void spi_flash_encryption_hal_prepare(uint32_t address, const uint32_t* buffer, uint32_t size);
+
+/**
+ * @brief flash data encryption operation is done.
+ */
+void spi_flash_encryption_hal_done(void);
+
+/**
+ * Destroy encrypted result
+ */
+void spi_flash_encryption_hal_destroy(void);
+
+/**
+ * Check if is qualified to encrypt the buffer
+ *
+ * @param address the address of written flash partition.
+ * @param length Buffer size.
+ */
+bool spi_flash_encryption_hal_check(uint32_t address, uint32_t length);

components/hal/include/hal/spi_flash_types.h

@@ -86,6 +86,44 @@ typedef struct {
     };
 } spi_flash_sus_cmd_conf;
 
+/// Structure for flash encryption operations.
+typedef struct
+{
+    /**
+     * @brief Enable the flash encryption
+    */
+    void (*flash_encryption_enable)(void);
+    /**
+     * @brief Disable the flash encryption
+    */
+    void (*flash_encryption_disable)(void);
+    /**
+     * Prepare flash encryption before operation.
+     *
+     * @param address The destination address in flash for the write operation.
+     * @param buffer Data for programming
+     * @param size Size to program.
+     *
+     * @note address and buffer must be 8-word aligned.
+     */
+    void (*flash_encryption_data_prepare)(uint32_t address, const uint32_t* buffer, uint32_t size);
+    /**
+     * @brief flash data encryption operation is done.
+     */
+    void (*flash_encryption_done)(void);
+    /**
+     * Destroy encrypted result
+    */
+    void (*flash_encryption_destroy)(void);
+    /**
+     * Check if is qualified to encrypt the buffer
+     *
+     * @param address the address of written flash partition.
+     * @param length Buffer size.
+     */
+    bool (*flash_encryption_check)(uint32_t address, uint32_t length);
+} spi_flash_encryption_t;
+
 ///Slowest io mode supported by ESP32, currently SlowRd
 #define SPI_FLASH_READ_MODE_MIN SPI_FLASH_SLOWRD
 

components/hal/linker.lf

@@ -8,6 +8,7 @@ entries:
     else:
         uart_hal_iram (default)
     spi_flash_hal_iram (noflash)
+    spi_flash_encrypt_hal_iram (noflash)
     ledc_hal_iram (noflash)
     i2c_hal_iram (noflash)
     cpu_hal (noflash)

components/hal/spi_flash_encrypt_hal_iram.c

@@ -0,0 +1,59 @@
+// Copyright 2021 Espressif Systems (Shanghai) PTE LTD
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// This part is put in iram.
+
+#include "hal/spi_flash_encrypted_ll.h"
+
+void spi_flash_encryption_hal_enable(void)
+{
+    spi_flash_encrypt_ll_enable();
+#if CONFIG_IDF_TARGET_ESP32S2
+    spi_flash_encrypt_ll_aes_accelerator_enable();
+#endif //CONFIG_IDF_TARGET_ESP32S2
+#if !CONFIG_IDF_TARGET_ESP32
+    spi_flash_encrypt_ll_type(FLASH_ENCRYPTION_MANU);
+#endif // !CONFIG_IDF_TARGET_ESP32
+}
+
+void spi_flash_encryption_hal_disable(void)
+{
+    spi_flash_encrypt_ll_disable();
+}
+
+void spi_flash_encryption_hal_prepare(uint32_t address, const uint32_t* buffer, uint32_t size)
+{
+#if !CONFIG_IDF_TARGET_ESP32
+    spi_flash_encrypt_ll_buffer_length(size);
+#endif // !CONFIG_IDF_TARGET_ESP32
+    spi_flash_encrypt_ll_address_save(address);
+    spi_flash_encrypt_ll_plaintext_save(address, buffer, size);
+    spi_flash_encrypt_ll_calculate_start();
+}
+
+void spi_flash_encryption_hal_done(void)
+{
+    spi_flash_encrypt_ll_calculate_wait_idle();
+    spi_flash_encrypt_ll_done();
+}
+
+void spi_flash_encryption_hal_destroy(void)
+{
+    spi_flash_encrypt_ll_destroy();
+}
+
+bool spi_flash_encryption_hal_check(uint32_t address, uint32_t length)
+{
+    return spi_flash_encrypt_ll_check(address, length);
+}

components/soc/esp32/include/soc/flash_encryption_reg.h

@@ -0,0 +1,20 @@
+// Copyright 2021 Espressif Systems (Shanghai) PTE LTD
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "soc.h"
+
+#define FLASH_ENCRYPTION_BUFFER_REG   (PERIPHS_SPI_ENCRYPT_BASEADDR)
+#define FLASH_ENCRYPTION_START_REG    (PERIPHS_SPI_ENCRYPT_BASEADDR + 0x20)
+#define FLASH_ENCRYPTION_ADDRESS_REG  (PERIPHS_SPI_ENCRYPT_BASEADDR + 0x24)
+#define FLASH_ENCRYPTION_DONE_REG     (PERIPHS_SPI_ENCRYPT_BASEADDR + 0x28)

components/soc/esp32/include/soc/soc_caps.h

@@ -264,6 +264,9 @@
 #define SOC_AES_SUPPORT_AES_192 (1)
 #define SOC_AES_SUPPORT_AES_256 (1)
 
+/*-------------------------- Flash Encryption CAPS----------------------------*/
+#define SOC_FLASH_ENCRYPTED_XTS_AES_BLOCK_MAX   (32)
+
 /*--------------- PHY REGISTER AND MEMORY SIZE CAPS --------------------------*/
 #define SOC_PHY_DIG_REGS_MEM_SIZE       (21*4)
 

components/soc/esp32c3/include/soc/soc_caps.h

@@ -97,6 +97,9 @@
 /*-------------------------- AES CAPS -----------------------------------------*/
 #define SOC_AES_SUPPORT_DMA     (1)
 
+/*-------------------------- Flash Encryption CAPS----------------------------*/
+#define SOC_FLASH_ENCRYPTED_XTS_AES_BLOCK_MAX   (32)
+
 /* Has a centralized DMA, which is shared with all peripherals */
 #define SOC_AES_GDMA            (1)
 

components/soc/esp32s2/include/soc/soc_caps.h

@@ -290,6 +290,9 @@
 #define SOC_AES_SUPPORT_DMA     (1)
 #define SOC_AES_SUPPORT_GCM     (1)
 
+/*-------------------------- Flash Encryption CAPS----------------------------*/
+#define SOC_FLASH_ENCRYPTED_XTS_AES_BLOCK_MAX   (64)
+
 /* Has "crypto DMA", which is shared with SHA */
 #define SOC_AES_CRYPTO_DMA      (1)
 

components/soc/esp32s3/include/soc/soc_caps.h

@@ -182,6 +182,9 @@
 #define SOC_AES_SUPPORT_AES_128 (1)
 #define SOC_AES_SUPPORT_AES_256 (1)
 
+/*-------------------------- Flash Encryption CAPS----------------------------*/
+#define SOC_FLASH_ENCRYPTED_XTS_AES_BLOCK_MAX   (64)
+
 // Attention: These fixed DMA channels are temporarily workaround before we have a centralized DMA controller API to help alloc the channel dynamically
 // Remove them when GDMA driver API is ready
 #define SOC_GDMA_SHA_DMA_CHANNEL  (3)

components/spi_flash/esp_flash_api.c

@@ -23,6 +23,9 @@
 #include "sdkconfig.h"
 #include "esp_flash_internal.h"
 #include "spi_flash_defs.h"
+#if CONFIG_IDF_TARGET_ESP32S2
+#include "esp_crypto_lock.h" // for locking flash encryption peripheral
+#endif //CONFIG_IDF_TARGET_ESP32S2
 
 static const char TAG[] = "spi_flash";
 
@@ -784,22 +787,126 @@ esp_err_t IRAM_ATTR esp_flash_write(esp_flash_t *chip, const void *buffer, uint3
     return rom_spiflash_api_funcs->flash_end_flush_cache(chip, err, bus_acquired, address, length);
 }
 
-//currently the legacy implementation is used, from flash_ops.c
-esp_err_t spi_flash_write_encrypted(size_t dest_addr, const void *src, size_t size);
-
 esp_err_t IRAM_ATTR esp_flash_write_encrypted(esp_flash_t *chip, uint32_t address, const void *buffer, uint32_t length)
 {
-    /*
-     * Since currently this feature is supported only by the hardware, there
-     * is no way to support non-standard chips. We use the legacy
-     * implementation and skip the chip and driver layers.
-     */
+    if (length == 0) {
+        return ESP_OK;
+    }
+
     esp_err_t err = rom_spiflash_api_funcs->chip_check(&chip);
+    // Flash encryption only support on main flash.
+    if (chip != esp_flash_default_chip) {
+        return ESP_ERR_NOT_SUPPORTED;
+    }
     if (err != ESP_OK) return err;
-    if (buffer == NULL || address > chip->size || address+length > chip->size) {
+    if (buffer == NULL || address + length > chip->size) {
+        return ESP_ERR_INVALID_ARG;
+    }
+
+    if ((address % 16) != 0) {
+        ESP_EARLY_LOGE(TAG, "flash encrypted write address must be 16 bytes aligned");
         return ESP_ERR_INVALID_ARG;
     }
-    return spi_flash_write_encrypted(address, buffer, length);
+
+    if ((length % 16) != 0) {
+        ESP_EARLY_LOGE(TAG, "flash encrypted write length must be multiple of 16");
+        return ESP_ERR_INVALID_SIZE;
+    }
+
+    bool bus_acquired = false;
+
+    const uint8_t *ssrc = (const uint8_t *)buffer;
+
+    /* On ESP32, write_encrypted encrypts data in RAM as it writes,
+       so copy to a temporary buffer - 32 bytes at a time.
+
+       Each call to write_encrypted 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).
+
+       On ESP32-S2 and later, the temporary buffer need to be
+       seperated into 16-bytes, 32-bytes, 64-bytes(if supported).
+
+       So, on ESP32-S2 and later, here has a totally different
+       data prepare implementation.
+    */
+    uint8_t encrypt_buf[64] __attribute__((aligned(4)));
+    uint32_t row_size_length;
+    for (size_t i = 0; i < length; i += row_size_length) {
+        uint32_t row_addr = address + i;
+        uint8_t row_size;
+        uint8_t encrypt_byte;
+#if CONFIG_IDF_TARGET_ESP32
+        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, row_size);
+            /* decrypt the first block from flash, will reencrypt to same bytes */
+            esp_flash_read_encrypted(chip, row_addr, encrypt_buf, 16);
+        } else if (length - 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, row_size);
+            /* decrypt the second block from flash, will reencrypt to same bytes */
+            esp_flash_read_encrypted(chip, row_addr + 16, encrypt_buf + 16, 16);
+        } else {
+            /* Writing a full 32 byte row (2 blocks) */
+            row_size = 32;
+            memcpy(encrypt_buf, ssrc + i, row_size);
+        }
+        encrypt_byte = 32;
+        row_size_length = row_size;
+#else // FOR ESP32-S2, ESP32-S3, ESP32-C3
+        if ((row_addr % 64) == 0 && (length - i) >= 64 && SOC_FLASH_ENCRYPTED_XTS_AES_BLOCK_MAX == 64) {
+            row_size = 64;
+            memcpy(encrypt_buf, ssrc + i, row_size);
+        } else if ((row_addr % 32) == 0 && (length - i) >= 32) {
+            row_size = 32;
+            memcpy(encrypt_buf, ssrc + i, row_size);
+        } else {
+            row_size = 16;
+            memcpy(encrypt_buf, ssrc + i, row_size);
+        }
+        encrypt_byte = row_size;
+        row_size_length = row_size;
+#endif //CONFIG_IDF_TARGET_ESP32
+
+#if CONFIG_IDF_TARGET_ESP32S2
+        esp_crypto_dma_lock_acquire();
+#endif //CONFIG_IDF_TARGET_ESP32S2
+        err = rom_spiflash_api_funcs->start(chip);
+
+        if (err != ESP_OK) {
+#if CONFIG_IDF_TARGET_ESP32S2
+            esp_crypto_dma_lock_release();
+#endif //CONFIG_IDF_TARGET_ESP32S2
+            break;
+        }
+        bus_acquired = true;
+
+        err = chip->chip_drv->write_encrypted(chip, (uint32_t *)encrypt_buf, row_addr, encrypt_byte);
+        if (err!= ESP_OK) {
+#if CONFIG_IDF_TARGET_ESP32S2
+            esp_crypto_dma_lock_release();
+#endif //CONFIG_IDF_TARGET_ESP32S2
+            bus_acquired = false;
+            assert(bus_acquired);
+            break;
+        }
+        err = rom_spiflash_api_funcs->end(chip, ESP_OK);
+#if CONFIG_IDF_TARGET_ESP32S2
+        esp_crypto_dma_lock_release();
+#endif //CONFIG_IDF_TARGET_ESP32S2
+        if (err != ESP_OK) {
+            bus_acquired = false;
+            break;
+        }
+        bus_acquired = false;
+    }
+    return rom_spiflash_api_funcs->flash_end_flush_cache(chip, err, bus_acquired, address, length);
 }
 
 inline static IRAM_ATTR bool regions_overlap(uint32_t a_start, uint32_t a_len,uint32_t b_start, uint32_t b_len)
@@ -889,6 +996,7 @@ static IRAM_ATTR esp_err_t spi_flash_translate_rc(esp_err_t err)
     switch (err) {
         case ESP_OK:
         case ESP_ERR_INVALID_ARG:
+        case ESP_ERR_INVALID_SIZE:
         case ESP_ERR_NO_MEM:
             return err;
 
@@ -928,4 +1036,10 @@ esp_err_t IRAM_ATTR spi_flash_read(size_t src, void *dstv, size_t size)
     return spi_flash_translate_rc(err);
 }
 
+esp_err_t IRAM_ATTR spi_flash_write_encrypted(size_t dest_addr, const void *src, size_t size)
+{
+    esp_err_t err = esp_flash_write_encrypted(NULL, dest_addr, src, size);
+    return spi_flash_translate_rc(err);
+}
+
 #endif // CONFIG_SPI_FLASH_USE_LEGACY_IMPL

components/spi_flash/flash_ops.c

@@ -91,7 +91,9 @@ static spi_flash_counters_t s_flash_stats;
 
 #endif //CONFIG_SPI_FLASH_ENABLE_COUNTERS
 
+#if CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 static esp_err_t spi_flash_translate_rc(esp_rom_spiflash_result_t rc);
+#endif //CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 static bool is_safe_write_address(size_t addr, size_t size);
 static void spi_flash_os_yield(void);
 
@@ -249,15 +251,6 @@ static esp_rom_spiflash_result_t IRAM_ATTR spi_flash_unlock(void)
     }
     return ESP_ROM_SPIFLASH_RESULT_OK;
 }
-#else
-static esp_rom_spiflash_result_t IRAM_ATTR spi_flash_unlock(void)
-{
-    esp_err_t err = esp_flash_set_chip_write_protect(NULL, false);
-    if (err != ESP_OK) {
-        return ESP_ROM_SPIFLASH_RESULT_ERR;
-    }
-    return ESP_ROM_SPIFLASH_RESULT_OK;
-}
 #endif // CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 
 esp_err_t IRAM_ATTR spi_flash_erase_sector(size_t sec)
@@ -567,7 +560,6 @@ void IRAM_ATTR flash_rom_init(void)
 {
     return;
 }
-#endif // !CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 
 esp_err_t IRAM_ATTR spi_flash_write_encrypted(size_t dest_addr, const void *src, size_t size)
 {
@@ -661,8 +653,6 @@ fail:
     return err;
 }
 
-
-#ifdef CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 esp_err_t IRAM_ATTR spi_flash_read(size_t src, void *dstv, size_t size)
 {
     // Out of bound reads are checked in ROM code, but we can give better
@@ -814,7 +804,7 @@ out:
     COUNTER_STOP(read);
     return spi_flash_translate_rc(rc);
 }
-#endif
+#endif // !CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 
 esp_err_t IRAM_ATTR spi_flash_read_encrypted(size_t src, void *dstv, size_t size)
 {
@@ -840,7 +830,7 @@ esp_err_t IRAM_ATTR spi_flash_read_encrypted(size_t src, void *dstv, size_t size
     return err;
 }
 
-
+#if CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 static esp_err_t IRAM_ATTR spi_flash_translate_rc(esp_rom_spiflash_result_t rc)
 {
     switch (rc) {
@@ -853,6 +843,7 @@ static esp_err_t IRAM_ATTR spi_flash_translate_rc(esp_rom_spiflash_result_t rc)
         return ESP_ERR_FLASH_OP_FAIL;
     }
 }
+#endif //CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 
 #if CONFIG_SPI_FLASH_ENABLE_COUNTERS
 

components/spi_flash/partition.c

@@ -395,7 +395,11 @@ esp_err_t esp_partition_write(const esp_partition_t* partition,
         if (partition->flash_chip != esp_flash_default_chip) {
             return ESP_ERR_NOT_SUPPORTED;
         }
+#ifndef CONFIG_SPI_FLASH_USE_LEGACY_IMPL
+        return esp_flash_write_encrypted(partition->flash_chip, dst_offset, src, size);
+#else
         return spi_flash_write_encrypted(dst_offset, src, size);
+#endif // CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 #else
         return ESP_ERR_NOT_SUPPORTED;
 #endif // CONFIG_SPI_FLASH_ENABLE_ENCRYPTED_READ_WRITE

components/spi_flash/spi_flash_chip_generic.c

@@ -17,6 +17,7 @@
 #include <sys/param.h> // For MIN/MAX
 #include "spi_flash_chip_generic.h"
 #include "spi_flash_defs.h"
+#include "hal/spi_flash_encrypt_hal.h"
 #include "esp_log.h"
 #include "esp_attr.h"
 
@@ -39,6 +40,16 @@ DRAM_ATTR const static flash_chip_dummy_t default_flash_chip_dummy = {
     .slowrd_dummy_bitlen = SPI_FLASH_SLOWRD_DUMMY_BITLEN,
 };
 
+// These are the pointer to HW flash encryption. Default using hardware encryption.
+DRAM_ATTR static spi_flash_encryption_t esp_flash_encryption_default __attribute__((__unused__)) = {
+    .flash_encryption_enable = spi_flash_encryption_hal_enable,
+    .flash_encryption_disable = spi_flash_encryption_hal_disable,
+    .flash_encryption_data_prepare = spi_flash_encryption_hal_prepare,
+    .flash_encryption_done = spi_flash_encryption_hal_done,
+    .flash_encryption_destroy = spi_flash_encryption_hal_destroy,
+    .flash_encryption_check = spi_flash_encryption_hal_check,
+};
+
 DRAM_ATTR flash_chip_dummy_t *rom_flash_chip_dummy = (flash_chip_dummy_t *)&default_flash_chip_dummy;
 
 #define SPI_FLASH_DEFAULT_IDLE_TIMEOUT_MS           200
@@ -265,7 +276,65 @@ esp_err_t spi_flash_chip_generic_write(esp_flash_t *chip, const void *buffer, ui
 
 esp_err_t spi_flash_chip_generic_write_encrypted(esp_flash_t *chip, const void *buffer, uint32_t address, uint32_t length)
 {
-    return ESP_ERR_FLASH_UNSUPPORTED_HOST; // TODO
+    spi_flash_encryption_t *esp_flash_encryption = &esp_flash_encryption_default;
+    esp_err_t err = ESP_OK;
+    // Encryption must happen on main flash.
+    if (chip != esp_flash_default_chip) {
+        return ESP_ERR_NOT_SUPPORTED;
+    }
+
+    /* Check if the buffer and length can qualify the requirments */
+    if (esp_flash_encryption->flash_encryption_check(address, length) != true) {
+        return ESP_ERR_NOT_SUPPORTED;
+    }
+
+    const uint8_t *data_bytes = (const uint8_t *)buffer;
+    esp_flash_encryption->flash_encryption_enable();
+    while (length > 0) {
+        int block_size;
+        /* Write the largest block if possible */
+        if (address % 64 == 0 && length >= 64) {
+            block_size = 64;
+        } else if (address % 32 == 0 && length >= 32) {
+            block_size = 32;
+        } else {
+            block_size = 16;
+        }
+        // Prepare the flash chip (same time as AES operation, for performance)
+        esp_flash_encryption->flash_encryption_data_prepare(address, (uint32_t *)data_bytes, block_size);
+        err = chip->chip_drv->set_chip_write_protect(chip, false);
+        if (err != ESP_OK) {
+            return err;
+        }
+        // Waiting for encrypting buffer to finish and making result visible for SPI1
+        esp_flash_encryption->flash_encryption_done();
+
+        // Note: For encryption function, after write flash command is sent. The hardware will write the encrypted buffer
+        // prepared in XTS_FLASH_ENCRYPTION register in function `flash_encryption_data_prepare`, instead of the origin
+        // buffer named `data_bytes`.
+
+        err = chip->chip_drv->write(chip, (uint32_t *)data_bytes, address, length);
+        if (err != ESP_OK) {
+            return err;
+        }
+        err = chip->chip_drv->wait_idle(chip, chip->chip_drv->timeout->page_program_timeout);
+        if (err != ESP_OK) {
+            return err;
+        }
+
+        // Note: we don't wait for idle status here, because this way
+        // the AES peripheral can start encrypting the next
+        // block while the SPI flash chip is busy completing the write
+
+        esp_flash_encryption->flash_encryption_destroy();
+
+        length -= block_size;
+        data_bytes += block_size;
+        address += block_size;
+    }
+
+    esp_flash_encryption->flash_encryption_disable();
+    return err;
 }
 
 esp_err_t spi_flash_chip_generic_set_write_protect(esp_flash_t *chip, bool write_protect)

components/spi_flash/test/test_flash_encryption.c

@@ -9,23 +9,42 @@
 #include <esp_attr.h>
 #include <esp_flash_encrypt.h>
 #include <string.h>
+#include "esp_log.h"
 
+/*-------------------- For running this test, some configurations are necessary -------------------*/
+/*     ESP32    |           CONFIG_SECURE_FLASH_ENC_ENABLED         |             SET              */
+/*    ESP32S2   |           CONFIG_SECURE_FLASH_ENC_ENABLED         |             SET              */
+/*              |                CONFIG_EFUSE_VIRTUAL               |           NOT SET            */
+/*              |    CONFIG_SECURE_FLASH_REQUIRE_ALREADY_ENABLED    |             SET              */
+/*    ESP32C3   |           CONFIG_SECURE_FLASH_ENC_ENABLED         |             SET              */
+/*              |                CONFIG_EFUSE_VIRTUAL               |           NOT SET            */
+/*              |    CONFIG_SECURE_FLASH_REQUIRE_ALREADY_ENABLED    |             SET              */
 
 #ifdef CONFIG_SECURE_FLASH_ENC_ENABLED
 
 static void test_encrypted_write(size_t offset, const uint8_t *data, size_t length);
-static void test_encrypted_write_new_impl(size_t offset, const uint8_t *data, size_t length);
 static void verify_erased_flash(size_t offset, size_t length);
 
 static size_t start;
 
 static void setup_tests(void)
 {
-    if (start == 0) {
-        const esp_partition_t *part = get_test_data_partition();
-        start = part->address;
-        printf("Test data partition @ 0x%x\n", start);
+    const esp_partition_t *part = get_test_data_partition();
+    start = part->address;
+    printf("Test data partition @ 0x%x\n", start);
+}
+
+static void verify_erased_flash(size_t offset, size_t length)
+{
+    uint8_t *readback = (uint8_t *)heap_caps_malloc(SPI_FLASH_SEC_SIZE, MALLOC_CAP_32BIT | MALLOC_CAP_8BIT | MALLOC_CAP_INTERNAL);
+    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];
+        TEST_ASSERT_EQUAL_HEX_MESSAGE(0xFF, readback[i], message);
     }
+    free(readback);
 }
 
 TEST_CASE("test 16 byte encrypted writes", "[flash_encryption][test_env=UT_T1_FlashEncryption]")
@@ -89,80 +108,6 @@ static void test_encrypted_write(size_t offset, const uint8_t *data, size_t leng
     TEST_ASSERT_EQUAL_HEX8_ARRAY(data, readback, length);
 }
 
-TEST_CASE("test 16 byte encrypted writes (esp_flash)", "[flash_encryption][esp_flash_enc][test_env=UT_T1_FlashEncryption]")
-{
-    setup_tests();
-
-    TEST_ASSERT_EQUAL_HEX(ESP_OK,
-                      spi_flash_erase_sector(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,
-                      esp_flash_write_encrypted(NULL, start+1, fortyeight_bytes, 32));
-
-    TEST_ASSERT_EQUAL_HEX(ESP_ERR_INVALID_SIZE,
-                      esp_flash_write_encrypted(NULL, start, fortyeight_bytes, 15));
-
-    /* ensure nothing happened to the flash yet */
-    verify_erased_flash(start, 0x20);
-
-    /* Write 32 byte block, this is the "normal" encrypted write */
-    test_encrypted_write_new_impl(start, fortyeight_bytes, 0x20);
-    verify_erased_flash(start + 0x20, 0x20);
-
-    /* Slip in an unaligned esp_flash_read_encrypted() test */
-    uint8_t buf[0x10];
-    esp_flash_read_encrypted(NULL, start+0x10, buf, 0x10);
-    TEST_ASSERT_EQUAL_HEX8_ARRAY(fortyeight_bytes+0x10, buf, 16);
-
-    /* Write 16 bytes unaligned */
-    test_encrypted_write_new_impl(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(start + 0x20, 0x10);
-    verify_erased_flash(start + 0x40, 0x10);
-
-    /* Write 48 bytes starting at a 32-byte aligned offset */
-    test_encrypted_write_new_impl(start + 0x40, fortyeight_bytes, 0x30);
-    /* 16 bytes after this write should still be 0xFF -unencrypted- */
-    verify_erased_flash(start + 0x70, 0x10);
-
-    /* Write 48 bytes starting at a 16-byte aligned offset */
-    test_encrypted_write_new_impl(start + 0x90, fortyeight_bytes, 0x30);
-    /* 16 bytes after this write should still be 0xFF -unencrypted- */
-    verify_erased_flash(start + 0x120, 0x10);
-}
-
-static void test_encrypted_write_new_impl(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,
-                          esp_flash_write_encrypted(NULL, offset, data, length));
-
-    TEST_ASSERT_EQUAL_HEX(ESP_OK,
-                          esp_flash_read_encrypted(NULL, 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);
-    }
-}
-
 TEST_CASE("test read & write random encrypted data", "[flash_encryption][test_env=UT_T1_FlashEncryption]")
 {
     const int MAX_LEN = 192;
@@ -229,4 +174,158 @@ TEST_CASE("test read & write random encrypted data", "[flash_encryption][test_en
     free(cmp_buf);
 }
 
+#ifndef CONFIG_SPI_FLASH_USE_LEGACY_IMPL
+
+static char TAG[] = "flash_encrypt_test";
+static const char plainttext_data[] = "$$$$#### Welcome! This is flash encryption test, ..., ..., hello_world. &&&&***";
+
+static void test_encrypted_write_new_impl(size_t offset, const uint8_t *data, size_t length)
+{
+    uint8_t *readback = (uint8_t *)heap_caps_malloc(SPI_FLASH_SEC_SIZE, MALLOC_CAP_32BIT | MALLOC_CAP_8BIT | MALLOC_CAP_INTERNAL);
+    printf("encrypt %d bytes at 0x%x\n", length, offset);
+    TEST_ASSERT_EQUAL_HEX(ESP_OK,
+                          esp_flash_write_encrypted(NULL, offset, data, length));
+
+    TEST_ASSERT_EQUAL_HEX(ESP_OK,
+                          esp_flash_read_encrypted(NULL, offset, readback, length));
+
+    TEST_ASSERT_EQUAL_HEX8_ARRAY(data, readback, length);
+    free(readback);
+}
+
+TEST_CASE("test 16 byte encrypted writes (esp_flash)", "[esp_flash_enc][flash_encryption][test_env=UT_T1_FlashEncryption]")
+{
+    setup_tests();
+
+    TEST_ASSERT_EQUAL_HEX(ESP_OK,
+                      spi_flash_erase_sector(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,
+                      esp_flash_write_encrypted(NULL, start+1, fortyeight_bytes, 32));
+
+    TEST_ASSERT_EQUAL_HEX(ESP_ERR_INVALID_SIZE,
+                      esp_flash_write_encrypted(NULL, start, fortyeight_bytes, 15));
+
+    /* ensure nothing happened to the flash yet */
+    verify_erased_flash(start, 0x20);
+
+    /* Write 32 byte block, this is the "normal" encrypted write */
+    test_encrypted_write_new_impl(start, fortyeight_bytes, 0x20);
+    verify_erased_flash(start + 0x20, 0x20);
+
+    /* Slip in an unaligned esp_flash_read_encrypted() test */
+    uint8_t buf[0x10];
+    esp_flash_read_encrypted(NULL, start+0x10, buf, 0x10);
+    TEST_ASSERT_EQUAL_HEX8_ARRAY(fortyeight_bytes+0x10, buf, 16);
+
+    /* Write 16 bytes unaligned */
+    test_encrypted_write_new_impl(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(start + 0x20, 0x10);
+    verify_erased_flash(start + 0x40, 0x10);
+
+    /* Write 48 bytes starting at a 32-byte aligned offset */
+    test_encrypted_write_new_impl(start + 0x40, fortyeight_bytes, 0x30);
+    /* 16 bytes after this write should still be 0xFF -unencrypted- */
+    verify_erased_flash(start + 0x70, 0x10);
+
+    /* Write 48 bytes starting at a 16-byte aligned offset */
+    test_encrypted_write_new_impl(start + 0x90, fortyeight_bytes, 0x30);
+    /* 16 bytes after this write should still be 0xFF -unencrypted- */
+    verify_erased_flash(start + 0x120, 0x10);
+}
+
+TEST_CASE("test read & write encrypted data(32 bytes alianed address)", "[esp_flash_enc][flash_encryption][test_env=UT_T1_FlashEncryption]")
+{
+    setup_tests();
+
+    TEST_ESP_OK(spi_flash_erase_sector(start / SPI_FLASH_SEC_SIZE));
+    start = (start + 31) & (~31); // round up to 32 byte boundary
+
+    ESP_LOG_BUFFER_HEXDUMP(TAG, plainttext_data, sizeof(plainttext_data), ESP_LOG_INFO);
+    printf("Encrypteed writting......\n");
+    TEST_ESP_OK(esp_flash_write_encrypted(NULL, start, plainttext_data, sizeof(plainttext_data)));
+
+    uint8_t *cmp_encrypt_buf = heap_caps_malloc(SPI_FLASH_SEC_SIZE, MALLOC_CAP_32BIT | MALLOC_CAP_8BIT | MALLOC_CAP_INTERNAL);
+    printf("Encrypted reading......\n");
+    TEST_ESP_OK(esp_flash_read_encrypted(NULL, start, cmp_encrypt_buf, SPI_FLASH_SEC_SIZE));
+    ESP_LOG_BUFFER_HEXDUMP(TAG, cmp_encrypt_buf, sizeof(plainttext_data), ESP_LOG_INFO);
+    TEST_ASSERT_EQUAL_HEX8_ARRAY(plainttext_data, cmp_encrypt_buf, sizeof(plainttext_data));
+
+    uint8_t *cmp_normal_buf = heap_caps_malloc(SPI_FLASH_SEC_SIZE, MALLOC_CAP_32BIT | MALLOC_CAP_8BIT | MALLOC_CAP_INTERNAL);
+    TEST_ESP_OK(esp_flash_read(NULL, cmp_normal_buf, start, SPI_FLASH_SEC_SIZE));
+    printf("Normal read(esp_flash_read)......\n");
+    ESP_LOG_BUFFER_HEXDUMP(TAG, cmp_normal_buf, sizeof(plainttext_data), ESP_LOG_INFO);
+
+    free(cmp_normal_buf);
+    free(cmp_encrypt_buf);
+}
+
+TEST_CASE("test read & write encrypted data(16 bytes alianed but 32 bytes unaligned)", "[esp_flash_enc][flash_encryption][test_env=UT_T1_FlashEncryption]")
+{
+    setup_tests();
+    TEST_ESP_OK(spi_flash_erase_sector(start/SPI_FLASH_SEC_SIZE));
+    do {
+        start++;
+    } while ((start % 16) != 0);
+
+    if (start % 32 == 0) {
+        start += 16;
+    }
+    printf("Write data partition @ 0x%x\n", start);
+
+    ESP_LOG_BUFFER_HEXDUMP(TAG, plainttext_data, sizeof(plainttext_data), ESP_LOG_INFO);
+    printf("Encrypteed writting......\n");
+    TEST_ESP_OK(esp_flash_write_encrypted(NULL, start, plainttext_data, sizeof(plainttext_data)));
+
+    uint8_t *cmp_encrypt_buf = heap_caps_malloc(SPI_FLASH_SEC_SIZE, MALLOC_CAP_32BIT | MALLOC_CAP_8BIT | MALLOC_CAP_INTERNAL);
+    printf("Encrypted reading......\n");
+    TEST_ESP_OK(esp_flash_read_encrypted(NULL, start, cmp_encrypt_buf, SPI_FLASH_SEC_SIZE));
+    ESP_LOG_BUFFER_HEXDUMP(TAG, cmp_encrypt_buf, sizeof(plainttext_data), ESP_LOG_INFO);
+    TEST_ASSERT_EQUAL_HEX8_ARRAY(plainttext_data, cmp_encrypt_buf, sizeof(plainttext_data));
+
+    uint8_t *cmp_normal_buf = heap_caps_malloc(SPI_FLASH_SEC_SIZE, MALLOC_CAP_32BIT | MALLOC_CAP_8BIT | MALLOC_CAP_INTERNAL);
+    TEST_ESP_OK(esp_flash_read(NULL, cmp_normal_buf, start, SPI_FLASH_SEC_SIZE));
+    printf("Normal read(esp_flash_read)......\n");
+    ESP_LOG_BUFFER_HEXDUMP(TAG, cmp_normal_buf, sizeof(plainttext_data), ESP_LOG_INFO);
+
+    free(cmp_normal_buf);
+    free(cmp_encrypt_buf);
+}
+
+static const uint8_t large_const_buffer[16432] = {
+    203, // first byte
+    1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
+    21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37,
+    [50 ... 99] = 2,
+    [108 ... 1520] = 0x9b,
+    [1600 ... 2000] = 0x3d,
+    [8000 ... 9000] = 0xf7,
+    [15000 ... 16398] = 0xe8,
+    43, 0x7f,
+    [16401 ... 16430] = 0xd1,
+    202, // last byte
+};
+
+TEST_CASE("test read & write encrypted data with large buffer(n*64+32+16)", "[esp_flash_enc][flash_encryption][test_env=UT_T1_FlashEncryption]")
+{
+    // The tested buffer should be n*64(or n*32)+16 bytes.
+    setup_tests();
+    TEST_ESP_OK(esp_flash_erase_region(NULL, start, 5 * 4096));
+    printf("Encrypteed writting......\n");
+
+    TEST_ESP_OK(esp_flash_write_encrypted(NULL, start, large_const_buffer, sizeof(large_const_buffer)));
+    uint8_t *buf = (uint8_t*)heap_caps_malloc(sizeof(large_const_buffer), MALLOC_CAP_32BIT | MALLOC_CAP_8BIT | MALLOC_CAP_INTERNAL);
+
+    TEST_ESP_OK(esp_flash_read_encrypted(NULL, start, buf, sizeof(large_const_buffer)));
+    TEST_ASSERT_EQUAL_HEX8_ARRAY(buf, large_const_buffer, sizeof(large_const_buffer));
+    free(buf);
+}
+#endif // CONFIG_SPI_FLASH_USE_LEGACY_IMPL
 #endif // CONFIG_SECURE_FLASH_ENC_ENABLED

examples/security/flash_encryption/main/flash_encrypt_main.c

@@ -27,7 +27,7 @@ static const char* TAG = "example";
 #if CONFIG_IDF_TARGET_ESP32
 #define TARGET_CRYPT_CNT_EFUSE  ESP_EFUSE_FLASH_CRYPT_CNT
 #define TARGET_CRYPT_CNT_WIDTH  7
-#elif CONFIG_IDF_TARGET_ESP32S2 || CONFIG_IDF_TARGET_ESP32C3
+#else
 #define TARGET_CRYPT_CNT_EFUSE ESP_EFUSE_SPI_BOOT_CRYPT_CNT
 #define TARGET_CRYPT_CNT_WIDTH  3
 #endif