From f4bb87a098bc68e1cac0232b1b0058ddc0163c3d Mon Sep 17 00:00:00 2001
From: Alex Hultman <alexhultman@gmail.com>
Date: Sat, 28 Dec 2019 22:47:15 +0100
Subject: [PATCH] Add WebSocket::cork

---
 src/App.h       |  2 +-
 src/WebSocket.h | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/App.h b/src/App.h
index e11f7b2..d6b1b06 100644
--- a/src/App.h
+++ b/src/App.h
@@ -223,7 +223,7 @@ public:
                             (us_socket_context_t *) webSocketContext, (us_socket_t *) res, sizeof(WebSocketData) + sizeof(UserData));
 
                 /* Update corked socket in case we got a new one (assuming we always are corked in handlers). */
-                webSocket->cork();
+                webSocket->AsyncSocket<SSL>::cork();
 
                 /* Initialize websocket with any moved backpressure intact */
                 httpContext->upgradeToWebSocket(
diff --git a/src/WebSocket.h b/src/WebSocket.h
index 9c36e32..eeb848f 100644
--- a/src/WebSocket.h
+++ b/src/WebSocket.h
@@ -139,6 +139,21 @@ public:
         webSocketData->subscriber = nullptr;
     }
 
+    /* Corks the response if possible. Leaves already corked socket be. */
+    void cork(fu2::unique_function<void()> &&handler) {
+        if (!Super::isCorked() && Super::canCork()) {
+            Super::cork();
+            handler();
+
+            /* There is no timeout when failing to uncork for WebSockets,
+             * as that is handled by idleTimeout */
+            auto [written, failed] = Super::uncork();
+        } else {
+            /* We are already corked, or can't cork so let's just call the handler */
+            handler();
+        }
+    }
+
     /* Subscribe to a topic according to MQTT rules and syntax */
     void subscribe(std::string_view topic) {
         WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) us_socket_context_ext(SSL,