From 2eaa3726e7ef5bbfb2ac6b0d9bd360db66c1c599 Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Fri, 6 Mar 2026 11:23:17 -0800
Subject: [PATCH] KEX Fix

1. When disabling mlkem768nistp256-sha256, there's a crash when doing a
   mlkem768x25519-sha256 KEX. It tries to do a DH key exchange instead
   of x25519 with ML-KEM. Modified the check.
2. Fixed a guard where the build was treating not having
   mlkem768nistp256-sha256 like a FIPS build disallowing using ML-KEM
   when generating sesssion keys. It shouldn't make that check.
3. Added a test case to the KEM test.
---
 src/internal.c | 12 ++----------
 tests/kex.c    |  4 ++++
 2 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index 61324e0d4..e40878584 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -2468,8 +2468,7 @@ int GenerateKey(byte hashId, byte keyId,
                 byte doKeyPad)
 #if defined(WOLFSSL_WOLFSSH) \
     && (LIBWOLFSSL_VERSION_HEX >= WOLFSSL_V5_0_0) \
-    && ((defined(HAVE_FIPS) && FIPS_VERSION_GE(5,2)) \
-        || defined(WOLFSSH_NO_NISTP256_MLKEM768_SHA256))
+    && (defined(HAVE_FIPS) && FIPS_VERSION_GE(5,2))
 /* Cannot use the SSH KDF with ML-KEM. With ML-KEM, doKeyPad must be false,
  * and the FIPS SSH KDF doesn't handle no-padding. Also, the ML-KEM algorithm
  * isn't in our FIPS boundary. */
@@ -13173,14 +13172,7 @@ int SendKexDhInit(WOLFSSH* ssh)
 
 
     if (ret == WS_SUCCESS) {
-        if (!ssh->handshake->useEcc
-#ifndef WOLFSSH_NO_NISTP256_MLKEM768_SHA256
-            && !ssh->handshake->useEccMlKem
-#endif
-#ifndef WOLFSSH_NO_CURVE25519_SHA256
-            && !ssh->handshake->useCurve25519
-#endif
-) {
+        if (ssh->handshake->useDh) {
 #ifndef WOLFSSH_NO_DH
             DhKey* privKey = &ssh->handshake->privKey.dh;
 
diff --git a/tests/kex.c b/tests/kex.c
index 1ce924879..26dc1262a 100644
--- a/tests/kex.c
+++ b/tests/kex.c
@@ -345,6 +345,10 @@ int wolfSSH_KexTest(int argc, char** argv)
     AssertIntEQ(wolfSSH_KexTest_Connect("mlkem768x25519-sha256"),
             EXIT_SUCCESS);
 #endif
+#if !defined(WOLFSSH_NO_NISTP256_MLKEM768_SHA256)
+    AssertIntEQ(wolfSSH_KexTest_Connect("mlkem768nistp256-sha256"),
+            EXIT_SUCCESS);
+#endif
 #if !defined(WOLFSSH_NO_NISTP384_MLKEM1024_SHA384)
     AssertIntEQ(wolfSSH_KexTest_Connect("mlkem1024nistp384-sha384"),
             EXIT_SUCCESS);