76 lines
2.3 KiB
Rust
76 lines
2.3 KiB
Rust
use crate::core::AppState;
|
|
use crate::http::context::CurrentUser;
|
|
use crate::models::user::Model as User;
|
|
use crate::routes::gateway::GatewayClient;
|
|
use axum::{
|
|
extract::{
|
|
State,
|
|
ws::{Message, WebSocket, WebSocketUpgrade},
|
|
},
|
|
response::IntoResponse,
|
|
};
|
|
use futures_util::{sink::SinkExt, stream::StreamExt};
|
|
use serde::Deserialize;
|
|
use tokio::sync::mpsc;
|
|
|
|
#[derive(Deserialize)]
|
|
pub struct WsQuery {
|
|
token: String,
|
|
}
|
|
|
|
pub async fn ws_handler(
|
|
ws: WebSocketUpgrade,
|
|
State(state): State<AppState>,
|
|
CurrentUser(user): CurrentUser,
|
|
) -> impl IntoResponse {
|
|
ws.on_upgrade(move |socket| handle_socket(socket, state, user))
|
|
|
|
// let token = query.token;
|
|
|
|
// Vérification du JWT
|
|
// match verify_jwt(&token, &state.config.jwt.secret) {
|
|
// Ok(claims) => ws.on_upgrade(move |socket| handle_socket(socket, state, claims.user_id)),
|
|
// Err(e) => {
|
|
// tracing::error!("WS auth error: {:?}", e);
|
|
// (axum::http::StatusCode::UNAUTHORIZED, "Invalid token").into_response()
|
|
// }
|
|
// }
|
|
}
|
|
|
|
async fn handle_socket(socket: WebSocket, state: AppState, user: User) {
|
|
let (mut sender, mut receiver) = socket.split();
|
|
let (tx, mut rx) = mpsc::unbounded_channel::<Message>();
|
|
|
|
let mut client = GatewayClient::new(user, tx);
|
|
client.on_connect().await;
|
|
state.gateway.add_client(client.clone());
|
|
|
|
// Task pour envoyer les messages du canal mpsc vers le WebSocket
|
|
let mut send_task = tokio::spawn(async move {
|
|
while let Some(message) = rx.recv().await {
|
|
if sender.send(message).await.is_err() {
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
|
|
// Task pour recevoir les messages du WebSocket
|
|
let client_clone = client.clone();
|
|
let state_clone = state.clone();
|
|
let mut recv_task = tokio::spawn(async move {
|
|
while let Some(Ok(message)) = receiver.next().await {
|
|
client_clone.on_message(message, &state_clone).await;
|
|
}
|
|
});
|
|
|
|
// Attente de la fin d'une des tâches (déconnexion)
|
|
tokio::select! {
|
|
_ = (&mut send_task) => recv_task.abort(),
|
|
_ = (&mut recv_task) => send_task.abort(),
|
|
};
|
|
|
|
state.gateway.remove_client(&client);
|
|
// // Déconnexion (Disconnect)
|
|
client.on_disconnect(&state).await;
|
|
}
|