use std::{ collections::HashMap, error::Error, sync::{ atomic::{AtomicUsize, Ordering}, mpsc::{self, Sender}, Arc, Mutex, }, thread, }; use headers::HeaderMapExt; use serde::{Deserialize, Serialize}; use touche::{upgrade::Upgrade, Body, Connection, Request, Response, Server, StatusCode}; use tungstenite::{protocol::Role, WebSocket}; #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type")] enum Event { #[serde(rename = "user")] User { id: usize }, #[serde(rename = "message", rename_all = "camelCase")] Message { user_id: usize, text: String }, } type Users = Arc>>>; fn main() -> std::io::Result<()> { static NEXT_USER_ID: AtomicUsize = AtomicUsize::new(1); let users: Users = Arc::new(Mutex::new(HashMap::new())); Server::bind("0.0.0.0:4444").serve(move |req: Request| { let users = users.clone(); if req.headers().typed_get::().is_some() { let res = tungstenite::handshake::server::create_response(&req.map(|_| ()))? .map(|_| Body::empty()); Ok::<_, Box>(res.upgrade(move |stream: Connection| { let users = users.clone(); let user_id = NEXT_USER_ID.fetch_add(1, Ordering::Relaxed); let (tx, rx) = mpsc::channel(); { let mut users = users.lock().unwrap(); users.insert(user_id, tx); let joined_msg = Event::User { id: user_id }; users.retain(|_id, tx| tx.send(joined_msg.clone()).is_ok()); }; let mut read_ws = WebSocket::from_raw_socket(stream.clone(), Role::Server, None); let mut write_ws = WebSocket::from_raw_socket(stream, Role::Server, None); let write_ws = thread::spawn(move || { for evt in rx { let msg = tungstenite::Message::Text(serde_json::to_string(&evt).unwrap()); if write_ws.write_message(msg).is_err() { break; } } }); let read_ws = thread::spawn(move || { while let Ok(msg) = read_ws.read_message() { match msg.to_text() { Ok(text) => { let text = text.to_owned(); let msg = Event::Message { user_id, text }; users .lock() .unwrap() .retain(|_id, tx| tx.send(msg.clone()).is_ok()); } Err(_) => break, } } }); write_ws.join().unwrap(); read_ws.join().unwrap(); })) } else { Ok(Response::builder() .status(StatusCode::OK) .header("content-type", "text/html") .body(include_str!("websocket-chat.html").into())?) } }) }