From 3fe7964e7488b09916d29fe7e0da8b69853be716 Mon Sep 17 00:00:00 2001
From: Alex Hultman <alexhultman@gmail.com>
Date: Sat, 9 Feb 2019 09:29:29 +0100
Subject: [PATCH] Properly unsubscribe from everywhere on close

---
 src/TopicTree.h        | 36 ++++++++++++++++++++++++++++++------
 src/WebSocket.h        |  4 ++++
 src/WebSocketContext.h |  3 +++
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/src/TopicTree.h b/src/TopicTree.h
index 6b9a19a..de207ee 100644
--- a/src/TopicTree.h
+++ b/src/TopicTree.h
@@ -44,13 +44,20 @@ private:
             }
         }
         /* Every subscriber should hold some backpressure cursor */
-        std::vector<std::pair<void *, bool *>> subscribers;
+        //std::vector<std::pair<void *, bool *>> subscribers;
+
+
+        std::set<void *> subscribers;
+
+
         std::string sharedMessage;
 
         /* We need backpressure stored */
         /* vector */
         std::string backpressure;
-    } *root = new Node;
+    } root; //topicToNode
+
+    std::map<void *, std::vector<Node *>> socketToNodeList;
 
     /* Nodes that hold something to send this iteration */
     std::set<Node *> pubNodes;
@@ -74,12 +81,14 @@ public:
             /* We say that all senders get their own message as well, for now being */
 
             for (Node *topicNode : pubNodes) {
-                for (auto [ws, valid] : topicNode->subscribers) {
+                for (auto /*[*/ws/*, valid]*/ : topicNode->subscribers) {
                     AsyncSocket<false> *asyncSocket = (AsyncSocket<false> *) ws; // assumes non-SSL for now
 
                     /* Writing optionally raw data */
                     auto [written, failed] = asyncSocket->write(topicNode->sharedMessage.data(), topicNode->sharedMessage.length(), true, 0);
 
+                    /* We should probably reset timeout for a WebSocket getting something sent */
+
 
                     /* Every subscriber to a topicNode will have int backpressure cursor to this room */
 
@@ -102,7 +111,7 @@ public:
 
     /* WebSocket.subscribe will lookup the Loop and subscribe in its tree */
     void subscribe(std::string topic, void *connection, bool *valid) {
-        Node *curr = root;
+        Node *curr = &root;
         for (int i = 0; i < topic.length(); i++) {
             int start = i;
             while (topic[i] != '/' && i < topic.length()) {
@@ -110,12 +119,27 @@ public:
             }
             curr = curr->get(topic.substr(start, i - start));
         }
-        curr->subscribers.push_back({connection, valid});
+        curr->subscribers.insert(connection);
+        /* Only do this if we did not aleady exist */
+        socketToNodeList[connection].push_back(curr);
+    }
+
+    /* Unsubscribe from all subscriptions */
+    void unsubscribeAll(void *connection) {
+
+        for (Node *node : socketToNodeList[connection]) {
+
+            /* Also make sure to update any backpressure here */
+
+            node->subscribers.erase(connection);
+        }
+
+        socketToNodeList.erase(connection);
     }
 
     /* WebSocket.publish looks up its tree and publishes to it */
     void publish(std::string topic, char *data, size_t length) {
-        Node *curr = root;
+        Node *curr = &root;
         for (int i = 0; i < topic.length(); i++) {
             int start = i;
             while (topic[i] != '/' && i < topic.length()) {
diff --git a/src/WebSocket.h b/src/WebSocket.h
index 16037d2..b0130fb 100644
--- a/src/WebSocket.h
+++ b/src/WebSocket.h
@@ -21,6 +21,7 @@
 #include "WebSocketData.h"
 #include "WebSocketProtocol.h"
 #include "AsyncSocket.h"
+#include "WebSocketContextData.h"
 
 #include <string_view>
 
@@ -124,6 +125,9 @@ public:
         if (webSocketContextData->closeHandler) {
             webSocketContextData->closeHandler(this, code, message);
         }
+
+        /* Make sure to unsubscribe from any pub/sub node at exit */
+        webSocketContextData->topicTree.unsubscribeAll(this);
     }
 
     /* Subscribe to a topic according to MQTT rules and syntax */
diff --git a/src/WebSocketContext.h b/src/WebSocketContext.h
index c7358df..a5c6011 100644
--- a/src/WebSocketContext.h
+++ b/src/WebSocketContext.h
@@ -246,6 +246,9 @@ private:
                 if (webSocketContextData->closeHandler) {
                     webSocketContextData->closeHandler((WebSocket<SSL, true> *) s, 1006, {});
                 }
+
+                /* Make sure to unsubscribe from any pub/sub node at exit */
+                webSocketContextData->topicTree.unsubscribeAll(s);
             }
 
             /* Destruct in-placed data struct */