From f2c97da66f7a7e47ba2d7136bdb79cf3e1cc4168 Mon Sep 17 00:00:00 2001
From: gojimmypi <jim@wolfssl.com>
Date: Fri, 16 Feb 2024 11:46:06 -0800
Subject: [PATCH] working single / multi threaded ESP32 SSH server signature
 w/SHA256 HW

---
 src/internal.c | 17 ++++++++++++++---
 src/ssh.c      |  8 ++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index 35aa4647c..a7af00a96 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -25,7 +25,7 @@
  */
 
 // #define MY_ZERO_FORCE
-#define MY_FIXED_VALUES
+// #define MY_FIXED_VALUES
 const char* newval = "";
 const char* TAG = "ssh";
 #ifdef HAVE_CONFIG_H
@@ -570,6 +570,10 @@ static HandshakeInfo* HandshakeInfoNew(void* heap)
         newHs->dhGexMaxSz = WOLFSSH_DEFAULT_GEXDH_MAX;
 #endif
     }
+    else {
+        WLOG(WS_LOG_DEBUG, "HandshakeInfoNew: Failed to allocation %d bytes\n",
+                           (int)sizeof(HandshakeInfo));
+    }
 
     return newHs;
 }
@@ -761,6 +765,7 @@ WOLFSSH* SshInit(WOLFSSH* ssh, WOLFSSH_CTX* ctx)
     HandshakeInfo* handshake;
     WC_RNG*        rng;
     void*          heap;
+    int            err;
 
     WLOG(WS_LOG_DEBUG, "Entering SshInit()");
 
@@ -770,10 +775,16 @@ WOLFSSH* SshInit(WOLFSSH* ssh, WOLFSSH_CTX* ctx)
 
     handshake = HandshakeInfoNew(heap);
     rng = (WC_RNG*)WMALLOC(sizeof(WC_RNG), heap, DYNTYPE_RNG);
+    if (handshake == NULL || rng == NULL) {
+        WLOG(WS_LOG_DEBUG, "SshInit: Cannot allocate memory.\n");
+    }
 
-    if (handshake == NULL || rng == NULL || wc_InitRng(rng) != 0) {
+    err = wc_InitRng(rng);
+    if (err != 0) {
+        WLOG(WS_LOG_DEBUG, "SshInit: wc_InitRng failed. err = %d\n", err);
+    }
 
-        WLOG(WS_LOG_DEBUG, "SshInit: Cannot allocate memory.\n");
+    if (handshake == NULL || rng == NULL || err != 0) {
         WFREE(handshake, heap, DYNTYPE_HS);
         WFREE(rng, heap, DYNTYPE_RNG);
         WFREE(ssh, heap, DYNTYPE_SSH);
diff --git a/src/ssh.c b/src/ssh.c
index a9c28cb56..28f054444 100644
--- a/src/ssh.c
+++ b/src/ssh.c
@@ -35,6 +35,10 @@
 #include <wolfssl/wolfcrypt/error-crypt.h>
 #include <wolfssl/wolfcrypt/random.h>
 
+#ifdef DEBUG_WOLFSSL
+    #include <wolfssl/wolfcrypt/logging.h>
+#endif
+
 #ifdef NO_INLINE
     #include <wolfssh/misc.h>
 #else
@@ -154,6 +158,10 @@ WOLFSSH* wolfSSH_new(WOLFSSH_CTX* ctx)
     }
 
     ssh = (WOLFSSH*)WMALLOC(sizeof(WOLFSSH), heap, DYNTYPE_SSH);
+    if (ssh == NULL) {
+        WOLFSSL_MSG_EX("wolfSSH_new failed to allocate %d bytes.",
+                    (int)sizeof(WOLFSSH));
+    }
     ssh = SshInit(ssh, ctx);
 
     WLOG(WS_LOG_DEBUG, "Leaving wolfSSH_new(), ssh = %p", ssh);