feat: Remove Groups and Direct Messages from server when nobody is online

Author: ByteMaster2003

server/Config.toml

@@ -1,5 +1,5 @@
 port = 8443
 
-[tls]
-cert_path = "/etc/letsencrypt/live/nulltalk.viveksahani.com/fullchain.pem"
-key_path = "/etc/letsencrypt/live/nulltalk.viveksahani.com/privkey.pem"
+# [tls]
+# cert_path = "/etc/letsencrypt/live/nulltalk.viveksahani.com/fullchain.pem"
+# key_path = "/etc/letsencrypt/live/nulltalk.viveksahani.com/privkey.pem"

server/src/handlers/client.rs

@@ -1,14 +1,15 @@
-use std::sync::Arc;
-
-use tokio::sync::{Mutex as AsyncMutex, mpsc::UnboundedSender};
-
 use crate::{
-    data::CLIENTS, handlers::task::start_reader_task, net::perform_handshake, types::Client,
+    data::{CLIENTS, CONVERSATIONS, GROUPS},
+    handlers::task::start_reader_task,
+    net::perform_handshake,
+    types::Client,
 };
 use common::{
     net::{AsyncStream, Packet, StreamReader, StreamWriter},
     utils::enc::public_key_to_user_id,
 };
+use std::sync::Arc;
+use tokio::sync::{Mutex as AsyncMutex, mpsc::UnboundedSender};
 
 /// Handle a new client connection
 pub async fn handle_client(
@@ -29,6 +30,8 @@ pub async fn handle_client(
         username: name.clone().to_string(),
         user_id: client_id.clone(),
         session_key: hex::encode(&session_key),
+        dms: Vec::new(),
+        groups: Vec::new(),
         writer: wt.clone(),
     };
 
@@ -44,10 +47,76 @@ pub async fn handle_client(
     let _ = read_task.await;
 
     println!("🔗 client disconnected: {}", &client_id[..8]);
+    drop(rd);
+    drop(wt);
+    cleanup_client_data(client_id.clone()).await;
+}
+
+/// Removes user from groups and update/remove DM Session, etc.
+async fn cleanup_client_data(client_id: String) {
+    // Get client from CLIENTS
+    let client = {
+        let clients_lock = CLIENTS.lock().await;
+        clients_lock.get(&client_id).cloned()
+    };
+
+    match client {
+        Some(client) => {
+            for dm in client.dms {
+                update_or_remove_dm_session(&dm, client_id.clone()).await;
+            }
+            for gp in client.groups {
+                update_or_remove_group_session(&gp, client_id.clone()).await;
+            }
+        }
+        None => {
+            println!("⚠️ Client not found for cleanup: {}", &client_id[..8]);
+        }
+    }
+    // Additional cleanup logic can be added here
 
-    // Remove client from map
+    // Remove client from CLIENTS
     {
         let mut clients_lock = CLIENTS.lock().await;
         clients_lock.remove(&client_id);
+    };
+}
+
+async fn update_or_remove_dm_session(session_id: &str, client_id: String) {
+    let mut conv = CONVERSATIONS.lock().await;
+    let need_to_remove = match conv.get_mut(session_id) {
+        Some(dm) => {
+            dm.members.insert(client_id.clone(), false);
+            // Check if we need to remove the DM session
+            dm.members.values().all(|is_active| !is_active)
+        }
+        None => {
+            println!("⚠️ DM session not found for ID: {}", session_id);
+            false
+        }
+    };
+
+    if need_to_remove {
+        conv.remove(session_id);
+    }
+}
+
+async fn update_or_remove_group_session(session_id: &str, client_id: String) {
+    let mut groups = GROUPS.lock().await;
+    let need_to_remove = match groups.get_mut(session_id) {
+        Some(group) => {
+            group.members.insert(client_id.clone(), false);
+
+            // Check if we need to remove the group session
+            group.members.values().all(|is_active| !is_active)
+        }
+        None => {
+            println!("⚠️ Group session not found for ID: {}", session_id);
+            false
+        }
+    };
+
+    if need_to_remove {
+        groups.remove(session_id);
     }
 }

server/src/handlers/cmd.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use crate::{
     data::{CLIENTS, CONVERSATIONS, GROUPS},
     types::{DmChat, GroupChat},
@@ -75,10 +77,23 @@ async fn create_new_session(payload: Vec<u8>, client_id: String) -> ServerRespon
                         session_key = dm.session_key.clone();
                     }
                     None => {
+                        let mut members: HashMap<String, bool> = HashMap::new();
+                        members.insert(client_id.clone(), true);
+                        members.insert(new_session.id.clone(), true);
+
+                        for member in members.keys() {
+                            let mut guard = CLIENTS.lock().await;
+                            if let Some(client) = guard.get_mut(member) {
+                                if !client.dms.contains(&session_id) {
+                                    client.dms.push(session_id.clone());
+                                }
+                            }
+                        }
+
                         let dm_chat = DmChat {
                             dm_id: session_id.clone(),
                             session_key: session_key.clone(),
-                            members: (client_id.clone(), new_session.id.clone()),
+                            members,
                         };
 
                         conversations
@@ -101,7 +116,7 @@ async fn create_new_session(payload: Vec<u8>, client_id: String) -> ServerRespon
                 }
             };
 
