├── src ├── config │ ├── mod.rs │ └── config_ws.rs ├── helpers │ ├── constants │ │ ├── mod.rs │ │ └── ws.rs │ ├── oncelock │ │ ├── mod.rs │ │ └── users.rs │ ├── enums │ │ ├── mod.rs │ │ └── ws_msg_type.rs │ ├── extension │ │ ├── mod.rs │ │ ├── ws_message.rs │ │ ├── ws_msg_type.rs │ │ └── user.rs │ ├── function │ │ ├── mod.rs │ │ ├── validation.rs │ │ ├── ws.rs │ │ ├── socket_addr.rs │ │ └── user.rs │ ├── mod.rs │ ├── types.rs │ └── ws_address.rs ├── events │ ├── mod.rs │ ├── sender │ │ ├── mod.rs │ │ ├── message.rs │ │ └── connection.rs │ └── handler.rs ├── models │ ├── mod.rs │ ├── message.rs │ └── user.rs └── main.rs ├── .gitignore ├── Cargo.toml └── readme.md /src/config/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod config_ws; -------------------------------------------------------------------------------- /src/helpers/constants/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod ws; 2 | -------------------------------------------------------------------------------- /src/helpers/oncelock/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod users; 2 | -------------------------------------------------------------------------------- /src/events/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod handler; 2 | pub mod sender; -------------------------------------------------------------------------------- /src/helpers/enums/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod ws_msg_type; 2 | -------------------------------------------------------------------------------- /src/models/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod user; 2 | pub mod message; -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | 3 | /target/ 4 | /Cargo.lock 5 | Cargo.lock 6 | -------------------------------------------------------------------------------- /src/events/sender/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod connection; 2 | pub mod message; 3 | -------------------------------------------------------------------------------- /src/helpers/constants/ws.rs: -------------------------------------------------------------------------------- 1 | pub const WS_AUTH: &str = "Authorization"; 2 | -------------------------------------------------------------------------------- /src/helpers/extension/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod user; 2 | pub mod ws_message; 3 | pub mod ws_msg_type; 4 | -------------------------------------------------------------------------------- /src/helpers/function/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod socket_addr; 2 | pub mod user; 3 | pub mod validation; 4 | pub mod ws; 5 | -------------------------------------------------------------------------------- /src/helpers/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod constants; 2 | pub mod enums; 3 | pub mod extension; 4 | pub mod function; 5 | pub mod oncelock; 6 | pub mod types; 7 | pub mod ws_address; 8 | -------------------------------------------------------------------------------- /src/models/message.rs: -------------------------------------------------------------------------------- 1 | use serde::{Deserialize, Serialize}; 2 | 3 | #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)] 4 | pub struct TextMessage { 5 | pub text: Option, 6 | pub to: Option, 7 | } 8 | -------------------------------------------------------------------------------- /src/helpers/extension/ws_message.rs: -------------------------------------------------------------------------------- 1 | use tokio_tungstenite::tungstenite::Message; 2 | 3 | pub trait WsMessageExt { 4 | fn to_ws_msg_text(&self) -> Message; 5 | } 6 | 7 | impl WsMessageExt for String { 8 | fn to_ws_msg_text(&self) -> Message { 9 | Message::Text(self.clone()) 10 | } 11 | } 12 | -------------------------------------------------------------------------------- /src/helpers/enums/ws_msg_type.rs: -------------------------------------------------------------------------------- 1 | use serde::{Deserialize, Serialize}; 2 | 3 | use crate::models::user::UserInfo; 4 | 5 | #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)] 6 | pub enum WsMsgType { 7 | NewConn(UserInfo), 8 | DisConn(UserInfo), 9 | ExistingConn(Vec), 10 | } 11 | -------------------------------------------------------------------------------- /src/helpers/extension/ws_msg_type.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | 3 | use crate::helpers::enums::ws_msg_type::WsMsgType; 4 | 5 | pub trait WsMsgTypeExt { 6 | fn to_json(&self) -> Result; 7 | } 8 | impl WsMsgTypeExt for WsMsgType { 9 | fn to_json(&self) -> Result { 10 | Ok(serde_json::to_string(&self)?) 11 | } 12 | } 13 | -------------------------------------------------------------------------------- /src/helpers/types.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | collections::HashMap, 3 | net::SocketAddr, 4 | sync::{Mutex, OnceLock}, 5 | }; 6 | use tokio::sync::mpsc::UnboundedSender; 7 | 8 | use tokio_tungstenite::tungstenite::protocol::Message; 9 | 10 | pub type USender = UnboundedSender; 11 | pub type SocketHashMap = OnceLock>>; 12 | -------------------------------------------------------------------------------- /src/helpers/oncelock/users.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | collections::HashSet, 3 | sync::{Mutex, OnceLock}, 4 | }; 5 | 6 | use crate::{helpers::types::SocketHashMap, models::user::User}; 7 | 8 | // pub static USERS: UsersHashMap = OnceLock::new(); 9 | pub static SOCKET_ADDRS: SocketHashMap = OnceLock::new(); 10 | pub static CURRENT_USER: OnceLock>> = OnceLock::new(); 11 | -------------------------------------------------------------------------------- /src/config/config_ws.rs: -------------------------------------------------------------------------------- 1 | use tokio::net::TcpListener; 2 | 3 | use crate::helpers::ws_address::ws_ip_address; 4 | use anyhow::{Ok, Result}; 5 | pub async fn ws_config() -> Result { 6 | let addr = ws_ip_address().await; 7 | let try_socket = TcpListener::bind(&addr).await; 8 | let listener = try_socket?; 9 | println!("Listening on: {}", addr); 10 | Ok(listener) 11 | } 12 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "socket_rust" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | tokio-tungstenite = "0.20" 10 | tokio = { version = "1.29", features = ["full"] } 11 | futures-util = "0.3" 12 | futures-channel = "0.3" 13 | async-recursion = "1.0" 14 | anyhow = "1.0" 15 | 16 | # json 17 | serde = { version = "1.0", features = ["derive"] } 18 | serde_json = "1.0" 19 | -------------------------------------------------------------------------------- /src/helpers/ws_address.rs: -------------------------------------------------------------------------------- 1 | use std::net::SocketAddr; 2 | 3 | use async_recursion::async_recursion; 4 | use tokio::net::TcpListener; 5 | 6 | pub async fn ws_ip_address() -> String { 7 | SocketAddr::from(([127, 0, 0, 1], free_port().await)).to_string() 8 | } 9 | 10 | #[async_recursion] 11 | async fn free_port() -> u16 { 12 | for port in 49152..=65535 { 13 | let addr = SocketAddr::from(([127, 0, 0, 1], port)); 14 | if let Ok(_) = TcpListener::bind(addr).await { 15 | return port; 16 | } 17 | } 18 | tokio::time::sleep(tokio::time::Duration::from_secs(2)).await; 19 | free_port().await 20 | } 21 | -------------------------------------------------------------------------------- /src/helpers/function/validation.rs: -------------------------------------------------------------------------------- 1 | use crate::models::message::TextMessage; 2 | use anyhow::Result; 3 | use serde_json::{from_str, from_value, Value}; 4 | use tokio_tungstenite::tungstenite::Message; 5 | 6 | pub fn validate_message_type(msg: Message) -> Result<(Option, Option)> { 7 | let msg = msg.to_string(); 8 | let parsed: Result = from_str(&msg.clone()); 9 | match parsed { 10 | Ok(value) => { 11 | let r: Result = from_value(value); 12 | match r { 13 | Ok(message) => Ok((Some(message), None)), 14 | Err(_) => Ok((None, Some(msg))), 15 | } 16 | } 17 | Err(_) => Ok((None, Some(msg))), 18 | } 19 | } 20 | -------------------------------------------------------------------------------- /src/models/user.rs: -------------------------------------------------------------------------------- 1 | use std::net::SocketAddr; 2 | 3 | use serde::{Deserialize, Serialize}; 4 | 5 | #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)] 6 | pub struct User { 7 | pub name: Option, 8 | pub addr: SocketAddr, 9 | } 10 | 11 | impl User { 12 | pub fn new(addr: SocketAddr) -> Self { 13 | Self { name: None, addr } 14 | } 15 | } 16 | 17 | #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)] 18 | pub struct UserInfo { 19 | pub name: Option, 20 | pub ws_id: Option, 21 | } 22 | 23 | impl UserInfo { 24 | pub fn new(n: String, w: String) -> Self { 25 | UserInfo { 26 | name: Some(n), 27 | ws_id: Some(w), 28 | } 29 | } 30 | } 31 | -------------------------------------------------------------------------------- /src/main.rs: -------------------------------------------------------------------------------- 1 | pub mod config; 2 | pub mod events; 3 | pub mod helpers; 4 | pub mod models; 5 | 6 | use anyhow::Result; 7 | use config::config_ws::ws_config; 8 | use events::handler::handle_connection; 9 | use helpers::function::{socket_addr::init_socket_addrs, user::init_current_user}; 10 | 11 | #[tokio::main] 12 | async fn main() { 13 | init_current_user(); 14 | init_socket_addrs(); 15 | match ws_run().await { 16 | Ok(_) => println!("ok"), 17 | Err(e) => println!("error: {:?}", e), 18 | } 19 | } 20 | 21 | pub async fn ws_run() -> Result<()> { 22 | let listener = ws_config().await?; 23 | while let Ok((stream, addr)) = listener.accept().await { 24 | tokio::task::spawn(async move { handle_connection(stream, addr).await }); 25 | } 26 | anyhow::Ok(()) 27 | } 28 | -------------------------------------------------------------------------------- /src/events/sender/message.rs: -------------------------------------------------------------------------------- 1 | use std::{net::SocketAddr, str::FromStr}; 2 | 3 | use anyhow::Result; 4 | use tokio::sync::mpsc::UnboundedSender; 5 | use tokio_tungstenite::tungstenite::Message; 6 | 7 | use crate::{ 8 | helpers::function::socket_addr::{get_all_other_u_sender, get_u_sender}, 9 | models::message::TextMessage, 10 | }; 11 | 12 | pub fn send_message( 13 | addr: SocketAddr, 14 | sender: UnboundedSender, 15 | message: TextMessage, 16 | ) -> Result<()> { 17 | if message.to.is_none() { 18 | send_message_to_all_other_users(addr.clone(), message.text.unwrap())?; 19 | return Ok(()); 20 | } 21 | let receiver = SocketAddr::from_str(&message.to.unwrap())?; 22 | match get_u_sender(receiver) { 23 | Some(recp) => { 24 | recp.send(Message::Text(message.text.unwrap()))?; 25 | } 26 | None => { 27 | sender.send(Message::Text(format!( 28 | "User not found: {}", 29 | receiver.to_string() 30 | )))?; 31 | } 32 | } 33 | Ok(()) 34 | } 35 | 36 | pub fn send_message_to_all_other_users(addr: SocketAddr, msg: String) -> Result<()> { 37 | let msg = Message::Text(msg); 38 | for recp in get_all_other_u_sender(addr.clone()) { 39 | recp.send(msg.clone())?; 40 | } 41 | Ok(()) 42 | } 43 | -------------------------------------------------------------------------------- /src/helpers/extension/user.rs: -------------------------------------------------------------------------------- 1 | use tokio_tungstenite::tungstenite::Message; 2 | 3 | use crate::{ 4 | helpers::enums::ws_msg_type::WsMsgType, 5 | models::user::{User, UserInfo}, 6 | }; 7 | pub trait UserExt { 8 | fn to_ws_user(&self) -> Message; 9 | } 10 | pub trait UsersExt { 11 | fn to_ws_users(&self) -> Message; 12 | } 13 | 14 | impl UserExt for User { 15 | fn to_ws_user(&self) -> Message { 16 | let users = serde_json::to_string(&self).unwrap(); 17 | Message::Text(users) 18 | } 19 | } 20 | 21 | impl UsersExt for Vec { 22 | fn to_ws_users(&self) -> Message { 23 | let users = serde_json::to_string(&self).unwrap(); 24 | Message::Text(users) 25 | } 26 | } 27 | 28 | pub trait UserInfoExt { 29 | fn to_new_conn(&self) -> WsMsgType; 30 | fn to_dis_conn(&self) -> WsMsgType; 31 | } 32 | 33 | impl UserInfoExt for UserInfo { 34 | fn to_new_conn(&self) -> WsMsgType { 35 | WsMsgType::NewConn(self.clone()) 36 | } 37 | 38 | fn to_dis_conn(&self) -> WsMsgType { 39 | WsMsgType::DisConn(self.clone()) 40 | } 41 | } 42 | pub trait UserInfoListExt { 43 | fn to_new_conn(&self) -> WsMsgType; 44 | } 45 | 46 | impl UserInfoListExt for Vec { 47 | fn to_new_conn(&self) -> WsMsgType { 48 | WsMsgType::ExistingConn(self.clone()) 49 | } 50 | } 51 | -------------------------------------------------------------------------------- /src/helpers/function/ws.rs: -------------------------------------------------------------------------------- 1 | use std::net::SocketAddr; 2 | 3 | use tokio_tungstenite::tungstenite::{ 4 | handshake::server::{ErrorResponse, Request, Response}, 5 | http::HeaderValue, 6 | Message, 7 | }; 8 | 9 | use crate::{events::sender::connection::close_connection_notify, helpers::constants::ws::WS_AUTH}; 10 | 11 | pub fn ws_disconnected(addr: SocketAddr, msg: Message) -> anyhow::Result { 12 | if msg.is_close() && msg.is_empty() { 13 | close_connection_notify(addr)?; 14 | return Ok(true); 15 | } 16 | return Ok(false); 17 | } 18 | 19 | pub fn ws_header_validation(req: &Request, res: Response) -> Result { 20 | for (ref header, value) in req.headers() { 21 | if header.as_str() == WS_AUTH.to_lowercase() { 22 | println!("* {}: {:?}", header, value); 23 | match validate_user_id(value) { 24 | Ok(v) => match v { 25 | true => return Ok(res), 26 | false => return ws_err(Some("Invalid user id".to_string())), 27 | }, 28 | Err(e) => return ws_err(Some(e.to_string())), 29 | } 30 | } 31 | } 32 | ws_err(None) 33 | } 34 | 35 | fn ws_err(msg: Option) -> Result { 36 | Err(Response::builder().status(401).body(msg).unwrap()) 37 | } 38 | 39 | fn validate_user_id(value: &HeaderValue) -> anyhow::Result { 40 | let id = value.to_str()?; 41 | Ok(id == "token") 42 | } 43 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # Rust WebSocket Example with tokio-tungstenite 2 | 3 | This is a simple example project demonstrating WebSocket implementation in Rust using the `tokio-tungstenite` crate. The project provides a basic WebSocket server written in Rust. 4 | 5 | ## Features 6 | 7 | - WebSocket server implementation using `tokio-tungstenite`. 8 | - Simple communication. 9 | 10 | ## Requirements 11 | 12 | - Rust: Make sure you have Rust installed. If not, you can install it from [https://www.rust-lang.org/tools/install](https://www.rust-lang.org/tools/install). 13 | 14 | ## Usage 15 | 16 | ### 1. Clone the Repository 17 | 18 | ```bash 19 | git clone https://github.com/blackshadowsoftwareltd/Websocket-Rust.git 20 | cd Websocket-Rust 21 | ``` 22 | 23 | ### 2. Build and Run the Server 24 | 25 | ```bash 26 | cargo run --release 27 | ``` 28 | 29 | #### After starting the server, you can see the IP address. now you can connect your Postman with this IP. 30 | 31 | ### 3. Headers. 32 | 33 | ### Before connecting your connection, you have to give the Header. 34 | 35 | ```bash 36 | key : Authorization 37 | value : token 38 | ``` 39 | 40 | #### Connect your multiple clients to send messages to other connections. after establishing a new connection you will see the IP of this connection. 41 | 42 | ### 4. To send a specific connection use this payload 43 | 44 | ```bash 45 | { 46 | "to":"127.0.0.1:48466", 47 | "text":"Hello!" 48 | } 49 | ``` 50 | 51 | ### Note : 52 | 53 | 1. If you do not have a provider ID, the message will be sent to all connections. 54 | 2. If you do not maintain this format (payload). or use any kind of String, the message will be sent to all of the connections. 55 | -------------------------------------------------------------------------------- /src/events/sender/connection.rs: -------------------------------------------------------------------------------- 1 | use std::net::SocketAddr; 2 | 3 | use anyhow::{Ok, Result}; 4 | use tokio::sync::mpsc::UnboundedSender; 5 | use tokio_tungstenite::tungstenite::Message; 6 | 7 | use crate::{ 8 | helpers::{ 9 | extension::{ 10 | user::{UserInfoExt, UserInfoListExt}, 11 | ws_message::WsMessageExt, 12 | ws_msg_type::WsMsgTypeExt, 13 | }, 14 | function::socket_addr::{ 15 | get_all_other_socket_addrs, get_all_other_u_sender, get_all_u_sender, 16 | remove_socket_addr, 17 | }, 18 | }, 19 | models::user::UserInfo, 20 | }; 21 | 22 | pub fn new_connection_notify(addr: SocketAddr) -> Result<()> { 23 | for recp in get_all_other_u_sender(addr) { 24 | let msg = UserInfo::new("Mr".to_string(), addr.to_string()) 25 | .to_new_conn() 26 | .to_json()? 27 | .to_ws_msg_text(); 28 | recp.send(msg)?; 29 | } 30 | Ok(()) 31 | } 32 | 33 | pub fn close_connection_notify(addr: SocketAddr) -> Result<()> { 34 | remove_socket_addr(addr); 35 | for recp in get_all_u_sender() { 36 | let msg = UserInfo::new("Mr".to_string(), addr.to_string()) 37 | .to_dis_conn() 38 | .to_json()? 39 | .to_ws_msg_text(); 40 | recp.send(msg)?; 41 | } 42 | Ok(()) 43 | } 44 | pub fn get_existing_connections(tx: UnboundedSender, addr: SocketAddr) -> Result<()> { 45 | let mut others: Vec = vec![]; 46 | for addr in get_all_other_socket_addrs(addr) { 47 | others.push(UserInfo::new("Mr".to_string(), addr.to_string())); 48 | } 49 | let msg = others.to_new_conn().to_json()?.to_ws_msg_text(); 50 | tx.send(msg)?; 51 | Ok(()) 52 | } 53 | -------------------------------------------------------------------------------- /src/helpers/function/socket_addr.rs: -------------------------------------------------------------------------------- 1 | use std::{collections::HashMap, net::SocketAddr, sync::Mutex}; 2 | 3 | use crate::helpers::{oncelock::users::SOCKET_ADDRS, types::USender}; 4 | 5 | // ---------- Users ---------- 6 | pub fn init_socket_addrs() { 7 | let mut set = SOCKET_ADDRS 8 | .get_or_init(|| Mutex::new(HashMap::new())) 9 | .lock() 10 | .unwrap(); 11 | set.clear(); 12 | } 13 | 14 | pub fn add_socket_addr(addr: SocketAddr, u_sender: USender) { 15 | let mut set = SOCKET_ADDRS 16 | .get_or_init(|| Mutex::new(HashMap::new())) 17 | .lock() 18 | .unwrap(); 19 | set.insert(addr, u_sender); 20 | } 21 | 22 | pub fn get_u_sender(addr: SocketAddr) -> Option { 23 | let set = SOCKET_ADDRS 24 | .get_or_init(|| Mutex::new(HashMap::new())) 25 | .lock() 26 | .unwrap(); 27 | set.get(&addr).cloned() 28 | } 29 | 30 | pub fn remove_socket_addr(addr: SocketAddr) { 31 | let mut set = SOCKET_ADDRS 32 | .get_or_init(|| Mutex::new(HashMap::new())) 33 | .lock() 34 | .unwrap(); 35 | set.remove(&addr); 36 | } 37 | 38 | pub fn get_all_socket_addrs() -> Vec { 39 | let set = SOCKET_ADDRS 40 | .get_or_init(|| Mutex::new(HashMap::new())) 41 | .lock() 42 | .unwrap(); 43 | set.iter().map(|(u, _)| u).cloned().collect() 44 | } 45 | 46 | pub fn get_all_other_socket_addrs(addr: SocketAddr) -> Vec { 47 | let set = SOCKET_ADDRS 48 | .get_or_init(|| Mutex::new(HashMap::new())) 49 | .lock() 50 | .unwrap(); 51 | set.clone() 52 | .into_iter() 53 | .map(|(u, _)| u) 54 | .filter(|x| x != &addr) 55 | .collect() 56 | } 57 | 58 | pub fn get_all_other_u_sender(addr: SocketAddr) -> Vec { 59 | let set = SOCKET_ADDRS 60 | .get_or_init(|| Mutex::new(HashMap::new())) 61 | .lock() 62 | .unwrap(); 63 | 64 | set.clone() 65 | .into_iter() 66 | .filter(|(u, _)| u != &addr) 67 | .map(|(_, s)| s) 68 | .collect() 69 | } 70 | 71 | pub fn get_all_u_sender() -> Vec { 72 | let set = SOCKET_ADDRS 73 | .get_or_init(|| Mutex::new(HashMap::new())) 74 | .lock() 75 | .unwrap(); 76 | 77 | set.clone().into_iter().map(|(_, s)| s).collect() 78 | } 79 | -------------------------------------------------------------------------------- /src/events/handler.rs: -------------------------------------------------------------------------------- 1 | use anyhow::Result; 2 | use futures_util::{SinkExt, StreamExt}; 3 | use std::net::SocketAddr; 4 | use tokio::sync::mpsc::unbounded_channel; 5 | use tokio_tungstenite::tungstenite::{ 6 | handshake::server::{Request, Response}, 7 | Message, 8 | }; 9 | 10 | use crate::{ 11 | events::sender::{ 12 | connection::{get_existing_connections, new_connection_notify}, 13 | message::{send_message, send_message_to_all_other_users}, 14 | }, 15 | helpers::function::{ 16 | socket_addr::add_socket_addr, 17 | validation::validate_message_type, 18 | ws::{ws_disconnected, ws_header_validation}, 19 | }, 20 | }; 21 | use tokio::net::TcpStream; 22 | 23 | pub async fn handle_connection(raw_stream: TcpStream, addr: SocketAddr) -> Result<()> { 24 | println!("Incoming TCP connection from: {}", addr); 25 | 26 | let ws_stream = 27 | tokio_tungstenite::accept_hdr_async(raw_stream, |req: &Request, res: Response| { 28 | ws_header_validation(req, res) 29 | }) 30 | .await 31 | .expect("Error during the websocket handshake occurred"); 32 | println!("WebSocket connection established: {}", addr); 33 | 34 | let (tx, mut rx) = unbounded_channel::(); 35 | add_socket_addr(addr.clone(), tx.clone()); 36 | 37 | let (mut ws_writer, mut ws_read) = ws_stream.split(); 38 | 39 | get_existing_connections(tx.clone(), addr.clone())?; 40 | new_connection_notify(addr.clone())?; 41 | loop { 42 | tokio::select! { 43 | Some(msg)=ws_read.next()=>{ 44 | match msg { 45 | Ok(msg) => { 46 | match ws_disconnected(addr.clone(),msg.clone())?{ 47 | true=>{ break; } 48 | false=>{ 49 | println!("Received a message from {}: {}", addr, msg.to_text()?); 50 | let (message,raw)= validate_message_type(msg.clone())?; 51 | if raw.is_some(){ 52 | send_message_to_all_other_users(addr.clone(),raw.unwrap())?; 53 | } else if message.is_some(){ 54 | send_message(addr.clone(),tx.clone(),message.unwrap())?; 55 | } 56 | } 57 | } 58 | } 59 | Err(e) => { 60 | println!("Error reading message from {}: {:?}", addr, e); 61 | break 62 | } 63 | } 64 | } 65 | Some(m)=rx.recv() =>{ 66 | ws_writer.send(m).await.expect("Failed to send msg"); 67 | } 68 | } 69 | } 70 | println!("Client {} disconnected", addr); 71 | anyhow::Ok(()) 72 | } 73 | -------------------------------------------------------------------------------- /src/helpers/function/user.rs: -------------------------------------------------------------------------------- 1 | use std::{collections::HashSet, sync::Mutex}; 2 | 3 | use crate::{helpers::oncelock::users::CURRENT_USER, models::user::User}; 4 | 5 | // // ---------- Users ---------- 6 | // pub fn init_users() { 7 | // let mut set = USERS 8 | // .get_or_init(|| Mutex::new(HashMap::new())) 9 | // .lock() 10 | // .unwrap(); 11 | // set.clear(); 12 | // } 13 | 14 | // pub fn add_user(user: User, u_sender: USender) { 15 | // let mut set = USERS 16 | // .get_or_init(|| Mutex::new(HashMap::new())) 17 | // .lock() 18 | // .unwrap(); 19 | // set.insert(user, u_sender); 20 | // } 21 | 22 | // pub fn get_u_sender(user: &User) -> Option { 23 | // let set = USERS 24 | // .get_or_init(|| Mutex::new(HashMap::new())) 25 | // .lock() 26 | // .unwrap(); 27 | // set.get(user).cloned() 28 | // } 29 | 30 | // pub fn remove_user(user: User) { 31 | // let mut set = USERS 32 | // .get_or_init(|| Mutex::new(HashMap::new())) 33 | // .lock() 34 | // .unwrap(); 35 | // set.remove(&user); 36 | // } 37 | 38 | // pub fn get_all_users() -> Vec { 39 | // let set = USERS 40 | // .get_or_init(|| Mutex::new(HashMap::new())) 41 | // .lock() 42 | // .unwrap(); 43 | // set.iter().map(|(u, _)| u).cloned().collect() 44 | // } 45 | 46 | // pub fn get_all_other_users(user: User) -> Vec { 47 | // let set = USERS 48 | // .get_or_init(|| Mutex::new(HashMap::new())) 49 | // .lock() 50 | // .unwrap(); 51 | // set.clone() 52 | // .into_iter() 53 | // .map(|(u, _)| u) 54 | // .filter(|x| x != &user) 55 | // .collect() 56 | // } 57 | 58 | // pub fn get_all_other_u_sender(user: User) -> Vec { 59 | // let set = USERS 60 | // .get_or_init(|| Mutex::new(HashMap::new())) 61 | // .lock() 62 | // .unwrap(); 63 | 64 | // set.clone() 65 | // .into_iter() 66 | // .filter(|(u, _)| u != &user) 67 | // .map(|(_, s)| s) 68 | // .collect() 69 | // } 70 | 71 | // pub fn get_all_u_sender() -> Vec { 72 | // let set = USERS 73 | // .get_or_init(|| Mutex::new(HashMap::new())) 74 | // .lock() 75 | // .unwrap(); 76 | 77 | // set.clone().into_iter().map(|(_, s)| s).collect() 78 | // } 79 | 80 | // pub fn add_addr_in_users(u: User, u_sender: USender) { 81 | // { 82 | // add_user(u, u_sender); 83 | // } 84 | // for x in get_all_users().iter() { 85 | // println!("user: {:?}", x); 86 | // } 87 | // } 88 | 89 | // ---------- Current User ---------- 90 | pub fn init_current_user() { 91 | let mut set = CURRENT_USER 92 | .get_or_init(|| Mutex::new(HashSet::new())) 93 | .lock() 94 | .unwrap(); 95 | set.clear(); 96 | } 97 | 98 | pub fn set_current_user(user: User) { 99 | let mut set = CURRENT_USER 100 | .get_or_init(|| Mutex::new(HashSet::new())) 101 | .lock() 102 | .unwrap(); 103 | { 104 | set.clear(); 105 | } 106 | { 107 | set.insert(user); 108 | } 109 | } 110 | 111 | pub fn get_current_user() -> Option { 112 | let set = CURRENT_USER 113 | .get_or_init(|| Mutex::new(HashSet::new())) 114 | .lock() 115 | .unwrap(); 116 | set.iter().cloned().last() 117 | } 118 | --------------------------------------------------------------------------------