From 72e62d4bd762fe197a29d4dbc67e6287b9affa73 Mon Sep 17 00:00:00 2001
From: szymon <szymon_jozef@proton.me>
Date: Sun, 16 Aug 2026 19:39:38 +0200
Subject: [PATCH] feat: add removing from connection list on connection broke

---
 src/network/connection.rs |  2 ++
 src/network/message.rs    |  4 ++++
 src/network/server.rs     | 30 +++++++++++++++++++++++++++---
 3 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/src/network/connection.rs b/src/network/connection.rs
index c6f56f3..7d58775 100644
--- a/src/network/connection.rs
+++ b/src/network/connection.rs
@@ -128,6 +128,7 @@ where
                             "Connection broken with {:?}: {}. Disconnecting...",
                             self.peer, e
                         );
+                        self.tcp = None;
                         break;
                     }
                 },
@@ -139,6 +140,7 @@ where
                         },
                         None => {
                             info!("Remote closed the channel for {:?}. Disconnecting...", self.peer);
+                            self.tcp = None;
                             break;
                         }
                     }
diff --git a/src/network/message.rs b/src/network/message.rs
index 9d175eb..460f6e6 100644
--- a/src/network/message.rs
+++ b/src/network/message.rs
@@ -1,3 +1,5 @@
+use std::net::SocketAddr;
+
 use serde::{Deserialize, Serialize, de::DeserializeOwned};
 use uuid::Uuid;
 
@@ -52,6 +54,8 @@ pub enum ClientMessage {
     AnswerHandshake { player_name: String },
     SignalReadiness,
 
+    ConnectionDead { addr: SocketAddr },
+
     RequestCapture { capture_path: CapturePath },
     RequestMove { from: Position, to: Position },
 
diff --git a/src/network/server.rs b/src/network/server.rs
index 00328fd..974ae24 100644
--- a/src/network/server.rs
+++ b/src/network/server.rs
@@ -105,22 +105,35 @@ impl Server {
 
         info!("Starting the server at: {:?}", listener.local_addr());
 
-        let sender = self.sender.clone();
+        let server_sender = self.sender.clone();
 
         tokio::spawn(async move {
             loop {
                 if let Ok((socket, addr)) = listener.accept().await {
                     info!("New connection from: {:?}", addr);
                     let mut conn: Connection<ClientMessage, ServerMessage> =
-                        Connection::new(ConnectionType::Client, sender.clone());
+                        Connection::new(ConnectionType::Client, server_sender.clone());
 
                     conn.delegate(socket, addr);
 
                     let sender = conn.get_sender();
+                    let disconnect_sender = server_sender.clone();
                     let _ = welcoming_tx.send((addr, sender)).await;
 
                     tokio::spawn(async move {
                         conn.start_listening().await;
+
+                        // TODO! Add waiting for reconnect
+
+                        let disconnect_msg = Message::new(ClientMessage::ConnectionDead { addr });
+                        let _ = disconnect_sender
+                            .send((
+                                addr,
+                                disconnect_msg.expect(
+                                    "Could not create disconnect msg. this shouldn't happen",
+                                ),
+                            ))
+                            .await;
                     });
                 }
             }
@@ -142,7 +155,6 @@ impl Server {
 
                             let msg = Message::new(ServerMessage::DeclineHandshake { reason: reason });
 
-
                             if let Ok(msg) = msg {
                             let _ = sender.send(msg).await; // we can't use send_message because we
                                                             // don't add addr to connections
@@ -166,6 +178,14 @@ impl Server {
                                     self.process_handshake(addr, player_name);
                                 }
 
+                                ClientMessage::ConnectionDead { addr } => {
+                                    warn!("Removing: {:?} from connection list", addr);
+                                    // TODO! Should decrement ready_counter if player was ready. But
+                                    // how to check that? Maybe create new struct that will hold
+                                    // sender and ready and maybe some other data.
+                                    self.connections.remove(&addr);
+                                },
+
                                 ClientMessage::SignalReadiness => {
                                     self.state.ready_count += 1;
                                     info!("{} signaled readiness. Currently ready: {}/{}", addr, self.state.ready_count, self.settings.max_connections);
@@ -236,4 +256,8 @@ impl Server {
             Ok(msg) => conn.send(msg).await,
         };
     }
+
+    fn remove_client(&mut self, addr: SocketAddr) {
+        self.connections.remove(&addr);
+    }
 }