From 48e59eaeb1ddbfbf474d85494460fbe6a61dfd39 Mon Sep 17 00:00:00 2001
From: Tesfa Mael <tesfa@wolfssl.com>
Date: Wed, 18 Dec 2019 17:39:53 -0800
Subject: [PATCH 1/3] Free EVP ctx pkey

---
 src/ssl.c           | 6 ++++--
 wolfcrypt/src/evp.c | 2 ++
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/ssl.c b/src/ssl.c
index fc2aa63bf..cdd515540 100644
--- a/src/ssl.c
+++ b/src/ssl.c
@@ -16132,7 +16132,7 @@ int wolfSSL_EVP_MD_type(const WOLFSSL_EVP_MD *md)
     {
         WOLFSSL_ENTER("EVP_MD_CTX_cleanup");
         if (ctx->pctx != NULL)
-            wolfSSL_EVP_PKEY_CTX_free(ctx->pctx);
+            XFREE(ctx->pctx, NULL, DYNAMIC_TYPE_PUBLIC_KEY);
 
         if (ctx->macType == (NID_hmac & 0xFF)) {
             wc_HmacFree(&ctx->hash.hmac);
@@ -23447,7 +23447,9 @@ void wolfSSL_EVP_PKEY_free(WOLFSSL_EVP_PKEY* key)
                 break;
             }
 
-            wc_FreeMutex(&key->refMutex);
+            if (wc_FreeMutex(&key->refMutex) != 0) {
+                WOLFSSL_MSG("Couldn't free pkey mutex");
+            }
             XFREE(key, key->heap, DYNAMIC_TYPE_PUBLIC_KEY);
         }
     }
diff --git a/wolfcrypt/src/evp.c b/wolfcrypt/src/evp.c
index dd8f857e3..df6fb050e 100644
--- a/wolfcrypt/src/evp.c
+++ b/wolfcrypt/src/evp.c
@@ -901,6 +901,8 @@ int wolfSSL_EVP_PKEY_CTX_free(WOLFSSL_EVP_PKEY_CTX *ctx)
 {
     if (ctx == NULL) return 0;
     WOLFSSL_ENTER("EVP_PKEY_CTX_free");
+    if (ctx->pkey != NULL)
+        wolfSSL_EVP_PKEY_free(ctx->pkey);
     XFREE(ctx, NULL, DYNAMIC_TYPE_PUBLIC_KEY);
     return WOLFSSL_SUCCESS;
 }

From 99a7aff31ee238106f605c7c20a1f132fbaa6247 Mon Sep 17 00:00:00 2001
From: Tesfa Mael <tesfa@wolfssl.com>
Date: Thu, 19 Dec 2019 13:49:46 -0800
Subject: [PATCH 2/3] Increment pkey references count

---
 src/ssl.c           |  2 +-
 wolfcrypt/src/evp.c | 12 ++++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/ssl.c b/src/ssl.c
index cdd515540..38d4fe484 100644
--- a/src/ssl.c
+++ b/src/ssl.c
@@ -16132,7 +16132,7 @@ int wolfSSL_EVP_MD_type(const WOLFSSL_EVP_MD *md)
     {
         WOLFSSL_ENTER("EVP_MD_CTX_cleanup");
         if (ctx->pctx != NULL)
-            XFREE(ctx->pctx, NULL, DYNAMIC_TYPE_PUBLIC_KEY);
+            wolfSSL_EVP_PKEY_CTX_free(ctx->pctx);
 
         if (ctx->macType == (NID_hmac & 0xFF)) {
             wc_HmacFree(&ctx->hash.hmac);
diff --git a/wolfcrypt/src/evp.c b/wolfcrypt/src/evp.c
index df6fb050e..288abbc0c 100644
--- a/wolfcrypt/src/evp.c
+++ b/wolfcrypt/src/evp.c
@@ -918,6 +918,7 @@ int wolfSSL_EVP_PKEY_CTX_free(WOLFSSL_EVP_PKEY_CTX *ctx)
 WOLFSSL_EVP_PKEY_CTX *wolfSSL_EVP_PKEY_CTX_new(WOLFSSL_EVP_PKEY *pkey, WOLFSSL_ENGINE *e)
 {
     WOLFSSL_EVP_PKEY_CTX* ctx;
+    int type = NID_undef;
 
     if (pkey == NULL) return 0;
     if (e != NULL) return 0;
@@ -931,7 +932,18 @@ WOLFSSL_EVP_PKEY_CTX *wolfSSL_EVP_PKEY_CTX_new(WOLFSSL_EVP_PKEY *pkey, WOLFSSL_E
 #if !defined(NO_RSA) && !defined(HAVE_USER_RSA)
     ctx->padding = RSA_PKCS1_PADDING;
 #endif
+    type = wolfSSL_EVP_PKEY_type(pkey->type);
 
+    if ((type == EVP_PKEY_RSA) ||
+        (type == EVP_PKEY_DSA) ||
+        (type == EVP_PKEY_EC)) {
+        if (wc_LockMutex(&pkey->refMutex) != 0) {
+            WOLFSSL_MSG("Couldn't lock pkey mutex");
+        }
+        pkey->references++;
+
+        wc_UnLockMutex(&pkey->refMutex);
+    }
     return ctx;
 }
 

From 4004963c6ac12da90deda1465f8f2b751f7ef2ae Mon Sep 17 00:00:00 2001
From: Tesfa Mael <tesfa@wolfssl.com>
Date: Mon, 30 Dec 2019 09:31:23 -0800
Subject: [PATCH 3/3] test pkey references count

---
 tests/api.c | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/tests/api.c b/tests/api.c
index 6d23cac3e..cae2578c4 100644
--- a/tests/api.c
+++ b/tests/api.c
@@ -26226,9 +26226,16 @@ static void test_wolfSSL_EVP_PKEY_encrypt(void)
     AssertIntEQ(EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_PADDING),
                 WOLFSSL_SUCCESS);
 
+    /* Test pkey references count is decremented. pkey shouldn't be destroyed
+     since ctx uses it.*/
+    AssertIntEQ(pkey->references, 2);
+    EVP_PKEY_free(pkey);
+    AssertIntEQ(pkey->references, 1);
+
     /* Encrypt data */
     AssertIntEQ(EVP_PKEY_encrypt(ctx, outEnc, &outEncLen,
                             (const unsigned char*)in, inlen), WOLFSSL_SUCCESS);
+
     /* Decrypt data */
     AssertIntEQ(EVP_PKEY_decrypt_init(ctx), WOLFSSL_SUCCESS);
 
@@ -26255,8 +26262,6 @@ static void test_wolfSSL_EVP_PKEY_encrypt(void)
                                  WOLFSSL_SUCCESS);
     AssertIntEQ(XMEMCMP(inTmp, outDecTmp, outDecLen), 0);
 #endif
-
-    EVP_PKEY_free(pkey);
     EVP_PKEY_CTX_free(ctx);
     XFREE(outEnc, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);
     XFREE(outDec, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER);