fix: don't allow sending typing notification to unknown channel

closes #177
This commit is contained in:
Paul Makles
2024-06-11 12:51:52 +01:00
parent 2cb20618da
commit 96fb0eecca
2 changed files with 46 additions and 25 deletions

View File

@@ -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 lru::LruCache;
use revolt_database::{Channel, Member, Server, User}; use revolt_database::{Channel, Member, Server, User};
@@ -60,8 +64,8 @@ pub struct State {
pub session_id: String, pub session_id: String,
pub private_topic: String, pub private_topic: String,
subscribed: HashSet<String>, pub state: SubscriptionStateChange,
state: SubscriptionStateChange, pub subscribed: Arc<RwLock<HashSet<String>>>,
} }
impl State { impl State {
@@ -81,7 +85,7 @@ impl State {
State { State {
cache, cache,
subscribed, subscribed: Arc::new(RwLock::new(subscribed)),
session_id, session_id,
private_topic, private_topic,
state: SubscriptionStateChange::Reset, state: SubscriptionStateChange::Reset,
@@ -89,15 +93,16 @@ impl State {
} }
/// Apply currently queued 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 state = std::mem::replace(&mut self.state, SubscriptionStateChange::None);
let mut subscribed = self.subscribed.write().await;
if let SubscriptionStateChange::Change { add, remove } = &state { if let SubscriptionStateChange::Change { add, remove } = &state {
for id in add { for id in add {
self.subscribed.insert(id.clone()); subscribed.insert(id.clone());
} }
for id in remove { 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() 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 /// Reset the current state
pub fn reset_state(&mut self) { pub async fn reset_state(&mut self) {
self.state = SubscriptionStateChange::Reset; self.state = SubscriptionStateChange::Reset;
self.subscribed.clear(); self.subscribed.write().await.clear();
} }
/// Add a new subscription /// Add a new subscription
pub fn insert_subscription(&mut self, subscription: String) { pub async fn insert_subscription(&mut self, subscription: String) {
if self.subscribed.contains(&subscription) { let mut subscribed = self.subscribed.write().await;
if subscribed.contains(&subscription) {
return; return;
} }
@@ -139,12 +140,13 @@ impl State {
SubscriptionStateChange::Reset => {} SubscriptionStateChange::Reset => {}
} }
self.subscribed.insert(subscription); subscribed.insert(subscription);
} }
/// Remove existing subscription /// Remove existing subscription
pub fn remove_subscription(&mut self, subscription: &str) { pub async fn remove_subscription(&mut self, subscription: &str) {
if !self.subscribed.contains(&subscription.to_string()) { let mut subscribed = self.subscribed.write().await;
if !subscribed.contains(&subscription.to_string()) {
return; return;
} }
@@ -161,6 +163,6 @@ impl State {
SubscriptionStateChange::Reset => panic!("Should not remove during a reset!"), SubscriptionStateChange::Reset => panic!("Should not remove during a reset!"),
} }
self.subscribed.remove(subscription); subscribed.remove(subscription);
} }
} }

View File

@@ -1,4 +1,4 @@
use std::net::SocketAddr; use std::{collections::HashSet, net::SocketAddr, sync::Arc};
use async_tungstenite::WebSocketStream; use async_tungstenite::WebSocketStream;
use authifier::AuthifierEvent; use authifier::AuthifierEvent;
@@ -19,7 +19,10 @@ use revolt_database::{
}; };
use revolt_presence::{create_session, delete_session}; 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 revolt_result::create_error;
use crate::config::{ProtocolConfiguration, WebsocketHandshakeCallback}; use crate::config::{ProtocolConfiguration, WebsocketHandshakeCallback};
@@ -44,6 +47,7 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr)
else { else {
return; return;
}; };
// Verify we've received a valid config, otherwise we should just drop the connection. // Verify we've received a valid config, otherwise we should just drop the connection.
let Ok(mut config) = receiver.await else { let Ok(mut config) = receiver.await else {
return; 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 { let Ok(ready_payload) = state.generate_ready_payload(db).await else {
return; return;
}; };
if write.send(config.encode(&ready_payload)).await.is_err() { if write.send(config.encode(&ready_payload)).await.is_err() {
return; return;
} }
@@ -116,10 +121,12 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr)
{ {
let write = Mutex::new(write); let write = Mutex::new(write);
let subscribed = state.subscribed.clone();
// Create a PubSub connection to poll on. // Create a PubSub connection to poll on.
let listener = listener(db, &mut state, addr, &config, &write).fuse(); let listener = listener(db, &mut state, addr, &config, &write).fuse();
// Read from WebSocket stream. // 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 both tasks.
pin_mut!(listener, worker); pin_mut!(listener, worker);
@@ -157,10 +164,12 @@ async fn listener(
let mut message_rx = subscriber.message_rx(); let mut message_rx = subscriber.message_rx();
loop { loop {
// Check for state changes for subscriptions. // Check for state changes for subscriptions.
match state.apply_state() { match state.apply_state().await {
SubscriptionStateChange::Reset => { SubscriptionStateChange::Reset => {
subscriber.unsubscribe_all().await.unwrap(); 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(); subscriber.subscribe(id).await.unwrap();
} }
@@ -255,6 +264,7 @@ async fn listener(
async fn worker( async fn worker(
addr: SocketAddr, addr: SocketAddr,
subscribed: Arc<RwLock<HashSet<String>>>,
user_id: String, user_id: String,
config: &ProtocolConfiguration, config: &ProtocolConfiguration,
mut read: WsReader, mut read: WsReader,
@@ -277,8 +287,13 @@ async fn worker(
let Ok(payload) = config.decode(&msg) else { let Ok(payload) = config.decode(&msg) else {
continue; continue;
}; };
match payload { match payload {
ClientMessage::BeginTyping { channel } => { ClientMessage::BeginTyping { channel } => {
if !subscribed.read().await.contains(&channel) {
break;
}
EventV1::ChannelStartTyping { EventV1::ChannelStartTyping {
id: channel.clone(), id: channel.clone(),
user: user_id.clone(), user: user_id.clone(),
@@ -287,6 +302,10 @@ async fn worker(
.await; .await;
} }
ClientMessage::EndTyping { channel } => { ClientMessage::EndTyping { channel } => {
if !subscribed.read().await.contains(&channel) {
break;
}
EventV1::ChannelStopTyping { EventV1::ChannelStopTyping {
id: channel.clone(), id: channel.clone(),
user: user_id.clone(), user: user_id.clone(),