HMC/handmade-revolt-backend
6
0
Fork 0
forked from jmug/stoatchat
Code Pull Requests Actions Packages Releases 2 Activity

Merge remote-tracking branch 'origin/main' into livekit

Browse Source
This commit is contained in:
Zomatree
2025-01-26 20:36:40 +00:00
parent 2f7294705e 71e7fe3086
commit 9de87a01ab
227 changed files with 16549 additions and 2943 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
crates/bonfire/Cargo.toml
View File

@@ -1,6 +1,6 @@
[package]
name = "revolt-bonfire"
version = "0.7.1"
version = "0.8.1"
license = "AGPL-3.0-or-later"
edition = "2021"
@@ -14,6 +14,8 @@ lru = "0.7.6"
ulid = "0.5.0"
once_cell = "1.9.0"
redis-kiss = "0.1.4"
lru_time_cache = "0.11.11"
async-channel = "2.3.1"
# parsing
querystring = "1.1.0"
@@ -34,11 +36,12 @@ async-std = { version = "1.8.0", features = [
] }
# core
authifier = { version = "1.0.9" }
revolt-result = { path = "../core/result" }
revolt-models = { path = "../core/models" }
revolt-config = { path = "../core/config" }
revolt-database = { path = "../core/database" }
revolt-permissions = { version = "0.7.1", path = "../core/permissions" }
revolt-permissions = { version = "0.8.1", path = "../core/permissions" }
revolt-presence = { path = "../core/presence", features = ["redis-is-patched"] }
revolt-voice = { path = "../core/voice" }

11
crates/bonfire/Dockerfile
View File

@@ -1,11 +1,12 @@
# Build Stage
FROM ghcr.io/revoltchat/base:latest AS builder
FROM debian:12 AS debian
# Bundle Stage
FROM debian:bullseye-slim
RUN apt-get update && \
apt-get install -y ca-certificates && \
apt-get clean
FROM gcr.io/distroless/cc-debian12:nonroot
COPY --from=builder /home/rust/src/target/release/revolt-bonfire ./
EXPOSE 9000
COPY --from=debian /usr/bin/uname /usr/bin/uname
EXPOSE 14703
USER nonroot
CMD ["./revolt-bonfire"]

12
crates/bonfire/src/config.rs
View File

@@ -1,5 +1,6 @@
use async_tungstenite::tungstenite::{handshake, Message};
use futures::channel::oneshot::Sender;
use revolt_database::events::client::ReadyPayloadFields;
use revolt_result::{create_error, Result};
use serde::{Deserialize, Serialize};
@@ -83,6 +84,17 @@ impl ProtocolConfiguration {
pub fn get_protocol_format(&self) -> &ProtocolFormat {
&self.format
}
/// Get ready payload fields
pub fn get_ready_payload_fields(&self) -> Vec<ReadyPayloadFields> {
vec![
ReadyPayloadFields::Users,
ReadyPayloadFields::Servers,
ReadyPayloadFields::Channels,
ReadyPayloadFields::Members,
ReadyPayloadFields::Emoji,
]
}
}
/// Object holding one side of a channel for receiving the parsed information

180
crates/bonfire/src/events/impl.rs
View File

