Fix iteration over remoteSockets during invalidation (#195)

When receiving INVL command from peer the connection socket is closed and
removed from remoteSockets map. However, since command processing is
performed inside an outer loop over remoteSockets, the removal also
invalidates the outer loop iterator which causes a segfault[0].

Fix the issue by refactoring the loop to assign the new iterator value
returned by map::erase() to the outer loop iterator.

Signed-off-by: Vlad Buslov <vladbu@nvidia.com>

Author: w1ldptr

src/core/nixl_listener.cpp

@@ -195,17 +195,22 @@ void nixlAgentData::commWorker(nixlAgent* myAgent){
         }
 
         // third, do remote commands
-        for (const auto& [sock_peer, socketClient] : remoteSockets ) {
+        auto socket_iter = remoteSockets.begin();
+        while (socket_iter != remoteSockets.end()) {
             std::string commands;
             std::vector<std::string> command_list;
             nixl_status_t ret;
 
-            ssize_t recv_bytes = recvCommMessage(socketClient, commands);
+            ssize_t recv_bytes = recvCommMessage(socket_iter->second, commands);
 
-            if(recv_bytes == 0 || recv_bytes == -1) continue;
+            if(recv_bytes == 0 || recv_bytes == -1) {
+                socket_iter++;
+                continue;
+            }
 
             command_list = str_split_substr(commands, "NIXLCOMM:");
 
+            bool invl = false;
             for(std::string command : command_list) {
 
                 if(command.size() < 4) continue;
@@ -225,14 +230,21 @@ void nixlAgentData::commWorker(nixlAgent* myAgent){
                     nixl_blob_t my_MD;
                     myAgent->getLocalMD(my_MD);
 
-                    sendCommMessage(socketClient, std::string("NIXLCOMM:LOAD" + my_MD));
+                    sendCommMessage(socket_iter->second, std::string("NIXLCOMM:LOAD" + my_MD));
                 } else if(header == "INVL") {
-                    close(socketClient);
-                    remoteSockets.erase(sock_peer);
+                    invl = true;
+                    break;
                 } else {
                     throw std::runtime_error("Received socket message with bad header" + header + ", critically failing\n");
                 }
             }
+
+            if (invl) {
+                close(socket_iter->second);
+                socket_iter = remoteSockets.erase(socket_iter);
+            } else {
+                socket_iter++;
+            }
         }
 
         nixlTime::us_t start = nixlTime::getUs();