From 5060dc2d2d6fc34fc123b3d8ac9ee0067780de06 Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Tue, 30 Apr 2024 16:38:54 -0700
Subject: [PATCH 1/3] Key Agree Update

1. In the key signature block, add flag for Ecc, and for the key
   allocation.
2. Add FreePubKey() to delete the pub key stored in the key signature
   block.
3. In DoKexDhReply(), break down the key agreement actions for the
   various supported key types into their own functions. Remove the
   redundant variables.
4. Using flags that are always present, reduce some of the complicated
   flag checks.
5. Fix a compile guard where the ECDH private key used by the client is
   disabled by ECDSA.
---
 src/internal.c     | 570 +++++++++++++++++++++++++++------------------
 wolfssh/internal.h |  13 +-
 2 files changed, 344 insertions(+), 239 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index b807a010..f8c689e7 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -4167,7 +4167,9 @@ static int DoKexDhInit(WOLFSSH* ssh, byte* buf, word32 len, word32* idx)
 
 
 struct wolfSSH_sigKeyBlock {
-    byte useRsa;
+    byte useRsa:1;
+    byte useEcc:1;
+    byte keyAllocated:1;
     word32 keySz;
     union {
 #ifndef WOLFSSH_NO_RSA
@@ -4222,8 +4224,10 @@ static int ParseRSAPubKey(WOLFSSH *ssh,
                                        &sigKeyBlock_ptr->sk.rsa.key);
     }
 
-    if (ret == 0)
+    if (ret == 0) {
         sigKeyBlock_ptr->keySz = (word32)sizeof(sigKeyBlock_ptr->sk.rsa.key);
+        sigKeyBlock_ptr->keyAllocated = 1;
+    }
     else
         ret = WS_RSA_E;
 #else
@@ -4284,6 +4288,7 @@ static int ParseECCPubKey(WOLFSSH *ssh,
         if (ret == 0) {
             sigKeyBlock_ptr->keySz =
                 (word32)sizeof(sigKeyBlock_ptr->sk.ecc.key);
+            sigKeyBlock_ptr->keyAllocated = 1;
         }
         else
             ret = WS_ECC_E;
@@ -4474,8 +4479,10 @@ static int ParseECCPubKeyCert(WOLFSSH *ssh,
         if (error == 0)
             error = wc_EccPublicKeyDecode(der, &idx,
                 &sigKeyBlock_ptr->sk.ecc.key, derSz);
-        if (error == 0)
+        if (error == 0) {
             sigKeyBlock_ptr->keySz = (word32)sizeof(sigKeyBlock_ptr->sk.ecc.key);
+            sigKeyBlock_ptr->keyAllocated = 1;
+        }
         if (error != 0)
             ret = error;
         WFREE(der, NULL, 0);
@@ -4512,6 +4519,7 @@ static int ParseRSAPubKeyCert(WOLFSSH *ssh,
         if (error == 0) {
             sigKeyBlock_ptr->keySz =
                 (word32)sizeof(sigKeyBlock_ptr->sk.rsa.key);
+            sigKeyBlock_ptr->keyAllocated = 1;
         }
         if (error != 0)
             ret = error;
@@ -4555,7 +4563,7 @@ static int ParsePubKey(WOLFSSH *ssh,
         case ID_ECDSA_SHA2_NISTP256:
         case ID_ECDSA_SHA2_NISTP384:
         case ID_ECDSA_SHA2_NISTP521:
-            sigKeyBlock_ptr->useRsa = 0;
+            sigKeyBlock_ptr->useEcc = 1;
             ret = ParseECCPubKey(ssh, sigKeyBlock_ptr, pubKey, pubKeySz);
             break;
 
@@ -4563,7 +4571,7 @@ static int ParsePubKey(WOLFSSH *ssh,
         case ID_X509V3_ECDSA_SHA2_NISTP256:
         case ID_X509V3_ECDSA_SHA2_NISTP384:
         case ID_X509V3_ECDSA_SHA2_NISTP521:
-            sigKeyBlock_ptr->useRsa = 0;
+            sigKeyBlock_ptr->useEcc = 1;
             ret = ParseECCPubKeyCert(ssh, sigKeyBlock_ptr, pubKey, pubKeySz);
             break;
     #endif
@@ -4576,6 +4584,320 @@ static int ParsePubKey(WOLFSSH *ssh,
 }
 
 
+static void FreePubKey(struct wolfSSH_sigKeyBlock *p)
+{
+    if (p && p->keyAllocated) {
+        if (p->useRsa) {
+        #ifndef WOLFSSH_NO_RSA
+            wc_FreeRsaKey(&p->sk.rsa.key);
+        #endif
+        }
+        else if (p->useEcc) {
+        #ifndef WOLFSSH_NO_ECDSA
+            wc_ecc_free(&p->sk.ecc.key);
+        #endif
+        }
+        p->keyAllocated = 0;
+    }
+}
+
+
+/* KeyAgreeDh_client
+ * f - peer public key
+ * fSz - peer public key size
+ */
+static int KeyAgreeDh_client(WOLFSSH* ssh, const byte* f, word32 fSz)
+#ifndef WOLFSSH_NO_DH
+{
+    int ret;
+
+    PRIVATE_KEY_UNLOCK();
+    ret = wc_DhAgree(&ssh->handshake->privKey.dh,
+                     ssh->k, &ssh->kSz,
+                     ssh->handshake->x, ssh->handshake->xSz,
+                     f, fSz);
+    PRIVATE_KEY_LOCK();
+    if (ret != 0) {
+        WLOG(WS_LOG_ERROR,
+                "Generate DH shared secret failed, %d", ret);
+        ret = WS_CRYPTO_FAILED;
+    }
+    ForceZero(ssh->handshake->x, ssh->handshake->xSz);
+    wc_FreeDhKey(&ssh->handshake->privKey.dh);
+    return ret;
+}
+#else /* WOLFSSH_NO_DH */
+{
+    WOLFSSH_UNUSED(ssh);
+    WOLFSSH_UNUSED(f);
+    WOLFSSH_UNUSED(fSz);
+    return WS_INVALID_ALGO_ID;
+}
+#endif /* WOLFSSH_NO_DH */
+
+
+/* KeyAgreeEcdh_client
+ * f - peer public key
+ * fSz - peer public key size
+ */
+static int KeyAgreeEcdh_client(WOLFSSH* ssh, const byte* f, word32 fSz)
+#ifndef WOLFSSH_NO_ECDH
+{
+    int ret = WS_SUCCESS;
+    ecc_key *key_ptr = NULL;
+    #ifndef WOLFSSH_SMALL_STACK
+        ecc_key key_s;
+    #endif
+    #ifdef WOLFSSH_SMALL_STACK
+        key_ptr = (ecc_key*)WMALLOC(sizeof(ecc_key),
+                ssh->ctx->heap, DYNTYPE_PRIVKEY);
+        if (key_ptr == NULL) {
+            ret = WS_MEMORY_E;
+        }
+    #else /* ! WOLFSSH_SMALL_STACK */
+        key_ptr = &key_s;
+    #endif /* WOLFSSH_SMALL_STACK */
+    ret = wc_ecc_init(key_ptr);
+    #ifdef HAVE_WC_ECC_SET_RNG
+    if (ret == 0)
+        ret = wc_ecc_set_rng(key_ptr, ssh->rng);
+    #endif
+    if (ret == 0)
+        ret = wc_ecc_import_x963(f, fSz, key_ptr);
+    if (ret == 0) {
+        PRIVATE_KEY_UNLOCK();
+        ret = wc_ecc_shared_secret(&ssh->handshake->privKey.ecc,
+                key_ptr, ssh->k, &ssh->kSz);
+        PRIVATE_KEY_LOCK();
+        if (ret != 0) {
+            WLOG(WS_LOG_ERROR,
+                    "Generate ECC shared secret failed, %d", ret);
+            ret = WS_CRYPTO_FAILED;
+        }
+    }
+    wc_ecc_free(key_ptr);
+    #ifdef WOLFSSH_SMALL_STACK
+    if (key_ptr) {
+        WFREE(key_ptr, ssh->ctx->heap, DYNTYPE_PRIVKEY);
+    }
+    #endif
+    wc_ecc_free(&ssh->handshake->privKey.ecc);
+    return ret;
+}
+#else /* WOLFSSH_NO_ECDH */
+{
+    WOLFSSH_UNUSED(ssh);
+    WOLFSSH_UNUSED(f);
+    WOLFSSH_UNUSED(fSz);
+    return WS_INVALID_ALGO_ID;
+}
+#endif /* WOLFSSH_NO_ECDH */
+
+
+/* KeyAgreeX25519_client
+ * f - peer public key
+ * fSz - peer public key size
+ */
+static int KeyAgreeX25519_client(WOLFSSH* ssh, const byte* f, word32 fSz)
+#ifndef WOLFSSH_NO_CURVE25519_SHA256
+{
+    int ret;
+    curve25519_key pub;
+
+    ret = wc_curve25519_init(&pub);
+    if (ret == 0) {
+        ret = wc_curve25519_check_public(f, fSz,
+                EC25519_LITTLE_ENDIAN);
+    }
+
+    if (ret == 0) {
+        ret = wc_curve25519_import_public_ex(f, fSz, &pub,
+                EC25519_LITTLE_ENDIAN);
+    }
+
+    if (ret == 0) {
+        PRIVATE_KEY_UNLOCK();
+        ret = wc_curve25519_shared_secret_ex(
+                  &ssh->handshake->privKey.curve25519, &pub,
+                  ssh->k, &ssh->kSz, EC25519_LITTLE_ENDIAN);
+        PRIVATE_KEY_LOCK();
+        if (ret != 0) {
+            WLOG(WS_LOG_ERROR,
+                    "Gen curve25519 shared secret failed, %d", ret);
+            ret = WS_CRYPTO_FAILED;
+        }
+    }
+
+    wc_curve25519_free(&pub);
+    wc_curve25519_free(&ssh->handshake->privKey.curve25519);
+
+    return ret;
+}
+#else /* WOLFSSH_NO_CURVE25519_SHA256 */
+{
+    WOLFSSH_UNUSED(ssh);
+    WOLFSSH_UNUSED(f);
+    WOLFSSH_UNUSED(fSz);
+    return WS_INVALID_ALGO_ID;
+}
+#endif /* WOLFSSH_NO_CURVE25519_SHA256 */
+
+
+/* KeyAgreeEcdhKyber1_client
+ * f - peer public key
+ * fSz - peer public key size
+ */
+static int KeyAgreeEcdhKyber1_client(WOLFSSH* ssh, const byte* f, word32 fSz)
+#ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
+{
+    int ret = WS_SUCCESS;
+    byte hashId = WC_HASH_TYPE_SHA256;
+    byte sharedSecretHashSz = 0;
+    byte *sharedSecretHash = NULL;
+    ecc_key *key_ptr = NULL;
+    #ifndef WOLFSSH_SMALL_STACK
+        ecc_key key_s;
+    #endif
+    #ifdef WOLFSSH_SMALL_STACK
+        key_ptr = (ecc_key*)WMALLOC(sizeof(ecc_key),
+                ssh->ctx->heap, DYNTYPE_PRIVKEY);
+        if (key_ptr == NULL) {
+            ret = WS_MEMORY_E;
+        }
+    #else /* ! WOLFSSH_SMALL_STACK */
+        key_ptr = &key_s;
+    #endif /* WOLFSSH_SMALL_STACK */
+    /* This is a a hybrid of ECDHE and a post-quantum KEM. In this
+     * case, I need to generated the ECC shared secret and
+     * decapsulate the ciphertext of the post-quantum KEM. */
+    OQS_KEM* kem = OQS_KEM_new(OQS_KEM_alg_kyber_512);
+    if (kem == NULL) {
+        ret = WS_MEMORY_E;
+    }
+
+    if ((ret == 0) && (fSz <= (word32)kem->length_ciphertext)) {
+        ret = WS_BUFFER_E;
+    }
+
+    if (ret == 0) {
+        ret = wc_ecc_init(key_ptr);
+    }
+    #ifdef HAVE_WC_ECC_SET_RNG
+    if (ret == 0) {
+        ret = wc_ecc_set_rng(key_ptr, ssh->rng);
+    }
+    #endif
+    if (ret == 0) {
+        ret = wc_ecc_import_x963(f + kem->length_ciphertext,
+                  fSz - (word32)kem->length_ciphertext,
+                  key_ptr);
+    }
+
+    if (ret == 0) {
+        PRIVATE_KEY_UNLOCK();
+        ret = wc_ecc_shared_secret(&ssh->handshake->privKey.ecc,
+                   key_ptr, ssh->k + kem->length_shared_secret,
+                   &ssh->kSz);
+        PRIVATE_KEY_LOCK();
+    }
+    wc_ecc_free(key_ptr);
+    #ifdef WOLFSSH_SMALL_STACK
+    if (key_ptr) {
+        WFREE(key_ptr, ssh->ctx->heap, DYNTYPE_PRIVKEY);
+    }
+    #endif
+    wc_ecc_free(&ssh->handshake->privKey.ecc);
+
+    if (ret == 0) {
+        if (OQS_KEM_decaps(kem, ssh->k, f, ssh->handshake->x)
+            != OQS_SUCCESS) {
+            ret = WS_ERROR;
+        }
+    }
+
+    if (ret == 0) {
+        ssh->kSz += kem->length_shared_secret;
+    } else {
+        ssh->kSz = 0;
+        WLOG(WS_LOG_ERROR,
+             "Generate ECC-kyber (decap) shared secret failed, %d",
+             ret);
+    }
+
+    if (kem != NULL) {
+        OQS_KEM_free(kem);
+    }
+
+    /* Replace the concatenated shared secrets with the hash. That
+     * will become the new shared secret. */
+    if (ret == 0) {
+        sharedSecretHashSz = wc_HashGetDigestSize(hashId);
+        sharedSecretHash = (byte *)WMALLOC(sharedSecretHashSz,
+                                           ssh->ctx->heap,
+                                           DYNTYPE_PRIVKEY);
+        if (sharedSecretHash == NULL) {
+            ret = WS_MEMORY_E;
+        }
+    }
+
+    if (ret == 0) {
+        ret = wc_Hash(hashId, ssh->k, ssh->kSz, sharedSecretHash,
+                      sharedSecretHashSz);
+    }
+
+    if (ret == 0) {
+        XMEMCPY(ssh->k, sharedSecretHash, sharedSecretHashSz);
+        ssh->kSz = sharedSecretHashSz;
+    }
+
+    if (sharedSecretHash) {
+        ForceZero(sharedSecretHash, sharedSecretHashSz);
+        WFREE(sharedSecretHash, ssh->ctx->heap, DYNTYPE_PRIVKEY);
+    }
+    return ret;
+}
+#else /* WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256 */
+{
+    WOLFSSH_UNUSED(ssh);
+    WOLFSSH_UNUSED(f);
+    WOLFSSH_UNUSED(fSz);
+    return WS_INVALID_ALGO_ID;
+}
+#endif /* WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256 */
+
+
+/* KeyAgree_client
+ * f - peer public key
+ * fSz - peer public key size
+ */
+static int KeyAgree_client(WOLFSSH* ssh, const byte* f, word32 fSz)
+{
+    int ret;
+
+    /* reset size here because a previous shared secret could
+     * potentially be smaller by a byte than usual and cause buffer
+     * issues with re-key */
+    ssh->kSz = MAX_KEX_KEY_SZ;
+
+    if (ssh->handshake->useDh) {
+        ret = KeyAgreeDh_client(ssh, f, fSz);
+    }
+    else if (ssh->handshake->useEcc) {
+        ret = KeyAgreeEcdh_client(ssh, f, fSz);
+    }
+    else if (ssh->handshake->useCurve25519) {
+        ret = KeyAgreeX25519_client(ssh, f, fSz);
+    }
+    else if (ssh->handshake->useEccKyber) {
+        ret = KeyAgreeEcdhKyber1_client(ssh, f, fSz);
+    }
+    else {
+        ret = WS_INVALID_ALGO_ID;
+    }
+    return ret;
+}
+
+
 static int DoKexDhReply(WOLFSSH* ssh, byte* buf, word32 len, word32* idx)
 {
     struct wolfSSH_sigKeyBlock *sigKeyBlock_ptr = NULL;
@@ -4592,17 +4914,6 @@ static int DoKexDhReply(WOLFSSH* ssh, byte* buf, word32 len, word32* idx)
     enum wc_HashType hashId;
     byte scratchLen[LENGTH_SZ];
     byte kPad = 0;
-    byte keyAllocated = 0;
-#ifndef WOLFSSH_NO_ECDH
-    ecc_key *key_ptr = NULL;
-    #ifndef WOLFSSH_SMALL_STACK
-        ecc_key key_s;
-    #endif
-#endif
-#ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
-    byte sharedSecretHashSz = 0;
-    byte *sharedSecretHash = NULL;
-#endif
 
     WLOG(WS_LOG_DEBUG, "Entering DoKexDhReply()");
 
@@ -4778,24 +5089,10 @@ static int DoKexDhReply(WOLFSSH* ssh, byte* buf, word32 len, word32* idx)
         if (sigKeyBlock_ptr == NULL) {
             ret = WS_MEMORY_E;
         }
-
-#ifdef WOLFSSH_SMALL_STACK
-#ifndef WOLFSSH_NO_ECDSA
-        key_ptr = (ecc_key*)WMALLOC(sizeof(ecc_key), ssh->ctx->heap,
-                DYNTYPE_PRIVKEY);
-        if (key_ptr == NULL) {
-            ret = WS_MEMORY_E;
-        }
-#endif /* WOLFSSH_NO_ECDSA */
-
-#else /* ! WOLFSSH_SMALL_STACK */
-#ifndef WOLFSSH_NO_ECDSA
-        key_ptr = &key_s;
-#endif
-#endif
     }
 
     if (ret == WS_SUCCESS) {
+        WMEMSET(sigKeyBlock_ptr, 0, sizeof(*sigKeyBlock_ptr));
         sig = buf + begin;
         begin += sigSz;
         *idx = begin;
@@ -4803,185 +5100,14 @@ static int DoKexDhReply(WOLFSSH* ssh, byte* buf, word32 len, word32* idx)
         ret = ParsePubKey(ssh, sigKeyBlock_ptr, pubKey, pubKeySz);
         /* Generate and hash in the shared secret */
         if (ret == WS_SUCCESS) {
-            /* Remember that the key needs to be freed */
-            keyAllocated = 1;
-            /* reset size here because a previous shared secret could
-             * potentially be smaller by a byte than usual and cause buffer
-             * issues with re-key */
-            ssh->kSz = MAX_KEX_KEY_SZ;
-            if (!ssh->handshake->useEcc
-#ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
-                && !ssh->handshake->useEccKyber
-#endif
-#ifndef WOLFSSH_NO_CURVE25519_SHA256
-                && !ssh->handshake->useCurve25519
-#endif
-               ) {
-#ifndef WOLFSSH_NO_DH
-                PRIVATE_KEY_UNLOCK();
-                ret = wc_DhAgree(&ssh->handshake->privKey.dh,
-                                 ssh->k, &ssh->kSz,
-                                 ssh->handshake->x, ssh->handshake->xSz,
-                                 f, fSz);
-                PRIVATE_KEY_LOCK();
-                ForceZero(ssh->handshake->x, ssh->handshake->xSz);
-                wc_FreeDhKey(&ssh->handshake->privKey.dh);
-                if (ret != 0) {
-                    WLOG(WS_LOG_ERROR,
-                            "Generate DH shared secret failed, %d", ret);
-                }
-#else
-                ret = WS_INVALID_ALGO_ID;
-#endif
-            }
-            else if (ssh->handshake->useEcc) {
-#ifndef WOLFSSH_NO_ECDH
-                ret = wc_ecc_init(key_ptr);
-#ifdef HAVE_WC_ECC_SET_RNG
-                if (ret == 0)
-                    ret = wc_ecc_set_rng(key_ptr, ssh->rng);
-#endif
-                if (ret == 0)
-                    ret = wc_ecc_import_x963(f, fSz, key_ptr);
-                if (ret == 0) {
-                    PRIVATE_KEY_UNLOCK();
-                    ret = wc_ecc_shared_secret(&ssh->handshake->privKey.ecc,
-                                               key_ptr, ssh->k, &ssh->kSz);
-                    PRIVATE_KEY_LOCK();
-                }
-                wc_ecc_free(key_ptr);
-                wc_ecc_free(&ssh->handshake->privKey.ecc);
-                if (ret != 0) {
-                    WLOG(WS_LOG_ERROR,
-                            "Generate ECC shared secret failed, %d", ret);
-                }
-#else
-                ret = WS_INVALID_ALGO_ID;
-#endif
-            }
-#ifndef WOLFSSH_NO_CURVE25519_SHA256
-            else if (ssh->handshake->useCurve25519) {
-                curve25519_key pub;
-                ret = wc_curve25519_init(&pub);
-
-                if (ret == 0)
-                    ret = wc_curve25519_check_public(f, fSz,
-                                                     EC25519_LITTLE_ENDIAN);
-
-                if (ret == 0) {
-                    ret = wc_curve25519_import_public_ex(f, fSz, &pub,
-                                                         EC25519_LITTLE_ENDIAN);
-                }
-
-                if (ret == 0) {
-                    PRIVATE_KEY_UNLOCK();
-                    ret = wc_curve25519_shared_secret_ex(
-                              &ssh->handshake->privKey.curve25519, &pub,
-                              ssh->k, &ssh->kSz, EC25519_LITTLE_ENDIAN);
-                    PRIVATE_KEY_LOCK();
-                }
-
-                wc_curve25519_free(&pub);
-                wc_curve25519_free(&ssh->handshake->privKey.curve25519);
-
-                if (ret != 0) {
-                    WLOG(WS_LOG_ERROR,
-                            "Gen curve25519 shared secret failed, %d", ret);
-                }
-            }
-#endif /* !WOLFSSH_NO_CURVE25519_SHA256 */
-#ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
-            else if (ssh->handshake->useEccKyber) {
-                /* This is a a hybrid of ECDHE and a post-quantum KEM. In this
-                 * case, I need to generated the ECC shared secret and
-                 * decapsulate the ciphertext of the post-quantum KEM. */
-                OQS_KEM* kem = OQS_KEM_new(OQS_KEM_alg_kyber_512);
-                if (kem == NULL) {
-                    ret = WS_INVALID_ALGO_ID;
-                }
-
-                if ((ret == 0) && (fSz <= (word32)kem->length_ciphertext)) {
-                    ret = WS_BUFFER_E;
-                }
-
-                if (ret == 0) {
-                    ret = wc_ecc_init(key_ptr);
-                }
-#ifdef HAVE_WC_ECC_SET_RNG
-                if (ret == 0) {
-                    ret = wc_ecc_set_rng(key_ptr, ssh->rng);
-                }
-#endif
-                if (ret == 0) {
-                    ret = wc_ecc_import_x963(f + kem->length_ciphertext,
-                              fSz - (word32)kem->length_ciphertext, key_ptr);
-                }
-
-                if (ret == 0) {
-                    PRIVATE_KEY_UNLOCK();
-                    ret = wc_ecc_shared_secret(&ssh->handshake->privKey.ecc,
-                               key_ptr, ssh->k + kem->length_shared_secret,
-                               &ssh->kSz);
-                    PRIVATE_KEY_LOCK();
-                }
-                wc_ecc_free(key_ptr);
-                wc_ecc_free(&ssh->handshake->privKey.ecc);
-
-                if (ret == 0) {
-                    if (OQS_KEM_decaps(kem, ssh->k, f, ssh->handshake->x)
-                        != OQS_SUCCESS) {
-                        ret = WS_ERROR;
-                    }
-                }
-
-                if (ret == 0) {
-                    ssh->kSz += kem->length_shared_secret;
-                } else {
-                    ssh->kSz = 0;
-                    WLOG(WS_LOG_ERROR,
-                         "Generate ECC-kyber (decap) shared secret failed, %d",
-                         ret);
-                }
-
-                if (kem != NULL) {
-                    OQS_KEM_free(kem);
-                }
-
-                /* Replace the concatenated shared secrets with the hash. That
-                 * will become the new shared secret. */
-                if (ret == 0) {
-                    sharedSecretHashSz = wc_HashGetDigestSize(hashId);
-                    sharedSecretHash = (byte *)WMALLOC(sharedSecretHashSz,
-                                                       ssh->ctx->heap,
-                                                       DYNTYPE_PRIVKEY);
-                    if (sharedSecretHash == NULL) {
-                        ret = WS_MEMORY_E;
-                    }
-                }
-
-                if (ret == 0) {
-                    ret = wc_Hash(hashId, ssh->k, ssh->kSz, sharedSecretHash,
-                                  sharedSecretHashSz);
-                }
-
-                if (ret == 0) {
-                    XMEMCPY(ssh->k, sharedSecretHash, sharedSecretHashSz);
-                    ssh->kSz = sharedSecretHashSz;
-                }
-            }
-#endif /* !WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256 */
-            else {
-                ret = WS_INVALID_ALGO_ID;
-            }
+            ret = KeyAgree_client(ssh, f, fSz);
         }
 
         /* Hash in the shared secret K. */
-        if (ret == 0
-#ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
-            && !ssh->handshake->useEccKyber
-#endif
-        ) {
-            ret = CreateMpint(ssh->k, &ssh->kSz, &kPad);
+        if (ret == WS_SUCCESS) {
+            if (!ssh->handshake->useEccKyber) {
+                ret = CreateMpint(ssh->k, &ssh->kSz, &kPad);
+            }
         }
 
         if (ret == 0) {
@@ -5115,27 +5241,12 @@ static int DoKexDhReply(WOLFSSH* ssh, byte* buf, word32 len, word32* idx)
                 }
             }
         }
-
-        if (keyAllocated) {
-            if (sigKeyBlock_ptr->useRsa) {
-#ifndef WOLFSSH_NO_RSA
-                wc_FreeRsaKey(&sigKeyBlock_ptr->sk.rsa.key);
-#endif
-            }
-            else {
-#ifndef WOLFSSH_NO_ECDSA
-                wc_ecc_free(&sigKeyBlock_ptr->sk.ecc.key);
-#endif
-            }
-        }
+        FreePubKey(sigKeyBlock_ptr);
     }
 
     if (ret == WS_SUCCESS) {
-        int useKeyPadding = 1;
-#if !defined(WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256)
-        useKeyPadding = !ssh->handshake->useEccKyber;
-#endif
-        ret = GenerateKeys(ssh, hashId, useKeyPadding);
+        /* If we aren't using EccKyber, use padding. */
+        ret = GenerateKeys(ssh, hashId, !ssh->handshake->useEccKyber);
     }
 
     if (ret == WS_SUCCESS)
@@ -5143,12 +5254,6 @@ static int DoKexDhReply(WOLFSSH* ssh, byte* buf, word32 len, word32* idx)
 
     if (sigKeyBlock_ptr)
         WFREE(sigKeyBlock_ptr, ssh->ctx->heap, DYNTYPE_PRIVKEY);
-#ifdef WOLFSSH_SMALL_STACK
-    #ifndef WOLFSSH_NO_ECDSA
-    if (key_ptr)
-        WFREE(key_ptr, ssh->ctx->heap, DYNTYPE_PRIVKEY);
-    #endif
-#endif
     WLOG(WS_LOG_DEBUG, "Leaving DoKexDhReply(), ret = %d", ret);
     return ret;
 }
@@ -11012,6 +11117,7 @@ int SendKexDhInit(WOLFSSH* ssh)
     switch (ssh->handshake->kexId) {
 #ifndef WOLFSSH_NO_DH_GROUP1_SHA1
         case ID_DH_GROUP1_SHA1:
+            ssh->handshake->useDh = 1;
             primeGroup = dhPrimeGroup1;
             primeGroupSz = dhPrimeGroup1Sz;
             generator = dhGenerator;
@@ -11020,6 +11126,7 @@ int SendKexDhInit(WOLFSSH* ssh)
 #endif
 #ifndef WOLFSSH_NO_DH_GROUP14_SHA1
         case ID_DH_GROUP14_SHA1:
+            ssh->handshake->useDh = 1;
             primeGroup = dhPrimeGroup14;
             primeGroupSz = dhPrimeGroup14Sz;
             generator = dhGenerator;
@@ -11028,6 +11135,7 @@ int SendKexDhInit(WOLFSSH* ssh)
 #endif
 #ifndef WOLFSSH_NO_DH_GEX_SHA256
         case ID_DH_GEX_SHA256:
+            ssh->handshake->useDh = 1;
             primeGroup = ssh->handshake->primeGroup;
             primeGroupSz = ssh->handshake->primeGroupSz;
             generator = ssh->handshake->generator;
diff --git a/wolfssh/internal.h b/wolfssh/internal.h
index 6ff2721c..c8aeefea 100644
--- a/wolfssh/internal.h
+++ b/wolfssh/internal.h
@@ -583,19 +583,16 @@ typedef struct HandshakeInfo {
     word32 generatorSz;
 #endif
 
-    byte useEcc;
-#ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
-    byte useEccKyber;
-#endif
-#ifndef WOLFSSH_NO_CURVE25519_SHA256
-    byte useCurve25519;
-#endif
+    byte useDh:1;
+    byte useEcc:1;
+    byte useEccKyber:1;
+    byte useCurve25519:1;
 
     union {
 #ifndef WOLFSSH_NO_DH
         DhKey dh;
 #endif
-#if !defined(WOLFSSH_NO_ECDSA) && !defined(WOLFSSH_NO_ECDH)
+#ifndef WOLFSSH_NO_ECDH
         ecc_key ecc;
 #endif
 #ifndef WOLFSSH_NO_CURVE25519_SHA256

From 49c420d59351dda2de416bfe4ac6ce12fb96364f Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Tue, 30 Apr 2024 16:50:18 -0700
Subject: [PATCH 2/3] Key Agree Update

1. In SendKexDhReply(), break down the key agreement actions for the
   various supported key types into their own functions. Remove
   the redundant variables.
2. In DoKexDhInit(), add flags for the various key agreement types, and
   set them as appropriate when checking the selected kexId. The flags
   are always present no matter the build options.
3. Simplify some of the flag checks for optional options.
---
 src/internal.c | 754 ++++++++++++++++++++++++++-----------------------
 1 file changed, 407 insertions(+), 347 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index f8c689e7..982b8245 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -10046,6 +10046,371 @@ int wolfSSH_RsaVerify(byte *sig, word32 sigSz,
 #endif /* WOLFSSH_NO_RSA */
 
 
+static int KeyAgreeDh_server(WOLFSSH* ssh, byte hashId, byte* f, word32* fSz)
+#ifndef WOLFSSH_NO_DH
+{
+    int ret = WS_SUCCESS;
+    byte *y_ptr = NULL;
+    const byte* primeGroup = NULL;
+    const byte* generator = NULL;
+    word32 ySz = MAX_KEX_KEY_SZ;
+    word32 primeGroupSz = 0;
+    word32 generatorSz = 0;
+    #ifdef WOLFSSH_SMALL_STACK
+    DhKey *privKey = (DhKey*)WMALLOC(sizeof(DhKey), heap,
+            DYNTYPE_PRIVKEY);
+    y_ptr = (byte*)WMALLOC(ySz, heap, DYNTYPE_PRIVKEY);
+    if (privKey == NULL || y_ptr == NULL)
+        ret = WS_MEMORY_E;
+    #else
+    DhKey privKey[1];
+    byte y_s[MAX_KEX_KEY_SZ];
+    y_ptr = y_s;
+    #endif
+
+    WOLFSSH_UNUSED(hashId);
+
+    if (ret == WS_SUCCESS) {
+        ret = GetDHPrimeGroup(ssh->handshake->kexId, &primeGroup,
+            &primeGroupSz, &generator, &generatorSz);
+
+        if (ret == WS_SUCCESS) {
+            ret = wc_InitDhKey(privKey);
+        }
+        if (ret == 0)
+            ret = wc_DhSetKey(privKey, primeGroup, primeGroupSz,
+                    generator, generatorSz);
+        if (ret == 0)
+            ret = wc_DhGenerateKeyPair(privKey, ssh->rng,
+                    y_ptr, &ySz, f, fSz);
+        if (ret == 0) {
+            PRIVATE_KEY_UNLOCK();
+            ret = wc_DhAgree(privKey, ssh->k, &ssh->kSz, y_ptr, ySz,
+                    ssh->handshake->e, ssh->handshake->eSz);
+            PRIVATE_KEY_LOCK();
+        }
+        ForceZero(y_ptr, ySz);
+        wc_FreeDhKey(privKey);
+    }
+    #ifdef WOLFSSH_SMALL_STACK
+    if (y_ptr)
+        WFREE(y_ptr, heap, DYNTYPE_PRIVKEY);
+    if (privKey) {
+        WFREE(privKey, heap, DYNTYPE_PRIVKEY);
+    }
+    #endif
+    return ret;
+}
+#else /* WOLFSSH_NO_DH */
+{
+    WOLFSSH_UNUSED(ssh);
+    WOLFSSH_UNUSED(hashId);
+    WOLFSSH_UNUSED(f);
+    WOLFSSH_UNUSED(fSz);
+    return WS_INVALID_ALGO_ID;
+}
+#endif /* WOLFSSH_NO_DH */
+
+
+static int KeyAgreeEcdh_server(WOLFSSH* ssh, byte hashId, byte* f, word32* fSz)
+#ifndef WOLFSSH_NO_ECDH
+{
+    int ret = WS_SUCCESS;
+    void* heap;
+#ifdef WOLFSSH_SMALL_STACK
+    ecc_key *pubKey = NULL, *privKey = NULL;
+    pubKey = (ecc_key*)WMALLOC(sizeof(ecc_key), heap,
+            DYNTYPE_PUBKEY);
+    privKey = (ecc_key*)WMALLOC(sizeof(ecc_key), heap,
+            DYNTYPE_PRIVKEY);
+    if (pubKey == NULL || privKey == NULL) {
+        ret = WS_MEMORY_E;
+    }
+#else
+    ecc_key pubKey[1];
+    ecc_key privKey[1];
+#endif
+    int primeId;
+
+    WOLFSSH_UNUSED(hashId);
+    heap = ssh->ctx->heap;
+    primeId  = wcPrimeForId(ssh->handshake->kexId);
+    if (primeId == ECC_CURVE_INVALID)
+        ret = WS_INVALID_PRIME_CURVE;
+
+    if (ret == 0)
+        ret = wc_ecc_init_ex(pubKey, heap, INVALID_DEVID);
+    if (ret == 0)
+        ret = wc_ecc_init_ex(privKey, heap, INVALID_DEVID);
+#ifdef HAVE_WC_ECC_SET_RNG
+    if (ret == 0)
+        ret = wc_ecc_set_rng(privKey, ssh->rng);
+#endif
+
+    if (ret == 0)
+        ret = wc_ecc_import_x963_ex(ssh->handshake->e,
+                                    ssh->handshake->eSz,
+                                    pubKey, primeId);
+
+    if (ret == 0)
+        ret = wc_ecc_make_key_ex(ssh->rng,
+                             wc_ecc_get_curve_size_from_id(primeId),
+                             privKey, primeId);
+    if (ret == 0) {
+        PRIVATE_KEY_UNLOCK();
+        ret = wc_ecc_export_x963(privKey, f, fSz);
+        PRIVATE_KEY_LOCK();
+    }
+    if (ret == 0) {
+        PRIVATE_KEY_UNLOCK();
+        ret = wc_ecc_shared_secret(privKey, pubKey,
+                                   ssh->k, &ssh->kSz);
+        PRIVATE_KEY_LOCK();
+    }
+    wc_ecc_free(privKey);
+    wc_ecc_free(pubKey);
+#ifdef WOLFSSH_SMALL_STACK
+    WFREE(pubKey, heap, DYNTYPE_PUBKEY);
+    WFREE(privKey, heap, DYNTYPE_PRIVKEY);
+    pubKey  = NULL;
+    privKey = NULL;
+#endif
+    return ret;
+}
+#else /* WOLFSSH_NO_ECDH */
+{
+    WOLFSSH_UNUSED(ssh);
+    WOLFSSH_UNUSED(hashId);
+    WOLFSSH_UNUSED(f);
+    WOLFSSH_UNUSED(fSz);
+    return WS_INVALID_ALGO_ID;
+}
+#endif /* WOLFSSH_NO_ECDH */
+
+
+static int KeyAgreeX25519_server(WOLFSSH* ssh, byte hashId,
+        byte* f, word32* fSz)
+#ifndef WOLFSSH_NO_CURVE25519_SHA256
+{
+    int ret = WS_SUCCESS;
+    void* heap = ssh->ctx->heap;
+#ifdef WOLFSSH_SMALL_STACK
+    curve25519_key *pubKey = NULL, *privKey = NULL;
+    pubKey = (curve25519_key*)WMALLOC(sizeof(curve25519_key),
+            heap, DYNTYPE_PUBKEY);
+    privKey = (curve25519_key*)WMALLOC(sizeof(curve25519_key),
+            heap, DYNTYPE_PRIVKEY);
+    if (pubKey == NULL || privKey == NULL) {
+        ret = WS_MEMORY_E;
+    }
+#else
+    curve25519_key pubKey[1], privKey[1];
+#endif
+
+    WOLFSSH_UNUSED(hashId);
+    if (ret == 0)
+        ret = wc_curve25519_init_ex(pubKey, heap, INVALID_DEVID);
+    if (ret == 0)
+        ret = wc_curve25519_init_ex(privKey, heap, INVALID_DEVID);
+    if (ret == 0)
+        ret = wc_curve25519_check_public(ssh->handshake->e,
+                ssh->handshake->eSz, EC25519_LITTLE_ENDIAN);
+    if (ret == 0)
+        ret = wc_curve25519_import_public_ex(
+                ssh->handshake->e, ssh->handshake->eSz,
+                pubKey, EC25519_LITTLE_ENDIAN);
+
+    if (ret == 0)
+        ret = wc_curve25519_make_key(ssh->rng, CURVE25519_KEYSIZE, privKey);
+
+    if (ret == 0) {
+        PRIVATE_KEY_UNLOCK();
+        ret = wc_curve25519_export_public_ex(privKey,
+                f, fSz, EC25519_LITTLE_ENDIAN);
+        PRIVATE_KEY_LOCK();
+    }
+
+    if (ret == 0) {
+        PRIVATE_KEY_UNLOCK();
+        ret = wc_curve25519_shared_secret_ex(privKey, pubKey,
+                  ssh->k, &ssh->kSz, EC25519_LITTLE_ENDIAN);
+        PRIVATE_KEY_LOCK();
+    }
+    wc_curve25519_free(privKey);
+    wc_curve25519_free(pubKey);
+#ifdef WOLFSSH_SMALL_STACK
+    WFREE(pubKey, heap, DYNTYPE_PUBKEY);
+    WFREE(privKey, heap, DYNTYPE_PRIVKEY);
+    pubKey  = NULL;
+    privKey = NULL;
+#endif
+    return ret;
+}
+#else /* WOLFSSH_NO_CURVE25519_SHA256 */
+{
+    WOLFSSH_UNUSED(ssh);
+    WOLFSSH_UNUSED(hashId);
+    WOLFSSH_UNUSED(f);
+    WOLFSSH_UNUSED(fSz);
+    return WS_INVALID_ALGO_ID;
+}
+#endif /* WOLFSSH_NO_CURVE25519_SHA256 */
+
+
+static int KeyAgreeEcdhKyber1_server(WOLFSSH* ssh, byte hashId,
+        byte* f, word32* fSz)
+#ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
+{
+    int ret = WS_SUCCESS;
+    void* heap = ssh->ctx->heap;
+    byte sharedSecretHashSz = 0;
+    byte *sharedSecretHash = NULL;
+    /* This is a hybrid KEM. In this case, I need to generate my ECC
+     * keypair, send the public one, use the private one to generate
+     * the shared secret, use the post-quantum public key to
+     * generate and encapsulate the shared secret and send the
+     * ciphertext. */
+    OQS_KEM* kem = NULL;
+    int primeId;
+    ret = 0;
+#ifdef WOLFSSH_SMALL_STACK
+    ecc_key *pubKey = NULL, *privKey = NULL;
+    pubKey = (ecc_key*)WMALLOC(sizeof(ecc_key), heap,
+            DYNTYPE_PUBKEY);
+    privKey = (ecc_key*)WMALLOC(sizeof(ecc_key), heap,
+            DYNTYPE_PRIVKEY);
+    if (pubKey == NULL || privKey == NULL) {
+        ret = WS_MEMORY_E;
+    }
+#else
+    ecc_key pubKey[1];
+    ecc_key privKey[1];
+#endif
+
+    if (ret == 0) {
+        XMEMSET(pubKey, 0, sizeof(*pubKey));
+        XMEMSET(privKey, 0, sizeof(*privKey));
+
+        primeId = wcPrimeForId(ssh->handshake->kexId);
+        if (primeId == ECC_CURVE_INVALID)
+            ret = WS_INVALID_PRIME_CURVE;
+    }
+
+    if (ret == 0) {
+        kem = OQS_KEM_new(OQS_KEM_alg_kyber_512);
+        if (kem == NULL) {
+            ret = WS_INVALID_ALGO_ID;
+        }
+    }
+
+    if ((ret == 0) &&
+        (ssh->handshake->eSz <= (word32)kem->length_public_key)) {
+        ret = WS_BUFFER_E;
+    }
+
+    if (ret == 0) {
+        if (OQS_KEM_encaps(kem, f, ssh->k,
+            ssh->handshake->e) != OQS_SUCCESS) {
+            ret = WS_PUBKEY_REJECTED_E;
+        }
+    }
+
+    if (ret == 0) {
+        *fSz -= kem->length_ciphertext;
+        ssh->kSz -= kem->length_shared_secret;
+    }
+    else {
+        *fSz = 0;
+        ssh->kSz = 0;
+        WLOG(WS_LOG_ERROR,
+             "Generate ECC-kyber (encap) shared secret failed, %d",
+             ret);
+    }
+
+    if (ret == 0) {
+        ret = wc_ecc_init_ex(pubKey, heap, INVALID_DEVID);
+    }
+    if (ret == 0) {
+        ret = wc_ecc_init_ex(privKey, heap, INVALID_DEVID);
+    }
+#ifdef HAVE_WC_ECC_SET_RNG
+    if (ret == 0) {
+        ret = wc_ecc_set_rng(privKey, ssh->rng);
+    }
+#endif
+    if (ret == 0) {
+        ret = wc_ecc_import_x963_ex(
+            ssh->handshake->e + kem->length_public_key,
+            ssh->handshake->eSz - (word32)kem->length_public_key,
+            pubKey, primeId);
+    }
+    if (ret == 0) {
+        ret = wc_ecc_make_key_ex(ssh->rng,
+                  wc_ecc_get_curve_size_from_id(primeId),
+                  privKey, primeId);
+    }
+    if (ret == 0) {
+        PRIVATE_KEY_UNLOCK();
+        ret = wc_ecc_export_x963(privKey,
+                  f + kem->length_ciphertext, fSz);
+        PRIVATE_KEY_LOCK();
+        *fSz += kem->length_ciphertext;
+    }
+    if (ret == 0) {
+        word32 tmp_kSz = ssh->kSz;
+        PRIVATE_KEY_UNLOCK();
+        ret = wc_ecc_shared_secret(privKey, pubKey,
+                  ssh->k + kem->length_shared_secret, &tmp_kSz);
+        PRIVATE_KEY_LOCK();
+        ssh->kSz = (word32)kem->length_shared_secret + tmp_kSz;
+    }
+    wc_ecc_free(privKey);
+    wc_ecc_free(pubKey);
+#ifdef WOLFSSH_SMALL_STACK
+    WFREE(pubKey, heap, DYNTYPE_PUBKEY);
+    WFREE(privKey, heap, DYNTYPE_PRIVKEY);
+    pubKey  = NULL;
+    privKey = NULL;
+#endif
+    if (kem != NULL) {
+        OQS_KEM_free(kem);
+        kem = NULL;
+    }
+
+    /* Replace the concatenated shared secrets with the hash. That
+     * will become the new shared secret.*/
+    if (ret == 0) {
+        sharedSecretHashSz = wc_HashGetDigestSize(hashId);
+        sharedSecretHash = (byte *)WMALLOC(sharedSecretHashSz, heap,
+                                           DYNTYPE_PRIVKEY);
+        if (sharedSecretHash == NULL) {
+            ret = WS_MEMORY_E;
+        }
+    }
+    if (ret == 0) {
+        ret = wc_Hash(hashId, ssh->k, ssh->kSz, sharedSecretHash,
+                      sharedSecretHashSz);
+    }
+    if (ret == 0) {
+        XMEMCPY(ssh->k, sharedSecretHash, sharedSecretHashSz);
+        ssh->kSz = sharedSecretHashSz;
+    }
+
+    WFREE(sharedSecretHash, heap, DYNTYPE_PRIVKEY);
+    sharedSecretHash = NULL;
+    return ret;
+}
+#else /* WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256 */
+{
+    WOLFSSH_UNUSED(ssh);
+    WOLFSSH_UNUSED(hashId);
+    WOLFSSH_UNUSED(f);
+    WOLFSSH_UNUSED(fSz);
+    return WS_INVALID_ALGO_ID;
+}
+#endif /* WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256 */
+
+
 /* SendKexDhReply()
  * It is also the funciton used for MSGID_KEXECDH_REPLY. The parameters
  * are analogous between the two messages. Where MSGID_KEXDH_REPLY has
@@ -10062,15 +10427,6 @@ int SendKexDhReply(WOLFSSH* ssh)
     byte scratchLen[LENGTH_SZ];
     word32 fSz = KEX_F_SIZE;
     word32 sigSz = KEX_SIG_SIZE;
-    byte useEcc = 0;
-#ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
-    byte useEccKyber = 0;
-    byte sharedSecretHashSz = 0;
-    byte *sharedSecretHash = NULL;
-#endif
-#ifndef WOLFSSH_NO_CURVE25519_SHA256
-    byte useCurve25519 = 0;
-#endif
     byte fPad = 0;
     byte kPad = 0;
     word32 sigBlockSz = 0;
@@ -10078,17 +10434,19 @@ int SendKexDhReply(WOLFSSH* ssh)
     byte* output = NULL;
     word32 idx = 0;
     word32 keyIdx = 0;
-    byte msgId = MSGID_KEXDH_REPLY;
     enum wc_HashType hashId = WC_HASH_TYPE_NONE;
     wc_HashAlg* hash = NULL;
     struct wolfSSH_sigKeyBlockFull *sigKeyBlock_ptr = NULL;
 #ifndef WOLFSSH_SMALL_STACK
     byte f_s[KEX_F_SIZE];
     byte sig_s[KEX_SIG_SIZE];
-
-    f_ptr = f_s;
-    sig_ptr = sig_s;
 #endif
+    byte msgId;
+    byte useDh = 0;
+    byte useEcc = 0;
+    byte useCurve25519 = 0;
+    byte useEccKyber = 0;
+
     WLOG(WS_LOG_DEBUG, "Entering SendKexDhReply()");
 
     if (ret == WS_SUCCESS) {
@@ -10106,6 +10464,9 @@ int SendKexDhReply(WOLFSSH* ssh)
     sig_ptr = (byte*)WMALLOC(KEX_SIG_SIZE, heap, DYNTYPE_BUFFER);
     if (f_ptr == NULL || sig_ptr == NULL)
         ret = WS_MEMORY_E;
+#else
+    f_ptr = f_s;
+    sig_ptr = sig_s;
 #endif
 
     sigKeyBlock_ptr = (struct wolfSSH_sigKeyBlockFull*)WMALLOC(
@@ -10135,6 +10496,24 @@ int SendKexDhReply(WOLFSSH* ssh)
                 (word32)WSTRLEN(sigKeyBlock_ptr->pubKeyFmtName);
 
         switch (ssh->handshake->kexId) {
+#ifndef WOLFSSH_NO_DH_GROUP1_SHA1
+            case ID_DH_GROUP1_SHA1:
+                useDh = 1;
+                msgId = MSGID_KEXDH_REPLY;
+                break;
+#endif
+#ifndef WOLFSSH_NO_DH_GROUP14_SHA1
+            case ID_DH_GROUP14_SHA1:
+                useDh = 1;
+                msgId = MSGID_KEXDH_REPLY;
+                break;
+#endif
+#ifndef WOLFSSH_NO_DH_GEX_SHA256
+            case ID_DH_GEX_SHA256:
+                useDh = 1;
+                msgId = MSGID_KEXDH_GEX_REPLY;
+                break;
+#endif
 #ifndef WOLFSSH_NO_ECDH_SHA2_NISTP256
             case ID_ECDH_SHA2_NISTP256:
                 useEcc = 1;
@@ -10165,13 +10544,15 @@ int SendKexDhReply(WOLFSSH* ssh)
                 msgId = MSGID_KEXKEM_REPLY;
                 break;
 #endif
+            default:
+                ret = WS_INVALID_ALGO_ID;
         }
-
-        hash = &ssh->handshake->kexHash;
-        hashId = (enum wc_HashType)ssh->handshake->kexHashId;
     }
 
     if (ret == WS_SUCCESS) {
+        hash = &ssh->handshake->kexHash;
+        hashId = (enum wc_HashType)ssh->handshake->kexHashId;
+
         for (keyIdx = 0; keyIdx < ssh->ctx->privateKeyCount; keyIdx++) {
             if (ssh->ctx->privateKey[keyIdx].publicKeyFmt
                     == sigKeyBlock_ptr->pubKeyFmtId) {
@@ -10193,341 +10574,27 @@ int SendKexDhReply(WOLFSSH* ssh)
     if (ret == WS_SUCCESS) {
         /* reset size here because a previous shared secret could potentially be
          * smaller by a byte than usual and cause buffer issues with re-key */
-        if (ret == 0)
-            ssh->kSz = MAX_KEX_KEY_SZ;
+        ssh->kSz = MAX_KEX_KEY_SZ;
 
         /* Make the server's DH f-value and the shared secret K. */
         /* Or make the server's ECDH private value, and the shared secret K. */
         if (ret == 0) {
-            if (!useEcc
-#ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
-                && !useEccKyber
-#endif
-#ifndef WOLFSSH_NO_CURVE25519_SHA256
-                && !useCurve25519
-#endif
-               ) {
-#ifndef WOLFSSH_NO_DH
-                byte *y_ptr = NULL;
-                const byte* primeGroup = NULL;
-                const byte* generator = NULL;
-                word32 ySz = MAX_KEX_KEY_SZ;
-                word32 primeGroupSz = 0;
-                word32 generatorSz = 0;
-            #ifdef WOLFSSH_SMALL_STACK
-                DhKey *privKey = (DhKey*)WMALLOC(sizeof(DhKey), heap,
-                        DYNTYPE_PRIVKEY);
-                y_ptr = (byte*)WMALLOC(ySz, heap, DYNTYPE_PRIVKEY);
-                if (privKey == NULL || y_ptr == NULL)
-                    ret = WS_MEMORY_E;
-            #else
-                DhKey privKey[1];
-                byte y_s[MAX_KEX_KEY_SZ];
-                y_ptr = y_s;
-            #endif
-                if (ret == WS_SUCCESS) {
-                    ret = GetDHPrimeGroup(ssh->handshake->kexId, &primeGroup,
-                        &primeGroupSz, &generator, &generatorSz);
-                    #ifndef WOLFSSH_NO_DH_GEX_SHA256
-                    if (ssh->handshake->kexId == ID_DH_GEX_SHA256)
-                        msgId = MSGID_KEXDH_GEX_REPLY;
-                    #endif
-
-                    if (ret == WS_SUCCESS) {
-                        ret = wc_InitDhKey(privKey);
-                    }
-                    if (ret == 0)
-                        ret = wc_DhSetKey(privKey, primeGroup, primeGroupSz,
-                                generator, generatorSz);
-                    if (ret == 0)
-                        ret = wc_DhGenerateKeyPair(privKey, ssh->rng,
-                                y_ptr, &ySz, f_ptr, &fSz);
-                    if (ret == 0) {
-                        PRIVATE_KEY_UNLOCK();
-                        ret = wc_DhAgree(privKey, ssh->k, &ssh->kSz, y_ptr, ySz,
-                                ssh->handshake->e, ssh->handshake->eSz);
-                        PRIVATE_KEY_LOCK();
-                    }
-                    ForceZero(y_ptr, ySz);
-                    wc_FreeDhKey(privKey);
-                }
-            #ifdef WOLFSSH_SMALL_STACK
-                if (y_ptr)
-                    WFREE(y_ptr, heap, DYNTYPE_PRIVKEY);
-                if (privKey) {
-                    WFREE(privKey, heap, DYNTYPE_PRIVKEY);
-                }
-            #endif
-#endif /* ! WOLFSSH_NO_DH */
+            if (useDh) {
+                ret = KeyAgreeDh_server(ssh, hashId, f_ptr, &fSz);
             }
             else if (useEcc) {
-#if !defined(WOLFSSH_NO_ECDH)
-            #ifdef WOLFSSH_SMALL_STACK
-                ecc_key *pubKey = NULL, *privKey = NULL;
-                pubKey = (ecc_key*)WMALLOC(sizeof(ecc_key), heap,
-                        DYNTYPE_PUBKEY);
-                privKey = (ecc_key*)WMALLOC(sizeof(ecc_key), heap,
-                        DYNTYPE_PRIVKEY);
-                if (pubKey == NULL || privKey == NULL) {
-                    ret = WS_MEMORY_E;
-                }
-            #else
-                ecc_key pubKey[1];
-                ecc_key privKey[1];
-            #endif
-                int primeId;
-
-                primeId  = wcPrimeForId(ssh->handshake->kexId);
-                if (primeId == ECC_CURVE_INVALID)
-                    ret = WS_INVALID_PRIME_CURVE;
-
-                if (ret == 0)
-                    ret = wc_ecc_init_ex(pubKey, heap, INVALID_DEVID);
-                if (ret == 0)
-                    ret = wc_ecc_init_ex(privKey, heap, INVALID_DEVID);
-#ifdef HAVE_WC_ECC_SET_RNG
-                if (ret == 0)
-                    ret = wc_ecc_set_rng(privKey, ssh->rng);
-#endif
-
-                if (ret == 0)
-                    ret = wc_ecc_import_x963_ex(ssh->handshake->e,
-                                                ssh->handshake->eSz,
-                                                pubKey, primeId);
-
-                if (ret == 0)
-                    ret = wc_ecc_make_key_ex(ssh->rng,
-                                         wc_ecc_get_curve_size_from_id(primeId),
-                                         privKey, primeId);
-                if (ret == 0) {
-                    PRIVATE_KEY_UNLOCK();
-                    ret = wc_ecc_export_x963(privKey, f_ptr, &fSz);
-                    PRIVATE_KEY_LOCK();
-                }
-                if (ret == 0) {
-                    PRIVATE_KEY_UNLOCK();
-                    ret = wc_ecc_shared_secret(privKey, pubKey,
-                                               ssh->k, &ssh->kSz);
-                    PRIVATE_KEY_LOCK();
-                }
-                wc_ecc_free(privKey);
-                wc_ecc_free(pubKey);
-            #ifdef WOLFSSH_SMALL_STACK
-                WFREE(pubKey, heap, DYNTYPE_PUBKEY);
-                WFREE(privKey, heap, DYNTYPE_PRIVKEY);
-                pubKey  = NULL;
-                privKey = NULL;
-            #endif
-#endif /* !defined(WOLFSSH_NO_ECDH) */
+                ret = KeyAgreeEcdh_server(ssh, hashId, f_ptr, &fSz);
             }
-#ifndef WOLFSSH_NO_CURVE25519_SHA256
             if (useCurve25519) {
-#ifdef WOLFSSH_SMALL_STACK
-                curve25519_key *pubKey = NULL, *privKey = NULL;
-                pubKey = (curve25519_key*)WMALLOC(sizeof(curve25519_key),
-                                                  heap, DYNTYPE_PUBKEY);
-                privKey = (curve25519_key*)WMALLOC(sizeof(curve25519_key),
-                                                   heap, DYNTYPE_PRIVKEY);
-                if (pubKey == NULL || privKey == NULL) {
-                    ret = WS_MEMORY_E;
-                }
-#else
-                curve25519_key pubKey[1], privKey[1];
-#endif
-
-                if (ret == 0)
-                    ret = wc_curve25519_init_ex(pubKey, heap,
-                                                INVALID_DEVID);
-                if (ret == 0)
-                    ret = wc_curve25519_init_ex(privKey, heap,
-                                                INVALID_DEVID);
-                if (ret == 0)
-                    ret = wc_curve25519_check_public(ssh->handshake->e,
-                              ssh->handshake->eSz, EC25519_LITTLE_ENDIAN);
-                if (ret == 0)
-                    ret = wc_curve25519_import_public_ex(
-                              ssh->handshake->e, ssh->handshake->eSz,
-                              pubKey, EC25519_LITTLE_ENDIAN);
-
-                if (ret == 0)
-                    ret = wc_curve25519_make_key(ssh->rng,
-                              CURVE25519_KEYSIZE, privKey);
-
-                if (ret == 0) {
-                    PRIVATE_KEY_UNLOCK();
-                    ret = wc_curve25519_export_public_ex(privKey,
-                              f_ptr, &fSz, EC25519_LITTLE_ENDIAN);
-                    PRIVATE_KEY_LOCK();
-                }
-
-                if (ret == 0) {
-                    PRIVATE_KEY_UNLOCK();
-                    ret = wc_curve25519_shared_secret_ex(privKey, pubKey,
-                              ssh->k, &ssh->kSz, EC25519_LITTLE_ENDIAN);
-                    PRIVATE_KEY_LOCK();
-                }
-                wc_curve25519_free(privKey);
-                wc_curve25519_free(pubKey);
-#ifdef WOLFSSH_SMALL_STACK
-                WFREE(pubKey, heap, DYNTYPE_PUBKEY);
-                WFREE(privKey, heap, DYNTYPE_PRIVKEY);
-                pubKey  = NULL;
-                privKey = NULL;
-#endif
+                ret = KeyAgreeX25519_server(ssh, hashId, f_ptr, &fSz);
             }
-#endif /* ! WOLFSSH_NO_CURVE25519_SHA256 */
-#ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
             else if (useEccKyber) {
-                /* This is a hybrid KEM. In this case, I need to generate my ECC
-                 * keypair, send the public one, use the private one to generate
-                 * the shared secret, use the post-quantum public key to
-                 * generate and encapsulate the shared secret and send the
-                 * ciphertext. */
-                OQS_KEM* kem = NULL;
-                int primeId;
-                ret = 0;
-            #ifdef WOLFSSH_SMALL_STACK
-                ecc_key *pubKey = NULL, *privKey = NULL;
-                pubKey = (ecc_key*)WMALLOC(sizeof(ecc_key), heap,
-                        DYNTYPE_PUBKEY);
-                privKey = (ecc_key*)WMALLOC(sizeof(ecc_key), heap,
-                        DYNTYPE_PRIVKEY);
-                if (pubKey == NULL || privKey == NULL) {
-                    ret = WS_MEMORY_E;
-                }
-            #else
-                ecc_key pubKey[1];
-                ecc_key privKey[1];
-            #endif
-
-                if (ret == 0) {
-                    XMEMSET(pubKey, 0, sizeof(*pubKey));
-                    XMEMSET(privKey, 0, sizeof(*privKey));
-
-                    primeId = wcPrimeForId(ssh->handshake->kexId);
-                    if (primeId == ECC_CURVE_INVALID)
-                        ret = WS_INVALID_PRIME_CURVE;
-                }
-
-                if (ret == 0) {
-                    kem = OQS_KEM_new(OQS_KEM_alg_kyber_512);
-                    if (kem == NULL) {
-                        ret = WS_INVALID_ALGO_ID;
-                    }
-                }
-
-                if ((ret == 0) &&
-                    (ssh->handshake->eSz <= (word32)kem->length_public_key)) {
-                    ret = WS_BUFFER_E;
-                }
-
-                if (ret == 0) {
-                    if (OQS_KEM_encaps(kem, f_ptr, ssh->k,
-                        ssh->handshake->e) != OQS_SUCCESS) {
-                        ret = WS_PUBKEY_REJECTED_E;
-                    }
-                }
-
-                if (ret == 0) {
-                    fSz -= kem->length_ciphertext;
-                    ssh->kSz -= kem->length_shared_secret;
-                }
-                else {
-                    fSz = 0;
-                    ssh->kSz = 0;
-                    WLOG(WS_LOG_ERROR,
-                         "Generate ECC-kyber (encap) shared secret failed, %d",
-                         ret);
-                }
-
-                if (ret == 0) {
-                    ret = wc_ecc_init_ex(pubKey, heap, INVALID_DEVID);
-                }
-                if (ret == 0) {
-                    ret = wc_ecc_init_ex(privKey, heap, INVALID_DEVID);
-                }
-#ifdef HAVE_WC_ECC_SET_RNG
-                if (ret == 0) {
-                    ret = wc_ecc_set_rng(privKey, ssh->rng);
-                }
-#endif
-                if (ret == 0) {
-                    ret = wc_ecc_import_x963_ex(
-                        ssh->handshake->e + kem->length_public_key,
-                        ssh->handshake->eSz - (word32)kem->length_public_key,
-                        pubKey, primeId);
-                }
-                if (ret == 0) {
-                    ret = wc_ecc_make_key_ex(ssh->rng,
-                              wc_ecc_get_curve_size_from_id(primeId),
-                              privKey, primeId);
-                }
-                if (ret == 0) {
-                    PRIVATE_KEY_UNLOCK();
-                    ret = wc_ecc_export_x963(privKey,
-                              f_ptr + kem->length_ciphertext, &fSz);
-                    fSz += kem->length_ciphertext;
-                    PRIVATE_KEY_LOCK();
-                }
-                if (ret == 0) {
-                    word32 tmp_kSz = ssh->kSz;
-                    PRIVATE_KEY_UNLOCK();
-                    ret = wc_ecc_shared_secret(privKey, pubKey,
-                              ssh->k + kem->length_shared_secret, &tmp_kSz);
-                    PRIVATE_KEY_LOCK();
-                    ssh->kSz = (word32)kem->length_shared_secret + tmp_kSz;
-                }
-                wc_ecc_free(privKey);
-                wc_ecc_free(pubKey);
-            #ifdef WOLFSSH_SMALL_STACK
-                WFREE(pubKey, heap, DYNTYPE_PUBKEY);
-                WFREE(privKey, heap, DYNTYPE_PRIVKEY);
-                pubKey  = NULL;
-                privKey = NULL;
-            #endif
-                if (kem != NULL) {
-                    OQS_KEM_free(kem);
-                    kem = NULL;
-                }
-
-                /* Replace the concatenated shared secrets with the hash. That
-                 * will become the new shared secret.*/
-                if (ret == 0) {
-                    sharedSecretHashSz = wc_HashGetDigestSize(hashId);
-                    sharedSecretHash = (byte *)WMALLOC(sharedSecretHashSz, heap,
-                                                       DYNTYPE_PRIVKEY);
-                    if (sharedSecretHash == NULL) {
-                        ret = WS_MEMORY_E;
-                    }
-                }
-                if (ret == 0) {
-                    ret = wc_Hash(hashId, ssh->k, ssh->kSz, sharedSecretHash,
-                                  sharedSecretHashSz);
-                }
-                if (ret == 0) {
-                    XMEMCPY(ssh->k, sharedSecretHash, sharedSecretHashSz);
-                    ssh->kSz = sharedSecretHashSz;
-                }
-
-                WFREE(sharedSecretHash, heap, DYNTYPE_PRIVKEY);
-                sharedSecretHash = NULL;
-            }
-#endif /* !WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256 */
-            else {
-                /* This should never happen */
-                ret = WS_ERROR;
+                ret = KeyAgreeEcdhKyber1_server(ssh, hashId, f_ptr, &fSz);
             }
         }
 
         /* Hash in the server's DH f-value. */
-        if (ret == 0
-#ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
-            && !useEccKyber
-#endif
-#ifndef WOLFSSH_NO_CURVE25519_SHA256
-            && !useCurve25519
-#endif
-           ) {
+        if (ret == 0 && (useDh || useEcc)) {
             ret = CreateMpint(f_ptr, &fSz, &fPad);
         }
         if (ret == 0) {
@@ -10543,11 +10610,7 @@ int SendKexDhReply(WOLFSSH* ssh)
         }
 
         /* Hash in the shared secret K. */
-        if (ret == 0
-#ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
-            && !useEccKyber
-#endif
-           ) {
+        if (ret == 0 && !useEccKyber) {
             ret = CreateMpint(ssh->k, &ssh->kSz, &kPad);
         }
         if (ret == 0) {
@@ -10742,11 +10805,8 @@ int SendKexDhReply(WOLFSSH* ssh)
     }
 
     if (ret == WS_SUCCESS) {
-        int doKeyPadding = 1;
-#if !defined(WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256)
-        doKeyPadding = !useEccKyber;
-#endif
-        ret = GenerateKeys(ssh, hashId, doKeyPadding);
+        /* If we aren't using EccKyber, use padding. */
+        ret = GenerateKeys(ssh, hashId, !useEccKyber);
     }
 
     /* Get the buffer, copy the packet data, once f is laid into the buffer,

From 6e93b92b22bcc72eb01f7ff7c8bbc0c816a8eb7c Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Wed, 1 May 2024 19:43:12 -0700
Subject: [PATCH 3/3] Key Agree Update

1. Add a parameter to the client key agree functions for the hashId.
   It's only really used for EcdhKyber1, but it keeps the functions
   parallel.
2. Add and update some top-of-function comments for the key agree
   functions.
3. Renamed the X25519 key agreement functions to Curve25519 to match the
   naming in the RFC.
4. Removed the temporary hashId local in the client EcdhKyber1
   function.
5. Messed around with some variable declarations in a few of the
   functions.
6. Fix a couple breaks for small stack build.
7. Fix where GEX-SHA2 key exchange wasn't allowed to work.
8. Disable EcdhKyber1 is ECDH-NISTP256 is disabled.
---
 src/internal.c     | 171 ++++++++++++++++++++++++++++++++-------------
 wolfssh/internal.h |   6 +-
 2 files changed, 127 insertions(+), 50 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index 982b8245..8dbc1d82 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -4603,14 +4603,19 @@ static void FreePubKey(struct wolfSSH_sigKeyBlock *p)
 
 
 /* KeyAgreeDh_client
+ * hashId - wolfCrypt hash type ID used
  * f - peer public key
  * fSz - peer public key size
  */
-static int KeyAgreeDh_client(WOLFSSH* ssh, const byte* f, word32 fSz)
+static int KeyAgreeDh_client(WOLFSSH* ssh, byte hashId,
+        const byte* f, word32 fSz)
 #ifndef WOLFSSH_NO_DH
 {
     int ret;
 
+    WLOG(WS_LOG_DEBUG, "Entering KeyAgreeDh_client()");
+    WOLFSSH_UNUSED(hashId);
+
     PRIVATE_KEY_UNLOCK();
     ret = wc_DhAgree(&ssh->handshake->privKey.dh,
                      ssh->k, &ssh->kSz,
@@ -4624,11 +4629,14 @@ static int KeyAgreeDh_client(WOLFSSH* ssh, const byte* f, word32 fSz)
     }
     ForceZero(ssh->handshake->x, ssh->handshake->xSz);
     wc_FreeDhKey(&ssh->handshake->privKey.dh);
+
+    WLOG(WS_LOG_DEBUG, "Leaving KeyAgreeDh_client(), ret = %d", ret);
     return ret;
 }
 #else /* WOLFSSH_NO_DH */
 {
     WOLFSSH_UNUSED(ssh);
+    WOLFSSH_UNUSED(hashId);
     WOLFSSH_UNUSED(f);
     WOLFSSH_UNUSED(fSz);
     return WS_INVALID_ALGO_ID;
@@ -4637,10 +4645,12 @@ static int KeyAgreeDh_client(WOLFSSH* ssh, const byte* f, word32 fSz)
 
 
 /* KeyAgreeEcdh_client
+ * hashId - wolfCrypt hash type ID used
  * f - peer public key
  * fSz - peer public key size
  */
-static int KeyAgreeEcdh_client(WOLFSSH* ssh, const byte* f, word32 fSz)
+static int KeyAgreeEcdh_client(WOLFSSH* ssh, byte hashId,
+        const byte* f, word32 fSz)
 #ifndef WOLFSSH_NO_ECDH
 {
     int ret = WS_SUCCESS;
@@ -4648,6 +4658,10 @@ static int KeyAgreeEcdh_client(WOLFSSH* ssh, const byte* f, word32 fSz)
     #ifndef WOLFSSH_SMALL_STACK
         ecc_key key_s;
     #endif
+
+    WLOG(WS_LOG_DEBUG, "Entering KeyAgreeEcdh_client()");
+    WOLFSSH_UNUSED(hashId);
+
     #ifdef WOLFSSH_SMALL_STACK
         key_ptr = (ecc_key*)WMALLOC(sizeof(ecc_key),
                 ssh->ctx->heap, DYNTYPE_PRIVKEY);
@@ -4682,11 +4696,14 @@ static int KeyAgreeEcdh_client(WOLFSSH* ssh, const byte* f, word32 fSz)
     }
     #endif
     wc_ecc_free(&ssh->handshake->privKey.ecc);
+
+    WLOG(WS_LOG_DEBUG, "Leaving KeyAgreeEcdh_client(), ret = %d", ret);
     return ret;
 }
 #else /* WOLFSSH_NO_ECDH */
 {
     WOLFSSH_UNUSED(ssh);
+    WOLFSSH_UNUSED(hashId);
     WOLFSSH_UNUSED(f);
     WOLFSSH_UNUSED(fSz);
     return WS_INVALID_ALGO_ID;
@@ -4694,16 +4711,21 @@ static int KeyAgreeEcdh_client(WOLFSSH* ssh, const byte* f, word32 fSz)
 #endif /* WOLFSSH_NO_ECDH */
 
 
-/* KeyAgreeX25519_client
+/* KeyAgreeCurve25519_client
+ * hashId - wolfCrypt hash type ID used
  * f - peer public key
  * fSz - peer public key size
  */
-static int KeyAgreeX25519_client(WOLFSSH* ssh, const byte* f, word32 fSz)
+static int KeyAgreeCurve25519_client(WOLFSSH* ssh, byte hashId,
+        const byte* f, word32 fSz)
 #ifndef WOLFSSH_NO_CURVE25519_SHA256
 {
     int ret;
     curve25519_key pub;
 
+    WLOG(WS_LOG_DEBUG, "Entering KeyAgreeCurve25519_client()");
+    WOLFSSH_UNUSED(hashId);
+
     ret = wc_curve25519_init(&pub);
     if (ret == 0) {
         ret = wc_curve25519_check_public(f, fSz,
@@ -4731,11 +4753,13 @@ static int KeyAgreeX25519_client(WOLFSSH* ssh, const byte* f, word32 fSz)
     wc_curve25519_free(&pub);
     wc_curve25519_free(&ssh->handshake->privKey.curve25519);
 
+    WLOG(WS_LOG_DEBUG, "Leaving KeyAgreeCurve25519_client(), ret = %d", ret);
     return ret;
 }
 #else /* WOLFSSH_NO_CURVE25519_SHA256 */
 {
     WOLFSSH_UNUSED(ssh);
+    WOLFSSH_UNUSED(hashId);
     WOLFSSH_UNUSED(f);
     WOLFSSH_UNUSED(fSz);
     return WS_INVALID_ALGO_ID;
@@ -4744,17 +4768,20 @@ static int KeyAgreeX25519_client(WOLFSSH* ssh, const byte* f, word32 fSz)
 
 
 /* KeyAgreeEcdhKyber1_client
+ * hashId - wolfCrypt hash type ID used
  * f - peer public key
  * fSz - peer public key size
  */
-static int KeyAgreeEcdhKyber1_client(WOLFSSH* ssh, const byte* f, word32 fSz)
+static int KeyAgreeEcdhKyber1_client(WOLFSSH* ssh, byte hashId,
+        const byte* f, word32 fSz)
 #ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
 {
     int ret = WS_SUCCESS;
-    byte hashId = WC_HASH_TYPE_SHA256;
     byte sharedSecretHashSz = 0;
     byte *sharedSecretHash = NULL;
     ecc_key *key_ptr = NULL;
+    OQS_KEM *kem;
+
     #ifndef WOLFSSH_SMALL_STACK
         ecc_key key_s;
     #endif
@@ -4767,10 +4794,13 @@ static int KeyAgreeEcdhKyber1_client(WOLFSSH* ssh, const byte* f, word32 fSz)
     #else /* ! WOLFSSH_SMALL_STACK */
         key_ptr = &key_s;
     #endif /* WOLFSSH_SMALL_STACK */
+
+    WLOG(WS_LOG_DEBUG, "Entering KeyAgreeEcdhKyber1_client()");
+
     /* This is a a hybrid of ECDHE and a post-quantum KEM. In this
      * case, I need to generated the ECC shared secret and
      * decapsulate the ciphertext of the post-quantum KEM. */
-    OQS_KEM* kem = OQS_KEM_new(OQS_KEM_alg_kyber_512);
+    kem = OQS_KEM_new(OQS_KEM_alg_kyber_512);
     if (kem == NULL) {
         ret = WS_MEMORY_E;
     }
@@ -4854,11 +4884,14 @@ static int KeyAgreeEcdhKyber1_client(WOLFSSH* ssh, const byte* f, word32 fSz)
         ForceZero(sharedSecretHash, sharedSecretHashSz);
         WFREE(sharedSecretHash, ssh->ctx->heap, DYNTYPE_PRIVKEY);
     }
+
+    WLOG(WS_LOG_DEBUG, "Leaving KeyAgreeEcdhKyber1_client(), ret = %d", ret);
     return ret;
 }
 #else /* WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256 */
 {
     WOLFSSH_UNUSED(ssh);
+    WOLFSSH_UNUSED(hashId);
     WOLFSSH_UNUSED(f);
     WOLFSSH_UNUSED(fSz);
     return WS_INVALID_ALGO_ID;
@@ -4867,10 +4900,11 @@ static int KeyAgreeEcdhKyber1_client(WOLFSSH* ssh, const byte* f, word32 fSz)
 
 
 /* KeyAgree_client
+ * hashId - wolfCrypt hash type ID used
  * f - peer public key
  * fSz - peer public key size
  */
-static int KeyAgree_client(WOLFSSH* ssh, const byte* f, word32 fSz)
+static int KeyAgree_client(WOLFSSH* ssh, byte hashId, const byte* f, word32 fSz)
 {
     int ret;
 
@@ -4880,16 +4914,16 @@ static int KeyAgree_client(WOLFSSH* ssh, const byte* f, word32 fSz)
     ssh->kSz = MAX_KEX_KEY_SZ;
 
     if (ssh->handshake->useDh) {
-        ret = KeyAgreeDh_client(ssh, f, fSz);
+        ret = KeyAgreeDh_client(ssh, hashId, f, fSz);
     }
     else if (ssh->handshake->useEcc) {
-        ret = KeyAgreeEcdh_client(ssh, f, fSz);
+        ret = KeyAgreeEcdh_client(ssh, hashId, f, fSz);
     }
     else if (ssh->handshake->useCurve25519) {
-        ret = KeyAgreeX25519_client(ssh, f, fSz);
+        ret = KeyAgreeCurve25519_client(ssh, hashId, f, fSz);
     }
     else if (ssh->handshake->useEccKyber) {
-        ret = KeyAgreeEcdhKyber1_client(ssh, f, fSz);
+        ret = KeyAgreeEcdhKyber1_client(ssh, hashId, f, fSz);
     }
     else {
         ret = WS_INVALID_ALGO_ID;
@@ -5100,7 +5134,7 @@ static int DoKexDhReply(WOLFSSH* ssh, byte* buf, word32 len, word32* idx)
         ret = ParsePubKey(ssh, sigKeyBlock_ptr, pubKey, pubKeySz);
         /* Generate and hash in the shared secret */
         if (ret == WS_SUCCESS) {
-            ret = KeyAgree_client(ssh, f, fSz);
+            ret = KeyAgree_client(ssh, hashId, f, fSz);
         }
 
         /* Hash in the shared secret K. */
@@ -10046,6 +10080,11 @@ int wolfSSH_RsaVerify(byte *sig, word32 sigSz,
 #endif /* WOLFSSH_NO_RSA */
 
 
+/* KeyAgreeDh_server
+ * hashId - wolfCrypt hash type ID used
+ * f - peer public key
+ * fSz - peer public key size
+ */
 static int KeyAgreeDh_server(WOLFSSH* ssh, byte hashId, byte* f, word32* fSz)
 #ifndef WOLFSSH_NO_DH
 {
@@ -10057,9 +10096,9 @@ static int KeyAgreeDh_server(WOLFSSH* ssh, byte hashId, byte* f, word32* fSz)
     word32 primeGroupSz = 0;
     word32 generatorSz = 0;
     #ifdef WOLFSSH_SMALL_STACK
-    DhKey *privKey = (DhKey*)WMALLOC(sizeof(DhKey), heap,
+    DhKey *privKey = (DhKey*)WMALLOC(sizeof(DhKey), ssh->ctx->heap,
             DYNTYPE_PRIVKEY);
-    y_ptr = (byte*)WMALLOC(ySz, heap, DYNTYPE_PRIVKEY);
+    y_ptr = (byte*)WMALLOC(ySz, ssh->ctx->heap, DYNTYPE_PRIVKEY);
     if (privKey == NULL || y_ptr == NULL)
         ret = WS_MEMORY_E;
     #else
@@ -10068,6 +10107,7 @@ static int KeyAgreeDh_server(WOLFSSH* ssh, byte hashId, byte* f, word32* fSz)
     y_ptr = y_s;
     #endif
 
+    WLOG(WS_LOG_DEBUG, "Entering KeyAgreeDh_server()");
     WOLFSSH_UNUSED(hashId);
 
     if (ret == WS_SUCCESS) {
@@ -10094,11 +10134,12 @@ static int KeyAgreeDh_server(WOLFSSH* ssh, byte hashId, byte* f, word32* fSz)
     }
     #ifdef WOLFSSH_SMALL_STACK
     if (y_ptr)
-        WFREE(y_ptr, heap, DYNTYPE_PRIVKEY);
+        WFREE(y_ptr, ssh->ctx->heap, DYNTYPE_PRIVKEY);
     if (privKey) {
-        WFREE(privKey, heap, DYNTYPE_PRIVKEY);
+        WFREE(privKey, ssh->ctx->heap, DYNTYPE_PRIVKEY);
     }
     #endif
+    WLOG(WS_LOG_DEBUG, "Leaving KeyAgreeDh_server(), ret = %d", ret);
     return ret;
 }
 #else /* WOLFSSH_NO_DH */
@@ -10112,6 +10153,11 @@ static int KeyAgreeDh_server(WOLFSSH* ssh, byte hashId, byte* f, word32* fSz)
 #endif /* WOLFSSH_NO_DH */
 
 
+/* KeyAgreeEcdh_server
+ * hashId - wolfCrypt hash type ID used
+ * f - peer public key
+ * fSz - peer public key size
+ */
 static int KeyAgreeEcdh_server(WOLFSSH* ssh, byte hashId, byte* f, word32* fSz)
 #ifndef WOLFSSH_NO_ECDH
 {
@@ -10132,7 +10178,9 @@ static int KeyAgreeEcdh_server(WOLFSSH* ssh, byte hashId, byte* f, word32* fSz)
 #endif
     int primeId;
 
+    WLOG(WS_LOG_DEBUG, "Entering KeyAgreeEcdh_server()");
     WOLFSSH_UNUSED(hashId);
+
     heap = ssh->ctx->heap;
     primeId  = wcPrimeForId(ssh->handshake->kexId);
     if (primeId == ECC_CURVE_INVALID)
@@ -10175,6 +10223,7 @@ static int KeyAgreeEcdh_server(WOLFSSH* ssh, byte hashId, byte* f, word32* fSz)
     pubKey  = NULL;
     privKey = NULL;
 #endif
+    WLOG(WS_LOG_DEBUG, "Leaving KeyAgreeEcdh_server(), ret = %d", ret);
     return ret;
 }
 #else /* WOLFSSH_NO_ECDH */
@@ -10188,7 +10237,12 @@ static int KeyAgreeEcdh_server(WOLFSSH* ssh, byte hashId, byte* f, word32* fSz)
 #endif /* WOLFSSH_NO_ECDH */
 
 
-static int KeyAgreeX25519_server(WOLFSSH* ssh, byte hashId,
+/* KeyAgreeCurve25519_server
+ * hashId - wolfCrypt hash type ID used
+ * f - peer public key
+ * fSz - peer public key size
+ */
+static int KeyAgreeCurve25519_server(WOLFSSH* ssh, byte hashId,
         byte* f, word32* fSz)
 #ifndef WOLFSSH_NO_CURVE25519_SHA256
 {
@@ -10207,7 +10261,9 @@ static int KeyAgreeX25519_server(WOLFSSH* ssh, byte hashId,
     curve25519_key pubKey[1], privKey[1];
 #endif
 
+    WLOG(WS_LOG_DEBUG, "Entering KeyAgreeCurve25519_server()");
     WOLFSSH_UNUSED(hashId);
+
     if (ret == 0)
         ret = wc_curve25519_init_ex(pubKey, heap, INVALID_DEVID);
     if (ret == 0)
@@ -10244,6 +10300,7 @@ static int KeyAgreeX25519_server(WOLFSSH* ssh, byte hashId,
     pubKey  = NULL;
     privKey = NULL;
 #endif
+    WLOG(WS_LOG_DEBUG, "Leaving KeyAgreeCurve25519_server(), ret = %d", ret);
     return ret;
 }
 #else /* WOLFSSH_NO_CURVE25519_SHA256 */
@@ -10257,34 +10314,45 @@ static int KeyAgreeX25519_server(WOLFSSH* ssh, byte hashId,
 #endif /* WOLFSSH_NO_CURVE25519_SHA256 */
 
 
+/* KeyAgreeEcdhKyber1_server
+ * hashId - wolfCrypt hash type ID used
+ * f - peer public key
+ * fSz - peer public key size
+ *
+ * This is a hybrid KEM. In this case, I need to generate my ECC
+ * keypair, send the public one, use the private one to generate
+ * the shared secret, use the post-quantum public key to
+ * generate and encapsulate the shared secret and send the
+ * ciphertext.
+ */
 static int KeyAgreeEcdhKyber1_server(WOLFSSH* ssh, byte hashId,
         byte* f, word32* fSz)
 #ifndef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
 {
     int ret = WS_SUCCESS;
-    void* heap = ssh->ctx->heap;
     byte sharedSecretHashSz = 0;
     byte *sharedSecretHash = NULL;
-    /* This is a hybrid KEM. In this case, I need to generate my ECC
-     * keypair, send the public one, use the private one to generate
-     * the shared secret, use the post-quantum public key to
-     * generate and encapsulate the shared secret and send the
-     * ciphertext. */
     OQS_KEM* kem = NULL;
+    ecc_key* pubKey = NULL;
+    ecc_key* privKey = NULL;
     int primeId;
-    ret = 0;
+#ifndef WOLFSSH_SMALL_STACK
+    ecc_key eccKeys[2];
+#endif
+
+    WLOG(WS_LOG_DEBUG, "Entering KeyAgreeEcdhKyber1_server()");
+
 #ifdef WOLFSSH_SMALL_STACK
-    ecc_key *pubKey = NULL, *privKey = NULL;
-    pubKey = (ecc_key*)WMALLOC(sizeof(ecc_key), heap,
-            DYNTYPE_PUBKEY);
-    privKey = (ecc_key*)WMALLOC(sizeof(ecc_key), heap,
-            DYNTYPE_PRIVKEY);
+    pubKey = (ecc_key*)WMALLOC(sizeof(ecc_key),
+            ssh->ctx->heap, DYNTYPE_PUBKEY);
+    privKey = (ecc_key*)WMALLOC(sizeof(ecc_key),
+            ssh->ctx->heap, DYNTYPE_PRIVKEY);
     if (pubKey == NULL || privKey == NULL) {
         ret = WS_MEMORY_E;
     }
 #else
-    ecc_key pubKey[1];
-    ecc_key privKey[1];
+    pubKey = &eccKeys[0];
+    privKey = &eccKeys[1];
 #endif
 
     if (ret == 0) {
@@ -10299,7 +10367,7 @@ static int KeyAgreeEcdhKyber1_server(WOLFSSH* ssh, byte hashId,
     if (ret == 0) {
         kem = OQS_KEM_new(OQS_KEM_alg_kyber_512);
         if (kem == NULL) {
-            ret = WS_INVALID_ALGO_ID;
+            ret = WS_MEMORY_E;
         }
     }
 
@@ -10309,9 +10377,13 @@ static int KeyAgreeEcdhKyber1_server(WOLFSSH* ssh, byte hashId,
     }
 
     if (ret == 0) {
-        if (OQS_KEM_encaps(kem, f, ssh->k,
-            ssh->handshake->e) != OQS_SUCCESS) {
+        if (OQS_KEM_encaps(kem, f, ssh->k, ssh->handshake->e) != OQS_SUCCESS) {
             ret = WS_PUBKEY_REJECTED_E;
+            WLOG(WS_LOG_ERROR,
+                 "Generate ECC-kyber (encap) shared secret failed, %d",
+                 ret);
+            *fSz = 0;
+            ssh->kSz = 0;
         }
     }
 
@@ -10328,10 +10400,10 @@ static int KeyAgreeEcdhKyber1_server(WOLFSSH* ssh, byte hashId,
     }
 
     if (ret == 0) {
-        ret = wc_ecc_init_ex(pubKey, heap, INVALID_DEVID);
+        ret = wc_ecc_init_ex(pubKey, ssh->ctx->heap, INVALID_DEVID);
     }
     if (ret == 0) {
-        ret = wc_ecc_init_ex(privKey, heap, INVALID_DEVID);
+        ret = wc_ecc_init_ex(privKey, ssh->ctx->heap, INVALID_DEVID);
     }
 #ifdef HAVE_WC_ECC_SET_RNG
     if (ret == 0) {
@@ -10351,8 +10423,7 @@ static int KeyAgreeEcdhKyber1_server(WOLFSSH* ssh, byte hashId,
     }
     if (ret == 0) {
         PRIVATE_KEY_UNLOCK();
-        ret = wc_ecc_export_x963(privKey,
-                  f + kem->length_ciphertext, fSz);
+        ret = wc_ecc_export_x963(privKey, f + kem->length_ciphertext, fSz);
         PRIVATE_KEY_LOCK();
         *fSz += kem->length_ciphertext;
     }
@@ -10367,10 +10438,10 @@ static int KeyAgreeEcdhKyber1_server(WOLFSSH* ssh, byte hashId,
     wc_ecc_free(privKey);
     wc_ecc_free(pubKey);
 #ifdef WOLFSSH_SMALL_STACK
-    WFREE(pubKey, heap, DYNTYPE_PUBKEY);
-    WFREE(privKey, heap, DYNTYPE_PRIVKEY);
-    pubKey  = NULL;
-    privKey = NULL;
+    if (pubKey)
+        WFREE(pubKey, ssh->ctx->heap, DYNTYPE_PUBKEY);
+    if (privKey)
+        WFREE(privKey, ssh->ctx->heap, DYNTYPE_PUBKEY);
 #endif
     if (kem != NULL) {
         OQS_KEM_free(kem);
@@ -10381,8 +10452,8 @@ static int KeyAgreeEcdhKyber1_server(WOLFSSH* ssh, byte hashId,
      * will become the new shared secret.*/
     if (ret == 0) {
         sharedSecretHashSz = wc_HashGetDigestSize(hashId);
-        sharedSecretHash = (byte *)WMALLOC(sharedSecretHashSz, heap,
-                                           DYNTYPE_PRIVKEY);
+        sharedSecretHash = (byte *)WMALLOC(sharedSecretHashSz,
+                ssh->ctx->heap, DYNTYPE_PRIVKEY);
         if (sharedSecretHash == NULL) {
             ret = WS_MEMORY_E;
         }
@@ -10396,8 +10467,12 @@ static int KeyAgreeEcdhKyber1_server(WOLFSSH* ssh, byte hashId,
         ssh->kSz = sharedSecretHashSz;
     }
 
-    WFREE(sharedSecretHash, heap, DYNTYPE_PRIVKEY);
-    sharedSecretHash = NULL;
+    if (sharedSecretHash) {
+        ForceZero(sharedSecretHash, sharedSecretHashSz);
+        WFREE(sharedSecretHash, ssh->ctx->heap, DYNTYPE_PRIVKEY);
+    }
+
+    WLOG(WS_LOG_DEBUG, "Leaving KeyAgreeEcdhKyber1_server(), ret = %d", ret);
     return ret;
 }
 #else /* WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256 */
@@ -10586,7 +10661,7 @@ int SendKexDhReply(WOLFSSH* ssh)
                 ret = KeyAgreeEcdh_server(ssh, hashId, f_ptr, &fSz);
             }
             if (useCurve25519) {
-                ret = KeyAgreeX25519_server(ssh, hashId, f_ptr, &fSz);
+                ret = KeyAgreeCurve25519_server(ssh, hashId, f_ptr, &fSz);
             }
             else if (useEccKyber) {
                 ret = KeyAgreeEcdhKyber1_server(ssh, hashId, f_ptr, &fSz);
diff --git a/wolfssh/internal.h b/wolfssh/internal.h
index c8aeefea..8bfc9563 100644
--- a/wolfssh/internal.h
+++ b/wolfssh/internal.h
@@ -157,7 +157,8 @@ extern "C" {
     #undef WOLFSSH_NO_ECDH_SHA2_ED25519
     #define WOLFSSH_NO_ECDH_SHA2_ED25519
 #endif
-#if !defined(WOLFSSH_HAVE_LIBOQS) || defined(NO_SHA256)
+#if !defined(WOLFSSH_HAVE_LIBOQS) || defined(NO_SHA256) \
+    || defined(WOLFSSH_NO_ECDH_SHA2_NISTP256)
     #undef WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
     #define WOLFSSH_NO_ECDH_NISTP256_KYBER_LEVEL1_SHA256
 #endif
@@ -1120,7 +1121,8 @@ enum WS_MessageIds {
 };
 
 
-#define MSGID_KEXDH_LIMIT 30
+/* Allows the server to receive up to KEXDH GEX Request during KEX. */
+#define MSGID_KEXDH_LIMIT MSGID_KEXDH_GEX_REQUEST
 
 /* The endpoints should not allow message IDs greater than or
  * equal to msgid 80 before user authentication is complete.