From 17b718a10be8aabf0800863e00624c9818f389b9 Mon Sep 17 00:00:00 2001
From: Alex Hultman <alexhultman@gmail.com>
Date: Wed, 10 Jun 2020 14:20:32 +0200
Subject: [PATCH] Stricter pub/sub grouping

---
 src/TopicTree.h            | 27 ++++++++++++++++++++++++---
 src/WebSocketContextData.h |  7 +++++++
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/src/TopicTree.h b/src/TopicTree.h
index 0f039bb..6b15704 100644
--- a/src/TopicTree.h
+++ b/src/TopicTree.h
@@ -77,7 +77,7 @@ private:
     Topic *triggeredTopics[64];
     int numTriggeredTopics = 0;
     Subscriber *min = (Subscriber *) UINTPTR_MAX;
-    
+
     /* Cull or trim unused Topic nodes from leaf to root */
     void trimTree(Topic *topic) {
         if (!topic->subs.size() && !topic->children.size() && !topic->terminatingWildcardChild && !topic->wildcardChild) {
@@ -197,7 +197,7 @@ public:
                 newTopic->terminatingWildcardChild = nullptr;
                 newTopic->wildcardChild = nullptr;
                 memcpy(newTopic->name, segment.data(), segment.length());
-                
+
                 /* For simplicity we do insert wildcards with text */
                 iterator->children.insert(lb, {std::string_view(newTopic->name, segment.length()), newTopic});
 
@@ -217,6 +217,11 @@ public:
             }
         }
 
+        /* If this topic is triggered, drain the tree before we join */
+        if (iterator->triggered) {
+            drain();
+        }
+
         /* Add socket to Topic's Set */
         auto [it, inserted] = iterator->subs.insert(subscriber);
 
@@ -253,6 +258,11 @@ public:
             /* Try and remove this topic from our list */
             for (auto it = subscriber->subscriptions.begin(); it != subscriber->subscriptions.end(); it++) {
                 if (*it == iterator) {
+                    /* If this topic is triggered, drain the tree before we leave */
+                    if (iterator->triggered) {
+                        drain();
+                    }
+
                     /* Remove topic ptr from our list */
                     subscriber->subscriptions.erase(it);
 
@@ -270,6 +280,17 @@ public:
     void unsubscribeAll(Subscriber *subscriber) {
         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! */
+
+                /* If this topic is triggered, drain the tree before we leave */
+                if (topic->triggered) {
+                    drain();
+                }
+
+                /* Remove us from the topic's set */
                 topic->subs.erase(subscriber);
                 trimTree(topic);
             }
@@ -325,7 +346,7 @@ public:
                 it[i] = triggeredTopics[i]->subs.begin();
                 end[i] = triggeredTopics[i]->subs.end();
             }
-            
+
             /* Empty all sets from unique subscribers */
             for (int nonEmpty = numTriggeredTopics; nonEmpty; ) {
 
diff --git a/src/WebSocketContextData.h b/src/WebSocketContextData.h
index 120e0eb..81f230b 100644
--- a/src/WebSocketContextData.h
+++ b/src/WebSocketContextData.h
@@ -65,6 +65,13 @@ 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 */