Merge pull request #118 from ejohnstown/nonblocking-fix

Non-blocking Fix

examples/client/client.c

@@ -104,6 +104,7 @@ static void ShowUsage(void)
     printf(" -P <password> password for username, prompted if omitted\n");
     printf(" -x            exit after successful connection without doing\n"
            "               read/write\n");
+    printf(" -N            use non-blocking sockets\n");
 }
 
 
@@ -150,6 +151,41 @@ static int wsUserAuth(byte authType,
 }
 
 
+static int NonBlockSSH_connect(WOLFSSH* ssh)
+{
+    int ret;
+    int error;
+    SOCKET_T sockfd;
+    int select_ret = 0;
+
+    ret = wolfSSH_connect(ssh);
+    error = wolfSSH_get_error(ssh);
+    sockfd = (SOCKET_T)wolfSSH_get_fd(ssh);
+
+    while (ret != WS_SUCCESS &&
+            (error == WS_WANT_READ || error == WS_WANT_WRITE))
+    {
+        if (error == WS_WANT_READ)
+            printf("... client would read block\n");
+        else if (error == WS_WANT_WRITE)
+            printf("... client would write block\n");
+
+        select_ret = tcp_select(sockfd, 1);
+        if (select_ret == WS_SELECT_RECV_READY ||
+            select_ret == WS_SELECT_ERROR_READY)
+        {
+            ret = wolfSSH_connect(ssh);
+        }
+        else if (select_ret == WS_SELECT_TIMEOUT)
+            error = WS_WANT_READ;
+        else
+            error = WS_FATAL_ERROR;
+    }
+
+    return ret;
+}
+
+
 THREAD_RETURN WOLFSSH_THREAD client_test(void* args)
 {
     WOLFSSH_CTX* ctx = NULL;
@@ -165,12 +201,13 @@ THREAD_RETURN WOLFSSH_THREAD client_test(void* args)
     const char* username = NULL;
     const char* password = NULL;
     byte imExit = 0;
+    byte nonBlock = 0;
 
     int     argc = ((func_args*)args)->argc;
     char**  argv = ((func_args*)args)->argv;
     ((func_args*)args)->return_code = 0;
 
-    while ((ch = mygetopt(argc, argv, "?h:p:u:P:x")) != -1) {
+    while ((ch = mygetopt(argc, argv, "?NP:h:p:u:x")) != -1) {
         switch (ch) {
             case 'h':
                 host = myoptarg;
@@ -197,6 +234,10 @@ THREAD_RETURN WOLFSSH_THREAD client_test(void* args)
                 imExit = 1;
                 break;
 
+            case 'N':
+                nonBlock = 1;
+                break;
+
             case '?':
                 ShowUsage();
                 exit(EXIT_SUCCESS);
@@ -237,13 +278,21 @@ THREAD_RETURN WOLFSSH_THREAD client_test(void* args)
     if (ret != 0)
         err_sys("Couldn't connect to server.");
 
+    if (nonBlock)
+        tcp_set_nonblocking(&sockFd);
+
     ret = wolfSSH_set_fd(ssh, (int)sockFd);
     if (ret != WS_SUCCESS)
         err_sys("Couldn't set the session's socket.");
 
+    if (!nonBlock)
         ret = wolfSSH_connect(ssh);
-    if (ret != WS_SUCCESS)
+    else
+        ret = NonBlockSSH_connect(ssh);
+    if (ret != WS_SUCCESS) {
+        printf("err = %s\n", wolfSSH_get_error_name(ssh));
         err_sys("Couldn't connect SSH stream.");
+    }
 
     if (!imExit) {
         ret = wolfSSH_stream_send(ssh, (byte*)testString,
@@ -251,9 +300,14 @@ THREAD_RETURN WOLFSSH_THREAD client_test(void* args)
         if (ret <= 0)
             err_sys("Couldn't send test string.");
 
+        do {
             ret = wolfSSH_stream_read(ssh, (byte*)rxBuf, sizeof(rxBuf) - 1);
-        if (ret <= 0)
+            if (ret <= 0) {
+                if (ret != WS_WANT_READ && ret != WS_WANT_WRITE)
                     err_sys("Stream read failed.");
+            }
+        } while (ret == WS_WANT_READ || ret == WS_WANT_WRITE);
+
         rxBuf[ret] = '\0';
         printf("Server said: %s\n", rxBuf);
     }

examples/server/server.c

@@ -46,6 +46,7 @@ typedef struct {
     WOLFSSH* ssh;
     SOCKET_T fd;
     word32 id;
+    char nonBlock;
 } thread_ctx_t;
 
 
@@ -97,12 +98,52 @@ static int dump_stats(thread_ctx_t* ctx)
 }
 
 
+static int NonBlockSSH_accept(WOLFSSH* ssh)
+{
+    int ret;
+    int error;
+    SOCKET_T sockfd;
+    int select_ret = 0;
+
+    ret = wolfSSH_accept(ssh);
+    error = wolfSSH_get_error(ssh);
+    sockfd = (SOCKET_T)wolfSSH_get_fd(ssh);
+
+    while (ret != WS_SUCCESS &&
+            (error == WS_WANT_READ || error == WS_WANT_WRITE))
+    {
+        if (error == WS_WANT_READ)
+            printf("... client would read block\n");
+        else if (error == WS_WANT_WRITE)
+            printf("... client would write block\n");
+
+        select_ret = tcp_select(sockfd, 1);
+        if (select_ret == WS_SELECT_RECV_READY ||
+            select_ret == WS_SELECT_ERROR_READY)
+        {
+            ret = wolfSSH_accept(ssh);
+        }
+        else if (select_ret == WS_SELECT_TIMEOUT)
+            error = WS_WANT_READ;
+        else
+            error = WS_FATAL_ERROR;
+    }
+
+    return ret;
+}
+
+
 static THREAD_RETURN WOLFSSH_THREAD server_worker(void* vArgs)
 {
     int ret;
     thread_ctx_t* threadCtx = (thread_ctx_t*)vArgs;
 
-    if ((ret = wolfSSH_accept(threadCtx->ssh)) == WS_SUCCESS) {
+    if (!threadCtx->nonBlock)
+        ret = wolfSSH_accept(threadCtx->ssh);
+    else
+        ret = NonBlockSSH_accept(threadCtx->ssh);
+
+    if (ret == WS_SUCCESS) {
         byte* buf = NULL;
         byte* tmpBuf;
         int bufSz, backlogSz = 0, rxSz, txSz, stop = 0, txSum;
@@ -117,9 +158,12 @@ static THREAD_RETURN WOLFSSH_THREAD server_worker(void* vArgs)
                 buf = tmpBuf;
 
             if (!stop) {
+                do {
                     rxSz = wolfSSH_stream_read(threadCtx->ssh,
                                                buf + backlogSz,
                                                EXAMPLE_BUFFER_SZ);
+                } while (rxSz == WS_WANT_READ || rxSz == WS_WANT_WRITE);
+
                 if (rxSz > 0) {
                     backlogSz += rxSz;
                     txSum = 0;
@@ -504,9 +548,10 @@ static int wsUserAuth(byte authType,
 static void ShowUsage(void)
 {
     printf("server %s\n", LIBWOLFSSH_VERSION_STRING);
-    printf("-h          Help, print this usage\n");
+    printf(" -h            display this help and exit\n");
-    printf("-m          Allow multiple connections\n");
+    printf(" -m            allow multiple connections\n");
-    printf("-e          Use ECC private key\n");
+    printf(" -e            use ECC private key\n");
+    printf(" -N            use non-blocking sockets\n");
 }
 
 
@@ -517,16 +562,17 @@ THREAD_RETURN WOLFSSH_THREAD server_test(void* args)
     SOCKET_T listenFd = 0;
     word32 defaultHighwater = EXAMPLE_HIGHWATER_MARK;
     word32 threadCount = 0;
-    int multipleConnections = 0;
-    int useEcc = 0;
-    char ch;
     word16 port = wolfSshPort;
+    char multipleConnections = 0;
+    char useEcc = 0;
+    char ch;
+    char nonBlock = 0;
 
     int     argc = ((func_args*)args)->argc;
     char**  argv = ((func_args*)args)->argv;
     ((func_args*)args)->return_code = 0;
 
-    while ((ch = mygetopt(argc, argv, "hme")) != -1) {
+    while ((ch = mygetopt(argc, argv, "hmeN")) != -1) {
         switch (ch) {
             case 'h' :
                 ShowUsage();
@@ -540,6 +586,10 @@ THREAD_RETURN WOLFSSH_THREAD server_test(void* args)
                 useEcc = 1;
                 break;
 
+            case 'N' :
+                nonBlock = 1;
+                break;
+
             default:
                 ShowUsage();
                 exit(MY_EX_USAGE);
@@ -626,11 +676,15 @@ THREAD_RETURN WOLFSSH_THREAD server_test(void* args)
         if (clientFd == -1)
             err_sys("tcp accept failed");
 
+        if (nonBlock)
+            tcp_set_nonblocking(&clientFd);
+
         wolfSSH_set_fd(ssh, (int)clientFd);
 
         threadCtx->ssh = ssh;
         threadCtx->fd = clientFd;
         threadCtx->id = threadCount++;
+        threadCtx->nonBlock = nonBlock;
 
 #ifndef SINGLE_THREADED
         ThreadStart(server_worker, threadCtx, &thread);

src/internal.c

@@ -7051,6 +7051,11 @@ int SendChannelData(WOLFSSH* ssh, word32 peerChannel,
             ret = WS_REKEYING;
     }
 
+    if (ret == WS_SUCCESS) {
+        if (ssh->outputBuffer.length != 0)
+            ret = wolfSSH_SendPacket(ssh);
+    }
+
     if (ret == WS_SUCCESS) {
         channel = ChannelFind(ssh, peerChannel, WS_CHANNEL_ID_PEER);
         if (channel == NULL) {
@@ -7069,7 +7074,8 @@ int SendChannelData(WOLFSSH* ssh, word32 peerChannel,
             dataSz = bound;
         }
 
-        ret = PreparePacket(ssh, MSG_ID_SZ + UINT32_SZ + LENGTH_SZ + dataSz);
+        ret = PreparePacket(ssh,
+                MSG_ID_SZ + UINT32_SZ + LENGTH_SZ + dataSz);
     }
 
     if (ret == WS_SUCCESS) {
@@ -7089,17 +7095,19 @@ int SendChannelData(WOLFSSH* ssh, word32 peerChannel,
         ret = BundlePacket(ssh);
     }
 
-    if (ret == WS_SUCCESS)
-        ret = wolfSSH_SendPacket(ssh);
-
     if (ret == WS_SUCCESS) {
         WLOG(WS_LOG_INFO, "  dataSz = %u", dataSz);
         WLOG(WS_LOG_INFO, "  peerWindowSz = %u", channel->peerWindowSz);
         channel->peerWindowSz -= dataSz;
         WLOG(WS_LOG_INFO, "  update peerWindowSz = %u", channel->peerWindowSz);
-        ret = dataSz;
     }
 
+    if (ret == WS_SUCCESS)
+        ret = wolfSSH_SendPacket(ssh);
+
+    if (ret == WS_SUCCESS)
+        ret = dataSz;
+
     WLOG(WS_LOG_DEBUG, "Leaving SendChannelData(), ret = %d", ret);
     return ret;
 }

wolfssh/test.h

@@ -539,6 +539,61 @@ static INLINE void tcp_listen(SOCKET_T* sockfd, word16* port, int useAnyAddr)
 
 #endif /* WOLFSSH_TEST_SERVER */
 
+static INLINE void tcp_set_nonblocking(SOCKET_T* sockfd)
+{
+    #ifdef USE_WINDOWS_API
+        unsigned long blocking = 1;
+        int ret = ioctlsocket(*sockfd, FIONBIO, &blocking);
+        if (ret == SOCKET_ERROR)
+            err_sys("ioctlsocket failed");
+    #elif defined(WOLFSSL_MDK_ARM) || defined(WOLFSSL_KEIL_TCP_NET) \
+        || defined (WOLFSSL_TIRTOS)|| defined(WOLFSSL_VXWORKS)
+         /* non blocking not supported, for now */
+    #else
+        int flags = fcntl(*sockfd, F_GETFL, 0);
+        if (flags < 0)
+            err_sys("fcntl get failed");
+        flags = fcntl(*sockfd, F_SETFL, flags | O_NONBLOCK);
+        if (flags < 0)
+            err_sys("fcntl set failed");
+    #endif
+}
+
+
+enum {
+    WS_SELECT_FAIL,
+    WS_SELECT_TIMEOUT,
+    WS_SELECT_RECV_READY,
+    WS_SELECT_ERROR_READY
+};
+
+static INLINE int tcp_select(SOCKET_T socketfd, int to_sec)
+{
+    fd_set recvfds, errfds;
+    SOCKET_T nfds = socketfd + 1;
+    struct timeval timeout = {(to_sec > 0) ? to_sec : 0, 0};
+    int result;
+
+    FD_ZERO(&recvfds);
+    FD_SET(socketfd, &recvfds);
+    FD_ZERO(&errfds);
+    FD_SET(socketfd, &errfds);
+
+    result = select(nfds, &recvfds, NULL, &errfds, &timeout);
+
+    if (result == 0)
+        return WS_SELECT_TIMEOUT;
+    else if (result > 0) {
+        if (FD_ISSET(socketfd, &recvfds))
+            return WS_SELECT_RECV_READY;
+        else if(FD_ISSET(socketfd, &errfds))
+            return WS_SELECT_ERROR_READY;
+    }
+
+    return WS_SELECT_FAIL;
+}
+
+
 /* Wolf Root Directory Helper */
 /* KEIL-RL File System does not support relative directory */
 #if !defined(WOLFSSL_MDK_ARM) && !defined(WOLFSSL_KEIL_FS) && !defined(WOLFSSL_TIRTOS) \