diff --git a/crates/bonfire/src/events/state.rs b/crates/bonfire/src/events/state.rs index 878eb8bb..988d5143 100644 --- a/crates/bonfire/src/events/state.rs +++ b/crates/bonfire/src/events/state.rs @@ -1,5 +1,9 @@ -use std::collections::{HashMap, HashSet}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; +use async_std::sync::RwLock; use lru::LruCache; use revolt_database::{Channel, Member, Server, User}; @@ -60,8 +64,8 @@ pub struct State { pub session_id: String, pub private_topic: String, - subscribed: HashSet, - state: SubscriptionStateChange, + pub state: SubscriptionStateChange, + pub subscribed: Arc>>, } impl State { @@ -81,7 +85,7 @@ impl State { State { cache, - subscribed, + subscribed: Arc::new(RwLock::new(subscribed)), session_id, private_topic, state: SubscriptionStateChange::Reset, @@ -89,15 +93,16 @@ impl State { } /// Apply currently queued state - pub fn apply_state(&mut self) -> SubscriptionStateChange { + pub async fn apply_state(&mut self) -> SubscriptionStateChange { let state = std::mem::replace(&mut self.state, SubscriptionStateChange::None); + let mut subscribed = self.subscribed.write().await; if let SubscriptionStateChange::Change { add, remove } = &state { for id in add { - self.subscribed.insert(id.clone()); + subscribed.insert(id.clone()); } for id in remove { - self.subscribed.remove(id); + subscribed.remove(id); } } @@ -109,20 +114,16 @@ impl State { self.cache.users.get(&self.cache.user_id).unwrap().clone() } - /// Iterate through all subscriptions - pub fn iter_subscriptions(&self) -> std::collections::hash_set::Iter<'_, std::string::String> { - self.subscribed.iter() - } - /// Reset the current state - pub fn reset_state(&mut self) { + pub async fn reset_state(&mut self) { self.state = SubscriptionStateChange::Reset; - self.subscribed.clear(); + self.subscribed.write().await.clear(); } /// Add a new subscription - pub fn insert_subscription(&mut self, subscription: String) { - if self.subscribed.contains(&subscription) { + pub async fn insert_subscription(&mut self, subscription: String) { + let mut subscribed = self.subscribed.write().await; + if subscribed.contains(&subscription) { return; } @@ -139,12 +140,13 @@ impl State { SubscriptionStateChange::Reset => {} } - self.subscribed.insert(subscription); + subscribed.insert(subscription); } /// Remove existing subscription - pub fn remove_subscription(&mut self, subscription: &str) { - if !self.subscribed.contains(&subscription.to_string()) { + pub async fn remove_subscription(&mut self, subscription: &str) { + let mut subscribed = self.subscribed.write().await; + if !subscribed.contains(&subscription.to_string()) { return; } @@ -161,6 +163,6 @@ impl State { SubscriptionStateChange::Reset => panic!("Should not remove during a reset!"), } - self.subscribed.remove(subscription); + subscribed.remove(subscription); } } diff --git a/crates/bonfire/src/websocket.rs b/crates/bonfire/src/websocket.rs index 0bfa666b..6745cfe9 100644 --- a/crates/bonfire/src/websocket.rs +++ b/crates/bonfire/src/websocket.rs @@ -1,4 +1,4 @@ -use std::net::SocketAddr; +use std::{collections::HashSet, net::SocketAddr, sync::Arc}; use async_tungstenite::WebSocketStream; use authifier::AuthifierEvent; @@ -19,7 +19,10 @@ use revolt_database::{ }; use revolt_presence::{create_session, delete_session}; -use async_std::{net::TcpStream, sync::Mutex}; +use async_std::{ + net::TcpStream, + sync::{Mutex, RwLock}, +}; use revolt_result::create_error; use crate::config::{ProtocolConfiguration, WebsocketHandshakeCallback}; @@ -44,6 +47,7 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr) else { return; }; + // Verify we've received a valid config, otherwise we should just drop the connection. let Ok(mut config) = receiver.await else { return; @@ -102,6 +106,7 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr) let Ok(ready_payload) = state.generate_ready_payload(db).await else { return; }; + if write.send(config.encode(&ready_payload)).await.is_err() { return; } @@ -116,10 +121,12 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr) { let write = Mutex::new(write); + let subscribed = state.subscribed.clone(); + // Create a PubSub connection to poll on. let listener = listener(db, &mut state, addr, &config, &write).fuse(); // Read from WebSocket stream. - let worker = worker(addr, user_id.clone(), &config, read, &write).fuse(); + let worker = worker(addr, subscribed, user_id.clone(), &config, read, &write).fuse(); // Pin both tasks. pin_mut!(listener, worker); @@ -157,10 +164,12 @@ async fn listener( let mut message_rx = subscriber.message_rx(); loop { // Check for state changes for subscriptions. - match state.apply_state() { + match state.apply_state().await { SubscriptionStateChange::Reset => { subscriber.unsubscribe_all().await.unwrap(); - for id in state.iter_subscriptions() { + + let subscribed = state.subscribed.read().await; + for id in subscribed.iter() { subscriber.subscribe(id).await.unwrap(); } @@ -255,6 +264,7 @@ async fn listener( async fn worker( addr: SocketAddr, + subscribed: Arc>>, user_id: String, config: &ProtocolConfiguration, mut read: WsReader, @@ -277,8 +287,13 @@ async fn worker( let Ok(payload) = config.decode(&msg) else { continue; }; + match payload { ClientMessage::BeginTyping { channel } => { + if !subscribed.read().await.contains(&channel) { + break; + } + EventV1::ChannelStartTyping { id: channel.clone(), user: user_id.clone(), @@ -287,6 +302,10 @@ async fn worker( .await; } ClientMessage::EndTyping { channel } => { + if !subscribed.read().await.contains(&channel) { + break; + } + EventV1::ChannelStopTyping { id: channel.clone(), user: user_id.clone(),