From e48633952378dc75fa7eba4793f183bad531bcae Mon Sep 17 00:00:00 2001
From: Gabi <gabrielalejandro7@gmail.com>
Date: Tue, 6 Feb 2024 20:18:49 -0300
Subject: [PATCH] fix(snownet): properly cleanup connections before adding a
 new one (#3587)

This was causing some problems in my tests with connlib, when the client
re-created the connection and recieved the candidates before the offer
response, this *seems* like it made the client hang waiting for the
connection to happen, without failing it.

I couldn't find out why exactly the hang happens since it seems like the
remote candidates are added to the new connection and stun packets are
directed to the new connection, so it may be unrelated but I think it's
still right to clean this hashmap up.

---------

Co-authored-by: Thomas Eizinger <thomas@eizinger.io>
---
 rust/connlib/snownet/src/node.rs | 32 ++++++++++++++++++++++++++++----
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/rust/connlib/snownet/src/node.rs b/rust/connlib/snownet/src/node.rs
index 93b519e13..1044ac37f 100644
--- a/rust/connlib/snownet/src/node.rs
+++ b/rust/connlib/snownet/src/node.rs
@@ -644,6 +644,14 @@ where
         allowed_stun_servers: HashSet<SocketAddr>,
         allowed_turn_servers: HashSet<(SocketAddr, String, String, String)>,
     ) -> Offer {
+        if self.initial_connections.remove(&id).is_some() {
+            tracing::info!(%id, "Replacing existing initial connection");
+        };
+
+        if self.negotiated_connections.remove(&id).is_some() {
+            tracing::info!(%id, "Replacing existing established connection");
+        };
+
         self.upsert_stun_servers(&allowed_stun_servers);
         self.upsert_turn_servers(&allowed_turn_servers);
 
@@ -674,7 +682,7 @@ where
             },
         };
 
-        self.initial_connections.insert(
+        let existing = self.initial_connections.insert(
             id,
             InitialConnection {
                 agent,
@@ -685,13 +693,16 @@ where
             },
         );
 
+        debug_assert!(existing.is_none());
+
         params
     }
 
     /// Accept an [`Answer`] from the remote for a connection previously created via [`Node::new_connection`].
     pub fn accept_answer(&mut self, id: TId, remote: PublicKey, answer: Answer) {
         let Some(initial) = self.initial_connections.remove(&id) else {
-            return; // TODO: Better error handling
+            debug_assert!(false, "No initial connection to accept answer for");
+            return;
         };
 
         let mut agent = initial.agent;
@@ -708,7 +719,9 @@ where
             initial.turn_servers,
         );
 
-        self.negotiated_connections.insert(id, connection);
+        let existing = self.negotiated_connections.insert(id, connection);
+
+        debug_assert!(existing.is_none());
     }
 }
 
@@ -728,6 +741,15 @@ where
         allowed_stun_servers: HashSet<SocketAddr>,
         allowed_turn_servers: HashSet<(SocketAddr, String, String, String)>,
     ) -> Answer {
+        debug_assert!(
+            !self.initial_connections.contains_key(&id),
+            "server to not use `initial_connections`"
+        );
+
+        if self.negotiated_connections.remove(&id).is_some() {
+            tracing::info!(%id, "Replacing existing established connection");
+        };
+
         self.upsert_stun_servers(&allowed_stun_servers);
         self.upsert_turn_servers(&allowed_turn_servers);
 
@@ -764,7 +786,9 @@ where
             allowed_stun_servers,
             allowed_turn_servers,
         );
-        self.negotiated_connections.insert(id, connection);
+        let existing = self.negotiated_connections.insert(id, connection);
+
+        debug_assert!(existing.is_none());
 
         answer
     }