From 38d191c159ebea44d98a4fddf874e5477590cb5f Mon Sep 17 00:00:00 2001
From: Andras Fekete <andras@wolfssl.com>
Date: Wed, 7 Aug 2024 16:56:57 -0400
Subject: [PATCH] More PR comment fixes

---
 wolfcrypt/src/dilithium.c |  6 ++++--
 wolfcrypt/src/rsa.c       | 12 +++++++-----
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/wolfcrypt/src/dilithium.c b/wolfcrypt/src/dilithium.c
index ad6a73f0d7..946289cebd 100644
--- a/wolfcrypt/src/dilithium.c
+++ b/wolfcrypt/src/dilithium.c
@@ -7934,8 +7934,10 @@ int wc_dilithium_check_key(dilithium_key* key)
         }
     }
 
-    /* Dispose of allocated memory. */
-    XFREE(s1, key->heap, DYNAMIC_TYPE_DILITHIUM);
+    if (key != NULL) {
+        /* Dispose of allocated memory. */
+        XFREE(s1, key->heap, DYNAMIC_TYPE_DILITHIUM);
+    }
 #else
     /* Validate parameter. */
     if (key == NULL) {
diff --git a/wolfcrypt/src/rsa.c b/wolfcrypt/src/rsa.c
index db458df0f9..4c7d3a0e89 100644
--- a/wolfcrypt/src/rsa.c
+++ b/wolfcrypt/src/rsa.c
@@ -5081,11 +5081,13 @@ int wc_MakeRsaKey(RsaKey* key, int size, long e, WC_RNG* rng)
 
 #if !defined(WOLFSSL_CRYPTOCELL) && !defined(WOLFSSL_SE050)
 #ifdef WOLFSSL_SMALL_STACK
-    XFREE(p, key->heap, DYNAMIC_TYPE_RSA);
-    XFREE(q, key->heap, DYNAMIC_TYPE_RSA);
-    XFREE(tmp1, key->heap, DYNAMIC_TYPE_RSA);
-    XFREE(tmp2, key->heap, DYNAMIC_TYPE_RSA);
-    XFREE(tmp3, key->heap, DYNAMIC_TYPE_RSA);
+    if (key != NULL) {
+        XFREE(p, key->heap, DYNAMIC_TYPE_RSA);
+        XFREE(q, key->heap, DYNAMIC_TYPE_RSA);
+        XFREE(tmp1, key->heap, DYNAMIC_TYPE_RSA);
+        XFREE(tmp2, key->heap, DYNAMIC_TYPE_RSA);
+        XFREE(tmp3, key->heap, DYNAMIC_TYPE_RSA);
+    }
 #elif defined(WOLFSSL_CHECK_MEM_ZERO)
     mp_memzero_check(p);
     mp_memzero_check(q);