From 50372b7033f475aa341ab557ba4add67c461d418 Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Tue, 31 Jul 2018 09:24:44 -0700
Subject: [PATCH] Memory Callback 1. Allow SetAllocators to use NULL as a
 callback. Memory wrappers will use default with NULL. 2. Remove the
 ResetAllocators function. Use SetAllocators with NULL. 3. Modify memory
 tracker to save the old allocators on init and restore them on cleanup.

---
 wolfcrypt/src/memory.c        | 30 +++---------------------------
 wolfcrypt/test/test.c         | 14 --------------
 wolfssl/wolfcrypt/mem_track.h | 14 ++++++++++++--
 wolfssl/wolfcrypt/memory.h    |  1 -
 4 files changed, 15 insertions(+), 44 deletions(-)

diff --git a/wolfcrypt/src/memory.c b/wolfcrypt/src/memory.c
index a586078ef..22ec298cd 100644
--- a/wolfcrypt/src/memory.c
+++ b/wolfcrypt/src/memory.c
@@ -66,33 +66,9 @@ int wolfSSL_SetAllocators(wolfSSL_Malloc_cb  mf,
                           wolfSSL_Free_cb    ff,
                           wolfSSL_Realloc_cb rf)
 {
-    int res = 0;
-
-    if (mf)
-        malloc_function = mf;
-    else
-        res = BAD_FUNC_ARG;
-
-    if (ff)
-        free_function = ff;
-    else
-        res = BAD_FUNC_ARG;
-
-    if (rf)
-        realloc_function = rf;
-    else
-        res = BAD_FUNC_ARG;
-
-    return res;
-}
-
-int wolfSSL_ResetAllocators(void)
-{
-    /* allow nulls to be set for callbacks to restore defaults */
-    malloc_function = NULL;
-    free_function = NULL;
-    realloc_function = NULL;
-
+    malloc_function = mf;
+    free_function = ff;
+    realloc_function = rf;
     return 0;
 }
 
diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c
index 941bec873..94dd18b64 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -19585,20 +19585,6 @@ int memcb_test(void)
     XFREE(b, NULL, DYNAMIC_TYPE_TMP_BUFFER);
     b = NULL;
 
-    /* Parameter Validation testing. */
-    if (wolfSSL_SetAllocators(NULL, (wolfSSL_Free_cb)&my_Free_cb,
-            (wolfSSL_Realloc_cb)&my_Realloc_cb) != BAD_FUNC_ARG) {
-        ERROR_OUT(-10002, exit_memcb);
-    }
-    if (wolfSSL_SetAllocators((wolfSSL_Malloc_cb)&my_Malloc_cb, NULL,
-            (wolfSSL_Realloc_cb)&my_Realloc_cb) != BAD_FUNC_ARG) {
-        ERROR_OUT(-10003, exit_memcb);
-    }
-    if (wolfSSL_SetAllocators((wolfSSL_Malloc_cb)&my_Malloc_cb,
-            (wolfSSL_Free_cb)&my_Free_cb, NULL) != BAD_FUNC_ARG) {
-        ERROR_OUT(-10004, exit_memcb);
-    }
-
     /* Use API. */
     if (wolfSSL_SetAllocators((wolfSSL_Malloc_cb)&my_Malloc_cb,
         (wolfSSL_Free_cb)&my_Free_cb, (wolfSSL_Realloc_cb)my_Realloc_cb) != 0) {
diff --git a/wolfssl/wolfcrypt/mem_track.h b/wolfssl/wolfcrypt/mem_track.h
index 8d22209b8..2fa50a5b8 100644
--- a/wolfssl/wolfcrypt/mem_track.h
+++ b/wolfssl/wolfcrypt/mem_track.h
@@ -308,9 +308,19 @@
     }
 
 #ifdef WOLFSSL_TRACK_MEMORY
+    static wolfSSL_Malloc_cb mfDefault = NULL;
+    static wolfSSL_Free_cb ffDefault = NULL;
+    static wolfSSL_Realloc_cb rfDefault = NULL;
+
     STATIC WC_INLINE int InitMemoryTracker(void)
     {
-        int ret = wolfSSL_SetAllocators(TrackMalloc, TrackFree, TrackRealloc);
+        int ret;
+
+        ret = wolfSSL_GetAllocators(&mfDefault, &ffDefault, &rfDefault);
+        if (ret < 0) {
+            printf("wolfSSL GetAllocators failed to get the defaults\n");
+        }
+        ret = wolfSSL_SetAllocators(TrackMalloc, TrackFree, TrackRealloc);
         if (ret < 0) {
             printf("wolfSSL SetAllocators failed for track memory\n");
             return ret;
@@ -380,7 +390,7 @@
     STATIC WC_INLINE int CleanupMemoryTracker(void)
     {
         /* restore default allocators */
-        return wolfSSL_ResetAllocators();
+        return wolfSSL_SetAllocators(mfDefault, ffDefault, rfDefault);
     }
 #endif
 
diff --git a/wolfssl/wolfcrypt/memory.h b/wolfssl/wolfcrypt/memory.h
index c69c093a4..a30978187 100644
--- a/wolfssl/wolfcrypt/memory.h
+++ b/wolfssl/wolfcrypt/memory.h
@@ -81,7 +81,6 @@
 WOLFSSL_API int wolfSSL_SetAllocators(wolfSSL_Malloc_cb,
                                       wolfSSL_Free_cb,
                                       wolfSSL_Realloc_cb);
-WOLFSSL_API int wolfSSL_ResetAllocators(void);
 WOLFSSL_API int wolfSSL_GetAllocators(wolfSSL_Malloc_cb*,
                                       wolfSSL_Free_cb*,
                                       wolfSSL_Realloc_cb*);