From 890eb415b1e326f7d0ac7974e1f4dc585fd38c9e Mon Sep 17 00:00:00 2001
From: Jacob Barthelmeh <jacob@wolfssl.com>
Date: Tue, 5 Nov 2019 15:13:26 -0700
Subject: [PATCH] fix for memory management on edge case with staticmemory

---
 examples/client/client.c | 11 ++++++++---
 src/ssl.c                | 11 +++++++++--
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/examples/client/client.c b/examples/client/client.c
index 0127aa61b..b2caca467 100644
--- a/examples/client/client.c
+++ b/examples/client/client.c
@@ -1470,6 +1470,7 @@ THREAD_RETURN WOLFSSL_THREAD client_test(void* args)
     #ifdef DEBUG_WOLFSSL
         WOLFSSL_MEM_STATS mem_stats;
     #endif
+    WOLFSSL_HEAP_HINT *heap = NULL;
 #endif
 
     ((func_args*)args)->return_code = -1; /* error state */
@@ -2162,20 +2163,24 @@ THREAD_RETURN WOLFSSL_THREAD client_test(void* args)
                                                   WOLFMEM_IO_POOL_FIXED));
     #endif /* DEBUG_WOLFSSL */
 
-    if (wolfSSL_CTX_load_static_memory(&ctx, method, memory, sizeof(memory),
-                                                     0, 1) != WOLFSSL_SUCCESS) {
+    if (wc_LoadStaticMemory(&heap, memory, sizeof(memory), WOLFMEM_GENERAL, 1)
+            != 0) {
         err_sys("unable to load static memory");
     }
 
+    ctx = wolfSSL_CTX_new_ex(method(heap), heap);
+    if (ctx == NULL)
+        err_sys("unable to get ctx");
+
     if (wolfSSL_CTX_load_static_memory(&ctx, NULL, memoryIO, sizeof(memoryIO),
            WOLFMEM_IO_POOL_FIXED | WOLFMEM_TRACK_STATS, 1) != WOLFSSL_SUCCESS) {
         err_sys("unable to load static memory");
     }
 #else
     ctx = wolfSSL_CTX_new(method(NULL));
-#endif
     if (ctx == NULL)
         err_sys("unable to get ctx");
+#endif
 
 #ifdef SINGLE_THREADED
     if (wolfSSL_CTX_new_rng(ctx) != WOLFSSL_SUCCESS) {
diff --git a/src/ssl.c b/src/ssl.c
index b31fedde4..bdd71e9a2 100644
--- a/src/ssl.c
+++ b/src/ssl.c
@@ -336,7 +336,15 @@ WOLFSSL_CTX* wolfSSL_CTX_new_ex(WOLFSSL_METHOD* method, void* heap)
 
     ctx = (WOLFSSL_CTX*) XMALLOC(sizeof(WOLFSSL_CTX), heap, DYNAMIC_TYPE_CTX);
     if (ctx) {
-        if (InitSSL_Ctx(ctx, method, heap) < 0) {
+        int ret;
+
+        ret = InitSSL_Ctx(ctx, method, heap);
+    #ifdef WOLFSSL_STATIC_MEMORY
+        if (heap != NULL) {
+            ctx->onHeap = 1; /* free the memory back to heap when done */
+        }
+    #endif
+        if (ret < 0) {
             WOLFSSL_MSG("Init CTX failed");
             wolfSSL_CTX_free(ctx);
             ctx = NULL;
@@ -1440,7 +1448,6 @@ int wolfSSL_CTX_load_static_memory(WOLFSSL_CTX** ctx, wolfSSL_method_func method
             WOLFSSL_MSG("Error creating ctx");
             return WOLFSSL_FAILURE;
         }
-        (*ctx)->onHeap = 1; /* free the memory back to heap when done */
     }
 
     /* determine what max applies too */