From 35554745bc07238662b5dd76e910ec0cbb50c105 Mon Sep 17 00:00:00 2001
From: John Safranek <john@wolfssl.com>
Date: Thu, 13 Apr 2023 16:17:59 -0700
Subject: [PATCH] Update sftpclient

1. Update the example SFTP client with changes made for the echoserver.
2. Better handling of rekeying status.
---
 examples/sftpclient/sftpclient.c | 8 +++++++-
 src/internal.c                   | 3 ++-
 src/ssh.c                        | 8 +++++++-
 src/wolfsftp.c                   | 6 +++++-
 4 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/examples/sftpclient/sftpclient.c b/examples/sftpclient/sftpclient.c
index 31292476..142f8d11 100644
--- a/examples/sftpclient/sftpclient.c
+++ b/examples/sftpclient/sftpclient.c
@@ -501,8 +501,14 @@ static int doCmds(func_args* args)
                 if (ret != WS_SUCCESS && ret == WS_FATAL_ERROR) {
                     ret = wolfSSH_get_error(ssh);
                 }
+                while (ret == WS_REKEYING || ssh->error == WS_REKEYING) {
+                    ret = wolfSSH_worker(ssh, NULL);
+                    if (ret != WS_SUCCESS && ret == WS_FATAL_ERROR) {
+                        ret = wolfSSH_get_error(ssh);
+                    }
+                }
             } while (ret == WS_WANT_READ || ret == WS_WANT_WRITE ||
-                    ret == WS_CHAN_RXD || ret == WS_REKEYING);
+                    ret == WS_CHAN_RXD);
 
 #ifndef WOLFSSH_NO_TIMESTAMP
             WMEMSET(currentFile, 0, WOLFSSH_MAX_FILENAME);
diff --git a/src/internal.c b/src/internal.c
index 1a469702..8df184f1 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -7411,7 +7411,8 @@ int DoReceive(WOLFSSH* ssh)
             ret = DoPacket(ssh);
             ssh->error = ret;
             if (ret < 0 && !(ret == WS_CHAN_RXD || ret == WS_EXTDATA ||
-                    ret == WS_CHANNEL_CLOSED || ret == WS_WANT_WRITE)) {
+                    ret == WS_CHANNEL_CLOSED || ret == WS_WANT_WRITE ||
+                    ret == WS_REKEYING)) {
                 return WS_FATAL_ERROR;
             }
             WLOG(WS_LOG_DEBUG, "PR3: peerMacSz = %u", peerMacSz);
diff --git a/src/ssh.c b/src/ssh.c
index 17feac85..55e13c17 100644
--- a/src/ssh.c
+++ b/src/ssh.c
@@ -1104,6 +1104,7 @@ int wolfSSH_stream_send(WOLFSSH* ssh, byte* buf, word32 bufSz)
         return WS_BAD_ARGUMENT;
 
     if (ssh->isKeying) {
+        ssh->error = WS_REKEYING;
         return WS_REKEYING;
     }
 
@@ -1786,11 +1787,16 @@ int wolfSSH_worker(WOLFSSH* ssh, word32* channelId)
         *channelId = ssh->lastRxId;
     }
 
+    if (ssh->isKeying) {
+        ssh->error = WS_REKEYING;
+        return WS_REKEYING;
+    }
+
     if (ret == WS_CHAN_RXD)
         WLOG(WS_LOG_DEBUG, "Leaving wolfSSH_worker(), "
                            "data received on channel %u", ssh->lastRxId);
     else {
-        if (ret == WS_SUCCESS && ssh->isKeying) {
+        if (ret != WS_SUCCESS && ssh->isKeying) {
             ret = WS_REKEYING;
         }
         WLOG(WS_LOG_DEBUG, "Leaving wolfSSH_worker(), ret = %d", ret);
diff --git a/src/wolfsftp.c b/src/wolfsftp.c
index a28fba34..47a2ac43 100644
--- a/src/wolfsftp.c
+++ b/src/wolfsftp.c
@@ -404,7 +404,8 @@ static INLINE int NoticeError(WOLFSSH* ssh)
 {
     return (ssh->error == WS_WANT_READ ||
             ssh->error == WS_WANT_WRITE ||
-            ssh->error == WS_CHAN_RXD);
+            ssh->error == WS_CHAN_RXD ||
+            ssh->error == WS_REKEYING);
 }
 
 
@@ -6927,6 +6928,9 @@ int wolfSSH_SFTP_SendReadPacket(WOLFSSH* ssh, byte* handle, word32 handleSz,
                 /* send header and type specific data */
                 ret = wolfSSH_SFTP_buffer_send(ssh, &state->buffer);
                 if (ret < 0) {
+                    if (ret == WS_REKEYING) {
+                        return ret;
+                    }
                     if (ssh->error != WS_WANT_READ &&
                             ssh->error != WS_WANT_WRITE) {
                         state->state = STATE_SEND_READ_CLEANUP;