fix(configs/mbedtls): patch Mbed TLS allocator to use app heap

The original Mbed TLS allocator uses a static variable as heap context
information. The library is used in apps, and each app should have its
own heap. This commit adds a patch to Mbed TLS so that the allocator
gets the heap context structure from the app heap area.

To make this work, the app heap area layout needs to be changed so that
the Mbed TLS heap context is the first. This is necessary so that any
app that uses the Mbed TLS allocator provides the heap context
structure on the same offset (0 in this case).

Change-Id: I4ca6a5882837f8fa9d9e9e2294eeec11808f2b2e
Signed-off-by: Mate Toth-Pal <mate.toth-pal@arm.com>
diff --git a/app/attestation/el0_app/src/attest_app.c b/app/attestation/el0_app/src/attest_app.c
index 8b42e99..6cc66c5 100644
--- a/app/attestation/el0_app/src/attest_app.c
+++ b/app/attestation/el0_app/src/attest_app.c
@@ -32,6 +32,14 @@
 
 /* The following static variables are common for all attest app instances */
 static bool attest_key_init_done;
+/*
+ * The initial heap is used during the global init call off the attestation app.
+ * This initialization happens once during RMM boot on the primary CPU, and
+ * this heap is going to hold the Realm Attestation Key. It is allocated in RMM
+ * core's .bss so that it cam be commonly used by all app instances.
+ */
+static buffer_alloc_ctx init_mbedtls_heap_ctx;
+static unsigned char init_mbedtls_heap_buf[4U * 1024U] __aligned(sizeof(unsigned long));
 
 /*
  * Variables that need to be allocated in the heap buffer
@@ -42,6 +50,7 @@
 	 * Variables that are not dynamically allocated, but need to be kept
 	 * local to app instance.
 	 */
+	buffer_alloc_ctx mbedtls_heap_ctx;
 	struct token_sign_cntxt token_sign_context;
 	uint8_t cca_attest_token_buf[REC_ATTEST_TOKEN_BUF_SIZE];
 	/* The realm token is never returned to RMM only used inside the app */
@@ -55,6 +64,16 @@
 /* Make sure that the buffer used by the allocator is aligned */
 COMPILER_ASSERT((UL(offsetof(struct attest_heap_t, mbedtls_heap_buf)) % sizeof(unsigned long)) == 0U);
 
+/* coverity[misra_c_2012_rule_5_8_violation:SUPPRESS] */
+void *mbedtls_app_get_heap(void)
+{
+	if (attest_key_init_done) {
+		return &((struct attest_heap_t *)get_heap_start())->mbedtls_heap_ctx;
+	} else {
+		return &init_mbedtls_heap_ctx;
+	}
+}
+
 /*
  * This function is used by Mbed TLS as a source of entropy. This means it is
  * called during RMM operation, to add entropy to the signing process.
@@ -88,14 +107,18 @@
 {
 	int ret;
 	psa_status_t psa_status;
-	struct attest_heap_t *heap = (struct attest_heap_t *)get_heap_start();
 
 	/* Enable Data Independent Timing feature */
 	write_dit(DIT_BIT);
 
