From 07970bd5df233c87608f8b5000bcbbe58ba15364 Mon Sep 17 00:00:00 2001
From: Laukik Hase <laukik.hase@espressif.com>
Date: Wed, 17 Jun 2026 13:40:52 +0530
Subject: [PATCH] fix(esp_tee): Prevent TEE from dispatching the REE heap
 poisoning callback

---
 components/esp_tee/include/esp_tee.h                 |  6 ++++--
 components/esp_tee/src/esp_tee_config.c              | 11 ++++++++++-
 .../esp_tee/subproject/main/common/multi_heap.c      | 12 ++++++++++++
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/components/esp_tee/include/esp_tee.h b/components/esp_tee/include/esp_tee.h
index a529928897a..63fee82cb56 100644
--- a/components/esp_tee/include/esp_tee.h
+++ b/components/esp_tee/include/esp_tee.h
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2024 Espressif Systems (Shanghai) CO LTD
+ * SPDX-FileCopyrightText: 2024-2026 Espressif Systems (Shanghai) CO LTD
  *
  * SPDX-License-Identifier: Apache-2.0
  */
@@ -43,7 +43,9 @@ typedef struct {
     uint32_t magic_word;
     uint32_t api_major_version;
     uint32_t api_minor_version;
-    uint32_t reserved[3];
+    uint32_t reserved[2];
+    /* Heap poisoning API */
+    void *ns_heap_poison_fill;
     /* TEE-related fields */
     void *s_int_handler;
     /* REE-related fields */
diff --git a/components/esp_tee/src/esp_tee_config.c b/components/esp_tee/src/esp_tee_config.c
index 22fd9f7f2ce..b4eeafdab83 100644
--- a/components/esp_tee/src/esp_tee_config.c
+++ b/components/esp_tee/src/esp_tee_config.c
@@ -1,11 +1,12 @@
 /*
- * SPDX-FileCopyrightText: 2021-2025 Espressif Systems (Shanghai) CO LTD
+ * SPDX-FileCopyrightText: 2021-2026 Espressif Systems (Shanghai) CO LTD
  *
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #include <stdarg.h>
 #include "esp_tee.h"
+#include "sdkconfig.h"
 
 /* U-mode interrupt handler */
 extern int _tee_interrupt_handler(void);
@@ -20,6 +21,13 @@ extern uint32_t _rodata_reserved_start;
 /* REE DROM end */
 extern uint32_t _rodata_reserved_end;
 
+#if CONFIG_HEAP_TLSF_USE_ROM_IMPL && (CONFIG_HEAP_POISONING_LIGHT || CONFIG_HEAP_POISONING_COMPREHENSIVE)
+extern void multi_heap_internal_poison_fill_region(void *start, size_t size, bool is_free);
+#define HEAP_POISON_FILL    ((void *)&multi_heap_internal_poison_fill_region)
+#else
+#define HEAP_POISON_FILL    NULL
+#endif
+
 esp_tee_config_t esp_tee_app_config __attribute__((section(".esp_tee_app_cfg"))) = {
     .magic_word = ESP_TEE_APP_CFG_MAGIC,
     .api_major_version = ESP_TEE_API_MAJOR_VER,
@@ -35,4 +43,5 @@ esp_tee_config_t esp_tee_app_config __attribute__((section(".esp_tee_app_cfg")))
     .ns_irom_end = &_instruction_reserved_end,
     .ns_drom_start = &_rodata_reserved_start,
     .ns_drom_end = &_rodata_reserved_end,
+    .ns_heap_poison_fill = HEAP_POISON_FILL,
 };
diff --git a/components/esp_tee/subproject/main/common/multi_heap.c b/components/esp_tee/subproject/main/common/multi_heap.c
index 536a4ae5cec..0eccfdd5c36 100644
--- a/components/esp_tee/subproject/main/common/multi_heap.c
+++ b/components/esp_tee/subproject/main/common/multi_heap.c
@@ -10,10 +10,16 @@
 #include "esp_rom_sys.h"
 #include "tlsf_block_functions.h"
 #include "multi_heap.h"
+#include "esp_tee.h"
 
 /* Handle to a registered TEE heap */
 static multi_heap_handle_t tee_heap;
 
+static inline void tee_heap_set_poison(bool enable)
+{
+    tlsf_poison_fill_pfunc_set(enable ? (poison_fill_pfunc_t)esp_tee_app_config.ns_heap_poison_fill : NULL);
+}
+
 inline static void multi_heap_assert(bool condition, const char *format, int line, intptr_t address)
 {
     /* Can't use libc assert() here as it calls printf() which can cause another malloc() for a newlib lock.
@@ -142,7 +148,10 @@ void esp_tee_heap_free(void *p)
 
     tee_heap->free_bytes += tlsf_block_size(p);
     tee_heap->free_bytes += tlsf_alloc_overhead();
+
+    tee_heap_set_poison(false);
     tlsf_free(tee_heap->heap_data, p);
+    tee_heap_set_poison(true);
 }
 
 void *malloc(size_t size)
@@ -166,7 +175,10 @@ void *realloc(void* ptr, size_t size)
     }
 
     size_t previous_block_size = tlsf_block_size(ptr);
+    tee_heap_set_poison(false);
     void *result = tlsf_realloc(tee_heap->heap_data, ptr, size);
+    tee_heap_set_poison(true);
+
     if (result) {
         /* No need to subtract the tlsf_alloc_overhead() as it has already
          * been subtracted when allocating the block at first with malloc */