Add tests for rekey behavior across highwater-mark

Author: yosuke-wolfssl

tests/api.c

@@ -1444,8 +1444,119 @@ static void test_wolfSSH_SFTP_SendReadPacket(void)
     ThreadJoin(serThread);
 }
 
+static void sftp_rekey_test(int nonBlock)
+{
+    func_args ser;
+    tcp_ready ready;
+    int argsCount;
+    int err;
+    WS_SOCKET_T clientFd;
+    WS_SFTPNAME* ls;
+    int i;
+
+    const char* args[10];
+    WOLFSSH_CTX* ctx = NULL;
+    WOLFSSH*     ssh = NULL;
+
+    THREAD_TYPE serThread;
+
+    WMEMSET(&ser, 0, sizeof(func_args));
+
+    argsCount = 0;
+    args[argsCount++] = ".";
+    args[argsCount++] = "-1";
+#ifndef USE_WINDOWS_API
+    args[argsCount++] = "-p";
+    args[argsCount++] = "0";
+#endif
+    ser.argv   = (char**)args;
+    ser.argc    = argsCount;
+    ser.signal = &ready;
+    InitTcpReady(ser.signal);
+    ThreadStart(echoserver_test, (void*)&ser, &serThread);
+    WaitTcpReady(&ready);
+
+    sftp_client_connect(&ctx, &ssh, ready.port);
+    AssertNotNull(ctx);
+    AssertNotNull(ssh);
+
+    /* Handshake completed in blocking mode; switch to non-blocking so the
+     * LS/rekey phase exercises the WS_WANT_READ/WS_WANT_WRITE early-return
+     * path in buffer_send/buffer_read. */
+    clientFd = wolfSSH_get_fd(ssh);
+    if (nonBlock) {
+        tcp_set_nonblocking(&clientFd);
+    }
+
+    /* Low threshold forces a rekey mid-SFTP-operation. In blocking mode the
+     * buffer_read/buffer_send fixes drive the rekey internally; in non-blocking
+     * mode the retry loop below advances it on
+     * WS_WANT_READ/WS_WANT_WRITE/WS_REKEYING. */
+    wolfSSH_SetHighwater(ssh, 256);
+
+    ls = NULL;
+    for (i = 0; i < 3; i++) {
+        /* Retry loop is a single pass in blocking mode; for non-blocking it
+         * drives the rekey forward while LS reports a want/rekey status. */
+        do {
+            ls  = wolfSSH_SFTP_LS(ssh, (char*)".");
+            err = wolfSSH_get_error(ssh);
+            if (ls == NULL && (err == WS_WANT_READ || err == WS_WANT_WRITE
+                        || err == WS_REKEYING)) {
+                tcp_select(clientFd, 1);
+            }
+        } while (ls == NULL && (err == WS_WANT_READ || err == WS_WANT_WRITE
+                    || err == WS_REKEYING));
+        AssertNotNull(ls);
+        wolfSSH_SFTPNAME_list_free(ls);
+        ls = NULL;
+    }
+
+    do {
+        argsCount = wolfSSH_shutdown(ssh);
+        err = wolfSSH_get_error(ssh);
+        if (argsCount != WS_SUCCESS && (err == WS_WANT_READ
+                    || err == WS_WANT_WRITE || err == WS_REKEYING)) {
+            tcp_select(clientFd, 1);
+        }
+    } while (argsCount != WS_SUCCESS && (err == WS_WANT_READ
+                || err == WS_WANT_WRITE || err == WS_REKEYING));
+    if (argsCount == WS_SOCKET_ERROR_E) {
+        argsCount = WS_SUCCESS;
+    }
+#if DEFAULT_HIGHWATER_MARK < 8000
+    if (argsCount == WS_REKEYING) {
+        /* in cases where highwater mark is really small a re-key could happen */
+        argsCount = WS_SUCCESS;
+    }
+#endif
+    AssertIntEQ(argsCount, WS_SUCCESS);
+
+    clientFd = wolfSSH_get_fd(ssh);
+    WCLOSESOCKET(clientFd);
+
+    wolfSSH_free(ssh);
+    wolfSSH_CTX_free(ctx);
+#ifdef WOLFSSH_ZEPHYR
+    k_sleep(Z_TIMEOUT_TICKS(100));
+#endif
+    ThreadJoin(serThread);
+}
+
+static void test_wolfSSH_SFTP_ReKey(void)
+{
+    sftp_rekey_test(0);
+}
+
+static void test_wolfSSH_SFTP_ReKey_NonBlock(void)
+{
+    sftp_rekey_test(1);
+}
+
 #else /* WOLFSSH_SFTP && !NO_WOLFSSH_CLIENT && !SINGLE_THREADED */
 static void test_wolfSSH_SFTP_SendReadPacket(void) { ; }
+static void test_wolfSSH_SFTP_ReKey(void) { ; }
+static void test_wolfSSH_SFTP_ReKey_NonBlock(void) { ; }
 #endif /* WOLFSSH_SFTP && !NO_WOLFSSH_CLIENT && !SINGLE_THREADED */
 
 
@@ -2159,6 +2270,8 @@ int wolfSSH_ApiTest(int argc, char** argv)
 
     /* SFTP tests */
     test_wolfSSH_SFTP_SendReadPacket();
+    test_wolfSSH_SFTP_ReKey();
+    test_wolfSSH_SFTP_ReKey_NonBlock();
 
     /* Either SCP or SFTP */
     test_wolfSSH_RealPath();