1 files changed, 31 insertions, 2 deletions
diff --git a/server-v1/source/EpistmoolServer/Session.cpp b/server-v1/source/EpistmoolServer/Session.cpp
index 6db51d8..cc95c6a 100644
--- a/server-v1/source/EpistmoolServer/Session.cpp
+++ b/server-v1/source/EpistmoolServer/Session.cpp
@@ -1,5 +1,7 @@
 #include "Session.hpp"
 
+#include <algorithm>
+
 using namespace Epistmool::Server;
 
 namespace {
@@ -63,14 +65,41 @@ Session* SessionManager::findSession(SessionId id)
 
 void SessionManager::addConnection(SessionId id, ConnectionId connId)
 {
+	// Add session -> connection links
 	auto session = findSession(id);
-	if (!session) return;
+	if (session) {
+		auto location = std::lower_bound(
+			session->mConnections.begin(),
+			session->mConnections.end(),
+			connId,
+			[](const ConnectionId& a, const ConnectionId& b) -> bool {
+				return a.index < b.index;
+			});
+		session->mConnections.insert(location, connId);
+	}
 
-	session->mConnections.push_back(connId);
+	// Add connection -> session links
 	mConnection2SessionMap.insert(connId.index, id.index);
 }
 
 void SessionManager::dropConnection(ConnectionId connId)
 {
+	// Remove session -> connection links
+	auto it = mConnection2SessionMap.find(connId.index);
+	while (it != mConnection2SessionMap.end() && it.key() == connId.index) {
+		auto& connList = mSessions.at(it.value()).mConnections;
+		auto range = std::equal_range(
+			connList.begin(),
+			connList.end(),
+			connId,
+			[](const ConnectionId& a, const ConnectionId& b) -> bool {
+				return a.index < b.index;
+			});
+		connList.erase(range.first, range.second);
+
+		++it;
+	}
+
+	// Remove connection -> sesion links
 	mConnection2SessionMap.remove(connId.index);
 }