added voice chat websocket route

src/routes/voice.rs

@@ -0,0 +1,126 @@
+use axum::{
+    Extension, Router,
+    extract::{
+        ConnectInfo, Path, Query, WebSocketUpgrade,
+        ws::{Message, WebSocket},
+    },
+    http::StatusCode,
+    response::IntoResponse,
+    routing::get,
+};
+use bytes::{BufMut, Bytes, BytesMut};
+use sqlx::PgPool;
+use std::{net::SocketAddr, time::Duration};
+use tokio::select;
+use uuid::Uuid;
+
+use crate::{
+    auth::verify_jwt_string,
+    db::{room_id_from_uuid, user_id_from_uuid},
+    realtime::RealTimeVoices,
+    routes::rooms::is_member,
+    routes::ws::WsAuthQuery,
+};
+
+pub fn routes() -> Router {
+    Router::new().route("/ws/voice/{room_uuid}", get(voice_ws_handler))
+}
+
+async fn voice_ws_handler(
+    ws: WebSocketUpgrade,
+    Path(room_uuid): Path<Uuid>,
+    Query(query): Query<WsAuthQuery>,
+    ConnectInfo(addr): ConnectInfo<SocketAddr>,
+    Extension(voice_manager): Extension<RealTimeVoices>,
+    Extension(db): Extension<PgPool>,
+) -> Result<impl IntoResponse, (StatusCode, String)> {
+    let claims = verify_jwt_string(&query.token)?;
+    let user_uuid = claims.sub;
+
+    let result = sqlx::query(
+        r#"
+        DELETE FROM ws_token_
+        WHERE token = $1
+          AND expires_at > now()
+        "#,
+    )
+    .bind(&query.token)
+    .execute(&db)
+    .await
+    .map_err(|e| {
+        tracing::error!("Failed to get WS token from DB: {e}");
+        (StatusCode::INTERNAL_SERVER_ERROR, "DB error".into())
+    })?;
+
+    if result.rows_affected() == 0 {
+        return Err((StatusCode::UNAUTHORIZED, "Invalid or expired token".into()));
+    }
+
+    let user_id = user_id_from_uuid(&db, user_uuid).await?;
+    let room_id = room_id_from_uuid(&db, room_uuid).await?;
+
+    if !is_member(user_id, room_id, &db).await {
+        return Err((StatusCode::FORBIDDEN, "Not a member of this room".into()));
+    }
+
+    tracing::info!("User {} joining voice in room {}", user_uuid, room_uuid);
+
+    let tx = voice_manager.get_or_create_room(room_uuid);
+    let rx = tx.subscribe();
+
+    Ok(ws.on_upgrade(move |socket| handle_voice_socket(socket, addr, user_uuid, tx, rx)))
+}
+
+async fn handle_voice_socket(
+    mut socket: WebSocket,
+    who: SocketAddr,
+    my_uuid: Uuid,
+    tx: tokio::sync::broadcast::Sender<(Uuid, Bytes)>,
+    mut rx: tokio::sync::broadcast::Receiver<(Uuid, Bytes)>,
+) {
+    let mut ping_interval = tokio::time::interval(Duration::from_secs(15));
+
+    loop {
+        select! {
+            // Receive audio from other users and send to client
+            voice_packet = rx.recv() => {
+                if let Ok((speaker_uuid, audio_data)) = voice_packet {
+                    if speaker_uuid != my_uuid {
+                        let mut msg = BytesMut::with_capacity(16 + audio_data.len());
+                        msg.put(speaker_uuid.as_bytes().as_slice());
+                        msg.put(audio_data);
+
+                        if socket.send(Message::Binary(msg.freeze().into())).await.is_err() {
+                            break;
+                        }
+                    }
+                }
+            }
+
+            // Receive audio from alient and broadcast to room
+            client_msg = socket.recv() => {
+                if let Some(Ok(msg)) = client_msg {
+                    match msg {
+                        Message::Binary(data) => {
+                            let _ = tx.send((my_uuid, Bytes::from(data)));
+                        }
+                        Message::Close(_) => {
+                            tracing::debug!("Voice client {} disconnected", who);
+                            break;
+                        }
+                        _ => {}
+                    }
+                } else {
+                    break;
+                }
+            }
+
+            // Keepalive
+            _ = ping_interval.tick() => {
+                if socket.send(Message::Ping(vec![].into())).await.is_err() {
+                    break;
+                }
+            }
+        }
+    }
+}