Merge pull request wolfSSL#7478 from JacobBarthelmeh/staticmemory

add global heap hint setter function
gojimmypi · May 15, 2024 · ac7aea9 · ac7aea9
2 parents c73e433 + 9aeef1d
commit ac7aea9
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 6 deletions.
diff --git a/wolfcrypt/src/memory.c b/wolfcrypt/src/memory.c
@@ -899,6 +899,31 @@ int wolfSSL_GetMemStats(WOLFSSL_HEAP* heap, WOLFSSL_MEM_STATS* stats)
 }
 
 
+/* global heap hint to fall back on when no heap hint is passed to
+ * XMALLOC/XFREE
+ * NOT thread safe, should be set once before any expected XMALLOC XFREE calls
+ */
+static void* globalHeapHint = NULL;
+
+
+/* Used to set a new global heap hint. Returns a pointer to the current global
+ * heap hint before being set. */
+void* wolfSSL_SetGlobalHeapHint(void* heap)
+{
+    void *oldHint = globalHeapHint;
+
+    globalHeapHint = heap;
+    return oldHint;
+}
+
+
+/* returns a pointer to the current global heap hint */
+void* wolfSSL_GetGlobalHeapHint(void)
+{
+    return globalHeapHint;
+}
+
+
 #ifdef WOLFSSL_DEBUG_MEMORY
 void* wolfSSL_Malloc(size_t size, void* heap, int type, const char* func, unsigned int line)
 #else
@@ -917,7 +942,7 @@ void* wolfSSL_Malloc(size_t size, void* heap, int type)
 #endif
 
     /* if no heap hint then use dynamic memory*/
-    if (heap == NULL) {
+    if (heap == NULL && globalHeapHint == NULL) {
         #ifdef WOLFSSL_HEAP_TEST
             /* allow using malloc for creating ctx and method */
             if (type == DYNAMIC_TYPE_CTX || type == DYNAMIC_TYPE_METHOD ||
@@ -952,7 +977,15 @@ void* wolfSSL_Malloc(size_t size, void* heap, int type)
     }
     else {
         WOLFSSL_HEAP_HINT* hint = (WOLFSSL_HEAP_HINT*)heap;
-        WOLFSSL_HEAP*      mem  = hint->memory;
+        WOLFSSL_HEAP*      mem;
+
+        if (hint == NULL) {
+            hint = (WOLFSSL_HEAP_HINT*)globalHeapHint;
+        #ifdef WOLFSSL_DEBUG_MEMORY
+            fprintf(stderr, "(Using global heap hint %p) ", hint);
+        #endif
+        }
+        mem = hint->memory;
 
         if (wc_LockMutex(&(mem->memory_mutex)) != 0) {
             WOLFSSL_MSG("Bad memory_mutex lock");
@@ -1073,7 +1106,7 @@ void wolfSSL_Free(void *ptr, void* heap, int type)
         }
     #endif
 
-        if (heap == NULL) {
+        if (heap == NULL && globalHeapHint == NULL) {
         #ifdef WOLFSSL_HEAP_TEST
             /* allow using malloc for creating ctx and method */
             if (type == DYNAMIC_TYPE_CTX || type == DYNAMIC_TYPE_METHOD ||
@@ -1098,9 +1131,17 @@ void wolfSSL_Free(void *ptr, void* heap, int type)
         }
         else {
             WOLFSSL_HEAP_HINT* hint = (WOLFSSL_HEAP_HINT*)heap;
-            WOLFSSL_HEAP*      mem  = hint->memory;
+            WOLFSSL_HEAP*      mem;
             word32 padSz = -(int)sizeof(wc_Memory) & (WOLFSSL_STATIC_ALIGN - 1);
 
+            if (hint == NULL) {
+                hint = (WOLFSSL_HEAP_HINT*)globalHeapHint;
+            #ifdef WOLFSSL_DEBUG_MEMORY
+                fprintf(stderr, "(Using global heap hint %p) ", hint);
+            #endif
+            }
+            mem = hint->memory;
+
             /* get memory struct and add it to available list */
             pt = (wc_Memory*)((byte*)ptr - sizeof(wc_Memory) - padSz);
             if (wc_LockMutex(&(mem->memory_mutex)) != 0) {
@@ -1181,7 +1222,7 @@ void* wolfSSL_Realloc(void *ptr, size_t size, void* heap, int type)
     }
 #endif
 
-    if (heap == NULL) {
+    if (heap == NULL && globalHeapHint == NULL) {
         #ifdef WOLFSSL_HEAP_TEST
             WOLFSSL_MSG("ERROR null heap hint passed in to XREALLOC");
         #endif
@@ -1193,9 +1234,17 @@ void* wolfSSL_Realloc(void *ptr, size_t size, void* heap, int type)
     }
     else {
         WOLFSSL_HEAP_HINT* hint = (WOLFSSL_HEAP_HINT*)heap;
-        WOLFSSL_HEAP*      mem  = hint->memory;
+        WOLFSSL_HEAP*      mem;
         word32 padSz = -(int)sizeof(wc_Memory) & (WOLFSSL_STATIC_ALIGN - 1);
 
+        if (hint == NULL) {
+            hint = (WOLFSSL_HEAP_HINT*)globalHeapHint;
+        #ifdef WOLFSSL_DEBUG_MEMORY
+            fprintf(stderr, "(Using global heap hint %p) ", hint);
+        #endif
+        }
+        mem = hint->memory;
+
         if (ptr == NULL) {
         #ifdef WOLFSSL_DEBUG_MEMORY
             return wolfSSL_Malloc(size, heap, type, func, line);

diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c
@@ -1059,6 +1059,9 @@ options: [-s max_relative_stack_bytes] [-m max_relative_heap_memory_bytes]\n\
         printf("unable to load static memory.\n");
         return(EXIT_FAILURE);
     }
+    #ifndef OPENSSL_EXTRA
+    wolfSSL_SetGlobalHeapHint(HEAP_HINT);
+    #endif
 #endif
 
 #if defined(DEBUG_WOLFSSL) && !defined(HAVE_VALGRIND)
@@ -2013,6 +2016,9 @@ options: [-s max_relative_stack_bytes] [-m max_relative_heap_memory_bytes]\n\
     #endif
 #endif
 
+#if defined(WOLFSSL_STATIC_MEMORY) && !defined(OPENSSL_EXTRA)
+    wolfSSL_SetGlobalHeapHint(NULL);
+#endif
     TEST_PASS("Test complete\n");
 
     EXIT_TEST(ret);

diff --git a/wolfssl/wolfcrypt/memory.h b/wolfssl/wolfcrypt/memory.h
@@ -240,6 +240,8 @@ WOLFSSL_API int wolfSSL_GetAllocators(wolfSSL_Malloc_cb* mf,
         byte        haFlag; /* flag used for checking handshake count */
     } WOLFSSL_HEAP_HINT;
 
+    WOLFSSL_API void* wolfSSL_SetGlobalHeapHint(void* heap);
+    WOLFSSL_API void* wolfSSL_GetGlobalHeapHint(void);
     WOLFSSL_API int wc_LoadStaticMemory_ex(WOLFSSL_HEAP_HINT** pHint,
             unsigned int listSz, const unsigned int *sizeList,
             const unsigned int *distList, unsigned char* buf, unsigned int sz,