fixes for SFTP want read/write cases

Author: JacobBarthelmeh

examples/echoserver/echoserver.c

@@ -1299,16 +1299,25 @@ static int sftp_worker(thread_ctx_t* threadCtx)
 {
     WOLFSSH* ssh = threadCtx->ssh;
     WS_SOCKET_T s;
-    int ret = WS_SUCCESS;
+    int ret;
     int error = -1;
     int selected;
     unsigned char peek_buf[1];
     int timeout = TEST_SFTP_TIMEOUT;
 
     s = (WS_SOCKET_T)wolfSSH_get_fd(ssh);
+    ret = error = wolfSSH_get_error(ssh);
+
+    /* there is an edge case where the last SFTP handshake message sent got a
+     * WANT_WRITE case, keep trying to send it here. */
+    while (error == WS_WANT_WRITE) {
+        ret = wolfSSH_worker(ssh, NULL);
+        error = wolfSSH_get_error(ssh);
+    }
 
     do {
-        if (ret == WS_WANT_WRITE || wolfSSH_SFTP_PendingSend(ssh)) {
+        if (ret == WS_WANT_WRITE || ret == WS_CHAN_RXD ||
+                wolfSSH_SFTP_PendingSend(ssh)) {
             /* Yes, process the SFTP data. */
             ret = wolfSSH_SFTP_read(ssh);
             error = wolfSSH_get_error(ssh);
@@ -1362,6 +1371,19 @@ static int sftp_worker(thread_ctx_t* threadCtx)
                 break;
             }
             if (ret != WS_SUCCESS && ret != WS_CHAN_RXD) {
+            #ifdef WOLFSSH_TEST_BLOCK
+                if (error == WS_WANT_READ) {
+                    while (error == WS_WANT_READ) {
+                        /* The socket had data but our test nonblocking code
+                        * returned want read. Loop over wolfSSH_worker here
+                        * until we get the data off the socket that select
+                        * indicated was available. */
+                        ret = wolfSSH_worker(ssh, NULL);
+                        error = wolfSSH_get_error(ssh);
+                    }
+                    continue;
+                }
+            #endif
                 if (ret == WS_WANT_WRITE) {
                     /* recall wolfSSH_worker here because is likely our custom
                      * highwater callback that returned up a WS_WANT_WRITE */
@@ -1433,20 +1455,45 @@ static int NonBlockSSH_accept(WOLFSSH* ssh)
             printf("... server would read block\n");
         else if (error == WS_WANT_WRITE)
             printf("... server would write block\n");
+        else if (error == WS_AUTH_PENDING)
+            printf("... server auth pending\n");
 
         select_ret = tcp_select(sockfd, 1);
-        if (select_ret == WS_SELECT_RECV_READY  ||
-            select_ret == WS_SELECT_ERROR_READY ||
-            error      == WS_WANT_WRITE ||
-            error      == WS_AUTH_PENDING)
-        {
+        if (select_ret == WS_SELECT_RECV_READY) {
             ret = wolfSSH_accept(ssh);
             error = wolfSSH_get_error(ssh);
+
+        #ifdef WOLFSSH_TEST_BLOCK
+            if (error == WS_WANT_READ) {
+                /* The socket had data but our test nonblocking code
+                 * returned want read. Loop over wolfSSH_accept here until
+                 * we get the data off the socket that select indicated was
+                 * available. */
+                while (error == WS_WANT_READ) {
+                    ret = wolfSSH_accept(ssh);
+                    error = wolfSSH_get_error(ssh);
+                }
+            }
+        #endif
         }
-        else if (select_ret == WS_SELECT_TIMEOUT)
-            error = WS_WANT_READ;
-        else
+        else if (select_ret == WS_SELECT_TIMEOUT) {
+            if (error == WS_WANT_WRITE || error == WS_AUTH_PENDING
+            #ifdef WOLFSSH_TEST_BLOCK
+                || error == WS_WANT_READ
+            #endif
+            ) {
+                /* For write or auth pending, we need to try again */
+                ret = wolfSSH_accept(ssh);
+                error = wolfSSH_get_error(ssh);
+            }
+            else {
+                error = WS_WANT_READ;
+            }
+        }
+        else {
             error = WS_FATAL_ERROR;
+            break;
+        }
     }
 
     return ret;

examples/sftpclient/sftpclient.c

@@ -205,7 +205,8 @@ static int NonBlockSSH_connect(void)
     sockfd = (SOCKET_T)wolfSSH_get_fd(ssh);
 
     while (ret != WS_SUCCESS &&
-            (error == WS_WANT_READ || error == WS_WANT_WRITE))
+            (error == WS_WANT_READ || error == WS_WANT_WRITE ||
+             error == WS_REKEYING || error == WS_AUTH_PENDING))
     {
         if (error == WS_WANT_READ)
             printf("... client would read block\n");
@@ -219,6 +220,19 @@ static int NonBlockSSH_connect(void)
         {
             ret = wolfSSH_SFTP_connect(ssh);
             error = wolfSSH_get_error(ssh);
+
+        #ifdef WOLFSSH_TEST_BLOCK
+            if (select_ret == WS_SELECT_RECV_READY && error == WS_WANT_READ) {
+                /* The socket has data to be read but the non blocking test
+                 * code returned want read. Loop over wolfSSH_SFTP_connect here
+                 * until we get the data off the socket that select indicated
+                 * was available. */
+                while (error == WS_WANT_READ) {
+                    ret = wolfSSH_SFTP_connect(ssh);
+                    error = wolfSSH_get_error(ssh);
+                }
+            }
+        #endif
         }
         else if (select_ret == WS_SELECT_TIMEOUT)
             error = WS_WANT_READ;

src/wolfsftp.c

@@ -543,6 +543,7 @@ static int wolfSSH_SFTP_buffer_read(WOLFSSH* ssh, WS_SFTP_BUFFER* buffer,
         int readSz)
 {
     int ret;
+    byte peekBuf[1];
 
     if (buffer == NULL || ssh == NULL) {
         return WS_FATAL_ERROR;
@@ -563,8 +564,18 @@ static int wolfSSH_SFTP_buffer_read(WOLFSSH* ssh, WS_SFTP_BUFFER* buffer,
     }
 
     do {
-        ret = wolfSSH_stream_read(ssh, buffer->data + buffer->idx,
+        if (!wolfSSH_stream_peek(ssh, peekBuf, 1)) {
+            /* poll more data off the wire */
+            ret = wolfSSH_worker(ssh, NULL);
+        }
+        else {
+            ret = WS_CHAN_RXD; /* existing data found with peek */
+        }
+
+        if (ret == WS_CHAN_RXD) {
+            ret = wolfSSH_stream_read(ssh, buffer->data + buffer->idx,
                 buffer->sz - buffer->idx);
+        }
         if (ret < 0) {
             return WS_FATAL_ERROR;
         }
@@ -5666,14 +5677,19 @@ int wolfSSH_SFTP_RecvFSetSTAT(WOLFSSH* ssh, int reqId, byte* data, word32 maxSz)
  * returns WS_SUCCESS on success
  */
 static int SFTP_ClientRecvInit(WOLFSSH* ssh) {
-    int  len;
+    int  len, ret;
     byte id;
     word32 sz = 0;
     word32 version = 0;
     byte buf[LENGTH_SZ + MSG_ID_SZ + UINT32_SZ];
 
     switch (ssh->sftpState) {
         case SFTP_RECV:
+            ret = wolfSSH_worker(ssh,NULL);
+            if (ret != WS_CHAN_RXD) {
+                return ret;
+            }
+
             if ((len = wolfSSH_stream_read(ssh, buf, sizeof(buf)))
                     != sizeof(buf)) {
                 /* @TODO partial read on small packet */
@@ -5785,7 +5801,7 @@ int wolfSSH_SFTP_connect(WOLFSSH* ssh)
 
     switch (ssh->sftpState) {
         case SFTP_BEGIN:
-            if ((ssh->error = SFTP_ClientSendInit(ssh)) != WS_SUCCESS) {
+            if (SFTP_ClientSendInit(ssh) != WS_SUCCESS) {
                 return WS_FATAL_ERROR;
             }
             ssh->sftpState = SFTP_RECV;
@@ -8174,7 +8190,7 @@ WS_SFTPNAME* wolfSSH_SFTP_RealPath(WOLFSSH* ssh, char* dir)
         case SFTP_REAL_GET_PACKET:
             /* read name response from Real Path packet */
             ret = wolfSSH_SFTP_DoName(ssh);
-            if (ret != NULL || (ret == NULL && ssh->error != WS_WANT_READ)) {
+            if (ret != NULL || (ret == NULL && !NoticeError(ssh))) {
                 wolfSSH_SFTP_ClearState(ssh, STATE_ID_NAME);
                 ssh->realState = SFTP_REAL_SEND_PACKET;
             }