@@ -1,8 +1,9 @@
use std::collections::HashSet;
use revolt_database::{
events::client::EventV1, util::permissions::DatabasePermissionQuery, Channel, Database, Member,
MemberCompositeKey, Presence, RelationshipStatus,
events::client::{EventV1, ReadyPayloadFields},
util::permissions::DatabasePermissionQuery,
Channel, Database, Member, MemberCompositeKey, Presence, RelationshipStatus,
};
use revolt_models::v0;
use revolt_permissions::{calculate_channel_permissions, ChannelPermission};
@@ -23,11 +24,11 @@ impl Cache {
let server = self.servers.get(server);
let mut query =
DatabasePermissionQuery::new(db, self.users.get(&self.user_id).unwrap())
.channel(&channel);
.channel(channel);
// let mut perms = perms(self.users.get(&self.user_id).unwrap()).channel(channel);
if let Some(member) = member {
query = query.member(&member);
query = query.member(member);
}
if let Some(server) = server {
@@ -92,8 +93,13 @@ impl Cache {
/// State Manager
impl State {
/// Generate a Ready packet for the current user
pub async fn generate_ready_payload(&mut self, db: &Database) -> Result<EventV1> {
pub async fn generate_ready_payload(
&mut self,
db: &Database,
fields: Vec<ReadyPayloadFields>,
) -> Result<EventV1> {
let user = self.clone_user();
self.cache.is_bot = user.bot.is_some();
// Find all relationships to the user.
let mut user_ids: HashSet<String> = user
@@ -151,14 +157,36 @@ impl State {
.await?;
// Fetch customisations.
let emojis = db
.fetch_emoji_by_parent_ids(
&servers
.iter()
.map(|x| x.id.to_string())
.collect::<Vec<String>>(),
let emojis = if fields.contains(&ReadyPayloadFields::Emoji) {
Some(
db.fetch_emoji_by_parent_ids(
&servers
.iter()
.map(|x| x.id.to_string())
.collect::<Vec<String>>(),
)
.await?,
)
.await?;
} else {
None
};
// Fetch user settings
let user_settings = if let Some(ReadyPayloadFields::UserSettings(keys)) = fields
.iter()
.find(|e| matches!(e, ReadyPayloadFields::UserSettings(_)))
{
Some(db.fetch_user_settings(&user.id, &keys).await?)
} else {
None
};
// Fetch channel unreads
let channel_unreads = if fields.contains(&ReadyPayloadFields::ChannelUnreads) {
Some(db.fetch_unreads(&user.id).await?)
} else {
None
};
// Copy data into local state cache.
self.cache.users = users.iter().cloned().map(|x| (x.id.clone(), x)).collect();
@@ -181,40 +209,69 @@ impl State {
.collect();
// Make sure we see our own user correctly.
users.push(user.into_self().await);
users.push(user.into_self(true).await);
// Set subscription state internally.
self.reset_state();
self.insert_subscription(self.private_topic.clone());
self.reset_state().await;
self.insert_subscription(self.private_topic.clone()).await;
for user in &users {
self.insert_subscription(user.id.clone());
self.insert_subscription(user.id.clone()).await;
}
for server in &servers {
self.insert_subscription(server.id.clone());
}
self.insert_subscription(server.id.clone()).await;
for channel in &channels {
self.insert_subscription(channel.id().to_string());
}
// fetch voice states for all the channels we can see
let mut voice_states = Vec::new();
for channel in &channels {
if let Ok(Some(voice_state)) = self.fetch_voice_state(channel).await {
voice_states.push(voice_state)
if self.cache.is_bot {
self.insert_subscription(format!("{}u", server.id)).await;
}
}
for channel in &channels {
self.insert_subscription(channel.id().to_string()).await;
}
let voice_states = if fields.contains(&ReadyPayloadFields::VoiceStates) {
// fetch voice states for all the channels we can see
let mut voice_states = Vec::new();
for channel in &channels {
if let Ok(Some(voice_state)) = self.fetch_voice_state(channel).await {
voice_states.push(voice_state)
}
}
Some(voice_states)
} else {
None
};
Ok(EventV1::Ready {
users,
servers: servers.into_iter().map(Into::into).collect(),
channels: channels.into_iter().map(Into::into).collect(),
members: members.into_iter().map(Into::into).collect(),
emojis: emojis.into_iter().map(Into::into).collect(),
users: if fields.contains(&ReadyPayloadFields::Users) {
Some(users)
} else {
None
},
servers: if fields.contains(&ReadyPayloadFields::Servers) {
Some(servers.into_iter().map(Into::into).collect())
} else {
None
},
channels: if fields.contains(&ReadyPayloadFields::Channels) {
Some(channels.into_iter().map(Into::into).collect())
} else {
None
},
members: if fields.contains(&ReadyPayloadFields::Members) {
Some(members.into_iter().map(Into::into).collect())
} else {
None
},
emojis: emojis.map(|vec| vec.into_iter().map(Into::into).collect()),
voice_states,
user_settings,
channel_unreads: channel_unreads.map(|vec| vec.into_iter().map(Into::into).collect()),
})
}
@@ -248,11 +305,11 @@ impl State {
let mut bulk_events = vec![];
for id in added_channels {
self.insert_subscription(id);
self.insert_subscription(id).await;
}
for id in removed_channels {
self.remove_subscription(&id);
self.remove_subscription(&id).await;
self.cache.channels.remove(&id);
bulk_events.push(EventV1::ChannelDelete { id });
@@ -275,7 +332,7 @@ impl State {
.channels
.insert(channel.id().to_string(), channel.clone());
self.insert_subscription(channel.id().to_string());
self.insert_subscription(channel.id().to_string()).await;
bulk_events.push(EventV1::ChannelCreate(channel.into()));
}
}
@@ -348,7 +405,7 @@ impl State {
match event {
EventV1::ChannelCreate(channel) => {
let id = channel.id().to_string();
self.insert_subscription(id.clone());
self.insert_subscription(id.clone()).await;
self.cache.channels.insert(id, channel.clone().into());
}
EventV1::ChannelUpdate {
@@ -388,17 +445,17 @@ impl State {
}
}
EventV1::ChannelDelete { id } => {
self.remove_subscription(id);
self.remove_subscription(id).await;
self.cache.channels.remove(id);
}
EventV1::ChannelGroupJoin { user, .. } => {
self.insert_subscription(user.clone());
self.insert_subscription(user.clone()).await;
}
EventV1::ChannelGroupLeave { id, user, .. } => {
if user == &self.cache.user_id {
self.remove_subscription(id);
self.remove_subscription(id).await;
} else if !self.cache.can_subscribe_to_user(user) {
self.remove_subscription(user);
self.remove_subscription(user).await;
}
}
@@ -408,7 +465,12 @@ impl State {
channels,
emojis: _,
} => {
self.insert_subscription(id.clone());
self.insert_subscription(id.clone()).await;
if self.cache.is_bot {
self.insert_subscription(format!("{}u", id)).await;
}
self.cache.servers.insert(id.clone(), server.clone().into());
let member = Member {
id: MemberCompositeKey {
@@ -445,13 +507,13 @@ impl State {
EventV1::ServerMemberJoin { .. } => {
// We will always receive ServerCreate when joining a new server.
}
EventV1::ServerMemberLeave { id, user } => {
EventV1::ServerMemberLeave { id, user, .. } => {
if user == &self.cache.user_id {
self.remove_subscription(id);
self.remove_subscription(id).await;
if let Some(server) = self.cache.servers.remove(id) {
for channel in &server.channels {
self.remove_subscription(channel);
self.remove_subscription(channel).await;
self.cache.channels.remove(channel);
}
}
@@ -459,11 +521,11 @@ impl State {
}
}
EventV1::ServerDelete { id } => {
self.remove_subscription(id);
self.remove_subscription(id).await;
if let Some(server) = self.cache.servers.remove(id) {
for channel in &server.channels {
self.remove_subscription(channel);
self.remove_subscription(channel).await;
self.cache.channels.remove(channel);
}
}
@@ -536,9 +598,23 @@ impl State {
self.cache.users.insert(id.clone(), user.clone().into());
if self.cache.can_subscribe_to_user(id) {
self.insert_subscription(id.clone());
self.insert_subscription(id.clone()).await;
} else {
self.remove_subscription(id);
self.remove_subscription(id).await;
}
}
EventV1::Message(message) => {
// Since Message events are fanned out to many clients,
// we must reconstruct the relationship value at this end.
if let Some(user) = &mut message.user {
user.relationship = self
.cache
.users
.get(&self.cache.user_id)
.expect("missing self?")
.relationship_with(&message.author)
.into();
}
}
@@ -552,11 +628,11 @@ impl State {
// Sub / unsub accordingly.
if let Some(id) = queue_add {
self.insert_subscription(id);
self.insert_subscription(id).await;
}
if let Some(id) = queue_remove {
self.remove_subscription(&id);
self.remove_subscription(&id).await;
}
true
@@ -582,7 +658,7 @@ impl State {
}
Ok(Some(v0::ChannelVoiceState {
id: channel.id(),
id: channel.id().to_string(),
participants,
}))
} else {

89
crates/bonfire/src/events/state.rs
View File

@@ -1,6 +1,12 @@
use std::collections::{HashMap, HashSet};
use std::{
collections::{HashMap, HashSet},
sync::Arc,
time::Duration,
};
use async_std::sync::{Mutex, RwLock};
use lru::LruCache;
use lru_time_cache::{LruCache as LruTimeCache, TimedEntry};
use revolt_database::{Channel, Member, Server, User};
/// Enumeration representing some change in subscriptions
@@ -30,6 +36,7 @@ pub enum SubscriptionStateChange {
#[derive(Debug)]
pub struct Cache {
pub user_id: String,
pub is_bot: bool,
pub users: HashMap<String, User>,
pub channels: HashMap<String, Channel>,
@@ -43,6 +50,7 @@ impl Default for Cache {
fn default() -> Self {
Cache {
user_id: Default::default(),
is_bot: false,
users: Default::default(),
channels: Default::default(),
@@ -58,14 +66,17 @@ impl Default for Cache {
pub struct State {
pub cache: Cache,
pub session_id: String,
pub private_topic: String,
subscribed: HashSet<String>,
state: SubscriptionStateChange,
pub state: SubscriptionStateChange,
pub subscribed: Arc<RwLock<HashSet<String>>>,
pub active_servers: Arc<Mutex<LruTimeCache<String, ()>>>,
}
impl State {
/// Create state from User
pub fn from(user: User) -> State {
pub fn from(user: User, session_id: String) -> State {
let mut subscribed = HashSet::new();
let private_topic = format!("{}!", user.id);
subscribed.insert(private_topic.clone());
@@ -80,22 +91,61 @@ impl State {
State {
cache,
subscribed,
subscribed: Arc::new(RwLock::new(subscribed)),
active_servers: Arc::new(Mutex::new(LruTimeCache::with_expiry_duration_and_capacity(
Duration::from_secs(900),
5,
))),
session_id,
private_topic,
state: SubscriptionStateChange::Reset,
}
}
/// Apply currently queued state
pub fn apply_state(&mut self) -> SubscriptionStateChange {
pub async fn apply_state(&mut self) -> SubscriptionStateChange {
// Check if we need to change subscriptions to member event topics
if !self.cache.is_bot {
enum Server {
Subscribe(String),
Unsubscribe(String),
}
let active_server_changes: Vec<Server> = {
let mut active_servers = self.active_servers.lock().await;
active_servers
.notify_iter()
.map(|e| match e {
TimedEntry::Valid(k, _) => Server::Subscribe(format!("{}u", k)),
TimedEntry::Expired(k, _) => Server::Unsubscribe(format!("{}u", k)),
})
.collect()
// It is bad practice to open more than one Mutex at once and could
// lead to a deadlock, so instead we choose to collect the changes.
};
for entry in active_server_changes {
match entry {
Server::Subscribe(k) => {
self.insert_subscription(k).await;
}
Server::Unsubscribe(k) => {
self.remove_subscription(&k).await;
}
}
}
}
// Flush changes to subscriptions
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);
}
}
@@ -107,20 +157,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;
}
@@ -137,12 +183,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;
}
@@ -159,6 +206,6 @@ impl State {
SubscriptionStateChange::Reset => panic!("Should not remove during a reset!"),
}
self.subscribed.remove(subscription);
subscribed.remove(subscription);
}
}

9
crates/bonfire/src/main.rs
View File

@@ -30,13 +30,16 @@ async fn main() {
database::connect().await;
// Clean up the current region information.
clear_region(None).await;
let no_clear_region = env::var("NO_CLEAR_PRESENCE").unwrap_or_else(|_| "0".into()) == "1";
if !no_clear_region {
clear_region(None).await;
}
VOICE_CLIENT.set(Arc::new(VoiceClient::from_revolt_config().await)).unwrap();
// Setup a TCP listener to accept WebSocket connections on.
// By default, we bind to port 9000 on all interfaces.
let bind = env::var("HOST").unwrap_or_else(|_| "0.0.0.0:9000".into());
// By default, we bind to port 14703 on all interfaces.
let bind = env::var("HOST").unwrap_or_else(|_| "0.0.0.0:14703".into());
info!("Listening on host {bind}");
let try_socket = TcpListener::bind(bind).await;
let listener = try_socket.expect("Failed to bind");

462
crates/bonfire/src/websocket.rs
View File

@@ -1,24 +1,33 @@
use std::net::SocketAddr;
use std::{collections::HashSet, net::SocketAddr, sync::Arc};
use async_tungstenite::WebSocketStream;
use authifier::AuthifierEvent;
use fred::{
error::{RedisError, RedisErrorKind},
interfaces::{ClientLike, EventInterface, PubsubInterface},
types::RedisConfig,
};
use futures::{
channel::oneshot,
pin_mut, select,
join, pin_mut, select,
stream::{SplitSink, SplitStream},
FutureExt, SinkExt, StreamExt, TryStreamExt,
};
use redis_kiss::{PayloadType, REDIS_PAYLOAD_TYPE, REDIS_URI};
use revolt_config::report_internal_error;
use revolt_database::{
events::{client::EventV1, server::ClientMessage},
Database, User, UserHint,
};
use revolt_presence::{create_session, delete_session};
use async_std::{net::TcpStream, sync::Mutex};
use async_std::{
net::TcpStream,
sync::{Mutex, RwLock},
task::spawn,
};
use revolt_result::create_error;
use sentry::Level;
use crate::config::{ProtocolConfiguration, WebsocketHandshakeCallback};
use crate::events::state::{State, SubscriptionStateChange};
@@ -42,10 +51,12 @@ 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;
};
info!(
"User {addr:?} provided protocol configuration (version = {}, format = {:?})",
config.get_protocol_version(),
@@ -57,10 +68,8 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr)
// If the user has not provided authentication, request information.
if config.get_session_token().is_none() {
while let Ok(message) = read.try_next().await {
if let Ok(ClientMessage::Authenticate { token }) =
config.decode(message.as_ref().unwrap())
{
while let Ok(Some(message)) = read.try_next().await {
if let Ok(ClientMessage::Authenticate { token }) = config.decode(&message) {
config.set_session_token(token);
break;
}
@@ -69,35 +78,48 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr)
// Try to authenticate the user.
let Some(token) = config.get_session_token().as_ref() else {
write
.send(config.encode(&EventV1::Error {
data: create_error!(InvalidSession),
}))
.await
.ok();
return;
};
let user = match User::from_token(db, token, UserHint::Any).await {
let (user, session_id) = match User::from_token(db, token, UserHint::Any).await {
Ok(user) => user,
Err(err) => {
write.send(config.encode(&err)).await.ok();
write
.send(config.encode(&EventV1::Error { data: err }))
.await
.ok();
return;
}
};
info!("User {addr:?} authenticated as @{}", user.username);
// Create local state.
let mut state = State::from(user);
let mut state = State::from(user, session_id);
let user_id = state.cache.user_id.clone();
// Notify socket we have authenticated.
if write
.send(config.encode(&EventV1::Authenticated))
.await
.is_err()
{
if report_internal_error!(write.send(config.encode(&EventV1::Authenticated)).await).is_err() {
return;
}
// Download required data to local cache and send Ready payload.
let Ok(ready_payload) = state.generate_ready_payload(db).await else {
return;
let ready_payload = match report_internal_error!(
state
.generate_ready_payload(db, config.get_ready_payload_fields())
.await
) {
Ok(ready_payload) => ready_payload,
Err(_) => return,
};
if write.send(config.encode(&ready_payload)).await.is_err() {
if report_internal_error!(write.send(config.encode(&ready_payload)).await).is_err() {
return;
}
@@ -110,20 +132,44 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr)
}
{
// Setup channels and mutexes
let write = Mutex::new(write);
let subscribed = state.subscribed.clone();
let active_servers = state.active_servers.clone();
let (topic_signal_s, topic_signal_r) = async_channel::unbounded();
// TODO: this needs to be rewritten
// Create channels through which the tasks can signal to each other they need to clean up
let (kill_signal_1_s, kill_signal_1_r) = async_channel::bounded(1);
let (kill_signal_2_s, kill_signal_2_r) = async_channel::bounded(1);
// 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();
// Pin both tasks.
pin_mut!(listener, worker);
// Wait for either disconnect or for listener to die.
select!(
() = listener => {},
() = worker => {}
let listener = listener_with_kill_signal(
db,
&mut state,
addr,
&config,
topic_signal_r,
kill_signal_1_r,
&write,
kill_signal_2_s,
);
// Read from WebSocket stream.
let worker = worker_with_kill_signal(
addr,
subscribed,
active_servers,
user_id.clone(),
&config,
topic_signal_s,
kill_signal_2_r,
read,
&write,
kill_signal_1_s,
);
join!(listener, worker);
}
// Clean up presence session.
let last_session = delete_session(&user_id, session_id).await;
@@ -134,29 +180,79 @@ pub async fn client(db: &'static Database, stream: TcpStream, addr: SocketAddr)
}
}
#[allow(clippy::too_many_arguments)]
async fn listener_with_kill_signal(
db: &'static Database,
state: &mut State,
addr: SocketAddr,
config: &ProtocolConfiguration,
topic_signal_r: async_channel::Receiver<()>,
kill_signal_r: async_channel::Receiver<()>,
write: &Mutex<WsWriter>,
kill_signal_s: async_channel::Sender<()>,
) {
listener(
db,
state,
addr,
config,
topic_signal_r,
kill_signal_r,
write,
)
.await;
kill_signal_s.send(()).await.ok();
}
async fn listener(
db: &'static Database,
state: &mut State,
addr: SocketAddr,
config: &ProtocolConfiguration,
topic_signal_r: async_channel::Receiver<()>,
kill_signal_r: async_channel::Receiver<()>,
write: &Mutex<WsWriter>,
) {
let redis_config = RedisConfig::from_url(&REDIS_URI).unwrap();
let Ok(subscriber) = fred::types::Builder::from_config(redis_config).build_subscriber_client()
else {
return;
let subscriber = match report_internal_error!(
fred::types::Builder::from_config(redis_config).build_subscriber_client()
) {
Ok(subscriber) => subscriber,
Err(_) => return,
};
if subscriber.init().await.is_err() {
if report_internal_error!(subscriber.init().await).is_err() {
return;
};
}
// Handle Redis connection dropping
let (clean_up_s, clean_up_r) = async_channel::bounded(1);
let clean_up_s = Arc::new(Mutex::new(clean_up_s));
subscriber.on_error(move |err| {
if let RedisErrorKind::Canceled = err.kind() {
let clean_up_s = clean_up_s.clone();
spawn(async move {
clean_up_s.lock().await.send(()).await.ok();
});
}
Ok(())
});
let mut message_rx = subscriber.message_rx();
loop {
'out: 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() {
subscriber.subscribe(id).await.unwrap();
if report_internal_error!(subscriber.unsubscribe_all().await).is_err() {
break 'out;
}
let subscribed = state.subscribed.read().await;
for id in subscribed.iter() {
if report_internal_error!(subscriber.subscribe(id).await).is_err() {
break 'out;
}
}
#[cfg(debug_assertions)]
@@ -167,112 +263,242 @@ async fn listener(
#[cfg(debug_assertions)]
info!("{addr:?} unsubscribing from {id}");
subscriber.unsubscribe(id).await.unwrap();
if report_internal_error!(subscriber.unsubscribe(id).await).is_err() {
break 'out;
}
}
for id in add {
#[cfg(debug_assertions)]
info!("{addr:?} subscribing to {id}");
subscriber.subscribe(id).await.unwrap();
if report_internal_error!(subscriber.subscribe(id).await).is_err() {
break 'out;
}
}
}
SubscriptionStateChange::None => {}
}
// Handle incoming events.
let Ok(message) = message_rx.recv().await.map_err(|e| {
warn!("Error while consuming pub/sub messages: {e:?}");
sentry::capture_error(&e);
}) else {
return;
};
let event = match *REDIS_PAYLOAD_TYPE {
PayloadType::Json => message
.value
.as_str()
.and_then(|s| serde_json::from_str::<EventV1>(s.as_ref()).ok()),
PayloadType::Msgpack => message
.value
.as_bytes()
.and_then(|b| rmp_serde::from_slice::<EventV1>(b).ok()),
PayloadType::Bincode => message
.value
.as_bytes()
.and_then(|b| bincode::deserialize::<EventV1>(b).ok()),
};
let Some(mut event) = event else {
warn!("Failed to deserialise an event for {}!", message.channel);
return;
};
let should_send = state.handle_incoming_event_v1(db, &mut event).await;
if !should_send {
continue;
}
let t1 = message_rx.recv().fuse();
let t2 = topic_signal_r.recv().fuse();
let t3 = kill_signal_r.recv().fuse();
let t4 = clean_up_r.recv().fuse();
let result = write.lock().await.send(config.encode(&event)).await;
if let Err(e) = result {
use async_tungstenite::tungstenite::Error;
if !matches!(e, Error::AlreadyClosed | Error::ConnectionClosed) {
warn!("Error while sending an event to {addr:?}: {e:?}");
pin_mut!(t1, t2, t3, t4);
select! {
_ = t4 => {
break 'out;
},
_ = t3 => {
break 'out;
},
_ = t2 => {},
message = t1 => {
// Handle incoming events.
let message = match report_internal_error!(message) {
Ok(message) => message,
Err(_) => break 'out
};
let event = match *REDIS_PAYLOAD_TYPE {
PayloadType::Json => message
.value
.as_str()
.and_then(|s| report_internal_error!(serde_json::from_str::<EventV1>(s.as_ref())).ok()),
PayloadType::Msgpack => message
.value
.as_bytes()
.and_then(|b| report_internal_error!(rmp_serde::from_slice::<EventV1>(b)).ok()),
PayloadType::Bincode => message
.value
.as_bytes()
.and_then(|b| report_internal_error!(bincode::deserialize::<EventV1>(b)).ok()),
};
let Some(mut event) = event else {
let err = format!(
"Failed to deserialise event for {}: `{:?}`",
message.channel,
message
.value
);
error!("{}", err);
sentry::capture_message(&err, Level::Error);
break 'out;
};
if let EventV1::Auth(auth) = &event {
if let AuthifierEvent::DeleteSession { session_id, .. } = auth {
if &state.session_id == session_id {
event = EventV1::Logout;
}
} else if let AuthifierEvent::DeleteAllSessions {
exclude_session_id, ..
} = auth
{
if let Some(excluded) = exclude_session_id {
if &state.session_id != excluded {
event = EventV1::Logout;
}
} else {
event = EventV1::Logout;
}
}
} else {
let should_send = state.handle_incoming_event_v1(db, &mut event).await;
if !should_send {
continue;
}
}
let result = write.lock().await.send(config.encode(&event)).await;
if let Err(e) = result {
use async_tungstenite::tungstenite::Error;
if !matches!(e, Error::AlreadyClosed | Error::ConnectionClosed) {
let err = format!("Error while sending an event to {addr:?}: {e:?}");
warn!("{}", err);
sentry::capture_message(&err, Level::Warning);
}
break 'out;
}
if let EventV1::Logout = event {
info!("User {addr:?} received log out event!");
break 'out;
}
}
return;
}
}
report_internal_error!(subscriber.quit().await).ok();
}
async fn worker(
#[allow(clippy::too_many_arguments)]
async fn worker_with_kill_signal(
addr: SocketAddr,
subscribed: Arc<RwLock<HashSet<String>>>,
active_servers: Arc<Mutex<lru_time_cache::LruCache<String, ()>>>,
user_id: String,
config: &ProtocolConfiguration,
topic_signal_s: async_channel::Sender<()>,
kill_signal_r: async_channel::Receiver<()>,
read: WsReader,
write: &Mutex<WsWriter>,
kill_signal_s: async_channel::Sender<()>,
) {
worker(
addr,
subscribed,
active_servers,
user_id,
config,
topic_signal_s,
kill_signal_r,
read,
write,
)
.await;
kill_signal_s.send(()).await.ok();
}
#[allow(clippy::too_many_arguments)]
async fn worker(
addr: SocketAddr,
subscribed: Arc<RwLock<HashSet<String>>>,
active_servers: Arc<Mutex<lru_time_cache::LruCache<String, ()>>>,
user_id: String,
config: &ProtocolConfiguration,
topic_signal_s: async_channel::Sender<()>,
kill_signal_r: async_channel::Receiver<()>,
mut read: WsReader,
write: &Mutex<WsWriter>,
) {
loop {
let result = read.try_next().await;
let msg = match result {
Ok(Some(msg)) => msg,
Ok(None) => return,
Err(e) => {
use async_tungstenite::tungstenite::Error;
if !matches!(e, Error::AlreadyClosed | Error::ConnectionClosed) {
warn!("Error while reading an event from {addr:?}: {e:?}");
}
return;
}
};
let t1 = read.try_next().fuse();
let t2 = kill_signal_r.recv().fuse();
let Ok(payload) = config.decode(&msg) else {
continue;
};
match payload {
ClientMessage::BeginTyping { channel } => {
EventV1::ChannelStartTyping {
id: channel.clone(),
user: user_id.clone(),
}
.p(channel.clone())
.await;
}
ClientMessage::EndTyping { channel } => {
EventV1::ChannelStopTyping {
id: channel.clone(),
user: user_id.clone(),
}
.p(channel.clone())
.await;
}
ClientMessage::Ping { data, responded } => {
if responded.is_none() {
write
.lock()
.await
.send(config.encode(&EventV1::Pong { data }))
.await
.ok();
pin_mut!(t1, t2);
select! {
_ = t2 => {
return;
},
result = t1 => {
let msg = match result {
Ok(Some(msg)) => msg,
Ok(None) => {
warn!("Received a None message!");
sentry::capture_message("Received a None message!", Level::Warning);
return;
}
Err(e) => {
use async_tungstenite::tungstenite::Error;
if !matches!(e, Error::AlreadyClosed | Error::ConnectionClosed) {
let err = format!("Error while reading an event from {addr:?}: {e:?}");
warn!("{}", err);
sentry::capture_message(&err, Level::Warning);
}
return;
}
};
let Ok(payload) = config.decode(&msg) else {
continue;
};
match payload {
ClientMessage::BeginTyping { channel } => {
if !subscribed.read().await.contains(&channel) {
continue;
}
EventV1::ChannelStartTyping {
id: channel.clone(),
user: user_id.clone(),
}
.p(channel.clone())
.await;
}
ClientMessage::EndTyping { channel } => {
if !subscribed.read().await.contains(&channel) {
continue;
}
EventV1::ChannelStopTyping {
id: channel.clone(),
user: user_id.clone(),
}
.p(channel.clone())
.await;
}
ClientMessage::Subscribe { server_id } => {
let mut servers = active_servers.lock().await;
let has_item = servers.contains_key(&server_id);
servers.insert(server_id, ());
if !has_item {
// Poke the listener to adjust subscriptions
topic_signal_s.send(()).await.ok();
}
}
ClientMessage::Ping { data, responded } => {
if responded.is_none() {
write
.lock()
.await
.send(config.encode(&EventV1::Pong { data }))
.await
.ok();
}
}
_ => {}
}
}
_ => {}
}
}
}