Merge branch 'feature/esp_tee_flash_prot' into 'master'

feat(esp_tee): Support for flash memory isolation and protection (SPI0)

See merge request espressif/esp-idf!35486

components/bootloader_support/bootloader_flash/include/bootloader_flash_priv.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2015-2024 Espressif Systems (Shanghai) CO LTD
+ * SPDX-FileCopyrightText: 2015-2025 Espressif Systems (Shanghai) CO LTD
  *
  * SPDX-License-Identifier: Apache-2.0
  */
@@ -128,7 +128,10 @@ esp_err_t bootloader_flash_read(size_t src_addr, void *dest, size_t size, bool a
  *
  * @note All of dest_addr, src and size have to be 4-byte aligned. If write_encrypted is set, dest_addr and size must be 32-byte aligned.
  *
- * Note: In bootloader, when write_encrypted == true, the src buffer is encrypted in place.
+ * @note In bootloader, when write_encrypted == true, the src buffer is encrypted in place.
+ *
+ * @note [ESP-TEE] Using this API from the TEE will return an error if the dest_addr lies
+ *       within the active TEE partition range.
  *
  * @param  dest_addr Destination address to write in Flash.
  * @param  src Pointer to the data to write to flash
@@ -152,6 +155,9 @@ esp_err_t bootloader_flash_erase_sector(size_t sector);
 /**
  * @brief  Erase the Flash range.
  *
+ * @note   [ESP-TEE] Using this API from the TEE will return an error if the start_addr lies
+ *         within the active TEE partition range.
+ *
  * @param  start_addr start address of flash offset
  * @param  size       sector aligned size to be erased
  *

components/bootloader_support/bootloader_flash/src/bootloader_flash.c

@@ -134,8 +134,11 @@ esp_err_t bootloader_flash_erase_range(uint32_t start_addr, uint32_t size)
 #include "spi_flash/spi_flash_defs.h"
 
 #if ESP_TEE_BUILD
+#include "esp_fault.h"
 #include "esp_flash_partitions.h"
 #include "esp32c6/rom/spi_flash.h"
+
+extern bool esp_tee_flash_check_paddr_in_active_tee_part(size_t paddr);
 #endif
 
 static const char *TAG = "bootloader_flash";
@@ -515,6 +518,19 @@ esp_err_t bootloader_flash_read(size_t src_addr, void *dest, size_t size, bool a
 
 esp_err_t bootloader_flash_write(size_t dest_addr, void *src, size_t size, bool write_encrypted)
 {
+    /* NOTE: [ESP-TEE] Flash operation address validation with anti-FI check
+     *
+     * Ensure that flash operations cannot be executed within forbidden memory ranges
+     * by validating the address before proceeding.
+     */
+#if ESP_TEE_BUILD
+    bool addr_chk = esp_tee_flash_check_paddr_in_active_tee_part(dest_addr);
+    if (addr_chk) {
+        ESP_EARLY_LOGE(TAG, "bootloader_flash_write invalid dest_addr");
+        return ESP_FAIL;
+    }
+    ESP_FAULT_ASSERT(!addr_chk);
+#endif
     size_t alignment = write_encrypted ? 32 : 4;
     if ((dest_addr % alignment) != 0) {
         ESP_EARLY_LOGE(TAG, "bootloader_flash_write dest_addr 0x%x not %d-byte aligned", dest_addr, alignment);
@@ -561,6 +577,13 @@ esp_err_t bootloader_flash_erase_sector(size_t sector)
 
 esp_err_t bootloader_flash_erase_range(uint32_t start_addr, uint32_t size)
 {
+#if ESP_TEE_BUILD
+    bool addr_chk = esp_tee_flash_check_paddr_in_active_tee_part(start_addr);
+    if (addr_chk) {
+        return ESP_ERR_INVALID_ARG;
+    }
+    ESP_FAULT_ASSERT(!addr_chk);
+#endif
     if (start_addr % FLASH_SECTOR_SIZE != 0) {
         return ESP_ERR_INVALID_ARG;
     }

components/bootloader_support/src/bootloader_utility.c

@@ -570,7 +570,7 @@ void bootloader_utility_load_tee_image(const bootloader_state_t *bs)
         ESP_LOGE(TAG, "Failed to load TEE app");
         bootloader_reset();
     }
-    tee_boot_part = tee_part_idx;
+    tee_boot_part = tee_active_part;
 
     ESP_LOGI(TAG, "Loaded TEE app from partition at offset 0x%"PRIx32, tee_active_part_pos->offset);
 }

components/esp_tee/scripts/esp32c6/secure_service.tbl

@@ -28,15 +28,19 @@
 26          IDF         esp_sha_dma                                   6
 27          IDF         esp_sha_read_digest_state                     2
 28          IDF         esp_sha_write_digest_state                    2
-29          custom      esp_tee_ota_begin                             0
+29          IDF         mmu_hal_map_region                            6
-30          custom      esp_tee_ota_write                             3
+30          IDF         mmu_hal_unmap_region                          3
-31          custom      esp_tee_ota_end                               0
+31          IDF         mmu_hal_vaddr_to_paddr                        4
-32          custom      esp_tee_sec_storage_init                      0
+32          IDF         mmu_hal_paddr_to_vaddr                        5
-33          custom      esp_tee_sec_storage_gen_key                   1
+33          custom      esp_tee_ota_begin                             0
-34          custom      esp_tee_sec_storage_get_signature             4
+34          custom      esp_tee_ota_write                             3
-35          custom      esp_tee_sec_storage_get_pubkey                2
+35          custom      esp_tee_ota_end                               0
-36          custom      esp_tee_sec_storage_encrypt                   8
+36          custom      esp_tee_sec_storage_init                      0
-37          custom      esp_tee_sec_storage_decrypt                   8
+37          custom      esp_tee_sec_storage_gen_key                   1
-38          custom      esp_tee_sec_storage_is_slot_empty             1
+38          custom      esp_tee_sec_storage_get_signature             4
-39          custom      esp_tee_sec_storage_clear_slot                1
+39          custom      esp_tee_sec_storage_get_pubkey                2
-40          custom      esp_tee_att_generate_token                    6
+40          custom      esp_tee_sec_storage_encrypt                   8
+41          custom      esp_tee_sec_storage_decrypt                   8
+42          custom      esp_tee_sec_storage_is_slot_empty             1
+43          custom      esp_tee_sec_storage_clear_slot                1
+44          custom      esp_tee_att_generate_token                    6

components/esp_tee/src/esp_secure_service_wrapper.c

@@ -4,20 +4,24 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 #include <stdarg.h>
-#include "secure_service_num.h"
+
-#include "hal/sha_types.h"
-#include "hal/sha_hal.h"
-#include "hal/wdt_hal.h"
-#include "esp_tee.h"
 #include "esp_err.h"
-#include "esp_hmac.h"
 #include "esp_efuse.h"
 #include "esp_random.h"
+
+#include "hal/sha_types.h"
+#include "hal/sha_hal.h"
+#include "hal/mmu_types.h"
+#include "hal/wdt_hal.h"
+
 #include "soc/soc_caps.h"
 
+#include "esp_tee.h"
+#include "secure_service_num.h"
+
 /* ---------------------------------------------- Interrupts ------------------------------------------------- */
 
-IRAM_ATTR void __wrap_esp_rom_route_intr_matrix(int cpu_no, uint32_t model_num, uint32_t intr_num)
+void IRAM_ATTR __wrap_esp_rom_route_intr_matrix(int cpu_no, uint32_t model_num, uint32_t intr_num)
 {
     esp_tee_service_call(4, SS_ESP_ROM_ROUTE_INTR_MATRIX, cpu_no, model_num, intr_num);
 }
@@ -216,3 +220,25 @@ void __wrap_esp_sha_write_digest_state(esp_sha_type sha_type, void *digest_state
 {
     esp_tee_service_call(3, SS_ESP_SHA_WRITE_DIGEST_STATE, sha_type, digest_state);
 }
+
+/* ---------------------------------------------- MMU HAL ------------------------------------------------- */
+
+void IRAM_ATTR __wrap_mmu_hal_map_region(uint32_t mmu_id, mmu_target_t mem_type, uint32_t vaddr, uint32_t paddr, uint32_t len, uint32_t *out_len)
+{
+    esp_tee_service_call(7, SS_MMU_HAL_MAP_REGION, mmu_id, mem_type, vaddr, paddr, len, out_len);
+}
+
+void IRAM_ATTR __wrap_mmu_hal_unmap_region(uint32_t mmu_id, uint32_t vaddr, uint32_t len)
+{
+    esp_tee_service_call(4, SS_MMU_HAL_UNMAP_REGION, mmu_id, vaddr, len);
+}
+
+bool IRAM_ATTR __wrap_mmu_hal_vaddr_to_paddr(uint32_t mmu_id, uint32_t vaddr, uint32_t *out_paddr, mmu_target_t *out_target)
+{
+    return esp_tee_service_call(5, SS_MMU_HAL_VADDR_TO_PADDR, mmu_id, vaddr, out_paddr, out_target);
+}
+
+bool IRAM_ATTR __wrap_mmu_hal_paddr_to_vaddr(uint32_t mmu_id, uint32_t paddr, mmu_target_t target, mmu_vaddr_t type, uint32_t *out_vaddr)
+{
+    return esp_tee_service_call(6, SS_MMU_HAL_PADDR_TO_VADDR, mmu_id, paddr, target, type, out_vaddr);
+}

components/esp_tee/src/esp_tee.c

@@ -19,10 +19,10 @@
 /* See esp_tee_u2m_switch.S */
 extern uint32_t _u2m_switch(int argc, va_list ap);
 
-static SemaphoreHandle_t s_tee_mutex;
+static DRAM_ATTR SemaphoreHandle_t s_tee_mutex;
-static StaticSemaphore_t s_tee_mutex_buf;
+static DRAM_ATTR StaticSemaphore_t s_tee_mutex_buf;
 
-static void init_mutex(void)
+static IRAM_ATTR void init_mutex(void)
 {
     static bool is_first_call = true;
     if (is_first_call) {
@@ -35,7 +35,7 @@ static void init_mutex(void)
  * TEE interface API used by untrusted side application
  * to call secure service in trusted side
  */
-uint32_t esp_tee_service_call(int argc, ...)
+uint32_t IRAM_ATTR esp_tee_service_call(int argc, ...)
 {
     init_mutex();
 
@@ -56,7 +56,7 @@ uint32_t esp_tee_service_call(int argc, ...)
     return val;
 }
 
-IRAM_ATTR uint32_t esp_tee_service_call_with_noniram_intr_disabled(int argc, ...)
+uint32_t IRAM_ATTR esp_tee_service_call_with_noniram_intr_disabled(int argc, ...)
 {
     uint32_t val = UINT32_MAX;
     va_list ap;

components/esp_tee/subproject/components/tee_flash_mgr/esp_tee_flash.c

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 #include <string.h>
+#include <sys/param.h>
 #include <sys/queue.h>
 
 #include "esp_err.h"
@@ -20,7 +21,10 @@ static const char *TAG = "esp_tee_flash";
 
 // Structure containing the valid flash address range for flash operations through TEE
 typedef struct {
-    uint32_t reserved;
+    uint32_t flash_reg_start_paddr;
+    uint32_t flash_reg_end_paddr;
+    uint32_t active_part_start_paddr;
+    uint32_t active_part_end_paddr;
 } tee_flash_protect_info_t;
 
 static tee_flash_protect_info_t tee_prot_ctx;
@@ -36,20 +40,6 @@ typedef struct _partition_list {
 
 SLIST_HEAD(partition_list, _partition_list) partition_table_list;
 
-esp_err_t esp_tee_flash_setup_prot_ctx(uint8_t tee_boot_part)
-{
-    static bool is_first_call = true;
-
-    if (is_first_call) {
-        // TODO: To-be-implemented for C6
-        (void)tee_boot_part;
-        tee_prot_ctx.reserved = UINT32_MAX;
-        is_first_call = false;
-    }
-
-    return ESP_OK;
-}
-
 static partition_list_t *create_partition_entry(const esp_partition_info_t* partition_info)
 {
     partition_list_t *partition_entry = calloc(1, sizeof(partition_list_t));
@@ -149,3 +139,65 @@ esp_partition_info_t *esp_tee_flash_get_running_ree_partition(void)
 {
     return &ree_running;
 }
+
+esp_err_t esp_tee_flash_setup_prot_ctx(uint8_t tee_boot_part)
+{
+    static bool is_first_call = true;
+    if (!is_first_call) {
+        return ESP_OK;
+    }
+
+    /* NOTE: Fetch the partition table and record TEE partitions range. */
+    esp_err_t err = get_partition_table();
+    if (err != ESP_OK) {
+        return err;
+    }
+
+    tee_prot_ctx.flash_reg_start_paddr = UINT32_MAX;
+
+    partition_list_t *partition_entry;
+    SLIST_FOREACH(partition_entry, &partition_table_list, next) {
+        const uint32_t start_addr = partition_entry->partition.pos.offset;
+        const uint32_t end_addr = start_addr + partition_entry->partition.pos.size;
+        const uint8_t type = partition_entry->partition.type;
+        const uint8_t subtype = partition_entry->partition.subtype;
+
+        bool needs_protection = false;
+        if (type == PART_TYPE_APP) {
+            needs_protection = (subtype == PART_SUBTYPE_TEE_0 || subtype == PART_SUBTYPE_TEE_1);
+        } else if (type == PART_TYPE_DATA) {
+            needs_protection = (subtype == PART_SUBTYPE_DATA_TEE_OTA || subtype == PART_SUBTYPE_DATA_TEE_SEC_STORAGE);
+        }
+
+        if (needs_protection) {
+            tee_prot_ctx.flash_reg_start_paddr = MIN(tee_prot_ctx.flash_reg_start_paddr, start_addr);
+            tee_prot_ctx.flash_reg_end_paddr = MAX(tee_prot_ctx.flash_reg_end_paddr, end_addr);
+        }
+
+        if (type == PART_TYPE_APP && subtype == tee_boot_part) {
+            tee_prot_ctx.active_part_start_paddr = start_addr;
+            tee_prot_ctx.active_part_end_paddr = end_addr;
+        }
+    }
+
+    is_first_call = false;
+    return ESP_OK;
+}
+
+bool esp_tee_flash_check_vaddr_in_tee_region(const size_t vaddr)
+{
+    bool prot_reg1_chk = (vaddr >= SOC_S_DROM_LOW) && (vaddr < SOC_S_IROM_HIGH);
+    bool prot_reg2_chk = (vaddr >= SOC_S_MMU_MMAP_RESV_START_VADDR) && (vaddr < SOC_MMU_END_VADDR);
+
+    return (prot_reg1_chk || prot_reg2_chk);
+}
+
+bool esp_tee_flash_check_paddr_in_tee_region(const size_t paddr)
+{
+    return ((paddr >= tee_prot_ctx.flash_reg_start_paddr) && (paddr < tee_prot_ctx.flash_reg_end_paddr));
+}
+
+bool esp_tee_flash_check_paddr_in_active_tee_part(const size_t paddr)
+{
+    return ((paddr >= tee_prot_ctx.active_part_start_paddr) && (paddr < tee_prot_ctx.active_part_end_paddr));
+}

components/esp_tee/subproject/components/tee_flash_mgr/include/esp_tee_flash.h

@@ -76,3 +76,30 @@ esp_err_t esp_tee_flash_set_running_ree_partition(uint32_t paddr);
  * @return esp_partition_info_t*  Partition table entry for the running REE (user) app partition
  */
 esp_partition_info_t *esp_tee_flash_get_running_ree_partition(void);
+
+/**
+ * @brief Check if the given virtual address falls within the TEE flash protected region
+ *
+ * @param addr Virtual address to check
+ *
+ * @return bool true if address is within protected region, false otherwise
+ */
+bool esp_tee_flash_check_vaddr_in_tee_region(const size_t vaddr);
+
+/**
+ * @brief Check if the given physical address falls within the TEE flash protected region
+ *
+ * @param addr Physical address to check
+ *
+ * @return bool true if address is within protected region, false otherwise
+ */
+bool esp_tee_flash_check_paddr_in_tee_region(const size_t paddr);
+
+/**
+ * @brief Check if the given physical address falls within the active TEE partition
+ *
+ * @param dest_addr Physical address to check
+ *
+ * @return bool true if address is within active TEE partition, false otherwise
+ */
+bool esp_tee_flash_check_paddr_in_active_tee_part(const size_t paddr);

components/esp_tee/subproject/main/common/panic/esp_tee_panic.c

@@ -10,6 +10,7 @@
 #include "esp_macros.h"
 #include "esp_rom_sys.h"
 #include "esp_rom_uart.h"
+#include "rom/cache.h"
 
 #include "riscv/rv_utils.h"
 #include "riscv/rvruntime-frames.h"
@@ -18,6 +19,7 @@
 
 #include "esp_tee.h"
 #include "esp_tee_apm_intr.h"
+#include "esp_tee_rv_utils.h"
 #include "panic_helper.h"
 
 #define RV_FUNC_STK_SZ    (32)
@@ -26,12 +28,22 @@
 
 static void tee_panic_end(void)
 {
-    // make sure all the panic handler output is sent from UART FIFO
+    // Disable interrupts
+    rv_utils_tee_intr_global_disable();
+
+    // Disable the cache
+    Cache_Disable_ICache();
+
+    // Clear the interrupt controller configurations
+    memset((void *)DR_REG_PLIC_MX_BASE, 0x00, (PLIC_MXINT_CLAIM_REG + 4 - DR_REG_PLIC_MX_BASE));
+    memset((void *)DR_REG_PLIC_UX_BASE, 0x00, (PLIC_UXINT_CLAIM_REG + 4 - DR_REG_PLIC_UX_BASE));
+
+    // Make sure all the panic handler output is sent from UART FIFO
     if (CONFIG_ESP_CONSOLE_UART_NUM >= 0) {
         esp_rom_output_tx_wait_idle(CONFIG_ESP_CONSOLE_UART_NUM);
     }
 
-    // generate core reset
+    // Generate system reset
     esp_rom_software_reset_system();
 }
 

components/esp_tee/subproject/main/common/panic/panic_helper_riscv.c

@@ -117,9 +117,7 @@ void panic_print_isrcause(const void *f, int core)
     static const char *pseudo_reason[] = {
         "Unknown reason",
         "Interrupt wdt timeout on CPU0",
-#if SOC_CPU_NUM > 1
         "Interrupt wdt timeout on CPU1",
-#endif
         "Cache error",
     };
 

components/esp_tee/subproject/main/core/esp_secure_service_table.c

@@ -3,7 +3,8 @@
  *
  * SPDX-License-Identifier: Apache-2.0
  */
-#include "stddef.h"
+#include <stddef.h>
+#include "esp_attr.h"
 #include "secure_service_num.h"
 #include "secure_service_dec.h"
 
@@ -14,7 +15,7 @@ typedef void (*secure_service_t)(void);
 #pragma GCC diagnostic ignored "-Woverride-init"
 #endif
 
-const secure_service_t tee_secure_service_table[] = {
+const DRAM_ATTR secure_service_t tee_secure_service_table[] = {
     [0 ... MAX_SECURE_SERVICES - 1] = (secure_service_t)NULL,
 
 #define __SECURE_SERVICE(nr, symbol, nargs)  [nr] = (secure_service_t)_ss_##symbol,

components/esp_tee/subproject/main/core/esp_secure_services.c

@@ -3,27 +3,31 @@
  *
  * SPDX-License-Identifier: Apache-2.0
  */
-
 #include <stdarg.h>
+
 #include "esp_cpu.h"
 #include "esp_efuse.h"
-#include "hal/efuse_hal.h"
+#include "esp_flash.h"
-#include "hal/wdt_hal.h"
-
-#include "esp_rom_efuse.h"
 #include "esp_flash_encrypt.h"
+#include "esp_rom_efuse.h"
+#include "esp_fault.h"
+
+#include "hal/efuse_hal.h"
+#include "hal/mmu_types.h"
+#include "hal/mmu_hal.h"
+#include "hal/wdt_hal.h"
 #include "hal/sha_hal.h"
+
 #include "soc/soc_caps.h"
+#include "aes/esp_aes.h"
+#include "sha/sha_dma.h"
 
 #include "esp_tee.h"
 #include "secure_service_num.h"
-
 #include "esp_tee_intr.h"
 #include "esp_tee_aes_intr.h"
 #include "esp_tee_rv_utils.h"
 
-#include "aes/esp_aes.h"
-#include "sha/sha_dma.h"
 #include "esp_tee_flash.h"
 #include "esp_tee_sec_storage.h"
 #include "esp_tee_ota_ops.h"
@@ -310,6 +314,52 @@ esp_err_t _ss_esp_tee_att_generate_token(const uint32_t nonce, const uint32_t cl
     return esp_att_generate_token(nonce, client_id, psa_cert_ref, token_buf, token_buf_size, token_len);
 }
 
+/* ---------------------------------------------- MMU HAL ------------------------------------------------- */
+
+void _ss_mmu_hal_map_region(uint32_t mmu_id, mmu_target_t mem_type, uint32_t vaddr,
+                            uint32_t paddr, uint32_t len, uint32_t *out_len)
+{
+    bool vaddr_chk = esp_tee_flash_check_vaddr_in_tee_region(vaddr);
+    bool paddr_chk = esp_tee_flash_check_paddr_in_tee_region(paddr);
+    if (vaddr_chk || paddr_chk) {
+        return;
+    }
+    ESP_FAULT_ASSERT(!vaddr_chk && !vaddr_chk);
+
+    mmu_hal_map_region(mmu_id, mem_type, vaddr, paddr, len, out_len);
+}
+
+void _ss_mmu_hal_unmap_region(uint32_t mmu_id, uint32_t vaddr, uint32_t len)
+{
+    bool vaddr_chk = esp_tee_flash_check_vaddr_in_tee_region(vaddr);
+    if (vaddr_chk) {
+        return;
+    }
+    ESP_FAULT_ASSERT(!vaddr_chk);
+
+    mmu_hal_unmap_region(mmu_id, vaddr, len);
+}
+
+bool _ss_mmu_hal_vaddr_to_paddr(uint32_t mmu_id, uint32_t vaddr, uint32_t *out_paddr, mmu_target_t *out_target)
+{
+    bool vaddr_chk = esp_tee_flash_check_vaddr_in_tee_region(vaddr);
+    if (vaddr_chk) {
+        return false;
+    }
+    ESP_FAULT_ASSERT(!vaddr_chk);
+    return mmu_hal_vaddr_to_paddr(mmu_id, vaddr, out_paddr, out_target);
+}
+
+bool _ss_mmu_hal_paddr_to_vaddr(uint32_t mmu_id, uint32_t paddr, mmu_target_t target, mmu_vaddr_t type, uint32_t *out_vaddr)
+{
+    bool paddr_chk = esp_tee_flash_check_paddr_in_tee_region(paddr);
+    if (paddr_chk) {
+        return false;
+    }
+    ESP_FAULT_ASSERT(!paddr_chk);
+    return mmu_hal_paddr_to_vaddr(mmu_id, paddr, target, type, out_vaddr);
+}
+
 /* ---------------------------------------------- Secure Service Dispatcher ------------------------------------------------- */
 
 /**

components/esp_tee/subproject/main/core/esp_tee_init.c

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 #include "esp_macros.h"
+#include "esp_fault.h"
 #include "esp_log.h"
 #include "esp_rom_sys.h"
 #include "riscv/rv_utils.h"
@@ -100,16 +101,16 @@ static void tee_mark_app_and_valid_cancel_rollback(void)
             return;
         }
         ESP_LOGE(TAG, "Failed to cancel rollback (0x%08x)", err);
-        esp_rom_software_reset_system();
+        abort();
     }
 
-    ESP_LOGW(TAG, "Rollback succeeded, erasing the passive TEE partition...");
+    ESP_LOGD(TAG, "Rollback succeeded, erasing the passive TEE partition...");
     uint8_t tee_next_part = bootloader_utility_tee_get_next_update_partition(&tee_ota_pos);
     esp_partition_info_t tee_next_part_info;
 
-    int ret = esp_tee_flash_find_partition(PART_TYPE_APP, tee_next_part, NULL, &tee_next_part_info);
+    err = esp_tee_flash_find_partition(PART_TYPE_APP, tee_next_part, NULL, &tee_next_part_info);
-    ret |= esp_tee_flash_erase_range(tee_next_part_info.pos.offset, tee_next_part_info.pos.size);
+    err |= esp_tee_flash_erase_range(tee_next_part_info.pos.offset, tee_next_part_info.pos.size);
-    if (ret != 0) {
+    if (err != ESP_OK) {
         ESP_LOGE(TAG, "Failed to find/erase the passive TEE partition!");
         return;
     }
@@ -148,7 +149,12 @@ void __attribute__((noreturn)) esp_tee_init(uint32_t ree_entry_addr, uint32_t re
              ((void *)&_tee_heap_start), TEE_HEAP_SIZE, TEE_HEAP_SIZE / 1024, "RAM");
 
     /* Setting up the permissible flash operation address range */
-    assert(esp_tee_flash_setup_prot_ctx(tee_boot_part) == ESP_OK);
+    esp_err_t err = esp_tee_flash_setup_prot_ctx(tee_boot_part);
+    if (err != ESP_OK) {
+        ESP_LOGE(TAG, "Failed to setup the TEE flash memory protection!");
+        abort();
+    }
+    ESP_FAULT_ASSERT(err == ESP_OK);
 
     /* Setting up the running non-secure app partition as per the address provided by the bootloader */
     assert(esp_tee_flash_set_running_ree_partition(ree_drom_addr) == ESP_OK);

components/esp_tee/subproject/main/ld/esp32c6/esp_tee.ld.in

@@ -105,6 +105,7 @@ SECTIONS
     *libtee_flash_mgr.a:*(.rodata .srodata .rodata.* .srodata.*)
     *libbootloader_support.a:bootloader_flash.*(.rodata .srodata .rodata.* .srodata.*)
     *libmain.a:panic_helper_riscv.*(.rodata .srodata .rodata.* .srodata.*)
+    *libmain.a:esp_tee_apm_intr.c.*(.rodata .srodata .rodata.* .srodata.*)
     _rodata_end = ABSOLUTE(.);
     _tee_dram_end = ABSOLUTE(.);
   } > dram_tee_seg

components/esp_tee/subproject/main/soc/esp32c6/esp_tee_apm_prot_cfg.c

@@ -35,11 +35,11 @@ static const char *TAG = "esp_tee_apm_prot_cfg";
 #endif
 #endif
 
-/*----------------HP APM Configuration-----------------------*/
+/*----------------------- HP APM range and filter configuration -----------------------*/
 
-/* HP APM Range and Filter configuration. */
+/* HP_APM: REE0 mode accessible regions */
 apm_ctrl_region_config_data_t hp_apm_pms_data[] = {
-    /* Region0: LP memory region access. (RWX)*/
+    /* Region 0: RTC memory (RWX) */
     {
         .regn_num = 0,
         .regn_start_addr = SOC_RTC_IRAM_LOW,
@@ -47,50 +47,61 @@ apm_ctrl_region_config_data_t hp_apm_pms_data[] = {
         .regn_pms        = 0x7,
         .filter_enable   = 1,
     },
-    /* Peripherals region access.(RW)*/
+    /* Region 1: Peripherals [Start - MMU] (RW) */
-    /* Region1: Next 2 entries for MMU peripheral protection. */
+    /* Protected: MMU */
     {
         .regn_num = 1,
         .regn_start_addr = SOC_PERIPHERAL_LOW,
+        .regn_end_addr   = (SPI_MEM_MMU_ITEM_CONTENT_REG(0) - 0x4),
+        .regn_pms        = 0x6,
+        .filter_enable   = 1,
+    },
+    /* Region 2: Peripherals [MMU - Interrupt Matrix] (RW) */
+    /* Protected: Interrupt Matrix */
+    {
+        .regn_num = 2,
+        .regn_start_addr = SPI_MEM_MMU_POWER_CTRL_REG(0),
         .regn_end_addr   = (DR_REG_INTMTX_BASE - 0x4),
         .regn_pms        = 0x6,
         .filter_enable   = 1,
     },
-    /* Region4: Interrupt Matrix protection. */
+    /* Region 3: Peripherals [H/W Lock - AES] (RW) */
+    /* Protected: AES, SHA */
     {
-        .regn_num = 2,
+        .regn_num = 3,
         .regn_start_addr = DR_REG_ATOMIC_BASE,
         .regn_end_addr   = (DR_REG_AES_BASE - 0x4),
         .regn_pms        = 0x6,
         .filter_enable   = 1,
     },
-    /* Region5: Peripherals region access. (RW)*/
+    /* Region 4: Peripherals [RSA - TEE Controller & APM] (RW) */
+    /* Protected: APM, TEE Controller */
     {
-        .regn_num = 3,
+        .regn_num = 4,
         .regn_start_addr = DR_REG_RSA_BASE,
         .regn_end_addr   = (DR_REG_TEE_BASE - 0x4),
         .regn_pms        = 0x6,
         .filter_enable   = 1,
     },
-    /* Region6: Peripherals region access. (RW)*/
+    /* Region 5: Peripherals [Miscellaneous - PMU] (RW) */
     {
-        .regn_num = 4,
+        .regn_num = 5,
         .regn_start_addr = DR_REG_MISC_BASE,
         .regn_end_addr   = (DR_REG_PMU_BASE - 0x04),
         .regn_pms        = 0x6,
         .filter_enable   = 1,
     },
-    /* Region7: Peripherals region access. (RW)*/
+    /* Region 6: Peripherals [DEBUG - PWDET] (RW) */
     {
-        .regn_num = 5,
+        .regn_num = 6,
         .regn_start_addr = DR_REG_OPT_DEBUG_BASE,
-        .regn_end_addr   = 0x600D0000, //PWDET_CONF_REG,
+        .regn_end_addr   = 0x600D0000,
         .regn_pms        = 0x6,
         .filter_enable   = 1,
     },
-    /* Region8: IRAM region access. (RW)*/
+    /* Region 7: REE SRAM region (RW) */
     {
-        .regn_num = 6,
+        .regn_num = 7,
         .regn_start_addr = SOC_NS_IRAM_START,
         .regn_end_addr   = SOC_IRAM_HIGH,
         .regn_pms        = 0x6,
@@ -124,23 +135,21 @@ apm_ctrl_region_config_data_t hp_apm_pms_data[] = {
  * +---------+-------------+
  */
 
-/* HP APM configuration for the Masters. */
+/* HP_APM: REE0 mode masters' configuration */
 apm_ctrl_secure_mode_config_t hp_apm_sec_mode_data = {
     .apm_ctrl   = HP_APM_CTRL,
     .apm_m_cnt  = HP_APM_MAX_ACCESS_PATH,
     .sec_mode   = APM_LL_SECURE_MODE_REE0,
-    /* Except crypto DMA (AES and SHA) and HP CPU, all other masters will be switch to REE0 mode - refer above note */
+    /* All masters except crypto DMA (AES/SHA) and HP CPU */
     .master_ids = 0xFF3FFFFE,
     .pms_data   = hp_apm_pms_data,
 };
 
-/*----------------HP APM TEE Configuration-----------------------*/
+/* HP_APM: TEE mode accessible regions */
-
-/* HP APM Range and Filter configuration. */
 apm_ctrl_region_config_data_t hp_apm_pms_data_tee[] = {
-    /* Region9: TEE All access enable. (RW)*/
+    /* Region 8: Entire memory region (RWX)*/
     {
-        .regn_num = 7,
+        .regn_num = 8,
         .regn_start_addr = 0x0,
         .regn_end_addr   = ~0x0,
         .regn_pms        = 0x7,
@@ -148,20 +157,21 @@ apm_ctrl_region_config_data_t hp_apm_pms_data_tee[] = {
     },
 };
 
+/* HP_APM: TEE mode masters' configuration */
 apm_ctrl_secure_mode_config_t hp_apm_sec_mode_data_tee = {
     .apm_ctrl   = HP_APM_CTRL,
     .apm_m_cnt  = HP_APM_MAX_ACCESS_PATH,
     .sec_mode   = APM_LL_SECURE_MODE_TEE,
-    /* Crypto DMA will be switch to TEE mode - refer above note*/
+    /* Crypto DMA (AES/SHA) and HP_CPU */
     .master_ids = 0xC00001,
     .pms_data   = hp_apm_pms_data_tee,
 };
 
-/*----------------LP APM0 Configuration-----------------------*/
+/*----------------------- LP_APM0 range and filter configuration -----------------------*/
 
-/* LP APM0 Range and Filter configuration. */
+/* LP_APM0: REE0 mode accessible regions */
 apm_ctrl_region_config_data_t lp_apm0_pms_data[] = {
-    /* Region0: LP memory. (RWX) */
+    /* Region 0: RTC memory (RWX) */
     {
         .regn_num = 0,
         .regn_start_addr = SOC_RTC_IRAM_LOW,
@@ -171,20 +181,21 @@ apm_ctrl_region_config_data_t lp_apm0_pms_data[] = {
     },
 };
 
-/* LP APM0 configuration for the Masters. */
+/* LP_APM0: REE0 mode masters' configuration */
 apm_ctrl_secure_mode_config_t lp_apm0_sec_mode_data = {
     .apm_ctrl   = LP_APM0_CTRL,
     .apm_m_cnt  = LP_APM0_MAX_ACCESS_PATH,
     .sec_mode   = APM_LL_SECURE_MODE_REE0,
-    .master_ids = 0x2, //Only LP_CPU here.
+    /* LP_CPU */
+    .master_ids = 0x2,
     .pms_data   = lp_apm0_pms_data,
 };
 
-/*----------------LP APM Configuration-----------------------*/
+/*----------------------- LP_APM range and filter configuration -----------------------*/
 
-/* LP APM Range and Filter configuration. */
+/* LP_APM: REE0 mode accessible regions */
 apm_ctrl_region_config_data_t lp_apm_pms_data[] = {
-    /* Region0: LP memory. (RWX) */
+    /* Region 0: RTC memory (RWX) */
     {
         .regn_num = 0,
         .regn_start_addr = SOC_RTC_IRAM_LOW,
@@ -192,7 +203,8 @@ apm_ctrl_region_config_data_t lp_apm_pms_data[] = {
         .regn_pms        = 0x7,
         .filter_enable   = 1,
     },
-    /* Region1: LP Peri 1. (RW) */
+    /* Region 1: LP Peripherals [PMU - eFuse BLK x] (RW) */
+    /* Protected: eFuse BLK x */
     {
         .regn_num = 1,
         .regn_start_addr = DR_REG_PMU_BASE,
@@ -200,7 +212,7 @@ apm_ctrl_region_config_data_t lp_apm_pms_data[] = {
         .regn_pms        = 0x6,
         .filter_enable   = 1,
     },
-    /* Region2: LP Peri 2. (RW) */
+    /* Region 2: LP Peripherals [eFuse - END] (RW) */
     {
         .regn_num = 2,
         .regn_start_addr = LP_APM_EFUSE_REG_END,
@@ -210,18 +222,19 @@ apm_ctrl_region_config_data_t lp_apm_pms_data[] = {
     },
 };
 
-/* LP APM configuration for the Masters. */
+/* LP_APM: REE0 mode masters' configuration */
 apm_ctrl_secure_mode_config_t lp_apm_sec_mode_data = {
     .apm_ctrl   = LP_APM_CTRL,
     .apm_m_cnt  = LP_APM_MAX_ACCESS_PATH,
     .sec_mode   = APM_LL_SECURE_MODE_REE0,
-    .master_ids = 0x3, //Only HP_CPU & LP_CPU here.
+    /* HP_CPU and LP_CPU */
+    .master_ids = 0x3,
     .pms_data   = lp_apm_pms_data,
 };
 
-/* LP APM Range and Filter configuration. */
+/* LP_APM: TEE mode accessible regions */
 apm_ctrl_region_config_data_t lp_apm_pms_data_tee[] = {
-    /* Region3: TEE All access enable. (RW)*/
+    /* Region 3: Entire memory region (RWX) */
     {
         .regn_num = 3,
         .regn_start_addr = 0x0,
@@ -231,12 +244,13 @@ apm_ctrl_region_config_data_t lp_apm_pms_data_tee[] = {
     },
 };
 
+/* LP_APM0: TEE mode masters' configuration */
 apm_ctrl_secure_mode_config_t lp_apm_sec_mode_data_tee = {
     .apm_ctrl   = LP_APM_CTRL,
     .apm_m_cnt  = LP_APM_MAX_ACCESS_PATH,
     .sec_mode   = APM_LL_SECURE_MODE_TEE,
-    /* HP CPU and LP CPU will be switch to TEE mode */
+    /* HP_CPU and LP_CPU */
-    .master_ids = 0x00003,
+    .master_ids = 0x3,
     .pms_data   = lp_apm_pms_data_tee,
 };
 

components/esp_tee/subproject/main/soc/esp32c6/include/esp_tee_rv_utils.h

@@ -21,9 +21,6 @@
 extern "C" {
 #endif
 
-#define SET_BIT(t, n)  (t |= (1UL << (n)))
-#define CLR_BIT(t, n)  (t &= ~(1UL << (n)))
-
 FORCE_INLINE_ATTR void rv_utils_tee_intr_global_enable(void)
 {
     /*
@@ -50,6 +47,13 @@ FORCE_INLINE_ATTR void rv_utils_tee_intr_global_enable(void)
     RV_SET_CSR(mstatus, MSTATUS_MIE);
 }
 
+FORCE_INLINE_ATTR void rv_utils_tee_intr_global_disable(void)
+{
+    RV_CLEAR_CSR(ustatus, USTATUS_UIE);
+    RV_CLEAR_CSR(mstatus, MSTATUS_UIE);
+    RV_CLEAR_CSR(mstatus, MSTATUS_MIE);
+}
+
 FORCE_INLINE_ATTR void rv_utils_tee_intr_enable(uint32_t intr_mask)
 {
     unsigned old_xstatus;

components/esp_tee/test_apps/tee_test_fw/main/CMakeLists.txt

@@ -2,7 +2,7 @@ idf_build_get_property(idf_path IDF_PATH)
 
 set(priv_requires bootloader_support driver esp_tee esp_timer mbedtls spi_flash)
 # Test FW related
-list(APPEND priv_requires cmock json test_utils unity)
+list(APPEND priv_requires cmock json nvs_flash test_utils unity)
 # TEE related
 list(APPEND priv_requires tee_sec_storage tee_attestation tee_ota_ops test_sec_srv)
 
@@ -13,7 +13,8 @@ list(APPEND srcs "test_esp_tee_ctx_switch.c"
                  "test_esp_tee_panic.c"
                  "test_esp_tee_sec_stg.c"
                  "test_esp_tee_ota.c"
-                 "test_esp_tee_att.c")
+                 "test_esp_tee_att.c"
+                 "test_esp_tee_flash_prot.c")
 
 set(mbedtls_test_srcs_dir "${idf_path}/components/mbedtls/test_apps/main")
 

components/esp_tee/test_apps/tee_test_fw/main/app_main.c

@@ -5,6 +5,8 @@
  */
 #include "freertos/FreeRTOS.h"
 #include "freertos/task.h"
+#include "nvs.h"
+#include "nvs_flash.h"
 #include "unity.h"
 #include "memory_checks.h"
 
@@ -42,5 +44,6 @@ static void test_task(void *pvParameters)
 
 void app_main(void)
 {
+    ESP_ERROR_CHECK(nvs_flash_init());
     xTaskCreatePinnedToCore(test_task, "testTask", CONFIG_UNITY_FREERTOS_STACK_SIZE, NULL, CONFIG_UNITY_FREERTOS_PRIORITY, NULL, CONFIG_UNITY_FREERTOS_CPU);
 }

components/esp_tee/test_apps/tee_test_fw/main/test_esp_tee_flash_prot.c

@@ -0,0 +1,151 @@
+/*
+ * SPDX-FileCopyrightText: 2024 Espressif Systems (Shanghai) CO LTD
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#include <string.h>
+
+#include "esp_log.h"
+#include "esp_system.h"
+#include "esp_heap_caps.h"
+
+#include "esp_flash.h"
+#include "esp_rom_spiflash.h"
+#include "spi_flash_mmap.h"
+#include "esp_private/cache_utils.h"
+#include "esp_partition.h"
+#include "nvs_flash.h"
+
+#include "esp_tee.h"
+#include "esp_tee_sec_storage.h"
+#include "secure_service_num.h"
+
+#include "unity.h"
+
+#define BOOT_COUNT_NAMESPACE "boot_count"
+
+static const char *TAG = "test_esp_tee_flash_prot";
+
+static void set_boot_count_in_nvs(uint8_t boot_count)
+{
+    nvs_handle_t boot_count_handle;
+    TEST_ESP_OK(nvs_open(BOOT_COUNT_NAMESPACE, NVS_READWRITE, &boot_count_handle));
+    TEST_ESP_OK(nvs_set_u8(boot_count_handle, "boot_count", boot_count));
+    TEST_ESP_OK(nvs_commit(boot_count_handle));
+    nvs_close(boot_count_handle);
+}
+
+static uint8_t get_boot_count_from_nvs(void)
+{
+    nvs_handle_t boot_count_handle;
+    esp_err_t err = nvs_open(BOOT_COUNT_NAMESPACE, NVS_READONLY, &boot_count_handle);
+    if (err == ESP_ERR_NVS_NOT_FOUND) {
+        set_boot_count_in_nvs(0);
+    }
+    uint8_t boot_count;
+    TEST_ESP_OK(nvs_get_u8(boot_count_handle, "boot_count", &boot_count));
+    nvs_close(boot_count_handle);
+    return boot_count;
+}
+
+static void test_initial_boot(void)
+{
+    ESP_LOGI(TAG, "Boot: 1");
+    set_boot_count_in_nvs(1);
+    esp_restart();
+}
+
+/* ---------------------------------------------- API family 1: esp_partition ------------------------------------------------- */
+
+static void test_esp_partition_mmap_api(void)
+{
+    uint8_t boot_count = get_boot_count_from_nvs();
+    boot_count++;
+    set_boot_count_in_nvs(boot_count);
+
+    const esp_partition_t *part = NULL;
+    esp_partition_mmap_handle_t out_handle;
+    const void *outptr = NULL;
+
+    switch (boot_count) {
+    case 2:
+        part = esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_TEE_0, NULL);
+        TEST_ASSERT_NOT_NULL(part);
+        TEST_ESP_OK(esp_partition_mmap(part, 0, part->size, ESP_PARTITION_MMAP_DATA, &outptr, &out_handle));
+        ESP_LOG_BUFFER_HEXDUMP(TAG, outptr, 0x20, ESP_LOG_INFO);
+        TEST_FAIL_MESSAGE("System fault should have been generated");
+        break;
+    case 3:
+        part = esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_TEE_1, NULL);
+        TEST_ASSERT_NOT_NULL(part);
+        TEST_ESP_OK(esp_partition_mmap(part, 0, part->size, ESP_PARTITION_MMAP_INST, &outptr, &out_handle));
+        ESP_LOG_BUFFER_HEXDUMP(TAG, outptr, 0x20, ESP_LOG_INFO);
+        TEST_FAIL_MESSAGE("System fault should have been generated");
+        break;
+    case 4:
+        part = esp_partition_find_first(ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_DATA_TEE_OTA, NULL);
+        TEST_ASSERT_NOT_NULL(part);
+        TEST_ESP_OK(esp_partition_mmap(part, 0, part->size, ESP_PARTITION_MMAP_DATA, &outptr, &out_handle));
+        ESP_LOG_BUFFER_HEXDUMP(TAG, outptr, 0x20, ESP_LOG_INFO);
+        TEST_FAIL_MESSAGE("System fault should have been generated");
+        break;
+    case 5:
+        part = esp_partition_find_first(ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_DATA_TEE_SEC_STORAGE, NULL);
+        TEST_ASSERT_NOT_NULL(part);
+        TEST_ESP_OK(esp_partition_mmap(part, 0, part->size, ESP_PARTITION_MMAP_DATA, &outptr, &out_handle));
+        ESP_LOG_BUFFER_HEXDUMP(TAG, outptr, 0x20, ESP_LOG_INFO);
+        TEST_FAIL_MESSAGE("System fault should have been generated");
+        break;
+    default:
+        TEST_FAIL_MESSAGE("Unexpected stage");
+        break;
+    }
+}
+
+TEST_CASE_MULTIPLE_STAGES("Test REE-TEE isolation: Flash - SPI0 (esp_partition_mmap)", "[flash_prot][timeout=60]",
+                          test_initial_boot, test_esp_partition_mmap_api, test_esp_partition_mmap_api,
+                          test_esp_partition_mmap_api, test_esp_partition_mmap_api);
+
+/* ---------------------------------------------- API family 2: spi_flash ------------------------------------------------- */
+
+static void test_spi_flash_mmap_api(void)
+{
+    uint8_t boot_count = get_boot_count_from_nvs();
+    boot_count++;
+    set_boot_count_in_nvs(boot_count);
+
+    const esp_partition_t *part = NULL;
+    spi_flash_mmap_handle_t handle;
+    const void *ptr = NULL;
+
+    switch (boot_count) {
+    case 2:
+        part = esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_TEE_0, NULL);
+        TEST_ASSERT_NOT_NULL(part);
+        TEST_ESP_OK(spi_flash_mmap(part->address, part->size, SPI_FLASH_MMAP_DATA, &ptr, &handle));
+        ESP_LOG_BUFFER_HEXDUMP(TAG, ptr, 0x20, ESP_LOG_INFO);
+        TEST_FAIL_MESSAGE("System fault should have been generated");
+        break;
+    case 3:
+        part = esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_TEE_1, NULL);
+        TEST_ASSERT_NOT_NULL(part);
+        TEST_ESP_OK(spi_flash_mmap(part->address, part->size, SPI_FLASH_MMAP_INST, &ptr, &handle));
+        ESP_LOG_BUFFER_HEXDUMP(TAG, ptr, 0x20, ESP_LOG_INFO);
+        TEST_FAIL_MESSAGE("System fault should have been generated");
+        break;
+    case 4:
+        part = esp_partition_find_first(ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_DATA_TEE_OTA, NULL);
+        TEST_ASSERT_NOT_NULL(part);
+        TEST_ESP_OK(spi_flash_mmap(part->address, part->size, SPI_FLASH_MMAP_DATA, &ptr, &handle));
+        ESP_LOG_BUFFER_HEXDUMP(TAG, ptr, 0x20, ESP_LOG_INFO);
+        TEST_FAIL_MESSAGE("System fault should have been generated");
+        break;
+    default:
+        TEST_FAIL_MESSAGE("Unexpected stage");
+        break;
+    }
+}
+
+TEST_CASE_MULTIPLE_STAGES("Test REE-TEE isolation: Flash - SPI0 (spi_flash_mmap)", "[flash_prot][timeout=60]",
+                          test_initial_boot, test_spi_flash_mmap_api, test_spi_flash_mmap_api,
+                          test_spi_flash_mmap_api);

components/esp_tee/test_apps/tee_test_fw/main/test_esp_tee_panic.c

@@ -8,6 +8,8 @@
 #include "soc/efuse_reg.h"
 #include "soc/lp_analog_peri_reg.h"
 #include "soc/lp_wdt_reg.h"
+#include "soc/spi_mem_reg.h"
+#include "soc/ext_mem_defs.h"
 
 #include "freertos/FreeRTOS.h"
 #include "freertos/task.h"
@@ -46,6 +48,15 @@ TEST_CASE("Test APM violation interrupt: eFuse", "[apm_violation]")
     TEST_FAIL_MESSAGE("APM violation interrupt should have been generated");
 }
 
+TEST_CASE("Test APM violation interrupt: MMU", "[apm_violation]")
+{
+    uint32_t val = UINT32_MAX;
+    REG_WRITE(SPI_MEM_MMU_ITEM_INDEX_REG(0), SOC_MMU_ENTRY_NUM - 2);
+    val = REG_READ(SPI_MEM_MMU_ITEM_CONTENT_REG(0));
+    TEST_ASSERT_EQUAL(0, val);
+    TEST_FAIL_MESSAGE("APM violation interrupt should have been generated");
+}
+
 /* TEE IRAM: Reserved/Vector-table boundary */
 TEST_CASE("Test TEE-TEE violation: IRAM (W1)", "[exception]")
 {

components/esp_tee/test_apps/tee_test_fw/pytest_esp_tee_ut.py

@@ -40,7 +40,7 @@ REE_ISOLATION_TEST_EXC_RSN: Dict[str, Any] = {
     }
 }
 
-TEE_APM_VIOLATION_EXC_CHK = ['AES', 'eFuse']
+TEE_APM_VIOLATION_EXC_CHK = ['AES', 'eFuse', 'MMU']
 
 # ---------------- TEE default tests ----------------
 
@@ -134,6 +134,53 @@ def test_esp_tee_isolation_checks(dut: IdfDut) -> None:
             raise RuntimeError('Incorrect exception received!')
         dut.expect('Exception origin: U-mode')
 
+
+def run_multiple_stages(dut: IdfDut, test_case_num: int, stages: int) -> None:
+    for stage in range(1, stages + 1):
+        dut.write(str(test_case_num))
+        dut.expect(r'\s+\((\d+)\)\s+"([^"]+)"\r?\n', timeout=30)
+        dut.write(str(stage))
+
+        if 1 < stage <= stages:
+            rst_rsn = dut.expect(r"Core ([01]) panic\'ed \(([^)]+)\)", timeout=30).group(2).decode()
+            if rst_rsn != 'Cache error':
+                raise RuntimeError('Incorrect reset reason observed after TEE image failure!')
+
+        if stage != stages:
+            dut.expect_exact('Press ENTER to see the list of tests.')
+
+
+@pytest.mark.generic
+@pytest.mark.parametrize('config', CONFIGS_OTA, indirect=True)
+@pytest.mark.parametrize('skip_autoflash', ['y'], indirect=True)
+def test_esp_tee_flash_prot_esp_partition_mmap(dut: IdfDut) -> None:
+    # Flash the bootloader, TEE and REE firmware
+    dut.serial.custom_flash()
+
+    # start test
+    extra_data = dut.parse_test_menu()
+    for test_case in extra_data:
+        if test_case.name == 'Test REE-TEE isolation: Flash - SPI1 (esp_partition_mmap)':
+            run_multiple_stages(dut, test_case.index, len(test_case.subcases))
+        else:
+            continue
+
+
+@pytest.mark.generic
+@pytest.mark.parametrize('config', CONFIGS_OTA, indirect=True)
+@pytest.mark.parametrize('skip_autoflash', ['y'], indirect=True)
+def test_esp_tee_flash_prot_spi_flash_mmap(dut: IdfDut) -> None:
+    # Flash the bootloader, TEE and REE firmware
+    dut.serial.custom_flash()
+
+    # start test
+    extra_data = dut.parse_test_menu()
+    for test_case in extra_data:
+        if test_case.name == 'Test REE-TEE isolation: Flash - SPI0 (spi_flash_mmap)':
+            run_multiple_stages(dut, test_case.index, len(test_case.subcases))
+        else:
+            continue
+
 # ---------------- TEE Local OTA tests ----------------