Removed user from subscriptions on disconnect

1 files changed, 16 insertions, 0 deletions

diff --git a/src/BootstrapNode.cpp b/src/BootstrapNode.cpp
index e62e48d..273abf0 100644
--- a/src/BootstrapNode.cpp
+++ b/src/BootstrapNode.cpp
@@ -17,6 +17,21 @@ namespace sibs
         connections(27130),
         socket(connections.createSocket(address, false, true))
     {
+        connections.onRemoveDisconnectedPeer([this](std::shared_ptr<DirectConnectionPeer> peer)
+        {
+            std::lock_guard<std::mutex> lock(subscribedPeersMutex);
+            for(auto &topicUsers : subscribedPeers)
+            {
+                for(auto it = topicUsers.second.begin(); it != topicUsers.second.end(); )
+                {
+                    if(peer->address == (*it)->address)
+                        it = topicUsers.second.erase(it);
+                    else
+                        ++it;
+                }
+            }
+        });
+
         if(UDT::listen(socket, 10) == UDT::ERROR)
         {
             std::string errMsg = "UDT: Failed to listen, error: ";
@@ -75,6 +90,7 @@ namespace sibs
         PubsubKey pubsubKey;
         deserializer.extract(pubsubKey.data.data(), PUBSUB_KEY_LENGTH);
         
+        std::lock_guard<std::mutex> lock(subscribedPeersMutex);
         auto &peers = subscribedPeers[pubsubKey];
         for(auto &peer : peers)
         {