From 15a0ae4244602f6fcc44c54a2ad4ed58c5b47178 Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Thu, 11 Apr 2024 13:59:28 -0700
Subject: [PATCH] Generic Memory Pools

1. Add generic function wolfSSL_StaticBufferSz_ex() where one specifies
   the memory bucket list sizes and distribution.
2. Rewrote wolfSSL_StaticBufferSz() in terms of the new function.
3. Changed the list pointers on wc_LoadStaticMemory_ex() and
   wc_init_memory_heap() to be pointers to const.
---
 wolfcrypt/src/memory.c     | 40 +++++++++++++++++++++++++-------------
 wolfssl/wolfcrypt/memory.h |  8 ++++++--
 2 files changed, 32 insertions(+), 16 deletions(-)

diff --git a/wolfcrypt/src/memory.c b/wolfcrypt/src/memory.c
index 5e37f6db1..872d6512f 100644
--- a/wolfcrypt/src/memory.c
+++ b/wolfcrypt/src/memory.c
@@ -633,7 +633,7 @@ static int wc_partition_static_memory(byte* buffer, word32 sz, int flag,
 }
 
 static int wc_init_memory_heap(WOLFSSL_HEAP* heap, unsigned int listSz,
-        unsigned int* sizeList, unsigned int* distList)
+        const unsigned int* sizeList, const unsigned int* distList)
 {
     if (heap == NULL || listSz > WOLFMEM_MAX_BUCKETS) {
         return BAD_FUNC_ARG;
@@ -653,9 +653,9 @@ static int wc_init_memory_heap(WOLFSSL_HEAP* heap, unsigned int listSz,
 }
 
 int wc_LoadStaticMemory_ex(WOLFSSL_HEAP_HINT** pHint,
-        unsigned int listSz,
-        unsigned int* memSzList, unsigned int* memDistList,
-        unsigned char* buf, unsigned int sz, int flag, int maxSz)
+        unsigned int listSz, const unsigned int* memSzList,
+        const unsigned int* memDistList, unsigned char* buf,
+        unsigned int sz, int flag, int maxSz)
 {
     WOLFSSL_HEAP*      heap = NULL;
     WOLFSSL_HEAP_HINT* hint = NULL;
@@ -743,17 +743,16 @@ int wolfSSL_MemoryPaddingSz(void)
 
 /* Used to calculate memory size for optimum use with buckets.
    returns the suggested size rounded down to the nearest bucket. */
-int wolfSSL_StaticBufferSz(byte* buffer, word32 sz, int flag)
+int wolfSSL_StaticBufferSz_ex(unsigned int listSz,
+        const unsigned int *sizeList, const unsigned int *distList,
+        byte* buffer, word32 sz, int flag)
 {
-    word32 bucketSz[WOLFMEM_MAX_BUCKETS] = {WOLFMEM_BUCKETS};
-    word32 distList[WOLFMEM_MAX_BUCKETS] = {WOLFMEM_DIST};
-
     word32 ava = sz;
     byte*  pt  = buffer;
     word32 memSz = (word32)sizeof(wc_Memory);
     word32 padSz = -(int)memSz & (WOLFSSL_STATIC_ALIGN - 1);
 
-    WOLFSSL_ENTER("wolfSSL_static_size");
+    WOLFSSL_ENTER("wolfSSL_StaticBufferSz_ex");
 
     if (buffer == NULL) {
         return BAD_FUNC_ARG;
@@ -776,26 +775,39 @@ int wolfSSL_StaticBufferSz(byte* buffer, word32 sz, int flag)
     else {
         int i, k;
 
-        if (ava < (bucketSz[0] + padSz + memSz)) {
+        if (ava < (sizeList[0] + padSz + memSz)) {
             return 0; /* not enough room for even one bucket */
         }
 
-        while ((ava >= (bucketSz[0] + padSz + memSz)) && (ava > 0)) {
+        while ((ava >= (sizeList[0] + padSz + memSz)) && (ava > 0)) {
             /* start at largest and move to smaller buckets */
-            for (i = (WOLFMEM_MAX_BUCKETS - 1); i >= 0; i--) {
+            for (i = (listSz- 1); i >= 0; i--) {
                 for (k = distList[i]; k > 0; k--) {
-                    if ((bucketSz[i] + padSz + memSz) <= ava) {
-                        ava -= bucketSz[i] + padSz + memSz;
+                    if ((sizeList[i] + padSz + memSz) <= ava) {
+                        ava -= sizeList[i] + padSz + memSz;
                     }
                 }
             }
         }
     }
 
+    WOLFSSL_LEAVE("wolfSSL_StaticBufferSz_ex", sz - ava);
     return sz - ava; /* round down */
 }
 
 
+/* Calls wolfSSL_StaticBufferSz_ex with the static memory pool config
+ * used by wolfSSL by default. */
+int wolfSSL_StaticBufferSz(byte* buffer, word32 sz, int flag)
+{
+    word32 bucketSz[WOLFMEM_MAX_BUCKETS] = {WOLFMEM_BUCKETS};
+    word32 distList[WOLFMEM_MAX_BUCKETS] = {WOLFMEM_DIST};
+
+    return wolfSSL_StaticBufferSz_ex(WOLFMEM_MAX_BUCKETS, bucketSz, distList,
+        buffer, sz, flag);
+}
+
+
 int FreeFixedIO(WOLFSSL_HEAP* heap, wc_Memory** io)
 {
     WOLFSSL_MSG("Freeing fixed IO buffer");
diff --git a/wolfssl/wolfcrypt/memory.h b/wolfssl/wolfcrypt/memory.h
index 9e3aec05c..f4c8248c2 100644
--- a/wolfssl/wolfcrypt/memory.h
+++ b/wolfssl/wolfcrypt/memory.h
@@ -217,8 +217,9 @@ WOLFSSL_API int wolfSSL_GetAllocators(wolfSSL_Malloc_cb* mf,
     } WOLFSSL_HEAP_HINT;
 
     WOLFSSL_API int wc_LoadStaticMemory_ex(WOLFSSL_HEAP_HINT** pHint,
-            unsigned int listSz, unsigned int *sizeList, unsigned int *distList,
-            unsigned char* buf, unsigned int sz, int flag, int max);
+            unsigned int listSz, const unsigned int *sizeList,
+            const unsigned int *distList, unsigned char* buf, unsigned int sz,
+            int flag, int max);
     WOLFSSL_API int wc_LoadStaticMemory(WOLFSSL_HEAP_HINT** pHint,
             unsigned char* buf, unsigned int sz, int flag, int max);
 
@@ -227,6 +228,9 @@ WOLFSSL_API int wolfSSL_GetAllocators(wolfSSL_Malloc_cb* mf,
     WOLFSSL_LOCAL int SetFixedIO(WOLFSSL_HEAP* heap, wc_Memory** io);
     WOLFSSL_LOCAL int FreeFixedIO(WOLFSSL_HEAP* heap, wc_Memory** io);
 
+    WOLFSSL_API int wolfSSL_StaticBufferSz_ex(unsigned int listSz,
+            const unsigned int *sizeList, const unsigned int *distList,
+            byte* buffer, word32 sz, int flag);
     WOLFSSL_API int wolfSSL_StaticBufferSz(byte* buffer, word32 sz, int flag);
     WOLFSSL_API int wolfSSL_MemoryPaddingSz(void);
 #endif /* WOLFSSL_STATIC_MEMORY */