From 411412f90296ed596340fe3928d6b9a46934f39f Mon Sep 17 00:00:00 2001
From: Wiebe Cazemier <wiebe@halfgaar.net>
Date: Sat, 16 Oct 2021 21:31:13 +0200
Subject: [PATCH] Fix length bug in sending websocket frames

---
 client.cpp    |  1 +
 iowrapper.cpp | 21 ++++++++++++++-------
 iowrapper.h   |  4 +++-
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/client.cpp b/client.cpp
index 0b845ce..bf2b25f 100644
--- a/client.cpp
+++ b/client.cpp
@@ -315,6 +315,7 @@ std::string Client::getKeepAliveInfoString() const
 void Client::resetBuffersIfEligible()
 {
     readbuf.resetSizeIfEligable(initialBufferSize);
+    ioWrapper.resetBuffersIfEligible();
 
     // Write buffers are written to from other threads, and this resetting takes place from the Client's own thread, so we need to lock.
     std::lock_guard<std::mutex> locker(writeBufMutex);
diff --git a/iowrapper.cpp b/iowrapper.cpp
index 6f61406..efaed66 100644
--- a/iowrapper.cpp
+++ b/iowrapper.cpp
@@ -579,27 +579,28 @@ ssize_t IoWrapper::websocketBytesToReadBuffer(void *buf, const size_t nbytes, Io
  * @param nbytes. The amount of bytes. Can be 0, for just an empty websocket frame.
  * @return
  */
-ssize_t IoWrapper::writeAsMuchOfBufAsWebsocketFrame(const void *buf, size_t nbytes, WebsocketOpcode opcode)
+ssize_t IoWrapper::writeAsMuchOfBufAsWebsocketFrame(const void *buf, const size_t nbytes, WebsocketOpcode opcode)
 {
     // We do allow pong frames to generate a zero payload packet, but for binary, that's not necessary.
     if (nbytes == 0 && opcode == WebsocketOpcode::Binary)
         return 0;
 
-    ssize_t nBytesReal = 0;
+    websocketWriteRemainder.ensureFreeSpace(nbytes + WEBSOCKET_MAX_SENDING_HEADER_SIZE, 131072);
+    const ssize_t nBytesReal = std::min<size_t>(nbytes, websocketWriteRemainder.freeSpace() - WEBSOCKET_MAX_SENDING_HEADER_SIZE);
 
     // We normally wrap each write in a frame, but if a previous one didn't fit in the system's write buffers, we're still working on it.
     if (websocketWriteRemainder.freeSpace() > WEBSOCKET_MAX_SENDING_HEADER_SIZE)
     {
         uint8_t extended_payload_length_num_bytes = 0;
         uint8_t payload_length = 0;
-        if (nbytes < 126)
-            payload_length = nbytes;
-        else if (nbytes >= 126 && nbytes <= 0xFFFF)
+        if (nBytesReal < 126)
+            payload_length = nBytesReal;
+        else if (nBytesReal >= 126 && nBytesReal <= 0xFFFF)
         {
             payload_length = 126;
             extended_payload_length_num_bytes = 2;
         }
-        else if (nbytes > 0xFFFF)
+        else if (nBytesReal > 0xFFFF)
         {
             payload_length = 127;
             extended_payload_length_num_bytes = 8;
@@ -613,7 +614,6 @@ ssize_t IoWrapper::writeAsMuchOfBufAsWebsocketFrame(const void *buf, size_t nbyt
         const int header_length = x + extended_payload_length_num_bytes;
 
         // This block writes the extended payload length.
-        nBytesReal = std::min<size_t>(nbytes, websocketWriteRemainder.freeSpace() - header_length);
         const uint64_t nbytes64 = nBytesReal;
         for (int z = extended_payload_length_num_bytes - 1; z >= 0; z--)
         {
@@ -667,3 +667,10 @@ ssize_t IoWrapper::writeWebsocketAndOrSsl(int fd, const void *buf, size_t nbytes
         return n;
     }
 }
+
+void IoWrapper::resetBuffersIfEligible()
+{
+    const size_t sz = websocket ? initialBufferSize : 0;
+    websocketPendingBytes.resetSizeIfEligable(sz),
+    websocketWriteRemainder.resetSizeIfEligable(sz);
+}
diff --git a/iowrapper.h b/iowrapper.h
index 00db122..2ffd804 100644
--- a/iowrapper.h
+++ b/iowrapper.h
@@ -118,7 +118,7 @@ class IoWrapper
     ssize_t websocketBytesToReadBuffer(void *buf, const size_t nbytes, IoWrapResult *error);
     ssize_t readOrSslRead(int fd, void *buf, size_t nbytes, IoWrapResult *error);
     ssize_t writeOrSslWrite(int fd, const void *buf, size_t nbytes, IoWrapResult *error);
-    ssize_t writeAsMuchOfBufAsWebsocketFrame(const void *buf, size_t nbytes, WebsocketOpcode opcode = WebsocketOpcode::Binary);
+    ssize_t writeAsMuchOfBufAsWebsocketFrame(const void *buf, const size_t nbytes, WebsocketOpcode opcode = WebsocketOpcode::Binary);
 public:
     IoWrapper(SSL *ssl, bool websocket, const size_t initialBufferSize, Client *parent);
     ~IoWrapper();
@@ -138,6 +138,8 @@ public:
 
     ssize_t readWebsocketAndOrSsl(int fd, void *buf, size_t nbytes, IoWrapResult *error);
     ssize_t writeWebsocketAndOrSsl(int fd, const void *buf, size_t nbytes, IoWrapResult *error);
+
+    void resetBuffersIfEligible();
 };
 
 #endif // IOWRAPPER_H
--
libgit2 0.21.4