feat: build out chat server websocket more, start identity websocket

chat/src/http/api/guilds/create_guild.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use axum::{Extension, Json};
 use foxchat::{
     http::ApiError,
-    model::{http::guild::CreateGuildParams, user::PartialUser, Guild, channel::PartialChannel},
+    model::{channel::PartialChannel, http::guild::CreateGuildParams, user::PartialUser, Guild},
     FoxError,
 };
 
@@ -39,8 +39,12 @@ pub async fn post_guilds(
             instance: user.instance.domain,
         },
         default_channel: PartialChannel {
+            id: channel.id.0.clone(),
+            name: channel.name.clone(),
+        },
+        channels: Some(vec![PartialChannel {
             id: channel.id.0,
             name: channel.name,
-        }
+        }]),
     }))
 }

chat/src/http/ws/mod.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 
 use axum::{
     extract::{
@@ -8,8 +8,9 @@ use axum::{
     response::Response,
     Extension,
 };
-use eyre::Result;
+use eyre::{eyre, Result};
 use foxchat::{
+    model::User,
     s2s::{Dispatch, Payload},
     signature::{parse_date, verify_signature},
 };
@@ -17,7 +18,11 @@ use futures::{
     stream::{SplitSink, SplitStream, StreamExt},
     SinkExt,
 };
-use tokio::sync::{broadcast, mpsc, RwLock};
+use rand::Rng;
+use tokio::{
+    sync::{broadcast, mpsc, RwLock},
+    time::timeout,
+};
 use tracing::error;
 
 use crate::{app_state::AppState, model::identity_instance::IdentityInstance};
@@ -28,6 +33,7 @@ pub async fn handler(ws: WebSocketUpgrade, Extension(state): Extension<Arc<AppSt
 
 struct SocketState {
     instance: Option<IdentityInstance>,
+    last_heartbeat: u64,
 }
 
 async fn handle_socket(app_state: Arc<AppState>, socket: WebSocket) {
@@ -36,16 +42,21 @@ async fn handle_socket(app_state: Arc<AppState>, socket: WebSocket) {
     let (tx, rx) = mpsc::channel::<Payload>(10);
     let dispatch = app_state.broadcast.subscribe();
     let socket_state = Arc::new(RwLock::new(SocketState {
-        instance: None, // Filled out after IDENTIFY
+        // These are filled out after IDENTIFY
+        instance: None,
+        last_heartbeat: 0,
     }));
 
+    // Function that merges the socket-local and broadcast streams for sending.
     tokio::spawn(merge_channels(
         app_state.clone(),
         socket_state.clone(),
         dispatch,
         tx.clone(),
     ));
+    // Function that writes all payloads to the socket.
     tokio::spawn(write(rx, sender));
+    // Function that reads from the socket.
     tokio::spawn(read(
         app_state.clone(),
         socket_state.clone(),
@@ -71,7 +82,8 @@ async fn read(
                         server,
                         signature,
                     } => {
-                        let instance = match IdentityInstance::get(app_state, &server).await {
+                        let instance = match IdentityInstance::get(app_state.clone(), &server).await
+                        {
                             Ok(i) => i,
                             Err(e) => {
                                 error!("getting instance {}: {}", server, e);
@@ -161,20 +173,105 @@ async fn read(
         return;
     }
 
-    // Send hello
-    tx.send(Payload::Hello {}).await.ok();
+    // Generate heartbeat interval and send it in the Hello payload
+    let heartbeat_interval = rand::thread_rng().gen_range(45_000..70_000);
+    tx.send(Payload::Hello { heartbeat_interval }).await.ok();
+    // Start the heartbeat loop
+    let (heartbeat_tx, heartbeat_rx) = mpsc::channel::<u64>(10);
+    tokio::spawn(heartbeat(
+        socket_state.clone(),
+        heartbeat_interval,
+        heartbeat_rx,
+        tx.clone(),
+    ));
 
     while let Some(msg) = receiver.next().await {
-        let msg = if let Ok(msg) = msg {
-            msg
-        } else {
+        let Ok(msg) = msg else {
             return;
         };
 
-        // TODO: handle incoming payloads
+        let Ok(msg) = msg.to_text() else {
+            tx.send(Payload::Error {
+                message: "Invalid message".into(),
+            })
+            .await
+            .ok();
+            return;
+        };
+
+        let Ok(msg) = serde_json::from_str::<Payload>(msg) else {
+            tx.send(Payload::Error {
+                message: "Invalid message".into(),
+            })
+            .await
+            .ok();
+            return;
+        };
+
+        match msg {
+            Payload::Connect { user_id } => {
+                match collect_ready(app_state.clone(), socket_state.clone(), &user_id).await {
+                    Ok(p) => {
+                        tx.send(p).await.ok();
+                    }
+                    Err(e) => {
+                        error!("Error collecting ready event data for {}: {}", user_id, e);
+                        tx.send(Payload::Error {
+                            message: "Internal server error".into(),
+                        })
+                        .await
+                        .ok();
+                    }
+                }
+            }
+            Payload::Heartbeat { timestamp } => {
+                // Heartbeats are handled in another function
+                heartbeat_tx.send(timestamp).await.ok();
+            }
+            _ => {
+                tx.send(Payload::Error {
+                    message: "Invalid send event".into(),
+                })
+                .await
+                .ok();
+                return;
+            }
+        }
     }
 }
 
+async fn heartbeat(
+    socket_state: Arc<RwLock<SocketState>>,
+    heartbeat_interval: u64,
+    mut rx: mpsc::Receiver<u64>,
+    tx: mpsc::Sender<Payload>,
+) {
+    // The timeout is twice the heartbeat interval. If no heartbeat is received by then, close the connection with an error.
+    while let Ok(i) = timeout(Duration::from_millis(heartbeat_interval * 2), rx.recv()).await {
+        match i {
+            // ACK the heartbeat with the same timestamp. TODO: validate the timestamp to make sure we aren't too out of sync with the identity server.
+            Some(timestamp) => {
+                socket_state.write().await.last_heartbeat = timestamp;
+                tx.send(Payload::HeartbeatAck { timestamp }).await.ok();
+            }
+            // If the channel returns None, that means it's been closed, which means the socket as a whole was closed
+            None => {
+                return;
+            }
+        };
+    }
+
+    // Send an error, which will automatically close the socket
+    tx.send(Payload::Error {
+        message: format!(
+            "Did not receive a heartbeat after {}ms",
+            heartbeat_interval * 2
+        ),
+    })
+    .await
+    .ok();
+}
+
 async fn write(mut rx: mpsc::Receiver<Payload>, mut sender: SplitSink<WebSocket, Message>) {
     loop {
         if let Some(ev) = rx.recv().await {
@@ -289,10 +386,78 @@ async fn filter_events(
     socket_state: Arc<RwLock<SocketState>>,
     evt: &Dispatch,
 ) -> Result<(bool, Vec<String>)> {
-    // If we're not authenticated yet, don't send anything
-    if socket_state.read().await.instance.is_none() {
+    let Some(instance) = &socket_state.read().await.instance else {
         return Ok((false, vec![]));
-    }
+    };
 
-    Ok((true, vec![]))
+    match evt {
+        Dispatch::MessageCreate {
+            id: _,
+            channel_id: _,
+            guild_id,
+            author: _,
+            content: _,
+            created_at: _,
+        } => {
+            let users = sqlx::query!(
+                r#"SELECT ARRAY(
+                SELECT u.remote_user_id FROM users u
+                JOIN guilds_users gu ON gu.user_id = u.id
+                WHERE u.instance_id = $1 AND gu.guild_id = $2
+            )"#,
+                instance.id,
+                guild_id
+            )
+            .fetch_one(&app_state.pool)
+            .await?;
+
+            if let Some(users) = users.array {
+                return Ok((users.len() > 0, users));
+            }
+            return Ok((false, vec![]));
+        }
+        Dispatch::Ready { user, guilds: _ } => {
+            let user = sqlx::query!(
+                "SELECT remote_user_id FROM users WHERE id = $1",
+                user.id.clone()
+            )
+            .fetch_one(&app_state.pool)
+            .await?;
+
+            return Ok((true, vec![user.remote_user_id]));
+        }
+    }
+}
+
+async fn collect_ready(
+    app_state: Arc<AppState>,
+    socket_state: Arc<RwLock<SocketState>>,
+    user_id: &str,
+) -> Result<Payload> {
+    let Some(instance) = &socket_state.read().await.instance else {
+        return Err(eyre!("instance was None when it shouldn't be"));
+    };
+
+    let user = sqlx::query!(
+        r#"SELECT u.*, i.domain FROM users u
+        JOIN identity_instances i ON i.id = u.instance_id
+        WHERE u.instance_id = $1 AND u.remote_user_id = $2"#,
+        instance.id,
+        user_id
+    )
+    .fetch_one(&app_state.pool)
+    .await?;
+
+    Ok(Payload::Dispatch {
+        event: Dispatch::Ready {
+            user: User {
+                id: user.id,
+                username: user.username,
+                instance: user.domain,
+                avatar_url: None,
+            },
+            guilds: vec![],
+        },
+        recipients: vec![user.remote_user_id],
+    })
 }