From abcfa3445b3675a68c681de7e094eb17c13d7aa7 Mon Sep 17 00:00:00 2001
From: dongresource <dongresource@protonmail.com>
Date: Thu, 9 Mar 2023 22:48:01 +0100
Subject: [PATCH] Move dead socket cleanup out of the poll() loop

This fixes a potential desync between the fds vector and the connections
map that could happen due to file descriptor number reuse.
---
 src/core/CNProtocol.cpp | 33 ++++++++++++++++++++-------------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/src/core/CNProtocol.cpp b/src/core/CNProtocol.cpp
index a54f35e..b5e7483 100644
--- a/src/core/CNProtocol.cpp
+++ b/src/core/CNProtocol.cpp
@@ -415,9 +415,9 @@ void CNServer::addPollFD(SOCKET s) {
     fds.push_back({s, POLLIN});
 }
 
-void CNServer::removePollFD(int i) {
+void CNServer::removePollFD(int fd) {
     auto it = fds.begin();
-    while (it != fds.end() && it->fd != fds[i].fd)
+    while (it != fds.end() && it->fd != fd)
         it++;
     assert(it != fds.end());
 
@@ -494,22 +494,29 @@ void CNServer::start() {
                 if (fds[i].revents & ~POLLIN)
                     cSock->kill();
 
-                if (cSock->isAlive()) {
+                if (cSock->isAlive())
                     cSock->step();
-                } else {
-                    killConnection(cSock);
-                    connections.erase(fds[i].fd);
-                    delete cSock;
-
-                    removePollFD(i);
-
-                    // a new entry was moved to this position, so we check it again
-                    i--;
-                }
             }
         }
 
         onStep();
+
+        // clean up dead connection sockets
+        auto it = connections.begin();
+        while (it != connections.end()) {
+            CNSocket *cSock = it->second;
+
+            if (!cSock->isAlive()) {
+                killConnection(cSock);
+                it = connections.erase(it);
+
+                removePollFD(cSock->sock);
+
+                delete cSock;
+            } else {
+                it++;
+            }
+        }
     }
 }