wolfssh/scp: complete rekey that starts mid-transfer

Author: yosuke-wolfssl

src/wolfscp.c

@@ -79,6 +79,94 @@ static int _DumpExtendedData(WOLFSSH* ssh)
 }
 
 
+/* Sends sz bytes from data, completing any rekey that blocks the send.
+ *
+ * Drives wolfSSH_worker while the connection is keying or the channel window is
+ * full, including a rekey triggered by a receive-side highwater check where
+ * ssh->error holds a status other than WS_REKEYING (caught via ssh->isKeying).
+ * In blocking mode the worker's DoReceive waits on the socket so progress
+ * depends on the peer finishing the rekey or adjusting the window; in
+ * non-blocking mode a WS_WANT_READ/WS_WANT_WRITE status is returned so a
+ * stalled rekey cannot spin forever. Returns the byte count from
+ * wolfSSH_stream_send (>= 0) or a negative error code, matching the wrapped
+ * call so callers keep their existing return handling.
+ */
+static int ScpStreamSend(WOLFSSH* ssh, byte* data, word32 sz)
+{
+    int ret = WS_SUCCESS;
+    int err;
+
+    if (ssh == NULL || data == NULL)
+        return WS_BAD_ARGUMENT;
+
+    /* Flush queued output before the drain loop. Otherwise a KEXINIT enqueued
+     * by a rekey triggered on the prior send sits unsent while wolfSSH_worker
+     * runs DoReceive before its own flush, blocking on the socket with the peer
+     * waiting for our KEXINIT. */
+    if (wolfSSH_OutputPending(ssh)) {
+        ret = wolfSSH_SendPacket(ssh);
+        if (ret < 0)
+            return ret;
+    }
+
+    err = wolfSSH_get_error(ssh);
+    while (ssh->isKeying || err == WS_WINDOW_FULL || err == WS_REKEYING) {
+        ret = wolfSSH_worker(ssh, NULL);
+        if (ret < 0 && ret != WS_REKEYING && ret != WS_WINDOW_FULL
+                && ret != WS_CHAN_RXD)
+            return ret;
+        err = wolfSSH_get_error(ssh);
+    }
+
+    return wolfSSH_stream_send(ssh, data, sz);
+}
+
+
+/* Reads up to sz bytes into data, completing any rekey that fires mid-read.
+ *
+ * Flushes queued output before reading so a KEXINIT enqueued by a receive-side
+ * highwater rekey is actually sent, otherwise the peer can wait for our KEXINIT
+ * while we block on the read. On a read that fails with WS_REKEYING the worker
+ * is driven to finish the rekey and the read is retried. The helper is
+ * error-code transparent: every other status (WS_EOF, WS_EXTDATA,
+ * WS_CHANNEL_CLOSED, WS_SOCKET_ERROR_E, WS_WANT_READ/WS_WANT_WRITE, byte count)
+ * is returned unchanged so each caller keeps its existing branch handling.
+ */
+static int ScpStreamRead(WOLFSSH* ssh, byte* data, word32 sz)
+{
+    int ret = WS_SUCCESS;
+    int done = 0;
+
+    if (ssh == NULL || data == NULL)
+        return WS_BAD_ARGUMENT;
+
+    do {
+        if (wolfSSH_OutputPending(ssh)) {
+            ret = wolfSSH_SendPacket(ssh);
+            if (ret < 0)
+                return ret;
+        }
+
+        ret = wolfSSH_stream_read(ssh, data, sz);
+        if (ret < 0 && wolfSSH_get_error(ssh) == WS_REKEYING) {
+            /* Drive the rekey to completion, then retry the read. A worker
+             * status that is not rekey or channel data means the rekey stalled
+             * or a non-blocking want occurred, so return it rather than
+             * spin. */
+            ret = wolfSSH_worker(ssh, NULL);
+            if (ret < 0 && ret != WS_CHAN_RXD
+                    && wolfSSH_get_error(ssh) != WS_REKEYING)
+                return ret;
+        }
+        else {
+            done = 1;
+        }
+    } while (!done);
+
+    return ret;
+}
+
+
 int DoScpSink(WOLFSSH* ssh)
 {
     int ret = WS_SUCCESS;
@@ -283,7 +371,7 @@ static int SendScpTimestamp(WOLFSSH* ssh)
 #endif
     bufSz = (int)WSTRLEN(buf);
 
-    ret = wolfSSH_stream_send(ssh, (byte*)buf, bufSz);
+    ret = ScpStreamSend(ssh, (byte*)buf, bufSz);
     if (ret != bufSz) {
         ret = WS_FATAL_ERROR;
     } else {
@@ -324,7 +412,7 @@ static int SendScpFileHeader(WOLFSSH* ssh)
         return WS_BAD_ARGUMENT;
 #endif
     bufSz = (int)WSTRLEN(filehdr);
-    ret = wolfSSH_stream_send(ssh, (byte*)filehdr, bufSz);
+    ret = ScpStreamSend(ssh, (byte*)filehdr, bufSz);
     if (ret != bufSz) {
         ret = WS_FATAL_ERROR;
     } else {
@@ -357,7 +445,7 @@ static int SendScpEnterDirectory(WOLFSSH* ssh)
 
     bufSz = (int)WSTRLEN(buf);
 
-    ret = wolfSSH_stream_send(ssh, (byte*)buf, bufSz);
+    ret = ScpStreamSend(ssh, (byte*)buf, bufSz);
     if (ret != bufSz) {
         ret = WS_FATAL_ERROR;
     } else {
@@ -383,7 +471,7 @@ static int SendScpExitDirectory(WOLFSSH* ssh)
     buf[0] = 'E';
     buf[1] = '\n';
 
-    ret = wolfSSH_stream_send(ssh, (byte*)buf, sizeof(buf));
+    ret = ScpStreamSend(ssh, (byte*)buf, sizeof(buf));
     if (ret != sizeof(buf)) {
         ret = WS_FATAL_ERROR;
     } else {
@@ -598,8 +686,8 @@ int DoScpSource(WOLFSSH* ssh)
             case SCP_SEND_FILE:
                 WLOG(WS_LOG_DEBUG, scpState, "SCP_SEND_FILE");
 
-                ret = wolfSSH_stream_send(ssh, ssh->scpFileBuffer,
-                                          ssh->scpBufferedSz);
+                ret = ScpStreamSend(ssh, ssh->scpFileBuffer,
+                                    ssh->scpBufferedSz);
                 if (ret == WS_WINDOW_FULL || ret == WS_REKEYING) {
                     ret = wolfSSH_worker(ssh, NULL);
                     if (ret == WS_SUCCESS || ssh->error == WS_WANT_READ)
@@ -736,7 +824,7 @@ int DoScpRequest(WOLFSSH* ssh)
 
         /* Peer MUST send back a SSH_MSG_CHANNEL_CLOSE unless already
             sent*/
-        ret = wolfSSH_stream_read(ssh, buf, 1);
+        ret = ScpStreamRead(ssh, buf, 1);
         if (ret == WS_SOCKET_ERROR_E || ret == WS_CHANNEL_CLOSED) {
             WLOG(WS_LOG_DEBUG, scpState, "Peer hung up, but SCP is done");
             ret = WS_SUCCESS;
@@ -1395,6 +1483,31 @@ int ReceiveScpMessage(WOLFSSH* ssh)
             return WS_BUFFER_E;
         }
 
+        /* Flush queued output before polling. A KEXINIT enqueued by a rekey
+         * would otherwise sit unsent while wolfSSH_worker runs DoReceive before
+         * its own flush, deadlocking against a peer that waits for our
+         * KEXINIT. */
+        if (wolfSSH_OutputPending(ssh)) {
+            ret = wolfSSH_SendPacket(ssh);
+            if (ret < 0)
+                return ret;
+        }
+
+        /* If channel data is already buffered, read it directly rather than
+         * polling the socket. A control message delivered into the channel
+         * buffer while a rekey was completing leaves wolfSSH_worker returning
+         * the rekey status (not WS_CHAN_RXD), so without this the buffered
+         * message is never read and the next worker blocks on the socket. */
+        if (wolfSSH_stream_peek(ssh, NULL, 1) > 0) {
+            sz = wolfSSH_stream_read(ssh, buf + ssh->scpRecvMsgSz,
+                    DEFAULT_SCP_MSG_SZ - ssh->scpRecvMsgSz);
+            if (sz < 0)
+                return sz;
+            ssh->scpRecvMsgSz += sz;
+            sz = ssh->scpRecvMsgSz;
+            continue;
+        }
+
         err = wolfSSH_worker(ssh, &lastChannel);
         if (err < 0) {
             int rc;
@@ -1517,7 +1630,7 @@ int ReceiveScpFile(WOLFSSH* ssh)
     }
 
     if (ret == WS_SUCCESS) {
-        ret = wolfSSH_stream_read(ssh, ssh->scpFileBuffer, partSz);
+        ret = ScpStreamRead(ssh, ssh->scpFileBuffer, partSz);
         if (ret > 0) {
             ssh->scpFileBufferSz = ret;
         }
@@ -1559,7 +1672,7 @@ int SendScpConfirmation(WOLFSSH* ssh)
 
     /* skip first byte for accurate strlen, may be 0 */
     msgSz = (int)XSTRLEN(msg + 1) + 1;
-    ret = wolfSSH_stream_send(ssh, (byte*)msg, msgSz);
+    ret = ScpStreamSend(ssh, (byte*)msg, msgSz);
     if (ret != msgSz || ssh->scpConfirm == WS_SCP_ABORT) {
         ret = WS_FATAL_ERROR;
 
@@ -1587,7 +1700,7 @@ int ReceiveScpConfirmation(WOLFSSH* ssh)
         return WS_BAD_ARGUMENT;
 
     WMEMSET(msg, 0, sizeof(msg));
-    msgSz = wolfSSH_stream_read(ssh, msg, DEFAULT_SCP_MSG_SZ);
+    msgSz = ScpStreamRead(ssh, msg, DEFAULT_SCP_MSG_SZ);
 
     if (msgSz < 0) {
         if (wolfSSH_get_error(ssh) == WS_EXTDATA)