refactor: clients now have a single websocket that handles all rooms the user is in

src/routes/messages.rs

@@ -18,6 +18,7 @@ pub struct MessageRow {
     pub uuid: Uuid,
     pub sender: String,
     pub sender_uuid: Uuid,
+    pub room_uuid: Uuid,
     pub message_type: String,
     pub content: String,
     pub sent_at: chrono::NaiveDateTime,
@@ -26,6 +27,7 @@ pub struct MessageRow {
 #[derive(serde::Serialize, Debug, Clone)]
 pub struct Message {
     pub uuid: Uuid,
+    pub room_uuid: Uuid,
     pub sender: String,
     pub sender_uuid: Uuid,
     pub message_type: String,
@@ -77,7 +79,7 @@ async fn list_messages(
             m.uuid,
             u.username AS sender,
             u.uuid AS sender_uuid,
-            r.uuid AS room,
+            r.uuid AS room_uuid,
             m.message_type,
             m.content,
             m.sent_at
@@ -106,6 +108,7 @@ async fn list_messages(
         .into_iter()
         .map(|m| Message {
             uuid: m.uuid,
+            room_uuid: m.room_uuid,
             sender: m.sender,
             sender_uuid: m.sender_uuid,
             message_type: m.message_type,
@@ -157,6 +160,7 @@ async fn create_message(
 
     let message = Message {
         uuid: uuid,
+        room_uuid,
         sender: sender_name,
         sender_uuid: claims.sub,
         message_type: payload.message_type,
@@ -164,8 +168,28 @@ async fn create_message(
         sent_at: sent_at.format("%Y-%m-%d %H:%M:%S").to_string(),
     };
 
-    let rt_sender = realtime.sender_for(room_id);
-    let _ = rt_sender.send(message.clone());
+    let recipients: Vec<Uuid> = sqlx::query_scalar(
+        r#"
+        SELECT u.uuid
+        FROM membership_ m
+        JOIN user_ u ON u.id = m.user_id
+        WHERE m.room = $1
+        "#,
+    )
+    .bind(room_id)
+    .fetch_all(&db)
+    .await
+    .map_err(|e| {
+        tracing::error!("Error fetching message recipients: {e}");
+        (StatusCode::INTERNAL_SERVER_ERROR, "DB error".into())
+    })?;
+
+    let rt = realtime.clone();
+    let msg_clone = message.clone();
+
+    tokio::spawn(async move {
+        rt.broadcast(recipients, msg_clone);
+    });
 
     Ok((StatusCode::CREATED, Json(message)))
 }

src/routes/ws.rs

@@ -3,23 +3,15 @@ use axum::extract::ws::{Message as WsMessage, WebSocket};
 use axum::extract::{ConnectInfo, Query};
 use axum::http::HeaderMap;
 use axum::routing::get;
-use axum::{
-    Extension,
-    extract::{Path, WebSocketUpgrade},
-    http::StatusCode,
-    response::IntoResponse,
-};
+use axum::{Extension, extract::WebSocketUpgrade, http::StatusCode, response::IntoResponse};
 use axum_extra::{TypedHeader, headers};
 use serde::Deserialize;
 use std::net::SocketAddr;
 use std::time::Duration;
 use tokio::select;
-use uuid::Uuid;
 
-use crate::auth::{create_jwt, verify_jwt};
-use crate::db::user_id_from_uuid;
-use crate::routes::rooms::is_member;
-use crate::{db::room_id_from_uuid, realtime::Realtime};
+use crate::auth::{create_jwt, verify_jwt, verify_jwt_string};
+use crate::realtime::Realtime;
 
 #[derive(sqlx::FromRow, serde::Serialize, Deserialize)]
 pub struct WsAuthQuery {
@@ -28,43 +20,27 @@ pub struct WsAuthQuery {
 
 pub fn routes() -> axum::Router {
     axum::Router::new()
-        .route("/ws/issue-token/rooms/{room_uuid}", get(issue_ws_token))
-        .route("/ws/rooms/{room_uuid}", get(ws_handler))
+        .route("/ws/messages/issue-token", get(issue_ws_token))
+        .route("/ws/messages", get(ws_handler))
 }
 
 pub async fn issue_ws_token(
     Extension(db): Extension<sqlx::PgPool>,
     headers: HeaderMap,
-    Path(room_uuid): Path<Uuid>,
 ) -> Result<(StatusCode, Json<WsAuthQuery>), (StatusCode, String)> {
     let claims = verify_jwt(headers)?;
 
-    let room_id = room_id_from_uuid(&db, room_uuid).await?;
-    let user_id = user_id_from_uuid(&db, claims.sub).await?;
-
-    if !is_member(user_id, room_id, &db).await {
-        return Err((
-            StatusCode::UNAUTHORIZED,
-            String::from("You are not a member of this room"),
-        ));
-    }
-
-    // tracing::info!(
-    //     "recieved token issue request from user {} for room {}",
-    //     claims.sub,
-    //     room_uuid
-    // );
+    tracing::debug!("Recieved token issue request from user {}", claims.sub);
 
     let token = create_jwt(claims.sub).map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
 
     sqlx::query(
         r#"
-        insert into ws_token_ (token, room_id, expires_at)
-        values ($1, $2, now() + interval '30 seconds')
+        insert into ws_token_ (token, expires_at)
+        values ($1, now() + interval '30 seconds')
         "#,
     )
     .bind(&token)
-    .bind(room_id)
     .execute(&db)
     .await
     .map_err(|_| {
@@ -80,46 +56,43 @@ pub async fn issue_ws_token(
 async fn ws_handler(
     ws: WebSocketUpgrade,
     user_agent: Option<TypedHeader<headers::UserAgent>>,
-    Path(room_uuid): Path<Uuid>,
     Query(query): Query<WsAuthQuery>,
     ConnectInfo(addr): ConnectInfo<SocketAddr>,
     Extension(realtime): Extension<Realtime>,
     Extension(db): Extension<sqlx::PgPool>,
-) -> Result<impl IntoResponse, axum::http::StatusCode> {
+) -> Result<impl IntoResponse, (StatusCode, String)> {
     // tracing::info!("recieved ws handshake: {}", room_uuid);
 
-    let room_id = room_id_from_uuid(&db, room_uuid)
-        .await
-        .map_err(|_| StatusCode::NOT_FOUND)?;
+    let claims = verify_jwt_string(&query.token)?;
+    let user_uuid = claims.sub;
 
-    let valid: Option<i32> = sqlx::query_scalar(
+    let result = sqlx::query(
         r#"
         delete from ws_token_
         where token = $1
-          and room_id = $2
           and expires_at > now()
-        returning room_id
         "#,
     )
     .bind(query.token)
-    .bind(room_id)
-    .fetch_optional(&db)
+    .execute(&db)
     .await
-    .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
+    .map_err(|e| {
+        tracing::error!("Failed to get WS token from DB: {e}");
+        (StatusCode::INTERNAL_SERVER_ERROR, "DB error".into())
+    })?;
 
-    if valid.is_none() {
-        return Err(StatusCode::UNAUTHORIZED);
+    if result.rows_affected() == 0 {
+        return Err((StatusCode::UNAUTHORIZED, "Wrong token".into()));
     }
 
-    let sender = realtime.sender_for(room_id);
-    let receiver = sender.subscribe();
+    let receiver = realtime.get_sender(user_uuid).subscribe();
 
     let user_agent = if let Some(TypedHeader(user_agent)) = user_agent {
         user_agent.to_string()
     } else {
         String::from("Unknown browser")
     };
-    tracing::debug!("`{user_agent}` at {addr} connected.");
+    tracing::debug!("`{user_agent}` {user_uuid} at {addr} connected.");
 
     Ok(ws.on_upgrade(move |socket| handle_socket(socket, addr, receiver)))
 }
@@ -133,7 +106,7 @@ async fn handle_socket(
 
     loop {
         select! {
-            // Receive broadcast messages and send to client
+            // Receive broadcast messages and send to client (any room)
             msg = receiver.recv() => {
                 if let Ok(msg) = msg {
                     if let Ok(json) = serde_json::to_string(&msg) {