-            if !group.members.contains(&client_id) {
+            if !group.members.contains_key(&client_id) {
                 response.success = false;
                 response.error = Some("You are not a member of this group".to_string());
 
@@ -110,6 +125,22 @@ async fn create_new_session(payload: Vec<u8>, client_id: String) -> ServerRespon
 
             let session_id = group.group_id.clone();
             let session_key = group.session_key.clone();
+
+            {
+                let mut guard = CLIENTS.lock().await;
+                if let Some(client) = guard.get_mut(&client_id) {
+                    if !client.groups.contains(&session_id) {
+                        client.groups.push(session_id.clone());
+                    }
+                }
+            }
+            {
+                let mut groups = GROUPS.lock().await;
+                if let Some(group) = groups.get_mut(&session_id) {
+                    group.members.insert(client_id.clone(), true);
+                }
+            }
+
             (session_id, session_key)
         }
     };
@@ -136,7 +167,7 @@ async fn create_new_group(payload: Vec<u8>, client_id: String) -> ServerResponse
         error: None,
     };
 
-    let mut members: Vec<String> = Vec::new();
+    let mut members: HashMap<String, bool> = HashMap::new();
 
     let (group_info, _): (NewGroupPayload, usize) =
         match bincode::decode_from_slice(&payload, bincode::config::standard()) {
@@ -149,10 +180,10 @@ async fn create_new_group(payload: Vec<u8>, client_id: String) -> ServerResponse
             }
         };
 
-    members.extend_from_slice(&group_info.members);
-    if !members.contains(&client_id) {
-        members.push(client_id.clone());
+    for member in &group_info.members {
+        members.insert(member.clone(), false);
     }
+    members.insert(client_id.clone(), true);
 
     let group_id = group_info
         .group_id
@@ -162,7 +193,7 @@ async fn create_new_group(payload: Vec<u8>, client_id: String) -> ServerResponse
         group_id: group_id.clone(),
         session_key: session_key.clone(),
         admin: client_id.clone(),
-        members,
+        members: members.clone(),
     };
 
     {
@@ -186,6 +217,15 @@ async fn create_new_group(payload: Vec<u8>, client_id: String) -> ServerResponse
         }
     }
 
+    for member in members.keys() {
+        let mut guard = CLIENTS.lock().await;
+        if let Some(client) = guard.get_mut(member) {
+            if !client.groups.contains(&group_id) {
+                client.groups.push(group_id.clone());
+            }
+        }
+    }
+
     let res_payload = NewGroupResponse {
         group_id,
         session_key,
@@ -231,19 +271,20 @@ async fn add_group_member(payload: Vec<u8>, client_id: String) -> ServerResponse
         return response;
     }
 
-    let member = match CLIENTS.lock().await.get(&data.member_id) {
-        Some(member) => member.clone(),
+    match CLIENTS.lock().await.get(&data.member_id).cloned() {
+        Some(mut member) => {
+            if !group.members.contains_key(&member.user_id) {
+                group.members.insert(member.user_id.clone(), true);
+            }
+            member.groups.push(group.group_id.clone());
+        }
         None => {
-            response.success = false;
-            response.error = Some("Member not found".to_string());
-            return response;
+            if !group.members.contains_key(&data.member_id) {
+                group.members.insert(data.member_id.clone(), false);
+            }
         }
     };
 
-    if !group.members.contains(&member.user_id) {
-        group.members.push(member.user_id.clone());
-    }
-
     response.payload = Some(
         bincode::encode_to_vec("Member Added successfully", bincode::config::standard()).unwrap(),
     );

server/src/handlers/msg.rs

@@ -18,7 +18,10 @@ pub async fn handle_group_message(packet: Packet, group_id: &str) {
         };
 
     // Broadcast the message to all clients in the group
-    for member_id in group.members {
+    for (member_id, is_active) in group.members {
+        if !is_active {
+            continue;
+        }
         if member_id == message.sender_id.clone() {
             continue;
         }
@@ -47,14 +50,13 @@ pub async fn handle_direct_message(packet: Packet, session_id: &str) {
         };
 
     // Find recipient
-    let member1 = dm.members.0;
-    let member2 = dm.members.1;
-    let recipient: String;
-
-    if message.sender_id.clone() == member1 {
-        recipient = member2;
-    } else {
-        recipient = member1;
+    let mut recipient = String::new();
+    for (member, _) in &dm.members {
+        if member == &message.sender_id {
+            continue;
+        } else {
+            recipient = member.clone();
+        }
     }
 
     // Check if client is online

server/src/types.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use common::net::StreamWriter;
 
 /// Represents a connected client
@@ -9,6 +11,10 @@ pub struct Client {
     pub user_id: String,
     /// session key of the client
     pub session_key: String,
+    /// direct message chats the client is part of
+    pub dms: Vec<String>,
+    /// group chats the client is part of
+    pub groups: Vec<String>,
     /// Stream writer for the client
     pub writer: StreamWriter,
 }
@@ -19,7 +25,7 @@ pub struct DmChat {
     /// unique identifier for the direct message chat
     pub dm_id: String,
     /// members of the direct message chat
-    pub members: (String, String),
+    pub members: HashMap<String, bool>,
     /// session key for the direct message chat
     pub session_key: Vec<u8>,
 }
@@ -32,7 +38,7 @@ pub struct GroupChat {
     /// unique identifier for the group chat
     pub group_id: String,
     /// members of the group chat
-    pub members: Vec<String>,
+    pub members: HashMap<String, bool>,
     /// session key for the group chat
     pub session_key: Vec<u8>,
     /// admin's user_id of the group chat