From fd2d36328d5ba3ee6d5837b40f2fc41404356e96 Mon Sep 17 00:00:00 2001
From: Alex Hultman <alexhultman@gmail.com>
Date: Wed, 10 Jun 2020 22:25:45 +0200
Subject: [PATCH] Never drain a closing socket from unsubscribeAll

---
 src/TopicTree.h            | 8 +++-----
 src/WebSocket.h            | 2 +-
 src/WebSocketContext.h     | 2 +-
 src/WebSocketContextData.h | 7 -------
 4 files changed, 5 insertions(+), 14 deletions(-)

diff --git a/src/TopicTree.h b/src/TopicTree.h
index 6b15704..29308ce 100644
--- a/src/TopicTree.h
+++ b/src/TopicTree.h
@@ -277,16 +277,14 @@ public:
     }
 
     /* Can be called with nullptr, ignore it then */
-    void unsubscribeAll(Subscriber *subscriber) {
+    void unsubscribeAll(Subscriber *subscriber, bool mayFlush = true) {
         if (subscriber) {
             for (Topic *topic : subscriber->subscriptions) {
 
-                /* This is questionable; we are called mostly from socket close, so we will
-                 * potentially call drain callback with a closed socket, make sure to check there!
-                 * Well it doesn't really matter since there are checks in uSockets but still! */
+                /* We do not want to flush when closing a socket, it makes no sense to do so */
 
                 /* If this topic is triggered, drain the tree before we leave */
-                if (topic->triggered) {
+                if (mayFlush && topic->triggered) {
                     drain();
                 }
 
diff --git a/src/WebSocket.h b/src/WebSocket.h
index dd32540..d28e912 100644
--- a/src/WebSocket.h
+++ b/src/WebSocket.h
@@ -154,7 +154,7 @@ public:
         }
 
         /* Make sure to unsubscribe from any pub/sub node at exit */
-        webSocketContextData->topicTree.unsubscribeAll(webSocketData->subscriber);
+        webSocketContextData->topicTree.unsubscribeAll(webSocketData->subscriber, false);
         delete webSocketData->subscriber;
         webSocketData->subscriber = nullptr;
     }
diff --git a/src/WebSocketContext.h b/src/WebSocketContext.h
index 73e76ec..7a740e4 100644
--- a/src/WebSocketContext.h
+++ b/src/WebSocketContext.h
@@ -249,7 +249,7 @@ private:
                 }
 
                 /* Make sure to unsubscribe from any pub/sub node at exit */
-                webSocketContextData->topicTree.unsubscribeAll(webSocketData->subscriber);
+                webSocketContextData->topicTree.unsubscribeAll(webSocketData->subscriber, false);
                 delete webSocketData->subscriber;
                 webSocketData->subscriber = nullptr;
             }
diff --git a/src/WebSocketContextData.h b/src/WebSocketContextData.h
index 81f230b..120e0eb 100644
--- a/src/WebSocketContextData.h
+++ b/src/WebSocketContextData.h
@@ -65,13 +65,6 @@ struct WebSocketContextData {
         /* We rely on writing to regular asyncSockets */
         auto *asyncSocket = (AsyncSocket<SSL> *) s->user;
 
-        /* We might be called from unsubscribeAll from close handler of a socket,
-         * so make sure to check if we are a closed socket before trying to send */
-        if (us_socket_is_closed(SSL, (struct us_socket_t *) s->user) || us_socket_is_shut_down(SSL, (struct us_socket_t *) s->user)) {
-            /* Return code means nothing, todo: remove it */
-            return 0;
-        }
-
         /* Check if we now have too much backpressure (todo: don't buffer up before check) */
         if (!maxBackpressure || (unsigned int) asyncSocket->getBufferedAmount() < maxBackpressure) {
             /* Pick uncompressed data track */