+	/*
+	 * First init the buffer with a global heap so that the memory that is
+	 * allocated from heap during initialisation is accessible later from
+	 * other app instances as well.
+	 */
 	mbedtls_memory_buffer_alloc_init(
-		heap->mem_buf,
-		sizeof(heap->mem_buf));
+		init_mbedtls_heap_buf,
+		sizeof(init_mbedtls_heap_buf));
 
 	psa_status = psa_crypto_init();
 	if (psa_status != PSA_SUCCESS) {
@@ -128,7 +151,7 @@
 	struct attest_heap_t *heap = (struct attest_heap_t *)get_heap_start();
 
 	enum attest_token_err_t ret =
-		attest_cca_token_create(&heap->token_sign_context,
+		attest_app_cca_token_create(&heap->token_sign_context,
 			heap->cca_attest_token_buf,
 			sizeof(heap->cca_attest_token_buf),
 			heap->realm_attest_token_buf,
@@ -177,7 +200,7 @@
 	struct attest_heap_t *heap = (struct attest_heap_t *)get_heap_start();
 
 	/* coverity[misra_c_2012_directive_4_7_violation:SUPPRESS] */
-	return (unsigned long)attest_realm_token_create(algorithm,
+	return (unsigned long)attest_app_realm_token_create(algorithm,
 			     params->measurements,
 			     MEASUREMENT_SLOT_NR,
 			     &(params->rpv),
@@ -225,7 +248,7 @@
 		return app_do_extend((enum hash_algo)arg_0, arg_1, (uint8_t *)shared);
 	case ATTESTATION_APP_FUNC_ID_TOKEN_SIGN:
 		/* coverity[misra_c_2012_directive_4_7_violation:SUPPRESS] */
-		token_ret = attest_realm_token_sign(
+		token_ret = attest_app_realm_token_sign(
 			&heap->token_sign_context, &heap->realm_token_len);
 		*((size_t *)get_shared_mem_start()) = heap->realm_token_len;
 		return (unsigned long)token_ret;
diff --git a/app/attestation/rmm_stub/src/attest.c b/app/attestation/rmm_stub/src/attest.c
index 685016a..f41c24a 100644
--- a/app/attestation/rmm_stub/src/attest.c
+++ b/app/attestation/rmm_stub/src/attest.c
@@ -97,7 +97,7 @@
 
 void attest_app_get_bss(uintptr_t *bss_pa, size_t *bss_size)
 {
-	static char attest_app_bss[2U * GRANULE_SIZE] __aligned(GRANULE_SIZE);
+	static char attest_app_bss[3U * GRANULE_SIZE] __aligned(GRANULE_SIZE);
 	*bss_pa = (uintptr_t)attest_app_bss;
 	*bss_size = sizeof(attest_app_bss);
 }
diff --git a/app/common/el0_app/CMakeLists.txt b/app/common/el0_app/CMakeLists.txt
index 59f90d8..91c08ec 100644
--- a/app/common/el0_app/CMakeLists.txt
+++ b/app/common/el0_app/CMakeLists.txt
@@ -15,8 +15,8 @@
 target_include_directories(rmm-app-common-el0_app
   PUBLIC "include"
          "include/${RMM_ARCH}"
-  PRIVATE "../../common/include"
-          "../../common/include/${RMM_ARCH}")
+         "../../common/include"
+         "../../common/include/${RMM_ARCH}")
 
 target_sources(rmm-app-common-el0_app
   PRIVATE
diff --git a/app/common/el0_app/include/el0_app_helpers.h b/app/common/el0_app/include/el0_app_helpers.h
index 6488b28..5281b51 100644
--- a/app/common/el0_app/include/el0_app_helpers.h
+++ b/app/common/el0_app/include/el0_app_helpers.h
@@ -12,6 +12,13 @@
 void *get_shared_mem_start(void);
 size_t get_shared_mem_size(void);
 
+/*
+ * This function must be defined by the app that uses Mbed TLS allocator
+ *
+ * The function returns the start address to a buffer_alloc_ctx object.
+ */
+void *mbedtls_app_get_heap(void);
+
 unsigned long el0_app_entry_func(
 	unsigned long func_id,
 	unsigned long arg_0,
diff --git a/app/random/el0_app/src/random_app.c b/app/random/el0_app/src/random_app.c
index c9c08ed..0d8b60c 100644
--- a/app/random/el0_app/src/random_app.c
+++ b/app/random/el0_app/src/random_app.c
@@ -21,12 +21,19 @@
 	/* Variables that are not dynamically allocated, but need to be kept
 	 * local to app instance.
 	 */
+	buffer_alloc_ctx mbedtls_heap_ctx;
 	bool prng_init_done;
 	mbedtls_hmac_drbg_context cpu_drbg_ctx;
 	/* The actual heap buffer */
 	unsigned char mbedtls_heap_buf[PRNG_HEAP_SIZE] __aligned(sizeof(unsigned long));
 };
 
+/* coverity[misra_c_2012_rule_5_8_violation:SUPPRESS] */
+void *mbedtls_app_get_heap(void)
+{
+	return &((struct random_heap_t *)get_heap_start())->mbedtls_heap_ctx;
+}
+
 /* Make sure that the buffer used by the allocator is 64 bit dword aligned. */
 COMPILER_ASSERT((U(offsetof(struct random_heap_t, mbedtls_heap_buf)) % sizeof(unsigned long)) == 0UL);
 
diff --git a/configs/mbedtls/0003-refactor-alloc-Use-the-heap-structure-provided-by-ap.patch b/configs/mbedtls/0003-refactor-alloc-Use-the-heap-structure-provided-by-ap.patch
new file mode 100644
index 0000000..413bf1a
--- /dev/null
+++ b/configs/mbedtls/0003-refactor-alloc-Use-the-heap-structure-provided-by-ap.patch
@@ -0,0 +1,554 @@
+From 90c2d7faeb675daa60eb704c31d64ec693a1f482 Mon Sep 17 00:00:00 2001
+From: Mate Toth-Pal <mate.toth-pal@arm.com>
+Date: Thu, 27 Feb 2025 13:00:13 +0000
+Subject: [PATCH 3/3] refactor: alloc: Use the heap structure provided by app
+
+---
+ include/mbedtls/memory_buffer_alloc.h |  36 +++++
+ library/memory_buffer_alloc.c         | 209 +++++++++++++-------------
+ 2 files changed, 138 insertions(+), 107 deletions(-)
+
+diff --git a/include/mbedtls/memory_buffer_alloc.h b/include/mbedtls/memory_buffer_alloc.h
+index b527d9b..1c7b627 100644
+--- a/include/mbedtls/memory_buffer_alloc.h
++++ b/include/mbedtls/memory_buffer_alloc.h
+@@ -34,6 +34,42 @@
+ #define MBEDTLS_MEMORY_VERIFY_ALWAYS       (MBEDTLS_MEMORY_VERIFY_ALLOC | \
+                                             MBEDTLS_MEMORY_VERIFY_FREE)
+ 
++typedef struct _memory_header memory_header;
++struct _memory_header {
++    size_t          magic1;
++    size_t          size;
++    size_t          alloc;
++    memory_header   *prev;
++    memory_header   *next;
++    memory_header   *prev_free;
++    memory_header   *next_free;
++#if defined(MBEDTLS_MEMORY_BACKTRACE)
++    char            **trace;
++    size_t          trace_count;
++#endif
++    size_t          magic2;
++};
++
++typedef struct {
++    unsigned char   *buf;
++    size_t          len;
++    memory_header   *first;
++    memory_header   *first_free;
++    int             verify;
++#if defined(MBEDTLS_MEMORY_DEBUG)
++    size_t          alloc_count;
++    size_t          free_count;
++    size_t          total_used;
++    size_t          maximum_used;
++    size_t          header_count;
++    size_t          maximum_header_count;
++#endif
++#if defined(MBEDTLS_THREADING_C)
++    mbedtls_threading_mutex_t   mutex;
++#endif
++}
++buffer_alloc_ctx;
++
+ #ifdef __cplusplus
+ extern "C" {
+ #endif
+diff --git a/library/memory_buffer_alloc.c b/library/memory_buffer_alloc.c
+index 79b0a8b..630120b 100644
+--- a/library/memory_buffer_alloc.c
++++ b/library/memory_buffer_alloc.c
+@@ -8,6 +8,7 @@
+ #include "common.h"
+ 
+ #if defined(MBEDTLS_MEMORY_BUFFER_ALLOC_C)
++#include <el0_app_helpers.h>
+ #include "mbedtls/memory_buffer_alloc.h"
+ 
+ /* No need for the header guard as MBEDTLS_MEMORY_BUFFER_ALLOC_C
+@@ -29,44 +30,6 @@
+ #define MAGIC2       0xEE119966
+ #define MAX_BT 20
+ 
+-typedef struct _memory_header memory_header;
+-struct _memory_header {
+-    size_t          magic1;
+-    size_t          size;
+-    size_t          alloc;
+-    memory_header   *prev;
+-    memory_header   *next;
+-    memory_header   *prev_free;
+-    memory_header   *next_free;
+-#if defined(MBEDTLS_MEMORY_BACKTRACE)
+-    char            **trace;
+-    size_t          trace_count;
+-#endif
+-    size_t          magic2;
+-};
+-
+-typedef struct {
+-    unsigned char   *buf;
+-    size_t          len;
+-    memory_header   *first;
+-    memory_header   *first_free;
+-    int             verify;
+-#if defined(MBEDTLS_MEMORY_DEBUG)
+-    size_t          alloc_count;
+-    size_t          free_count;
+-    size_t          total_used;
+-    size_t          maximum_used;
+-    size_t          header_count;
+-    size_t          maximum_header_count;
+-#endif
+-#if defined(MBEDTLS_THREADING_C)
+-    mbedtls_threading_mutex_t   mutex;
+-#endif
+-}
+-buffer_alloc_ctx;
+-
+-static buffer_alloc_ctx heap;
+-
+ #if defined(MBEDTLS_MEMORY_DEBUG)
+ static void debug_header(memory_header *hdr)
+ {
+@@ -92,7 +55,9 @@ static void debug_header(memory_header *hdr)
+ 
+ static void debug_chain(void)
+ {
+-    memory_header *cur = heap.first;
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
++    memory_header *cur = heap->first;
+ 
+     mbedtls_fprintf(stderr, "\nBlock list\n");
+     while (cur != NULL) {
+@@ -101,7 +66,7 @@ static void debug_chain(void)
+     }
+ 
+     mbedtls_fprintf(stderr, "Free list\n");
+-    cur = heap.first_free;
++    cur = heap->first_free;
+ 
+     while (cur != NULL) {
+         debug_header(cur);
+@@ -152,7 +117,9 @@ static int verify_header(memory_header *hdr)
+ 
+ static int verify_chain(void)
+ {
+-    memory_header *prv = heap.first, *cur;
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
++    memory_header *prv = heap->first, *cur;
+ 
+     if (prv == NULL || verify_header(prv) != 0) {
+ #if defined(MBEDTLS_MEMORY_DEBUG)
+@@ -162,7 +129,7 @@ static int verify_chain(void)
+         return 1;
+     }
+ 
+-    if (heap.first->prev != NULL) {
++    if (heap->first->prev != NULL) {
+ #if defined(MBEDTLS_MEMORY_DEBUG)
+         mbedtls_fprintf(stderr, "FATAL: verification failed: "
+                                 "first->prev != NULL\n");
+@@ -170,7 +137,7 @@ static int verify_chain(void)
+         return 1;
+     }
+ 
+-    cur = heap.first->next;
++    cur = heap->first->next;
+ 
+     while (cur != NULL) {
+         if (verify_header(cur) != 0) {
+@@ -198,7 +165,9 @@ static int verify_chain(void)
+ 
+ static void *buffer_alloc_calloc(size_t n, size_t size)
+ {
+-    memory_header *new, *cur = heap.first_free;
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
++    memory_header *new, *cur = heap->first_free;
+     unsigned char *p;
+     void *ret;
+     size_t original_len, len;
+@@ -207,7 +176,7 @@ static void *buffer_alloc_calloc(size_t n, size_t size)
+     size_t trace_cnt;
+ #endif
+ 
+-    if (heap.buf == NULL || heap.first == NULL) {
++    if (heap->buf == NULL || heap->first == NULL) {
+         return NULL;
+     }
+ 
+@@ -247,7 +216,7 @@ static void *buffer_alloc_calloc(size_t n, size_t size)
+     }
+ 
+ #if defined(MBEDTLS_MEMORY_DEBUG)
+-    heap.alloc_count++;
++    heap->alloc_count++;
+ #endif
+ 
+     // Found location, split block if > memory_header + 4 room left
+@@ -261,7 +230,7 @@ static void *buffer_alloc_calloc(size_t n, size_t size)
+         if (cur->prev_free != NULL) {
+             cur->prev_free->next_free = cur->next_free;
+         } else {
+-            heap.first_free = cur->next_free;
++            heap->first_free = cur->next_free;
+         }
+ 
+         if (cur->next_free != NULL) {
+@@ -272,9 +241,9 @@ static void *buffer_alloc_calloc(size_t n, size_t size)
+         cur->next_free = NULL;
+ 
+ #if defined(MBEDTLS_MEMORY_DEBUG)
+-        heap.total_used += cur->size;
+-        if (heap.total_used > heap.maximum_used) {
+-            heap.maximum_used = heap.total_used;
++        heap->total_used += cur->size;
++        if (heap->total_used > heap->maximum_used) {
++            heap->maximum_used = heap->total_used;
+         }
+ #endif
+ #if defined(MBEDTLS_MEMORY_BACKTRACE)
+@@ -283,7 +252,7 @@ static void *buffer_alloc_calloc(size_t n, size_t size)
+         cur->trace_count = trace_cnt;
+ #endif
+ 
+-        if ((heap.verify & MBEDTLS_MEMORY_VERIFY_ALLOC) && verify_chain() != 0) {
++        if ((heap->verify & MBEDTLS_MEMORY_VERIFY_ALLOC) && verify_chain() != 0) {
+             mbedtls_exit(1);
+         }
+ 
+@@ -318,7 +287,7 @@ static void *buffer_alloc_calloc(size_t n, size_t size)
+     if (new->prev_free != NULL) {
+         new->prev_free->next_free = new;
+     } else {
+-        heap.first_free = new;
++        heap->first_free = new;
+     }
+ 
+     if (new->next_free != NULL) {
+@@ -332,13 +301,13 @@ static void *buffer_alloc_calloc(size_t n, size_t size)
+     cur->next_free = NULL;
+ 
+ #if defined(MBEDTLS_MEMORY_DEBUG)
+-    heap.header_count++;
+-    if (heap.header_count > heap.maximum_header_count) {
+-        heap.maximum_header_count = heap.header_count;
++    heap->header_count++;
++    if (heap->header_count > heap->maximum_header_count) {
++        heap->maximum_header_count = heap->header_count;
+     }
+-    heap.total_used += cur->size;
+-    if (heap.total_used > heap.maximum_used) {
+-        heap.maximum_used = heap.total_used;
++    heap->total_used += cur->size;
++    if (heap->total_used > heap->maximum_used) {
++        heap->maximum_used = heap->total_used;
+     }
+ #endif
+ #if defined(MBEDTLS_MEMORY_BACKTRACE)
+@@ -347,7 +316,7 @@ static void *buffer_alloc_calloc(size_t n, size_t size)
+     cur->trace_count = trace_cnt;
+ #endif
+ 
+-    if ((heap.verify & MBEDTLS_MEMORY_VERIFY_ALLOC) && verify_chain() != 0) {
++    if ((heap->verify & MBEDTLS_MEMORY_VERIFY_ALLOC) && verify_chain() != 0) {
+         mbedtls_exit(1);
+     }
+ 
+@@ -359,14 +328,16 @@ static void *buffer_alloc_calloc(size_t n, size_t size)
+ 
+ static void buffer_alloc_free(void *ptr)
+ {
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
+     memory_header *hdr, *old = NULL;
+     unsigned char *p = (unsigned char *) ptr;
+ 
+-    if (ptr == NULL || heap.buf == NULL || heap.first == NULL) {
++    if (ptr == NULL || heap->buf == NULL || heap->first == NULL) {
+         return;
+     }
+ 
+-    if (p < heap.buf || p >= heap.buf + heap.len) {
++    if (p < heap->buf || p >= heap->buf + heap->len) {
+ #if defined(MBEDTLS_MEMORY_DEBUG)
+         mbedtls_fprintf(stderr, "FATAL: mbedtls_free() outside of managed "
+                                 "space\n");
+@@ -392,8 +363,8 @@ static void buffer_alloc_free(void *ptr)
+     hdr->alloc = 0;
+ 
+ #if defined(MBEDTLS_MEMORY_DEBUG)
+-    heap.free_count++;
+-    heap.total_used -= hdr->size;
++    heap->free_count++;
++    heap->total_used -= hdr->size;
+ #endif
+ 
+ #if defined(MBEDTLS_MEMORY_BACKTRACE)
+@@ -406,7 +377,7 @@ static void buffer_alloc_free(void *ptr)
+     //
+     if (hdr->prev != NULL && hdr->prev->alloc == 0) {
+ #if defined(MBEDTLS_MEMORY_DEBUG)
+-        heap.header_count--;
++        heap->header_count--;
+ #endif
+         hdr->prev->size += sizeof(memory_header) + hdr->size;
+         hdr->prev->next = hdr->next;
+@@ -424,7 +395,7 @@ static void buffer_alloc_free(void *ptr)
+     //
+     if (hdr->next != NULL && hdr->next->alloc == 0) {
+ #if defined(MBEDTLS_MEMORY_DEBUG)
+-        heap.header_count--;
++        heap->header_count--;
+ #endif
+         hdr->size += sizeof(memory_header) + hdr->next->size;
+         old = hdr->next;
+@@ -434,7 +405,7 @@ static void buffer_alloc_free(void *ptr)
+             if (hdr->prev_free != NULL) {
+                 hdr->prev_free->next_free = hdr->next_free;
+             } else {
+-                heap.first_free = hdr->next_free;
++                heap->first_free = hdr->next_free;
+             }
+ 
+             if (hdr->next_free != NULL) {
+@@ -448,7 +419,7 @@ static void buffer_alloc_free(void *ptr)
+         if (hdr->prev_free != NULL) {
+             hdr->prev_free->next_free = hdr;
+         } else {
+-            heap.first_free = hdr;
++            heap->first_free = hdr;
+         }
+ 
+         if (hdr->next_free != NULL) {
+@@ -466,21 +437,23 @@ static void buffer_alloc_free(void *ptr)
+     // (Does not have to stay in same order as prev / next list)
+     //
+     if (old == NULL) {
+-        hdr->next_free = heap.first_free;
+-        if (heap.first_free != NULL) {
+-            heap.first_free->prev_free = hdr;
++        hdr->next_free = heap->first_free;
++        if (heap->first_free != NULL) {
++            heap->first_free->prev_free = hdr;
+         }
+-        heap.first_free = hdr;
++        heap->first_free = hdr;
+     }
+ 
+-    if ((heap.verify & MBEDTLS_MEMORY_VERIFY_FREE) && verify_chain() != 0) {
++    if ((heap->verify & MBEDTLS_MEMORY_VERIFY_FREE) && verify_chain() != 0) {
+         mbedtls_exit(1);
+     }
+ }
+ 
+ void mbedtls_memory_buffer_set_verify(int verify)
+ {
+-    heap.verify = verify;
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
++    heap->verify = verify;
+ }
+ 
+ int mbedtls_memory_buffer_alloc_verify(void)
+@@ -491,16 +464,18 @@ int mbedtls_memory_buffer_alloc_verify(void)
+ #if defined(MBEDTLS_MEMORY_DEBUG)
+ void mbedtls_memory_buffer_alloc_status(void)
+ {
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
+     mbedtls_fprintf(stderr,
+                     "Current use: %zu blocks / %zu bytes, max: %zu blocks / "
+                     "%zu bytes (total %zu bytes), alloc / free: %zu / %zu\n",
+-                    heap.header_count, heap.total_used,
+-                    heap.maximum_header_count, heap.maximum_used,
+-                    heap.maximum_header_count * sizeof(memory_header)
+-                    + heap.maximum_used,
+-                    heap.alloc_count, heap.free_count);
++                    heap->header_count, heap->total_used,
++                    heap->maximum_header_count, heap->maximum_used,
++                    heap->maximum_header_count * sizeof(memory_header)
++                    + heap->maximum_used,
++                    heap->alloc_count, heap->free_count);
+ 
+-    if (heap.first->next == NULL) {
++    if (heap->first->next == NULL) {
+         mbedtls_fprintf(stderr, "All memory de-allocated in stack buffer\n");
+     } else {
+         mbedtls_fprintf(stderr, "Memory currently allocated:\n");
+@@ -510,38 +485,48 @@ void mbedtls_memory_buffer_alloc_status(void)
+ 
+ void mbedtls_memory_buffer_alloc_count_get(size_t *alloc_count, size_t *free_count)
+ {
+-    *alloc_count = heap.alloc_count;
+-    *free_count = heap.free_count;
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
++    *alloc_count = heap->alloc_count;
++    *free_count = heap->free_count;
+ }
+ 
+ void mbedtls_memory_buffer_alloc_max_get(size_t *max_used, size_t *max_blocks)
+ {
+-    *max_used   = heap.maximum_used;
+-    *max_blocks = heap.maximum_header_count;
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
++    *max_used   = heap->maximum_used;
++    *max_blocks = heap->maximum_header_count;
+ }
+ 
+ void mbedtls_memory_buffer_alloc_max_reset(void)
+ {
+-    heap.maximum_used = 0;
+-    heap.maximum_header_count = 0;
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
++    heap->maximum_used = 0;
++    heap->maximum_header_count = 0;
+ }
+ 
+ void mbedtls_memory_buffer_alloc_cur_get(size_t *cur_used, size_t *cur_blocks)
+ {
+-    *cur_used   = heap.total_used;
+-    *cur_blocks = heap.header_count;
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
++    *cur_used   = heap->total_used;
++    *cur_blocks = heap->header_count;
+ }
+ #endif /* MBEDTLS_MEMORY_DEBUG */
+ 
+ #if defined(MBEDTLS_THREADING_C)
+ static void *buffer_alloc_calloc_mutexed(size_t n, size_t size)
+ {
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
+     void *buf;
+-    if (mbedtls_mutex_lock(&heap.mutex) != 0) {
++    if (mbedtls_mutex_lock(&heap->mutex) != 0) {
+         return NULL;
+     }
+     buf = buffer_alloc_calloc(n, size);
+-    if (mbedtls_mutex_unlock(&heap.mutex)) {
++    if (mbedtls_mutex_unlock(&heap->mutex)) {
+         return NULL;
+     }
+     return buf;
+@@ -549,22 +534,26 @@ static void *buffer_alloc_calloc_mutexed(size_t n, size_t size)
+ 
+ static void buffer_alloc_free_mutexed(void *ptr)
+ {
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
+     /* We have no good option here, but corrupting the heap seems
+      * worse than losing memory. */
+-    if (mbedtls_mutex_lock(&heap.mutex)) {
++    if (mbedtls_mutex_lock(&heap->mutex)) {
+         return;
+     }
+     buffer_alloc_free(ptr);
+-    (void) mbedtls_mutex_unlock(&heap.mutex);
++    (void) mbedtls_mutex_unlock(&heap->mutex);
+ }
+ #endif /* MBEDTLS_THREADING_C */
+ 
+ void mbedtls_memory_buffer_alloc_init(unsigned char *buf, size_t len)
+ {
+-    memset(&heap, 0, sizeof(buffer_alloc_ctx));
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
++    memset(heap, 0, sizeof(buffer_alloc_ctx));
+ 
+ #if defined(MBEDTLS_THREADING_C)
+-    mbedtls_mutex_init(&heap.mutex);
++    mbedtls_mutex_init(&heap->mutex);
+     mbedtls_platform_set_calloc_free(buffer_alloc_calloc_mutexed,
+                                      buffer_alloc_free_mutexed);
+ #else
+@@ -583,22 +572,24 @@ void mbedtls_memory_buffer_alloc_init(unsigned char *buf, size_t len)
+ 
+     memset(buf, 0, len);
+ 
+-    heap.buf = buf;
+-    heap.len = len;
++    heap->buf = buf;
++    heap->len = len;
+ 
+-    heap.first = (memory_header *) buf;
+-    heap.first->size = len - sizeof(memory_header);
+-    heap.first->magic1 = MAGIC1;
+-    heap.first->magic2 = MAGIC2;
+-    heap.first_free = heap.first;
++    heap->first = (memory_header *) buf;
++    heap->first->size = len - sizeof(memory_header);
++    heap->first->magic1 = MAGIC1;
++    heap->first->magic2 = MAGIC2;
++    heap->first_free = heap->first;
+ }
+ 
+ void mbedtls_memory_buffer_alloc_free(void)
+ {
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
+ #if defined(MBEDTLS_THREADING_C)
+-    mbedtls_mutex_free(&heap.mutex);
++    mbedtls_mutex_free(&heap->mutex);
+ #endif
+-    mbedtls_platform_zeroize(&heap, sizeof(buffer_alloc_ctx));
++    mbedtls_platform_zeroize(heap, sizeof(buffer_alloc_ctx));
+ }
+ 
+ #if defined(MBEDTLS_SELF_TEST)
+@@ -617,12 +608,14 @@ static int check_pointer(void *p)
+ 
+ static int check_all_free(void)
+ {
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
+     if (
+ #if defined(MBEDTLS_MEMORY_DEBUG)
+-        heap.total_used != 0 ||
++        heap->total_used != 0 ||
+ #endif
+-        heap.first != heap.first_free ||
+-        (void *) heap.first != (void *) heap.buf) {
++        heap->first != heap->first_free ||
++        (void *) heap->first != (void *) heap->buf) {
+         return -1;
+     }
+ 
+@@ -641,6 +634,8 @@ static int check_all_free(void)
+ 
+ int mbedtls_memory_buffer_alloc_self_test(int verbose)
+ {
++    buffer_alloc_ctx *heap = mbedtls_app_get_heap();
++
+     unsigned char buf[1024];
+     unsigned char *p, *q, *r, *end;
+     int ret = 0;
+@@ -666,7 +661,7 @@ int mbedtls_memory_buffer_alloc_self_test(int verbose)
+     TEST_ASSERT(check_all_free() == 0);
+ 
+     /* Memorize end to compare with the next test */
+-    end = heap.buf + heap.len;
++    end = heap->buf + heap->len;
+ 
+     mbedtls_memory_buffer_alloc_free();
+ 
+@@ -680,7 +675,7 @@ int mbedtls_memory_buffer_alloc_self_test(int verbose)
+ 
+     mbedtls_memory_buffer_alloc_init(buf + 1, sizeof(buf) - 1);
+ 
+-    TEST_ASSERT(heap.buf + heap.len == end);
++    TEST_ASSERT(heap->buf + heap->len == end);
+ 
+     p = mbedtls_calloc(1, 1);
+     q = mbedtls_calloc(1, 128);
+-- 
+2.34.1
+
diff --git a/lib/mbedtls/CMakeLists.txt b/lib/mbedtls/CMakeLists.txt
index dbd4887..707757b 100644
--- a/lib/mbedtls/CMakeLists.txt
+++ b/lib/mbedtls/CMakeLists.txt
@@ -17,7 +17,8 @@
 set(MBEDTLS_PATCH_DIR "${RMM_SOURCE_DIR}/configs/mbedtls")
 set(MBEDTLS_PATCH_FILES
     "${MBEDTLS_PATCH_DIR}/0001-Remove-compiler-options-for-clang.patch"
-    "${MBEDTLS_PATCH_DIR}/0002-Disable-explicit_bzero.patch")
+    "${MBEDTLS_PATCH_DIR}/0002-Disable-explicit_bzero.patch"
+    "${MBEDTLS_PATCH_DIR}/0003-refactor-alloc-Use-the-heap-structure-provided-by-ap.patch")
 
 Git_Apply_Patches(${MBEDTLS_SRC_DIR} "${MBEDTLS_PATCH_FILES}")