From 73839469ede89768da40b68080f78e59b8423189 Mon Sep 17 00:00:00 2001
From: Jacob Barthelmeh <jacob@wolfssl.com>
Date: Fri, 18 Jan 2019 10:19:18 -0700
Subject: [PATCH] adjust stream_send behavior with WANT_WRITE

---
 src/internal.c     | 4 ++++
 src/ssh.c          | 7 +++----
 wolfssh/internal.h | 1 +
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/internal.c b/src/internal.c
index fcb4b2b..67908a6 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -1410,6 +1410,7 @@ int wolfSSH_SendPacket(WOLFSSH* ssh)
     }
 
     ssh->outputBuffer.idx = 0;
+    ssh->outputBuffer.plainSz = 0;
 
     WLOG(WS_LOG_DEBUG, "SB: Shrinking output buffer");
     ShrinkBuffer(&ssh->outputBuffer, 0);
@@ -7147,6 +7148,9 @@ int SendChannelData(WOLFSSH* ssh, word32 peerChannel,
     if (ret == WS_SUCCESS)
         ret = dataSz;
 
+    if (ssh->error == WS_WANT_WRITE)
+        ssh->outputBuffer.plainSz = dataSz;
+
     WLOG(WS_LOG_DEBUG, "Leaving SendChannelData(), ret = %d", ret);
     return ret;
 }
diff --git a/src/ssh.c b/src/ssh.c
index 5430a40..34d8346 100644
--- a/src/ssh.c
+++ b/src/ssh.c
@@ -766,14 +766,13 @@ int wolfSSH_stream_send(WOLFSSH* ssh, byte* buf, word32 bufSz)
     if (ssh->error == WS_WANT_WRITE && ssh->outputBuffer.length != 0) {
         int ret;
 
-        bytesTxd = ssh->outputBuffer.length;
-        WLOG(WS_LOG_DEBUG, "Trying to resend %d bytes\n", bytesTxd);
+        bytesTxd = ssh->outputBuffer.plainSz;
+        WLOG(WS_LOG_DEBUG, "\n\nTrying to resend %d bytes\n\n\n", bytesTxd);
         ssh->error = WS_SUCCESS;
         ret = wolfSSH_SendPacket(ssh);
 
         /* return the amount sent on success otherwise return error found */
-        return (ret == WS_SUCCESS)?
-            bytesTxd - (int)ssh->outputBuffer.length : ret;
+        return (ret == WS_SUCCESS)? bytesTxd : ret;
     }
 
     bytesTxd = SendChannelData(ssh, ssh->channelList->peerChannel, buf, bufSz);
diff --git a/wolfssh/internal.h b/wolfssh/internal.h
index 6379464..a94b236 100644
--- a/wolfssh/internal.h
+++ b/wolfssh/internal.h
@@ -155,6 +155,7 @@ WOLFSSH_LOCAL const char* IdToName(byte);
 
 typedef struct Buffer {
     void* heap;       /* Heap for allocations */
+    int   plainSz;    /* amount of plain text bytes to send with WANT_WRITE */
     word32 length;    /* total buffer length used */
     word32 idx;       /* idx to part of length already consumed */
     byte* buffer;     /* place holder for actual buffer */