simple websocket implementation (without auth use)

src/utils/db_pool.rs

@@ -2,13 +2,25 @@ use std::sync::Arc;
 use dashmap::DashMap;
 use r2d2::{Pool};
 use r2d2_sqlite::SqliteConnectionManager;
+use tokio::sync::mpsc;
+use axum::extract::{ws::{Message, WebSocket, WebSocketUpgrade}, State};
 
 type HotelId = i32; // or i32 if you want numeric ids
 
+    /// Type alias: user_id → sender to that user
+type UserMap = DashMap<i32, mpsc::UnboundedSender<Message>>;
+/// hotel_id → users
+type HotelMap = DashMap<i32, Arc<UserMap>>;
+/// global map of all hotels
+type WsMap = Arc<HotelMap>;
+    /// Type alias: user_id → sender to that user
+
+
 #[derive(Clone)]
 pub struct AppState {
     pub hotel_pools: HotelPool,
     pub logs_pool: Pool<SqliteConnectionManager>,
+    pub ws_map: WsMap,
    
 }
 

src/utils/mod.rs

@@ -1,3 +1,4 @@
 pub mod db_pool;
 pub mod auth;
-pub mod routes;
+pub mod routes;
+pub mod websocket;

src/utils/routes.rs

@@ -5,7 +5,7 @@ use axum::{
 
 use crate::utils::auth::*;
 use crate::utils::db_pool::{HotelPool, AppState, };
-
+use crate::utils::websocket::ws_handler;
 
 // ROOTS
 pub fn utils_routes() -> Router<AppState> {
@@ -13,6 +13,8 @@ pub fn utils_routes() -> Router<AppState> {
     Router::new()
         .route("/login", put(clean_auth_loging))
         .route("/register", put(register_user))
+        .route("/ws/{hotel_id}/{user_id}", get(ws_handler))
         .route("/tokentest", put(token_tester))
+
         //.with_state(state)
     }

src/utils/websocket.rs

@@ -0,0 +1,108 @@
+use dashmap::DashMap;
+use std::sync::Arc;
+use axum::extract::{ws::{Message, WebSocket, WebSocketUpgrade}, State};
+use tokio::sync::mpsc;
+use axum::extract::Path;
+use axum::response::IntoResponse;
+//use futures_util::stream::stream::StreamExt;
+use futures_util::{StreamExt, SinkExt};
+
+use crate::utils::db_pool::{HotelPool,AppState};
+
+
+
+    /// Type alias: user_id → sender to that user
+type UserMap = DashMap<i32, mpsc::UnboundedSender<Message>>;
+/// hotel_id → users
+type HotelMap = DashMap<i32, Arc<UserMap>>;
+/// global map of all hotels
+type WsMap = Arc<HotelMap>;
+    /// Type alias: user_id → sender to that user
+
+
+async fn handle_socket(
+    mut socket: WebSocket,
+    state: AppState,
+    hotel_id: i32,
+    user_id: i32,
+) {
+    // channel for sending messages TO this client
+    let (tx, mut rx) = mpsc::unbounded_channel::<Message>();
+
+    // insert into hotel → user map
+    let user_map = state
+        .ws_map
+        .entry(hotel_id)
+        .or_insert_with(|| Arc::new(UserMap::new()))
+        .clone();
+
+    user_map.insert(user_id, tx);
+
+    // ✅ print after upgrading
+    print_ws_state(&state);
+
+    // split socket into sender/receiver
+    let (mut sender, mut receiver) = socket.split();
+
+    
+
+    // task for sending messages from server to client
+    let mut rx_task = tokio::spawn(async move {
+        while let Some(msg) = rx.recv().await {
+            if sender.send(msg).await.is_err() {
+                break;
+            }
+        }
+    });
+
+    // task for receiving messages from client
+    let state_clone = state.clone();
+    let mut recv_task = tokio::spawn(async move {
+        while let Some(Ok(msg)) = receiver.next().await {
+            match msg {
+                Message::Text(text) => {
+                    println!("Hotel {hotel_id}, User {user_id} said: {text}");
+                    // echo back just as an example
+                    if let Some(hotel_entry) = state_clone.ws_map.get(&hotel_id) {
+                        if let Some(sender) = hotel_entry.get(&user_id) {
+                            let _ = sender.send(Message::Text(format!("echo: {text}").into()));
+                        }
+                    }
+                }
+                Message::Close(_) => break,
+                _ => {}
+            }
+        }
+    });
+
+    // wait for either side to finish
+    tokio::select! {
+        _ = (&mut rx_task) => recv_task.abort(),
+        _ = (&mut recv_task) => rx_task.abort(),
+    }
+
+    // cleanup
+    user_map.remove(&user_id);
+    if user_map.is_empty() {
+        state.ws_map.remove(&hotel_id);
+    }
+}
+
+pub async fn ws_handler(
+    ws: WebSocketUpgrade,
+    State(state): State<AppState>,
+    Path((hotel_id, user_id)): Path<(i32, i32)>,
+) -> impl IntoResponse {
+    ws.on_upgrade(move |socket| handle_socket(socket, state, hotel_id, user_id))
+}
+
+fn print_ws_state(state: &AppState) {
+    println!("--- Current WebSocket state ---");
+    for hotel_entry in state.ws_map.iter() {
+        let hotel_id = *hotel_entry.key();
+        let user_map = hotel_entry.value();
+        let users: Vec<_> = user_map.iter().map(|u| *u.key()).collect();
+        println!("Hotel {hotel_id}: users {:?}", users);
+    }
+    println!("--------------------------------");
+}