From f52d81603205aa449ef1a43bf54185e91f555ea0 Mon Sep 17 00:00:00 2001
From: Alex Hultman <alexhultman@gmail.com>
Date: Mon, 22 Jun 2020 07:10:43 +0200
Subject: [PATCH] Never compress 0 bytes, finish pub/sub ded. compr.

---
 examples/BroadcastingEchoServer.cpp |  2 +-
 examples/EchoServer.cpp             |  2 +-
 src/PerMessageDeflate.h             |  4 +++-
 src/WebSocket.h                     |  6 +++---
 src/WebSocketContextData.h          | 17 +++++++++++++----
 5 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/examples/BroadcastingEchoServer.cpp b/examples/BroadcastingEchoServer.cpp
index 8b4fda6..a7049a0 100644
--- a/examples/BroadcastingEchoServer.cpp
+++ b/examples/BroadcastingEchoServer.cpp
@@ -11,7 +11,7 @@ int main() {
     /* Very simple WebSocket broadcasting echo server */
     uWS::App().ws<PerSocketData>("/*", {
         /* Settings */
-        .compression = uWS::SHARED_COMPRESSOR,
+        .compression = uWS::DEDICATED_COMPRESSOR_3KB,
         .maxPayloadLength = 16 * 1024 * 1024,
         .idleTimeout = 10,
         .maxBackpressure = 1 * 1024 * 1024,
diff --git a/examples/EchoServer.cpp b/examples/EchoServer.cpp
index 10a4c5b..a6d231b 100644
--- a/examples/EchoServer.cpp
+++ b/examples/EchoServer.cpp
@@ -28,7 +28,7 @@ int main() {
             /* Open event here, you may access ws->getUserData() which points to a PerSocketData struct */
         },
         .message = [](auto *ws, std::string_view message, uWS::OpCode opCode) {
-            ws->send(message, opCode);
+            ws->send(message, opCode, true);
         },
         .drain = [](auto *ws) {
             /* Check ws->getBufferedAmount() here */
diff --git a/src/PerMessageDeflate.h b/src/PerMessageDeflate.h
index f5a75d5..94d309c 100644
--- a/src/PerMessageDeflate.h
+++ b/src/PerMessageDeflate.h
@@ -131,7 +131,7 @@ struct DeflationStream {
         deflateInit2(&deflationStream, 1, Z_DEFLATED, windowBits, memLevel, Z_DEFAULT_STRATEGY);
     }
 
-    /* Deflate and optionally reset */
+    /* Deflate and optionally reset. You must not deflate an empty string. */
     std::string_view deflate(ZlibContext *zlibContext, std::string_view raw, bool reset) {
         /* Odd place to clear this one, fix */
         zlibContext->dynamicDeflationBuffer.clear();
@@ -167,6 +167,8 @@ struct DeflationStream {
             return {(char *) zlibContext->dynamicDeflationBuffer.data(), zlibContext->dynamicDeflationBuffer.length() - 4};
         }
 
+        /* Note: We will get an interger overflow resulting in heap buffer overflow if Z_BUF_ERROR is returned
+         * from passing 0 as avail_in. Therefore we must not deflate an empty string */
         return {
             zlibContext->deflationBuffer,
             DEFLATE_OUTPUT_CHUNK - deflationStream.avail_out - 4
diff --git a/src/WebSocket.h b/src/WebSocket.h
index 51b27d7..fd518c3 100644
--- a/src/WebSocket.h
+++ b/src/WebSocket.h
@@ -71,8 +71,8 @@ public:
         if (compress) {
             WebSocketData *webSocketData = (WebSocketData *) Super::getAsyncSocketData();
 
-            /* Check and correct the compress hint */
-            if (opCode < 3 && webSocketData->compressionStatus == WebSocketData::ENABLED) {
+            /* Check and correct the compress hint. It is never valid to compress 0 bytes */
+            if (message.length() && opCode < 3 && webSocketData->compressionStatus == WebSocketData::ENABLED) {
                 LoopData *loopData = Super::getLoopData();
                 /* Compress using either shared or dedicated deflationStream */
                 if (webSocketData->deflationStream) {
@@ -94,7 +94,7 @@ public:
 
         /* Get size, alloate size, write if needed */
         size_t messageFrameSize = protocol::messageFrameSize(message.length());
-        auto[sendBuffer, requiresWrite] = Super::getSendBuffer(messageFrameSize);
+        auto [sendBuffer, requiresWrite] = Super::getSendBuffer(messageFrameSize);
         protocol::formatMessage<isServer>(sendBuffer, message.data(), message.length(), opCode, message.length(), compress);
         /* This is the slow path, when we couldn't cork for the user */
         if (requiresWrite) {
diff --git a/src/WebSocketContextData.h b/src/WebSocketContextData.h
index 6081d39..e5a9b60 100644
--- a/src/WebSocketContextData.h
+++ b/src/WebSocketContextData.h
@@ -92,15 +92,22 @@ public:
                     WebSocket<SSL, true> *ws = (WebSocket<SSL, true> *) asyncSocket;
 
                     /* We need to handle being corked, and corking here */
-
+                    bool needsUncorking = false;
+                    if (!ws->isCorked() && ws->canCork()) {
+                        asyncSocket->cork();
+                        needsUncorking = true;
+                    }
 
                     while (selectedData.length()) {
-                        /* Interpret the data like so */
+                        /* Interpret the data like so, because this is how we shoved it in */
                         MessageMetadata mm;
                         memcpy((char *) &mm, selectedData.data(), sizeof(MessageMetadata));
                         std::string_view unframedMessage(selectedData.data() + sizeof(MessageMetadata), mm.length);
 
-                        //std::cout << "<" << unframedMessage << ">" << std::endl;
+                        /* Skip this message if our backpressure is too high */
+                        if (maxBackpressure && ws->getBufferedAmount() > maxBackpressure) {
+                            break;
+                        }
 
                         /* Here we perform the actual compression and framing */
                         ws->send(unframedMessage, mm.opCode, mm.compress);
@@ -110,7 +117,9 @@ public:
                     }
 
                     /* Here we need to uncork or keep it as was */
-
+                    if (needsUncorking) {
+                        asyncSocket->uncork();
+                    }
 
                     /* See below */
                     return 0;