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

