src/main.rs

@@ -87,7 +87,7 @@ async fn handle_connection(
                         if parts.len() > 1 {
                             let requested_username = parts[1].to_string();
                             if let Err(e) =
-                                sd.handle_connect(addr, tx.clone(), requested_username).await
+                                sd.handle_connect_cmd(addr, tx.clone(), requested_username).await
                             {
                                 error!("handle_connect: {}", e);
                                 let _ = send(&tx, e.to_string().as_str());
@@ -96,6 +96,12 @@ async fn handle_connection(
                             let _ = send(&tx, "ERROR:INVALID:ID:");
                         }
                     }
+					"DISCONNECT" => {
+						if let Err(e) = sd.handle_disconnect_cmd(addr, tx.clone()).await {
+							error!("handle_disconnect: {}", e);
+							let _ = send(&tx, e.to_string().as_str());
+						}
+					}
                     "READY" => {
                         if let Err(e) = sd.handle_ready(addr, tx.clone()).await {
                             error!("handle_ready: {}", e);

src/server.rs

@@ -30,7 +30,7 @@ impl Server {
     }
 
     // Handler for CONNECT:<username>
-    pub async fn handle_connect(
+    pub async fn handle_connect_cmd(
         &self,
         addr: SocketAddr,
         tx: UnboundedSender<Message>,
@@ -55,6 +55,8 @@ impl Server {
 
         drop(clients_guard);
 
+        self.remove_observer_from_all_matches(addr).await;
+
         // not taken
         self.observers.write().await.remove(&addr);
         self.usernames.write().await.insert(requested_username.clone(), addr);
@@ -71,6 +73,35 @@ impl Server {
         Ok(())
     }
 
+    pub async fn handle_disconnect_cmd(&self, addr: SocketAddr, tx: UnboundedSender<Message>) -> Result<(), anyhow::Error> {
+        let clients_guard = self.clients.read().await;
+        let client_opt = clients_guard.get(&addr).cloned();
+
+        if client_opt.is_none() {
+            return Err(anyhow::anyhow!("ERROR:INVALID:DISCONNECT"));
+        }
+
+        drop(clients_guard);
+
+        let mut client = client_opt.as_ref().unwrap().write().await;
+        self.usernames.write().await.remove(&client.username);
+        client.ready = false;
+        client.color = Color::None;
+
+        if client.current_match.is_some() {
+            let match_id = client.current_match.unwrap();
+            drop(client);
+
+            self.terminate_match(match_id).await;
+        }
+
+        self.clients.write().await.remove(&addr);
+        self.observers.write().await.insert(addr, tx.clone());
+        let _ = send(&tx, "DISCONNECT:ACK");
+
+        Ok(())
+    }
+
     // Handler for READY
     pub async fn handle_ready(
         &self,
@@ -636,6 +667,26 @@ impl Server {
         Ok(())
     }
 
+    pub async fn remove_observer_from_all_matches(&self, addr: SocketAddr) {
+        let matches_guard = self.matches.read().await;
+
+        for match_guard in matches_guard.values() {
+            let mut found = false;
+            let mut a_match = match_guard.write().await;
+            for i in 0..a_match.viewers.len() {
+                if a_match.viewers[i] == addr {
+                    a_match.viewers.remove(i);
+                    found = true;
+                    break;
+                }
+            }
+
+            if found {
+                break;
+            }
+        }
+    }
+
     pub async fn terminate_match(&self, match_id: u32) {
         let matches_guard = self.matches.read().await;
         let the_match = matches_guard.get(&match_id);