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 revolt_database::{Channel, Member, Server, User};
@@ -60,8 +64,8 @@ pub struct State {
pub session_id: String,
pub private_topic: String,
subscribed: HashSet<String>,
state: SubscriptionStateChange,
pub state: SubscriptionStateChange,
pub subscribed: Arc<RwLock<HashSet<String>>>,
}
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);
}
}

View File

@@ -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<RwLock<HashSet<String>>>,
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(),