118 lines
3.2 KiB
Rust
118 lines
3.2 KiB
Rust
use crate::AppState;
|
|
use axum::extract::ws::{Message, WebSocket};
|
|
use futures_util::{SinkExt, StreamExt};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use tokio::sync::{RwLock, mpsc};
|
|
use uuid::Uuid;
|
|
|
|
#[derive(Default)]
|
|
pub struct ChatHub {
|
|
// user_id -> sender
|
|
clients: RwLock<HashMap<Uuid, mpsc::UnboundedSender<ServerEvent>>>,
|
|
}
|
|
|
|
#[derive(Serialize, Clone)]
|
|
#[serde(tag = "type", rename_all = "snake_case")]
|
|
pub enum ServerEvent {
|
|
MessageCreated {
|
|
channel_id: Uuid,
|
|
message: serde_json::Value,
|
|
},
|
|
DmCreated {
|
|
other_user_id: Uuid,
|
|
message: serde_json::Value,
|
|
},
|
|
UserPresence {
|
|
user_id: Uuid,
|
|
online: bool,
|
|
},
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
#[serde(tag = "type", rename_all = "snake_case")]
|
|
enum ClientEvent {
|
|
// Currently no interactive client events for the general chat WS
|
|
Ping,
|
|
}
|
|
|
|
impl ChatHub {
|
|
pub async fn add_client(&self, user_id: Uuid, tx: mpsc::UnboundedSender<ServerEvent>) {
|
|
let mut clients = self.clients.write().await;
|
|
clients.insert(user_id, tx);
|
|
}
|
|
|
|
pub async fn remove_client(&self, user_id: Uuid) {
|
|
let mut clients = self.clients.write().await;
|
|
clients.remove(&user_id);
|
|
}
|
|
|
|
pub async fn get_online_users(&self) -> Vec<Uuid> {
|
|
let clients = self.clients.read().await;
|
|
clients.keys().cloned().collect()
|
|
}
|
|
|
|
pub async fn broadcast_all(&self, event: ServerEvent) {
|
|
let clients = self.clients.read().await;
|
|
for tx in clients.values() {
|
|
let _ = tx.send(event.clone());
|
|
}
|
|
}
|
|
|
|
pub async fn broadcast_to_user(&self, user_id: Uuid, event: ServerEvent) {
|
|
let clients = self.clients.read().await;
|
|
if let Some(tx) = clients.get(&user_id) {
|
|
let _ = tx.send(event);
|
|
}
|
|
}
|
|
|
|
pub async fn broadcast_to_many(&self, user_ids: Vec<Uuid>, event: ServerEvent) {
|
|
let clients = self.clients.read().await;
|
|
for user_id in user_ids {
|
|
if let Some(tx) = clients.get(&user_id) {
|
|
let _ = tx.send(event.clone());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub async fn handle_socket(state: AppState, socket: WebSocket, user_id: Uuid) {
|
|
let (mut ws_sender, mut ws_receiver) = socket.split();
|
|
let (tx, mut rx) = mpsc::unbounded_channel::<ServerEvent>();
|
|
|
|
state.chat.add_client(user_id, tx).await;
|
|
state
|
|
.chat
|
|
.broadcast_all(ServerEvent::UserPresence {
|
|
user_id,
|
|
online: true,
|
|
})
|
|
.await;
|
|
|
|
let send_task = tokio::spawn(async move {
|
|
while let Some(event) = rx.recv().await {
|
|
let Ok(payload) = serde_json::to_string(&event) else {
|
|
continue;
|
|
};
|
|
if ws_sender.send(Message::Text(payload.into())).await.is_err() {
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
|
|
while let Some(Ok(msg)) = ws_receiver.next().await {
|
|
if let Message::Close(_) = msg {
|
|
break;
|
|
}
|
|
}
|
|
|
|
send_task.abort();
|
|
state.chat.remove_client(user_id).await;
|
|
state
|
|
.chat
|
|
.broadcast_all(ServerEvent::UserPresence {
|
|
user_id,
|
|
online: false,
|
|
})
|
|
.await;
|
|
}
|