From fe87f161145c8ebe281a5874e76e78f22e0f5be4 Mon Sep 17 00:00:00 2001
From: Anthony Hu <anthony@wolfssl.com>
Date: Fri, 26 Jan 2024 14:51:51 -0500
Subject: [PATCH] Fixes that prevent memory leaks when using OQS.

Fixes ZD 17177.
---
 src/tls.c                              | 28 +++++++++++++++++++++-----
 wolfcrypt/src/port/liboqs/liboqs.c     |  5 +++++
 wolfcrypt/src/wc_port.c                |  4 ++++
 wolfssl/internal.h                     |  1 +
 wolfssl/wolfcrypt/port/liboqs/liboqs.h |  2 ++
 5 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/src/tls.c b/src/tls.c
index 04794ae944..c32aa31062 100644
--- a/src/tls.c
+++ b/src/tls.c
@@ -7658,6 +7658,13 @@ static int TLSX_KeyShare_GenPqcKey(WOLFSSL *ssl, KeyShareEntry* kse)
     word32 privSz = 0;
     word32 pubSz = 0;
 
+    /* This gets called twice. Once during parsing of the key share and once
+     * during the population of the extension. No need to do work the second
+     * time. Just return success if its already been done. */
+    if (kse->pubKey != NULL) {
+        return ret;
+    }
+
     findEccPqc(&ecc_group, &oqs_group, kse->group);
     ret = kyber_id2type(oqs_group, &type);
     if (ret == NOT_COMPILED_IN) {
@@ -7735,10 +7742,11 @@ static int TLSX_KeyShare_GenPqcKey(WOLFSSL *ssl, KeyShareEntry* kse)
 
         /* Note we are saving the OQS private key and ECC private key
          * separately. That's because the ECC private key is not simply a
-         * buffer. Its is an ecc_key struct.
-         */
+         * buffer. Its is an ecc_key struct. Typically do not need the private
+         * key size, but will need to zero it out upon freeing. */
         kse->privKey = privKey;
         privKey = NULL;
+        kse->privKeyLen = privSz;
 
         kse->key = ecc_kse->key;
         ecc_kse->key = NULL;
@@ -7814,9 +7822,19 @@ static void TLSX_KeyShare_FreeAll(KeyShareEntry* list, void* heap)
 #endif
         }
 #ifdef HAVE_PQC
-        else if (WOLFSSL_NAMED_GROUP_IS_PQC(current->group) &&
-                 current->key != NULL) {
-            ForceZero((byte*)current->key, current->keyLen);
+        else if (WOLFSSL_NAMED_GROUP_IS_PQC(current->group)) {
+            if (current->key != NULL) {
+                ForceZero((byte*)current->key, current->keyLen);
+            }
+            if (current->pubKey != NULL) {
+                XFREE(current->pubKey, heap, DYNAMIC_TYPE_PUBLIC_KEY);
+                current->pubKey = NULL;
+            }
+            if (current->privKey != NULL) {
+                ForceZero(current->privKey, current->privKeyLen);
+                XFREE(current->privKey, heap, DYNAMIC_TYPE_PRIVATE_KEY);
+                current->privKey = NULL;
+            }
         }
 #endif
         else {
diff --git a/wolfcrypt/src/port/liboqs/liboqs.c b/wolfcrypt/src/port/liboqs/liboqs.c
index 2568330529..f04cab02c7 100644
--- a/wolfcrypt/src/port/liboqs/liboqs.c
+++ b/wolfcrypt/src/port/liboqs/liboqs.c
@@ -99,6 +99,11 @@ int wolfSSL_liboqsInit(void)
     return ret;
 }
 
+void wolfSSL_liboqsClose(void)
+{
+    wc_FreeRng(&liboqsDefaultRNG);
+}
+
 int wolfSSL_liboqsRngMutexLock(WC_RNG* rng)
 {
     int ret = wolfSSL_liboqsInit();
diff --git a/wolfcrypt/src/wc_port.c b/wolfcrypt/src/wc_port.c
index a6ebe7e91c..8cb8a0c189 100644
--- a/wolfcrypt/src/wc_port.c
+++ b/wolfcrypt/src/wc_port.c
@@ -493,6 +493,10 @@ int wolfCrypt_Cleanup(void)
     #endif
     }
 
+#if defined(HAVE_LIBOQS)
+    wolfSSL_liboqsClose();
+#endif
+
     return ret;
 }
 
diff --git a/wolfssl/internal.h b/wolfssl/internal.h
index 5baeb93b0f..572114e52a 100644
--- a/wolfssl/internal.h
+++ b/wolfssl/internal.h
@@ -3365,6 +3365,7 @@ typedef struct KeyShareEntry {
     word32                pubKeyLen; /* Public key length                 */
 #if !defined(NO_DH) || defined(HAVE_PQC)
     byte*                 privKey;   /* Private key - DH and PQ KEMs only */
+    word32                privKeyLen;/* Only for PQ KEMs. */
 #endif
 #ifdef WOLFSSL_ASYNC_CRYPT
     int                   lastRet;
diff --git a/wolfssl/wolfcrypt/port/liboqs/liboqs.h b/wolfssl/wolfcrypt/port/liboqs/liboqs.h
index d93ec2f2b7..58da9ba2be 100644
--- a/wolfssl/wolfcrypt/port/liboqs/liboqs.h
+++ b/wolfssl/wolfcrypt/port/liboqs/liboqs.h
@@ -47,6 +47,8 @@ implementations for Post-Quantum cryptography algorithms.
 
 int wolfSSL_liboqsInit(void);
 
+void wolfSSL_liboqsClose(void);
+
 int wolfSSL_liboqsRngMutexLock(WC_RNG* rng);
 
 int wolfSSL_liboqsRngMutexUnlock(void);