From c5844b6281586e9e025e12e2fd9fa113a94bb475 Mon Sep 17 00:00:00 2001 From: Zomatree Date: Thu, 5 Mar 2026 04:24:54 +0000 Subject: [PATCH] fix/hosepipe --- Cargo.lock | 2 + crates/bonfire/Cargo.toml | 4 ++ crates/bonfire/src/main.rs | 88 +++++++++++++++++++++++ crates/bonfire/src/websocket.rs | 44 ++++++------ crates/core/database/src/amqp/amqp.rs | 2 + crates/core/database/src/events/client.rs | 31 ++++++-- crates/core/database/src/util/mod.rs | 1 + crates/core/database/src/util/rabbit.rs | 48 +++++++++++++ crates/delta/src/main.rs | 50 ++++--------- crates/delta/src/util/test.rs | 12 +++- 10 files changed, 215 insertions(+), 67 deletions(-) create mode 100644 crates/core/database/src/util/rabbit.rs diff --git a/Cargo.lock b/Cargo.lock index ac4ed100..5659cae3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6838,8 +6838,10 @@ dependencies = [ name = "revolt-bonfire" version = "0.11.5" dependencies = [ + "amqprs", "async-channel 2.5.0", "async-std", + "async-trait", "async-tungstenite", "authifier", "bincode", diff --git a/crates/bonfire/Cargo.toml b/crates/bonfire/Cargo.toml index 93ac8d4b..3f529d82 100644 --- a/crates/bonfire/Cargo.toml +++ b/crates/bonfire/Cargo.toml @@ -36,6 +36,7 @@ async-std = { version = "1.8.0", features = [ "tokio02", "attributes", ] } +async-trait = "0.1.89" # core authifier = { version = "1.0.16" } @@ -48,3 +49,6 @@ revolt-presence = { path = "../core/presence", features = ["redis-is-patched"] } # redis fred = { version = "8.0.1", features = ["subscriber-client"] } + +# Redis +amqprs = { version = "1.7.0" } \ No newline at end of file diff --git a/crates/bonfire/src/main.rs b/crates/bonfire/src/main.rs index f4e8a387..b3400682 100644 --- a/crates/bonfire/src/main.rs +++ b/crates/bonfire/src/main.rs @@ -1,6 +1,18 @@ use std::env; +use amqprs::{ + channel::{ + BasicConsumeArguments, Channel, ExchangeDeclareArguments, QueueBindArguments, + QueueDeclareArguments, + }, + connection::{Connection, OpenConnectionArguments}, + consumer::AsyncConsumer, + BasicProperties, Deliver, +}; use async_std::net::TcpListener; +use async_trait::async_trait; +use redis_kiss::AsyncCommands; +use revolt_database::util::rabbit::set_rabbitmq_connection; use revolt_presence::clear_region; #[macro_use] @@ -31,6 +43,53 @@ async fn main() { let try_socket = TcpListener::bind(bind).await; let listener = try_socket.expect("Failed to bind"); + let config = revolt_config::config().await; + + let rmq_conn = Connection::open(&OpenConnectionArguments::new( + &config.rabbit.host, + config.rabbit.port, + &config.rabbit.username, + &config.rabbit.password, + )) + .await + .expect("Failed to connect to RabbitMQ"); + + set_rabbitmq_connection(rmq_conn.clone()); + + let channel = rmq_conn + .open_channel(None) + .await + .expect("Failed to open RabbitMQ channel."); + + channel + .exchange_declare( + ExchangeDeclareArguments::new("events", "fanout") + .durable(true) + .finish(), + ) + .await + .expect("wires"); + + channel + .queue_declare(QueueDeclareArguments::new("events").durable(true).finish()) + .await + .expect("wires 2"); + + channel + .queue_bind(QueueBindArguments::new("events", "events", "events")) + .await + .expect("wires 3"); + + channel + .basic_consume( + RabbitToRedisConsumer, + BasicConsumeArguments::new("events", "") + .manual_ack(false) + .finish(), + ) + .await + .expect("wires 4"); + // Start accepting new connections and spawn a client for each connection. while let Ok((stream, addr)) = listener.accept().await { async_std::task::spawn(async move { @@ -40,3 +99,32 @@ async fn main() { }); } } + +struct RabbitToRedisConsumer; + +#[async_trait] +impl AsyncConsumer for RabbitToRedisConsumer { + async fn consume( + &mut self, + _channel: &Channel, + _deliver: Deliver, + basic_properties: BasicProperties, + content: Vec, + ) { + let mut redis_conn = redis_kiss::get_connection() + .await + .expect("Failed to connect to Redis."); + + let pubsub_channel = basic_properties + .headers() + .expect("No headers") + .get(&"c".try_into().unwrap()) + .expect("No channel header") + .to_string(); + + redis_conn + .publish(pubsub_channel, content) + .await + .expect("failed to publish") + } +} diff --git a/crates/bonfire/src/websocket.rs b/crates/bonfire/src/websocket.rs index 4a8cfe4c..9ec562c3 100644 --- a/crates/bonfire/src/websocket.rs +++ b/crates/bonfire/src/websocket.rs @@ -464,30 +464,30 @@ async fn worker( match payload { // disabled events - // ClientMessage::BeginTyping { channel } => { - // if !subscribed.read().await.contains(&channel) { - // continue; - // } + 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::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; - // } + 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); diff --git a/crates/core/database/src/amqp/amqp.rs b/crates/core/database/src/amqp/amqp.rs index b29105ed..6a643683 100644 --- a/crates/core/database/src/amqp/amqp.rs +++ b/crates/core/database/src/amqp/amqp.rs @@ -11,6 +11,8 @@ use revolt_presence::filter_online; use serde_json::to_string; +// TODO: move away from storing Connection and Channel to using thread local singletons - #659 + #[derive(Clone)] pub struct AMQP { #[allow(unused)] diff --git a/crates/core/database/src/events/client.rs b/crates/core/database/src/events/client.rs index 4615273d..ff1e2529 100644 --- a/crates/core/database/src/events/client.rs +++ b/crates/core/database/src/events/client.rs @@ -1,12 +1,20 @@ +use amqprs::{ + channel::{BasicPublishArguments}, + BasicProperties, FieldTable, +}; use authifier::AuthifierEvent; use revolt_result::Error; use serde::{Deserialize, Serialize}; use revolt_models::v0::{ - AppendMessage, Channel, ChannelUnread, ChannelVoiceState, Emoji, FieldsChannel, FieldsMember, FieldsMessage, FieldsRole, FieldsServer, FieldsUser, FieldsWebhook, Member, MemberCompositeKey, Message, PartialChannel, PartialMember, PartialMessage, PartialRole, PartialServer, PartialUser, PartialUserVoiceState, PartialWebhook, PolicyChange, RemovalIntention, Report, Server, User, UserSettings, UserVoiceState, Webhook + AppendMessage, Channel, ChannelUnread, ChannelVoiceState, Emoji, FieldsChannel, FieldsMember, + FieldsMessage, FieldsRole, FieldsServer, FieldsUser, FieldsWebhook, Member, MemberCompositeKey, + Message, PartialChannel, PartialMember, PartialMessage, PartialRole, PartialServer, + PartialUser, PartialUserVoiceState, PartialWebhook, PolicyChange, RemovalIntention, Report, + Server, User, UserSettings, UserVoiceState, Webhook, }; -use crate::Database; +use crate::{util::rabbit::get_channel, Database}; /// Ping Packet #[derive(Serialize, Deserialize, Debug, Clone)] @@ -304,14 +312,23 @@ pub enum EventV1 { impl EventV1 { /// Publish helper wrapper pub async fn p(self, channel: String) { - #[cfg(not(debug_assertions))] - redis_kiss::p(channel, self).await; - #[cfg(debug_assertions)] info!("Publishing event to {channel}: {self:?}"); - #[cfg(debug_assertions)] - redis_kiss::publish(channel, self).await.unwrap(); + let rmq = get_channel().await; + + let mut headers = FieldTable::new(); + headers.insert("c".try_into().unwrap(), channel.clone().into()); + let mut properties = BasicProperties::default(); + properties.with_headers(headers); + + rmq.basic_publish( + properties, + serde_json::to_string(&self).unwrap().into_bytes(), + BasicPublishArguments::new("events", "events"), + ) + .await + .unwrap(); } /// Publish user event diff --git a/crates/core/database/src/util/mod.rs b/crates/core/database/src/util/mod.rs index 1a03a176..87b1ca70 100644 --- a/crates/core/database/src/util/mod.rs +++ b/crates/core/database/src/util/mod.rs @@ -5,5 +5,6 @@ pub mod idempotency; pub mod permissions; pub mod reference; pub mod test_fixtures; +pub mod rabbit; pub use funcs::*; diff --git a/crates/core/database/src/util/rabbit.rs b/crates/core/database/src/util/rabbit.rs new file mode 100644 index 00000000..0119df05 --- /dev/null +++ b/crates/core/database/src/util/rabbit.rs @@ -0,0 +1,48 @@ +use amqprs::{channel::Channel, connection::Connection}; +use once_cell::sync::OnceCell; +use std::{ + collections::HashMap, future::ready, sync::{LazyLock, RwLock}, thread::{ThreadId, current} +}; + +static RABBIT_CONNECTION: OnceCell = OnceCell::new(); +static RABBIT_CHANNELS: LazyLock>> = + LazyLock::new(|| RwLock::new(HashMap::new())); + +pub async fn get_channel_with_init Channel>(init: F) -> Channel { + let conn = RABBIT_CONNECTION + .get() + .expect("Rabbit connection is not initialised."); + + let thread_id = current().id(); + + let channel = RABBIT_CHANNELS + .read() + .expect("Channels poisioned") + .get(&thread_id) + .cloned(); + + if let Some(channel) = channel { + channel + } else { + let mut channel = + conn.open_channel(None) + .await + .expect("Failed to open rabbitmq channel"); + + channel = init(channel).await; + + RABBIT_CHANNELS + .write() + .expect("Channels poisioned") + .insert(thread_id, channel.clone()); + channel + } +} + +pub async fn get_channel() -> Channel { + get_channel_with_init(ready).await +} + +pub fn set_rabbitmq_connection(connection: Connection) -> bool { + RABBIT_CONNECTION.set(connection).is_ok() +} diff --git a/crates/delta/src/main.rs b/crates/delta/src/main.rs index 882609cc..b7be572a 100644 --- a/crates/delta/src/main.rs +++ b/crates/delta/src/main.rs @@ -9,8 +9,7 @@ pub mod routes; pub mod util; use revolt_config::config; -use revolt_database::events::client::EventV1; -use revolt_database::AMQP; +use revolt_database::{AMQP, util::rabbit::{get_channel, get_channel_with_init, set_rabbitmq_connection}}; use revolt_ratelimits::rocket as ratelimiter; use rocket::{Build, Rocket}; use rocket_cors::{AllowedOrigins, CorsOptions}; @@ -19,11 +18,9 @@ use std::net::Ipv4Addr; use std::str::FromStr; use amqprs::{ - channel::ExchangeDeclareArguments, + channel::{Channel, ExchangeDeclareArguments}, connection::{Connection, OpenConnectionArguments}, }; -use async_std::channel::unbounded; -use authifier::AuthifierEvent; use rocket::data::ToByteUnit; use revolt_database::voice::VoiceClient; @@ -39,28 +36,9 @@ pub async fn web() -> Rocket { log::info!("database_here {db:?}"); db.migrate_database().await.unwrap(); - // Setup Authifier event channel - let (_, receiver) = unbounded(); - // Setup Authifier let authifier = db.clone().to_authifier().await; - // Launch a listener for Authifier events - async_std::task::spawn(async move { - while let Ok(event) = receiver.recv().await { - match &event { - AuthifierEvent::CreateSession { .. } | AuthifierEvent::CreateAccount { .. } => { - EventV1::Auth(event).global().await - } - AuthifierEvent::DeleteSession { user_id, .. } - | AuthifierEvent::DeleteAllSessions { user_id, .. } => { - let id = user_id.to_string(); - EventV1::Auth(event).private(id).await - } - } - } - }); - // Configure CORS let cors = CorsOptions { allowed_origins: AllowedOrigins::All, @@ -121,19 +99,19 @@ pub async fn web() -> Rocket { .await .expect("Failed to connect to RabbitMQ"); - let channel = connection - .open_channel(None) - .await - .expect("Failed to open RabbitMQ channel"); + set_rabbitmq_connection(connection.clone()); + let channel = get_channel_with_init(|channel: Channel| async { + channel + .exchange_declare( + ExchangeDeclareArguments::new(&config.pushd.exchange, "direct") + .durable(true) + .finish(), + ) + .await + .expect("Failed to declare exchange"); - channel - .exchange_declare( - ExchangeDeclareArguments::new(&config.pushd.exchange, "direct") - .durable(true) - .finish(), - ) - .await - .expect("Failed to declare exchange"); + channel + }).await; let amqp = AMQP::new(connection, channel); diff --git a/crates/delta/src/util/test.rs b/crates/delta/src/util/test.rs index a6251509..6d97e671 100644 --- a/crates/delta/src/util/test.rs +++ b/crates/delta/src/util/test.rs @@ -6,7 +6,9 @@ use futures::StreamExt; use rand::Rng; use redis_kiss::redis::aio::PubSub; use revolt_database::{ - events::client::EventV1, Channel, Database, Member, Message, PartialRole, Server, User, AMQP, + events::client::EventV1, + util::rabbit::{get_channel, set_rabbitmq_connection}, + Channel, Database, Member, Message, PartialRole, Server, User, AMQP, }; use revolt_database::{util::idempotency::IdempotencyKey, Role}; use revolt_models::v0; @@ -59,10 +61,16 @@ impl TestHarness { ) .await .unwrap(); - let channel = connection.open_channel(None).await.unwrap(); + + set_rabbitmq_connection(connection.clone()); + let channel = get_channel().await; let amqp = AMQP::new(connection, channel); + async_std::task::spawn(async { + + }); + TestHarness { client, authifier,