From e4bb851c104a24e2eb32c6b9e5b4dd4fbf591b9f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Mon, 11 May 2020 14:04:13 +0200
Subject: [PATCH] aquatic_ws: in listener, remove previous connection if it
 exists

---
 aquatic_ws/src/lib/network.rs | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/aquatic_ws/src/lib/network.rs b/aquatic_ws/src/lib/network.rs
index 842f179..5838653 100644
--- a/aquatic_ws/src/lib/network.rs
+++ b/aquatic_ws/src/lib/network.rs
@@ -131,6 +131,35 @@ pub fn run_socket_worker(
 }
 
 
+fn remove_connection_if_exists(
+    poll: &mut Poll,
+    connections: &mut ConnectionMap,
+    token: Token,
+){
+    if let Some(connection) = connections.remove(&token){
+        match connection.stage {
+            ConnectionStage::Stream(mut stream) => {
+                poll.registry()
+                    .deregister(&mut stream)
+                    .unwrap();
+            },
+            ConnectionStage::MidHandshake(mut handshake) => {
+                poll.registry()
+                    .deregister(handshake.get_mut().get_mut())
+                    .unwrap();
+            }
+            ConnectionStage::Established(mut peer_connection) => {
+                poll.registry()
+                    .deregister(peer_connection.ws.get_mut())
+                    .unwrap();
+            }
+        };
+
+        connections.remove(&token);
+    }
+}
+
+
 fn accept_new_streams(
     listener: &mut TcpListener,
     poll: &mut Poll,
@@ -149,6 +178,8 @@ fn accept_new_streams(
 
                 let token = *poll_token_counter;
 
+                remove_connection_if_exists(poll, connections, token);
+
                 poll.registry()
                     .register(&mut stream, token, Interest::READABLE)
                     .unwrap();