First Kex Packet Follows Test

Add a regression for checking the `first_kex_packet_follows` flag versus
the guesses for KEX algorithm and public key algorithm.

Author: ejohnstown

src/internal.c

@@ -854,6 +854,19 @@ int wolfSSH_TestIsMessageAllowed(WOLFSSH* ssh, byte msg, byte state)
 {
     return IsMessageAllowed(ssh, msg, state);
 }
+
+static int DoKexInit(WOLFSSH* ssh, byte* buf, word32 len, word32* idx);
+static int DoKexDhInit(WOLFSSH* ssh, byte* buf, word32 len, word32* idx);
+
+int wolfSSH_TestDoKexInit(WOLFSSH* ssh, byte* buf, word32 len, word32* idx)
+{
+    return DoKexInit(ssh, buf, len, idx);
+}
+
+int wolfSSH_TestDoKexDhInit(WOLFSSH* ssh, byte* buf, word32 len, word32* idx)
+{
+    return DoKexDhInit(ssh, buf, len, idx);
+}
 #endif
 
 

tests/regress.c

@@ -1619,6 +1619,149 @@ static void TestKeyboardResponseNullCtx(WOLFSSH* ssh)
 #endif /* WOLFSSH_KEYBOARD_INTERACTIVE */
 
 
+#if !defined(WOLFSSH_NO_ECDH_SHA2_NISTP256) \
+    && !defined(WOLFSSH_NO_RSA) \
+    && !defined(WOLFSSH_NO_CURVE25519_SHA256) \
+    && !defined(WOLFSSH_NO_RSA_SHA2_256)
+
+#define FPF_KEX_GOOD "ecdh-sha2-nistp256"
+#define FPF_KEX_BAD  "curve25519-sha256"
+#define FPF_KEY_GOOD "ssh-rsa"
+#define FPF_KEY_BAD  "rsa-sha2-256"
+
+static word32 BuildKexInitPayload(const char* kexList, const char* keyList,
+        byte firstPacketFollows, byte* out, word32 outSz)
+{
+    word32 idx = 0;
+
+    /* cookie */
+    AssertTrue(idx + COOKIE_SZ <= outSz);
+    WMEMSET(out + idx, 0, COOKIE_SZ);
+    idx += COOKIE_SZ;
+
+    idx = AppendString(out, outSz, idx, kexList);
+    idx = AppendString(out, outSz, idx, keyList);
+    idx = AppendString(out, outSz, idx, "aes128-ctr");
+    idx = AppendString(out, outSz, idx, "aes128-ctr");
+    idx = AppendString(out, outSz, idx, "hmac-sha2-256");
+    idx = AppendString(out, outSz, idx, "hmac-sha2-256");
+    idx = AppendString(out, outSz, idx, "none");
+    idx = AppendString(out, outSz, idx, "none");
+    idx = AppendString(out, outSz, idx, "");
+    idx = AppendString(out, outSz, idx, "");
+
+    idx = AppendByte(out, outSz, idx, firstPacketFollows);
+    idx = AppendUint32(out, outSz, idx, 0); /* reserved */
+
+    return idx;
+}
+
+typedef struct {
+    const char* description;
+    const char* kexList;
+    const char* keyList;
+    byte firstPacketFollows;
+    byte expectIgnore;
+} FirstPacketFollowsCase;
+
+static const FirstPacketFollowsCase firstPacketFollowsCases[] = {
+    { "follows=0, guesses irrelevant: flag stays off",
+      FPF_KEX_BAD "," FPF_KEX_GOOD, FPF_KEY_BAD "," FPF_KEY_GOOD, 0, 0 },
+    { "follows=1, both guesses match: do not skip",
+      FPF_KEX_GOOD, FPF_KEY_GOOD, 1, 0 },
+    { "follows=1, KEX guess wrong: skip",
+      FPF_KEX_BAD "," FPF_KEX_GOOD, FPF_KEY_GOOD, 1, 1 },
+    { "follows=1, host-key guess wrong: skip", /* regression case */
+      FPF_KEX_GOOD, FPF_KEY_BAD "," FPF_KEY_GOOD, 1, 1 },
+    { "follows=1, both guesses wrong: skip",
+      FPF_KEX_BAD "," FPF_KEX_GOOD, FPF_KEY_BAD "," FPF_KEY_GOOD, 1, 1 },
+};
+
+static void RunFirstPacketFollowsCase(const FirstPacketFollowsCase* tc)
+{
+    WOLFSSH_CTX* ctx;
+    WOLFSSH* ssh;
+    byte payload[512];
+    word32 payloadSz;
+    word32 idx = 0;
+
+    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_SERVER, NULL);
+    AssertNotNull(ctx);
+
+    ssh = wolfSSH_new(ctx);
+    AssertNotNull(ssh);
+
+    AssertIntEQ(wolfSSH_SetAlgoListKex(ssh, FPF_KEX_GOOD), WS_SUCCESS);
+    AssertIntEQ(wolfSSH_SetAlgoListKey(ssh, FPF_KEY_GOOD), WS_SUCCESS);
+
+    payloadSz = BuildKexInitPayload(tc->kexList, tc->keyList,
+            tc->firstPacketFollows, payload, sizeof(payload));
+
+    /* DoKexInit's tail hashes and sends a response; on a stripped-down
+     * WOLFSSH without a loaded host key or a primed peer proto id, that
+     * tail errors. We only care about the parse path up through
+     * first_packet_follows, where ignoreNextKexMsg is set. */
+    (void)wolfSSH_TestDoKexInit(ssh, payload, payloadSz, &idx);
+
+    AssertNotNull(ssh->handshake);
+    if (ssh->handshake->ignoreNextKexMsg != tc->expectIgnore) {
+        Fail(("ignoreNextKexMsg == %u (%s)",
+                    tc->expectIgnore, tc->description),
+             ("%u", ssh->handshake->ignoreNextKexMsg));
+    }
+
+    wolfSSH_free(ssh);
+    wolfSSH_CTX_free(ctx);
+}
+
+/* With ignoreNextKexMsg set, DoKexDhInit must consume the packet, clear the
+ * flag, and not advance clientState to CLIENT_KEXDH_INIT_DONE. */
+static void TestFirstPacketFollowsDhInitSkipped(void)
+{
+    WOLFSSH_CTX* ctx;
+    WOLFSSH* ssh;
+    byte dhPayload[8];
+    word32 idx = 0;
+    int ret;
+
+    ctx = wolfSSH_CTX_new(WOLFSSH_ENDPOINT_SERVER, NULL);
+    AssertNotNull(ctx);
+
+    ssh = wolfSSH_new(ctx);
+    AssertNotNull(ssh);
+
+    ssh->handshake = AllocHandshake(ssh);
+    ssh->handshake->ignoreNextKexMsg = 1;
+    ssh->clientState = CLIENT_KEXINIT_DONE;
+
+    /* Garbage payload — must never be parsed when skipped. */
+    WMEMSET(dhPayload, 0xAB, sizeof(dhPayload));
+
+    ret = wolfSSH_TestDoKexDhInit(ssh, dhPayload, sizeof(dhPayload), &idx);
+    AssertIntEQ(ret, WS_SUCCESS);
+    AssertIntEQ(idx, sizeof(dhPayload));
+    AssertIntEQ(ssh->handshake->ignoreNextKexMsg, 0);
+    AssertIntEQ(ssh->clientState, CLIENT_KEXINIT_DONE);
+
+    wolfSSH_free(ssh);
+    wolfSSH_CTX_free(ctx);
+}
+
+static void TestFirstPacketFollows(void)
+{
+    size_t i;
+    size_t n = sizeof(firstPacketFollowsCases)
+            / sizeof(firstPacketFollowsCases[0]);
+
+    for (i = 0; i < n; i++) {
+        RunFirstPacketFollowsCase(&firstPacketFollowsCases[i]);
+    }
+    TestFirstPacketFollowsDhInitSkipped();
+}
+
+#endif /* first_packet_follows coverage guard */
+
+
 int main(int argc, char** argv)
 {
     WOLFSSH_CTX* ctx;
@@ -1653,6 +1796,11 @@ int main(int argc, char** argv)
     TestAgentChannelNullAgentSendsOpenFail();
 #endif
     TestKexInitRejectedWhenKeying(ssh);
+#if !defined(WOLFSSH_NO_ECDH_SHA2_NISTP256) && !defined(WOLFSSH_NO_RSA) \
+    && !defined(WOLFSSH_NO_CURVE25519_SHA256) \
+    && !defined(WOLFSSH_NO_RSA_SHA2_256)
+    TestFirstPacketFollows();
+#endif
     TestDisconnectSetsDisconnectError();
     TestClientBuffersIdempotent();
     TestPasswordEofNoCrash();

wolfssh/internal.h

@@ -1322,6 +1322,10 @@ enum WS_MessageIdLimits {
 #ifdef WOLFSSH_TEST_INTERNAL
     WOLFSSH_API int wolfSSH_TestIsMessageAllowed(WOLFSSH* ssh, byte msg,
             byte state);
+    WOLFSSH_API int wolfSSH_TestDoKexInit(WOLFSSH* ssh, byte* buf,
+            word32 len, word32* idx);
+    WOLFSSH_API int wolfSSH_TestDoKexDhInit(WOLFSSH* ssh, byte* buf,
+            word32 len, word32* idx);
 #endif
 
 /* dynamic memory types */