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*);