From 367dd5a8f28073a779af52337349fee4b06ae4d3 Mon Sep 17 00:00:00 2001
From: Yosuke Shimizu <yosuke@wolfssl.com>
Date: Wed, 15 Apr 2026 10:43:15 +0900
Subject: [PATCH] Add regress test for DoKexDhReply

---
 tests/regress.c | 42 +++++++++++++++++++++++++++++++++++++++---
 1 file changed, 39 insertions(+), 3 deletions(-)

diff --git a/tests/regress.c b/tests/regress.c
index 321151b8..f5747a38 100644
--- a/tests/regress.c
+++ b/tests/regress.c
@@ -749,8 +749,8 @@ static void FreeKexReplyHarness(KexReplyHarness* harness)
     }
 }
 
-static void InitKexReplyHarness(KexReplyHarness* harness,
-        const char* keyAlgo, byte mutateReply)
+static void InitKexReplyHarnessEx(KexReplyHarness* harness,
+        const char* keyAlgo, byte mutateReply, byte skipPublicKeyCheck)
 {
     byte keyBuf[2048];
     word32 keySz;
@@ -781,7 +781,9 @@ static void InitKexReplyHarness(KexReplyHarness* harness,
 
     wolfSSH_SetUserAuth(harness->clientCtx, RegressionClientUserAuth);
     wolfSSH_SetUserAuth(harness->serverCtx, RegressionServerUserAuth);
-    wolfSSH_CTX_SetPublicKeyCheck(harness->clientCtx, AcceptAnyServerHostKey);
+    if (!skipPublicKeyCheck) {
+        wolfSSH_CTX_SetPublicKeyCheck(harness->clientCtx, AcceptAnyServerHostKey);
+    }
 
     keySz = LoadFileBuffer(REGRESS_SERVER_KEY_PATH, keyBuf, sizeof(keyBuf));
     AssertTrue(keySz > 0);
@@ -802,6 +804,12 @@ static void InitKexReplyHarness(KexReplyHarness* harness,
             WS_SUCCESS);
 }
 
+static void InitKexReplyHarness(KexReplyHarness* harness,
+        const char* keyAlgo, byte mutateReply)
+{
+    InitKexReplyHarnessEx(harness, keyAlgo, mutateReply, 0);
+}
+
 static int IsHandshakeRetryable(int err)
 {
     return err == WS_WANT_READ || err == WS_WANT_WRITE ||
@@ -903,6 +911,33 @@ static void TestKexDhReplyRejectsRsaSha2_512SigNameDowngrade(void)
 }
 #endif
 
+static void AssertHandshakeRejectsWithNoPublicKeyCheck(const char* keyAlgo)
+{
+    KexReplyHarness harness;
+    KexReplyRunResult result;
+
+    InitKexReplyHarnessEx(&harness, keyAlgo, 0, 1 /* skipPublicKeyCheck */);
+    RunKexReplyHandshake(&harness, &result);
+
+    AssertFalse(result.clientSuccess);
+    AssertTrue(result.clientRet == WS_FATAL_ERROR);
+    AssertTrue(result.clientErr != WS_WANT_READ && result.clientErr != WS_WANT_WRITE);
+    AssertIntEQ(result.clientErr, WS_PUBKEY_REJECTED_E);
+    AssertFalse(harness.client->connectState >= CONNECT_KEYED);
+
+    FreeKexReplyHarness(&harness);
+}
+
+static void TestKexDhReplyRejectsNoPublicKeyCheck(void)
+{
+#ifndef WOLFSSH_NO_RSA_SHA2_256
+    AssertHandshakeRejectsWithNoPublicKeyCheck("rsa-sha2-256");
+#endif
+#ifndef WOLFSSH_NO_RSA_SHA2_512
+    AssertHandshakeRejectsWithNoPublicKeyCheck("rsa-sha2-512");
+#endif
+}
+
 #endif /* KEXDH_REPLY_REGRESS_KEX_ALGO */
 
 static void AssertChannelOpenFailResponse(const ChannelOpenHarness* harness,
@@ -1667,6 +1702,7 @@ int main(int argc, char** argv)
     #ifndef WOLFSSH_NO_RSA_SHA2_512
     TestKexDhReplyRejectsRsaSha2_512SigNameDowngrade();
     #endif
+    TestKexDhReplyRejectsNoPublicKeyCheck();
 #endif
 
 #ifdef WOLFSSH_SFTP