From 5b1985381ae829a92c80a19e91a414cd9dc4de93 Mon Sep 17 00:00:00 2001 From: Angelo Kontaxis Date: Mon, 18 May 2026 23:46:17 +0100 Subject: [PATCH] chore: switch to lapin (#767) * chore: begin switching to lapin fully Signed-off-by: Zomatree * chore: update rest of pushd to lapin Signed-off-by: Zomatree * chore: cleanup code Signed-off-by: Zomatree * chore: cleanup code Signed-off-by: Zomatree * fix: github webui sucks Signed-off-by: IAmTomahawkx --------- Signed-off-by: Zomatree Signed-off-by: Tom Signed-off-by: IAmTomahawkx Co-authored-by: Tom Release-As: 0.13.6 --- Cargo.lock | 41 +-- Cargo.toml | 1 - crates/core/database/Cargo.toml | 2 +- crates/core/database/src/amqp/amqp.rs | 301 ++++++++---------- crates/daemons/crond/src/tasks/acks.rs | 38 ++- crates/daemons/pushd/Cargo.toml | 2 +- .../pushd/src/consumers/inbound/ack.rs | 137 +++----- .../pushd/src/consumers/inbound/dm_call.rs | 140 +++----- .../src/consumers/inbound/fr_accepted.rs | 136 +++----- .../src/consumers/inbound/fr_received.rs | 136 +++----- .../pushd/src/consumers/inbound/generic.rs | 136 +++----- .../pushd/src/consumers/inbound/internal.rs | 53 --- .../src/consumers/inbound/mass_mention.rs | 149 +++------ .../pushd/src/consumers/inbound/message.rs | 143 +++------ .../pushd/src/consumers/inbound/mod.rs | 1 - .../pushd/src/consumers/outbound/apn.rs | 159 +++++---- .../pushd/src/consumers/outbound/fcm.rs | 111 +++---- .../pushd/src/consumers/outbound/vapid.rs | 143 ++++----- crates/daemons/pushd/src/main.rs | 251 +++++++++------ crates/daemons/pushd/src/utils/consumer.rs | 91 ++++++ crates/daemons/pushd/src/utils/mod.rs | 3 + crates/daemons/voice-ingress/Cargo.toml | 3 - crates/daemons/voice-ingress/src/main.rs | 2 +- crates/delta/Cargo.toml | 2 +- crates/delta/src/main.rs | 35 +- crates/delta/src/util/test.rs | 18 +- 26 files changed, 911 insertions(+), 1323 deletions(-) delete mode 100644 crates/daemons/pushd/src/consumers/inbound/internal.rs create mode 100644 crates/daemons/pushd/src/utils/consumer.rs diff --git a/Cargo.lock b/Cargo.lock index 9bccaccd..c65f954e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -174,31 +174,6 @@ dependencies = [ "url", ] -[[package]] -name = "amqp_serde" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f5f450f572a1ec4cdb4af7af09cbd0c7c3e1b9da2bfc7414c059a780993a8e16" -dependencies = [ - "bytes 1.11.1", - "serde", - "serde_bytes", -] - -[[package]] -name = "amqprs" -version = "1.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f1b4afcbd862e16c272b7625b6b057930b052d63c720bc90f6afab0d9abe8a8" -dependencies = [ - "amqp_serde", - "async-trait", - "bytes 1.11.1", - "serde", - "serde_bytes_ng", - "tokio 1.51.0", -] - [[package]] name = "android_system_properties" version = "0.1.5" @@ -7633,7 +7608,6 @@ dependencies = [ name = "revolt-database" version = "0.13.5" dependencies = [ - "amqprs", "async-lock 2.8.0", "async-recursion", "async-std", @@ -7648,6 +7622,7 @@ dependencies = [ "indexmap 2.13.1", "isahc", "iso8601-timestamp", + "lapin", "linkify", "livekit-api", "livekit-protocol", @@ -7684,7 +7659,6 @@ dependencies = [ name = "revolt-delta" version = "0.13.5" dependencies = [ - "amqprs", "async-channel 2.5.0", "async-std", "authifier", @@ -7694,6 +7668,7 @@ dependencies = [ "futures", "impl_ops", "iso8601-timestamp", + "lapin", "lettre", "linkify", "livekit-api", @@ -7867,7 +7842,6 @@ dependencies = [ name = "revolt-pushd" version = "0.13.5" dependencies = [ - "amqprs", "anyhow", "async-trait", "authifier", @@ -7875,6 +7849,7 @@ dependencies = [ "fcm_v1", "isahc", "iso8601-timestamp", + "lapin", "log", "pretty_env_logger", "redis-kiss", @@ -7931,7 +7906,6 @@ dependencies = [ name = "revolt-voice-ingress" version = "0.13.5" dependencies = [ - "amqprs", "async-std", "chrono", "futures", @@ -8987,15 +8961,6 @@ dependencies = [ "serde_core", ] -[[package]] -name = "serde_bytes_ng" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bdb0ebce8684e2253f964e8b6ce51f0ccc6666bbb448fb4a6788088bda6544b6" -dependencies = [ - "serde", -] - [[package]] name = "serde_core" version = "1.0.228" diff --git a/Cargo.toml b/Cargo.toml index 3b497784..3b2b92fa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -159,7 +159,6 @@ opentelemetry-appender-tracing = "0.31.1" authifier = "1.0.16" # RabbitMQ -amqprs = "1.7.0" lapin = "4.7.1" # Voice diff --git a/crates/core/database/Cargo.toml b/crates/core/database/Cargo.toml index 5031009d..dc213fd0 100644 --- a/crates/core/database/Cargo.toml +++ b/crates/core/database/Cargo.toml @@ -95,7 +95,7 @@ revolt_rocket_okapi = { workspace = true, optional = true } authifier = { workspace = true } # RabbitMQ -amqprs = { workspace = true } +lapin = { workspace = true, features = ["tokio"] } # Voice livekit-api = { workspace = true, features = ["rustls-tls-native-roots"], optional = true } diff --git a/crates/core/database/src/amqp/amqp.rs b/crates/core/database/src/amqp/amqp.rs index e0c644e5..b0e5e175 100644 --- a/crates/core/database/src/amqp/amqp.rs +++ b/crates/core/database/src/amqp/amqp.rs @@ -1,103 +1,77 @@ use std::collections::HashSet; +use std::sync::Arc; use crate::events::rabbit::*; use crate::User; -use amqprs::channel::{ - BasicPublishArguments, ExchangeDeclareArguments, ExchangeType, QueueBindArguments, - QueueDeclareArguments, +use lapin::{ + options::BasicPublishOptions, + protocol::basic::AMQPProperties, + types::{AMQPValue, FieldTable}, + Channel, Connection, ConnectionProperties, Error as AMQPError, }; -use amqprs::connection::OpenConnectionArguments; -use amqprs::{channel::Channel, connection::Connection, error::Error as AMQPError}; -use amqprs::{BasicProperties, FieldTable}; use revolt_models::v0::PushNotification; use revolt_presence::filter_online; +use revolt_result::Result; use serde_json::to_string; #[derive(Clone)] pub struct AMQP { + friend_request_accepted: Arc, + friend_request_received: Arc, + generic_message: Arc, + message_sent: Arc, + mass_mention_message_sent: Arc, + ack_notification_message: Arc, + dm_call_updated: Arc, + process_ack: Arc, #[allow(unused)] - connection: Connection, - channel: Channel, + connection: Arc, } impl AMQP { - pub fn new(connection: Connection, channel: Channel) -> AMQP { - AMQP { + pub async fn new(connection: Arc) -> Self { + Self { + friend_request_accepted: Self::create_channel(&connection).await, + friend_request_received: Self::create_channel(&connection).await, + generic_message: Self::create_channel(&connection).await, + message_sent: Self::create_channel(&connection).await, + mass_mention_message_sent: Self::create_channel(&connection).await, + ack_notification_message: Self::create_channel(&connection).await, + dm_call_updated: Self::create_channel(&connection).await, + process_ack: Self::create_channel(&connection).await, connection, - channel, } } - pub async fn new_auto() -> revolt_result::Result { + pub async fn new_auto() -> Self { let config = revolt_config::config().await; - let connection = Connection::open(&OpenConnectionArguments::new( - &config.rabbit.host, - config.rabbit.port, - &config.rabbit.username, - &config.rabbit.password, - )) - .await - .expect("Failed to connect to RabbitMQ"); - - let channel = connection - .open_channel(None) - .await - .expect("Failed to open RabbitMQ channel"); - - let mut resp = AMQP::new(connection, channel); - //resp.configure_channels().await?; - Ok(resp) - } - - pub async fn repoen_channel(&mut self) { - self.channel = self - .connection - .open_channel(None) - .await - .expect("Failed to open RabbitMQ channel"); - } - - pub async fn configure_channels(&mut self) -> revolt_result::Result<()> { - let config = revolt_config::config().await; - - if !self.channel.is_open() { - self.repoen_channel().await; - } - - self.channel - .exchange_declare( - ExchangeDeclareArguments::new( - &config.rabbit.default_exchange, - &ExchangeType::Topic.to_string(), - ) - .durable(true) - .finish(), + let connection = Arc::new( + Connection::connect( + &format!( + "amqp://{}:{}@{}:{}", + &config.rabbit.username, + &config.rabbit.password, + &config.rabbit.host, + &config.rabbit.port, + ), + ConnectionProperties::default(), ) .await - .expect("Failed to declare exchange"); + .expect("Failed to connect to RabbitMQ"), + ); - // Configure acks channel & routing - self.channel - .queue_declare( - QueueDeclareArguments::new(&config.rabbit.queues.acks) - .durable(true) - .no_wait(true) - .finish(), - ) - .await - .expect("Failed to bind queue"); + Self::new(connection).await + } - self.channel - .queue_bind(QueueBindArguments::new( - &config.rabbit.queues.acks, - &config.rabbit.default_exchange, - &config.rabbit.queues.acks, - )) - .await - .expect("Failed to bind channel"); - Ok(()) + async fn create_channel(connection: &Connection) -> Arc { + Arc::new( + connection + .create_channel() + .await + .expect("Failed to create channel"), + ) } pub async fn friend_request_accepted( @@ -117,19 +91,20 @@ impl AMQP { config.pushd.get_fr_accepted_routing_key(), payload ); - self.channel + + self.friend_request_accepted .basic_publish( - BasicProperties::default() - .with_content_type("application/json") - .with_persistence(true) - .finish(), - payload.into(), - BasicPublishArguments::new( - &config.pushd.exchange, - &config.pushd.get_fr_accepted_routing_key(), - ), + config.pushd.exchange.clone().into(), + config.pushd.get_fr_accepted_routing_key().into(), + BasicPublishOptions::default(), + payload.as_bytes(), + AMQPProperties::default() + .with_content_type("application/json".into()) + .with_delivery_mode(2), ) - .await + .await?; + + Ok(()) } pub async fn friend_request_received( @@ -150,19 +125,19 @@ impl AMQP { payload ); - self.channel + self.friend_request_received .basic_publish( - BasicProperties::default() - .with_content_type("application/json") - .with_persistence(true) - .finish(), - payload.into(), - BasicPublishArguments::new( - &config.pushd.exchange, - &config.pushd.get_fr_received_routing_key(), - ), + config.pushd.exchange.clone().into(), + config.pushd.get_fr_received_routing_key().into(), + BasicPublishOptions::default(), + payload.as_bytes(), + AMQPProperties::default() + .with_content_type("application/json".into()) + .with_delivery_mode(2), ) - .await + .await?; + + Ok(()) } pub async fn generic_message( @@ -187,19 +162,19 @@ impl AMQP { payload ); - self.channel + self.generic_message .basic_publish( - BasicProperties::default() - .with_content_type("application/json") - .with_persistence(true) - .finish(), - payload.into(), - BasicPublishArguments::new( - &config.pushd.exchange, - &config.pushd.get_generic_routing_key(), - ), + config.pushd.exchange.clone().into(), + config.pushd.get_generic_routing_key().into(), + BasicPublishOptions::default(), + payload.as_bytes(), + AMQPProperties::default() + .with_content_type("application/json".into()) + .with_delivery_mode(2), ) - .await + .await?; + + Ok(()) } pub async fn message_sent( @@ -230,19 +205,19 @@ impl AMQP { payload ); - self.channel + self.message_sent .basic_publish( - BasicProperties::default() - .with_content_type("application/json") - .with_persistence(true) - .finish(), - payload.into(), - BasicPublishArguments::new( - &config.pushd.exchange, - &config.pushd.get_message_routing_key(), - ), + config.pushd.exchange.clone().into(), + config.pushd.get_message_routing_key().into(), + BasicPublishOptions::default(), + payload.as_bytes(), + AMQPProperties::default() + .with_content_type("application/json".into()) + .with_delivery_mode(2), ) - .await + .await?; + + Ok(()) } pub async fn mass_mention_message_sent( @@ -265,16 +240,19 @@ impl AMQP { routing_key, payload ); - self.channel + self.mass_mention_message_sent .basic_publish( - BasicProperties::default() - .with_content_type("application/json") - .with_persistence(true) - .finish(), - payload.into(), - BasicPublishArguments::new(&config.pushd.exchange, routing_key.as_str()), + config.pushd.exchange.clone().into(), + routing_key.into(), + BasicPublishOptions::default(), + payload.as_bytes(), + AMQPProperties::default() + .with_content_type("application/json".into()) + .with_delivery_mode(2), ) - .await + .await?; + + Ok(()) } /// # Sends an ack to pushd to update badges on iPhones. @@ -299,23 +277,25 @@ impl AMQP { config.pushd.ack_queue, payload ); - let mut headers = FieldTable::new(); + let mut headers = FieldTable::default(); headers.insert( - "x-deduplication-header".try_into().unwrap(), - format!("{}-{}", &user_id, &channel_id).into(), + "x-deduplication-header".into(), + AMQPValue::LongString(format!("{}-{}", &user_id, &channel_id).into()), ); - self.channel + self.ack_notification_message .basic_publish( - BasicProperties::default() - .with_content_type("application/json") - .with_persistence(true) - //.with_headers(headers) - .finish(), - payload.into(), - BasicPublishArguments::new(&config.pushd.exchange, &config.pushd.ack_queue), + config.pushd.exchange.clone().into(), + config.pushd.ack_queue.into(), + BasicPublishOptions::default(), + payload.as_bytes(), + AMQPProperties::default() + .with_content_type("application/json".into()) + .with_delivery_mode(2), ) - .await + .await?; + + Ok(()) } /// # DM Call Update @@ -349,19 +329,19 @@ impl AMQP { payload ); - self.channel + self.dm_call_updated .basic_publish( - BasicProperties::default() - .with_content_type("application/json") - .with_persistence(true) - .finish(), - payload.into(), - BasicPublishArguments::new( - &config.pushd.exchange, - &config.pushd.get_dm_call_routing_key(), - ), + config.pushd.exchange.clone().into(), + config.pushd.get_dm_call_routing_key().into(), + BasicPublishOptions::default(), + payload.as_bytes(), + AMQPProperties::default() + .with_content_type("application/json".into()) + .with_delivery_mode(2), ) - .await + .await?; + + Ok(()) } /// # Send an ack to crond for processing @@ -385,19 +365,18 @@ impl AMQP { config.rabbit.default_exchange, config.rabbit.queues.acks, payload ); - self.channel + self.process_ack .basic_publish( - BasicProperties::default() - .with_content_type("application/json") - .with_persistence(true) - //.with_headers(headers) - .finish(), - payload.into(), - BasicPublishArguments::new( - &config.rabbit.default_exchange, - &config.rabbit.queues.acks, - ), + config.rabbit.default_exchange.clone().into(), + config.rabbit.queues.acks.into(), + BasicPublishOptions::default(), + payload.as_bytes(), + AMQPProperties::default() + .with_content_type("application/json".into()) + .with_delivery_mode(2), ) - .await + .await?; + + Ok(()) } } diff --git a/crates/daemons/crond/src/tasks/acks.rs b/crates/daemons/crond/src/tasks/acks.rs index 68b13144..2a3dc3a6 100644 --- a/crates/daemons/crond/src/tasks/acks.rs +++ b/crates/daemons/crond/src/tasks/acks.rs @@ -3,7 +3,7 @@ use lapin::{ options::*, types::FieldTable, uri::{AMQPAuthority, AMQPQueryString, AMQPUri, AMQPUserInfo}, - ConnectionBuilder, ConnectionProperties, + ConnectionBuilder, ConnectionProperties, ExchangeKind, }; use log::info; use redis_kiss::{get_connection, AsyncCommands, Conn as RedisConnection}; @@ -46,6 +46,42 @@ pub async fn task(db: Database) -> Result<()> { .await .expect("Failed to create channel"); + reader_channel + .exchange_declare( + config.rabbit.default_exchange.clone().into(), + ExchangeKind::Topic, + ExchangeDeclareOptions { + durable: true, + ..Default::default() + }, + FieldTable::default(), + ) + .await + .expect("Failed to declare exchange"); + + reader_channel + .queue_declare( + config.rabbit.queues.acks.clone().into(), + QueueDeclareOptions { + durable: true, + ..Default::default() + }, + FieldTable::default(), + ) + .await + .expect("Failed to bind queue"); + + reader_channel + .queue_bind( + config.rabbit.queues.acks.clone().into(), + config.rabbit.default_exchange.into(), + config.rabbit.queues.acks.clone().into(), + QueueBindOptions::default(), + FieldTable::default(), + ) + .await + .expect("Failed to bind channel"); + let mut consumer = reader_channel .basic_consume( config.rabbit.queues.acks.into(), diff --git a/crates/daemons/pushd/Cargo.toml b/crates/daemons/pushd/Cargo.toml index c47eb11a..8fb22999 100644 --- a/crates/daemons/pushd/Cargo.toml +++ b/crates/daemons/pushd/Cargo.toml @@ -15,7 +15,7 @@ revolt-parser = { workspace = true } anyhow = { workspace = true } -amqprs = { workspace = true } +lapin = { workspace = true } fcm_v1 = { workspace = true } web-push = { workspace = true } isahc = { workspace = true, features = ["json"], optional = true } diff --git a/crates/daemons/pushd/src/consumers/inbound/ack.rs b/crates/daemons/pushd/src/consumers/inbound/ack.rs index 77dd5d44..77df8b27 100644 --- a/crates/daemons/pushd/src/consumers/inbound/ack.rs +++ b/crates/daemons/pushd/src/consumers/inbound/ack.rs @@ -1,96 +1,69 @@ -use crate::consumers::inbound::internal::*; -use amqprs::{ - channel::{BasicPublishArguments, Channel}, - connection::Connection, - consumer::AsyncConsumer, - BasicProperties, Deliver, -}; +use std::sync::Arc; + +use crate::utils::Consumer; +use anyhow::Result; use async_trait::async_trait; +use lapin::{message::Delivery, Channel, Connection}; use revolt_database::{events::rabbit::*, Database}; +#[derive(Clone)] +#[allow(unused)] pub struct AckConsumer { - #[allow(dead_code)] db: Database, authifier_db: authifier::Database, - conn: Option, - channel: Option, + connection: Arc, + channel: Arc, } -impl Channeled for AckConsumer { - fn get_connection(&self) -> Option<&Connection> { - if self.conn.is_none() { - None - } else { - Some(self.conn.as_ref().unwrap()) - } - } - - fn get_channel(&self) -> Option<&Channel> { - if self.channel.is_none() { - None - } else { - Some(self.channel.as_ref().unwrap()) - } - } - - fn set_connection(&mut self, conn: Connection) { - self.conn = Some(conn); - } - - fn set_channel(&mut self, channel: Channel) { - self.channel = Some(channel) - } -} - -impl AckConsumer { - pub fn new(db: Database, authifier_db: authifier::Database) -> AckConsumer { - AckConsumer { +#[async_trait] +impl Consumer for AckConsumer { + async fn create( + db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, + ) -> Self { + Self { db, authifier_db, - conn: None, - channel: None, + connection, + channel, } } -} -#[allow(unused_variables)] -#[async_trait] -impl AsyncConsumer for AckConsumer { + fn channel(&self) -> &Arc { + &self.channel + } + /// This consumer processes all acks the platform receives, and sends relevant badge updates to apple platforms. - async fn consume( - &mut self, - channel: &Channel, - deliver: Deliver, - basic_properties: BasicProperties, - content: Vec, - ) { - let content = String::from_utf8(content).unwrap(); - let payload: AckPayload = serde_json::from_str(content.as_str()).unwrap(); + async fn consume(&self, delivery: Delivery) -> Result<()> { + let payload: AckPayload = serde_json::from_slice(&delivery.data)?; // Step 1: fetch unreads and don't continue if there's no unreads - #[allow(clippy::disallowed_methods)] - let unreads = self.db.fetch_unread_mentions(&payload.user_id).await; + // #[allow(clippy::disallowed_methods)] debug!("Processing unreads for {:}", &payload.user_id); - if let Ok(u) = &unreads { + let unreads = if let Ok(u) = self.db.fetch_unread_mentions(&payload.user_id).await { if u.is_empty() { debug!( "Discarding unread task (no mentions found) for {:}", &payload.user_id ); - return; - } + return Ok(()); + }; + + u } else { - return; - } + return Ok(()); + }; if let Ok(sessions) = self.authifier_db.find_sessions(&payload.user_id).await { let config = revolt_config::config().await; // Step 2: find any apple sessions, since we don't need to calculate this for anything else. // If there's no apple sessions, we can return early - let apple_sessions: Vec<&authifier::models::Session> = sessions - .iter() + let mut apple_sessions = sessions + .into_iter() .filter(|session| { if let Some(sub) = &session.subscription { sub.endpoint == "apn" @@ -98,19 +71,19 @@ impl AsyncConsumer for AckConsumer { false } }) - .collect(); + .peekable(); - if apple_sessions.is_empty() { + if apple_sessions.peek().is_none() { debug!( "Discarding unread task (no apn sessions found) for {:}", &payload.user_id ); - return; + return Ok(()); } // Step 3: calculate the actual mention count, since we have to send it out let mut mention_count = 0; - for u in &unreads.unwrap() { + for u in &unreads { mention_count += u.mentions.as_ref().unwrap().len() } @@ -123,26 +96,22 @@ impl AsyncConsumer for AckConsumer { token: session.subscription.as_ref().unwrap().auth.clone(), extras: Default::default(), }; - let raw_service_payload = serde_json::to_string(&service_payload); + let payload = serde_json::to_string(&service_payload)?; - if let Ok(p) = raw_service_payload { - let args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.apn.queue.as_str(), - ) - .finish(); + log::debug!( + "Publishing ack to apn session {}", + session.subscription.as_ref().unwrap().auth + ); - log::debug!( - "Publishing ack to apn session {}", - session.subscription.as_ref().unwrap().auth - ); - - publish_message(self, p.into(), args).await; - } else { - log::warn!("Failed to serialize ack badge update payload!"); - revolt_config::capture_error(&raw_service_payload.unwrap_err()); - } + self.publish_message( + payload.as_bytes(), + &config.pushd.exchange, + &config.pushd.apn.queue, + ) + .await?; } } + + Ok(()) } } diff --git a/crates/daemons/pushd/src/consumers/inbound/dm_call.rs b/crates/daemons/pushd/src/consumers/inbound/dm_call.rs index 873b55cc..3175d8d8 100644 --- a/crates/daemons/pushd/src/consumers/inbound/dm_call.rs +++ b/crates/daemons/pushd/src/consumers/inbound/dm_call.rs @@ -1,70 +1,44 @@ -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; -use crate::consumers::inbound::internal::*; -use amqprs::{ - channel::{BasicPublishArguments, Channel}, - connection::Connection, - consumer::AsyncConsumer, - BasicProperties, Deliver, -}; +use crate::utils::Consumer; use anyhow::Result; use async_trait::async_trait; +use lapin::{message::Delivery, Channel, Connection}; use log::debug; use revolt_database::{events::rabbit::*, Database}; +#[derive(Clone)] +#[allow(unused)] pub struct DmCallConsumer { - #[allow(dead_code)] db: Database, authifier_db: authifier::Database, - conn: Option, - channel: Option, + connection: Arc, + channel: Arc, } -impl Channeled for DmCallConsumer { - fn get_connection(&self) -> Option<&Connection> { - if self.conn.is_none() { - None - } else { - Some(self.conn.as_ref().unwrap()) - } - } - - fn get_channel(&self) -> Option<&Channel> { - if self.channel.is_none() { - None - } else { - Some(self.channel.as_ref().unwrap()) - } - } - - fn set_connection(&mut self, conn: Connection) { - self.conn = Some(conn); - } - - fn set_channel(&mut self, channel: Channel) { - self.channel = Some(channel) - } -} - -impl DmCallConsumer { - pub fn new(db: Database, authifier_db: authifier::Database) -> DmCallConsumer { - DmCallConsumer { +#[async_trait] +impl Consumer for DmCallConsumer { + async fn create( + db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, + ) -> Self { + Self { db, authifier_db, - conn: None, - channel: None, + connection, + channel, } } - async fn consume_event( - &mut self, - _channel: &Channel, - _deliver: Deliver, - _basic_properties: BasicProperties, - content: Vec, - ) -> Result<()> { - let content = String::from_utf8(content)?; - let _p: InternalDmCallPayload = serde_json::from_str(content.as_str())?; + fn channel(&self) -> &Arc { + &self.channel + } + + /// This consumer handles delegating messages into their respective platform queues. + async fn consume(&self, delivery: Delivery) -> Result<()> { + let _p: InternalDmCallPayload = serde_json::from_slice(&delivery.data)?; let payload = _p.payload; debug!("Received dm call start/stop event"); @@ -107,36 +81,27 @@ impl DmCallConsumer { extras: HashMap::new(), }; - let args: BasicPublishArguments; + let routing_key = match sub.endpoint.as_str() { + "apn" => &config.pushd.apn.queue, + "fcm" => &config.pushd.fcm.queue, + endpoint => { + sendable.extras.insert("p256dh".to_string(), sub.p256dh); + sendable + .extras + .insert("endpoint".to_string(), endpoint.to_string()); - if sub.endpoint == "apn" { - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.apn.queue.as_str(), - ) - .finish(); - } else if sub.endpoint == "fcm" { - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.fcm.queue.as_str(), - ) - .finish(); - } else { - // web push (vapid) - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.vapid.queue.as_str(), - ) - .finish(); - sendable.extras.insert("p256dh".to_string(), sub.p256dh); - sendable - .extras - .insert("endpoint".to_string(), sub.endpoint.clone()); - } + &config.pushd.vapid.queue + } + }; let payload = serde_json::to_string(&sendable)?; - publish_message(self, payload.into(), args).await; + self.publish_message( + payload.as_bytes(), + &config.pushd.exchange, + routing_key, + ) + .await?; } } } @@ -145,24 +110,3 @@ impl DmCallConsumer { Ok(()) } } - -#[allow(unused_variables)] -#[async_trait] -impl AsyncConsumer for DmCallConsumer { - /// This consumer handles delegating messages into their respective platform queues. - async fn consume( - &mut self, - channel: &Channel, - deliver: Deliver, - basic_properties: BasicProperties, - content: Vec, - ) { - if let Err(err) = self - .consume_event(channel, deliver, basic_properties, content) - .await - { - revolt_config::capture_anyhow(&err); - warn!("Failed to process dm call start/stop event: {err:?}"); - } - } -} diff --git a/crates/daemons/pushd/src/consumers/inbound/fr_accepted.rs b/crates/daemons/pushd/src/consumers/inbound/fr_accepted.rs index 115e4a26..bcfa1c2f 100644 --- a/crates/daemons/pushd/src/consumers/inbound/fr_accepted.rs +++ b/crates/daemons/pushd/src/consumers/inbound/fr_accepted.rs @@ -1,70 +1,44 @@ -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; -use crate::consumers::inbound::internal::*; -use amqprs::{ - channel::{BasicPublishArguments, Channel}, - connection::Connection, - consumer::AsyncConsumer, - BasicProperties, Deliver, -}; +use crate::utils::Consumer; use anyhow::Result; use async_trait::async_trait; +use lapin::{message::Delivery, Channel, Connection}; use log::debug; use revolt_database::{events::rabbit::*, Database}; +#[derive(Clone)] +#[allow(unused)] pub struct FRAcceptedConsumer { - #[allow(dead_code)] db: Database, authifier_db: authifier::Database, - conn: Option, - channel: Option, + connection: Arc, + channel: Arc, } -impl Channeled for FRAcceptedConsumer { - fn get_connection(&self) -> Option<&Connection> { - if self.conn.is_none() { - None - } else { - Some(self.conn.as_ref().unwrap()) - } - } - - fn get_channel(&self) -> Option<&Channel> { - if self.channel.is_none() { - None - } else { - Some(self.channel.as_ref().unwrap()) - } - } - - fn set_connection(&mut self, conn: Connection) { - self.conn = Some(conn); - } - - fn set_channel(&mut self, channel: Channel) { - self.channel = Some(channel) - } -} - -impl FRAcceptedConsumer { - pub fn new(db: Database, authifier_db: authifier::Database) -> FRAcceptedConsumer { - FRAcceptedConsumer { +#[async_trait] +impl Consumer for FRAcceptedConsumer { + async fn create( + db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, + ) -> Self { + Self { db, authifier_db, - conn: None, - channel: None, + connection, + channel, } } - async fn consume_event( - &mut self, - _channel: &Channel, - _deliver: Deliver, - _basic_properties: BasicProperties, - content: Vec, - ) -> Result<()> { - let content = String::from_utf8(content)?; - let payload: FRAcceptedPayload = serde_json::from_str(content.as_str())?; + fn channel(&self) -> &Arc { + &self.channel + } + + /// This consumer handles delegating messages into their respective platform queues. + async fn consume(&self, delivery: Delivery) -> Result<()> { + let payload: FRAcceptedPayload = serde_json::from_slice(&delivery.data)?; debug!("Received FR accept event"); @@ -80,36 +54,23 @@ impl FRAcceptedConsumer { extras: HashMap::new(), }; - let args: BasicPublishArguments; + let routing_key = match sub.endpoint.as_str() { + "apn" => &config.pushd.apn.queue, + "fcm" => &config.pushd.fcm.queue, + endpoint => { + sendable.extras.insert("p256dh".to_string(), sub.p256dh); + sendable + .extras + .insert("endpoint".to_string(), endpoint.to_string()); - if sub.endpoint == "apn" { - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.apn.queue.as_str(), - ) - .finish(); - } else if sub.endpoint == "fcm" { - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.fcm.queue.as_str(), - ) - .finish(); - } else { - // web push (vapid) - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.vapid.queue.as_str(), - ) - .finish(); - sendable.extras.insert("p256dh".to_string(), sub.p256dh); - sendable - .extras - .insert("endpoint".to_string(), sub.endpoint.clone()); - } + &config.pushd.vapid.queue + } + }; let payload = serde_json::to_string(&sendable)?; - publish_message(self, payload.into(), args).await; + self.publish_message(payload.as_bytes(), &config.pushd.exchange, routing_key) + .await?; } } } @@ -117,24 +78,3 @@ impl FRAcceptedConsumer { Ok(()) } } - -#[allow(unused_variables)] -#[async_trait] -impl AsyncConsumer for FRAcceptedConsumer { - /// This consumer handles delegating messages into their respective platform queues. - async fn consume( - &mut self, - channel: &Channel, - deliver: Deliver, - basic_properties: BasicProperties, - content: Vec, - ) { - if let Err(err) = self - .consume_event(channel, deliver, basic_properties, content) - .await - { - revolt_config::capture_anyhow(&err); - eprintln!("Failed to process friend request accepted event: {err:?}"); - } - } -} diff --git a/crates/daemons/pushd/src/consumers/inbound/fr_received.rs b/crates/daemons/pushd/src/consumers/inbound/fr_received.rs index c611dfaf..66fce72d 100644 --- a/crates/daemons/pushd/src/consumers/inbound/fr_received.rs +++ b/crates/daemons/pushd/src/consumers/inbound/fr_received.rs @@ -1,70 +1,44 @@ -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; -use crate::consumers::inbound::internal::*; -use amqprs::{ - channel::{BasicPublishArguments, Channel}, - connection::Connection, - consumer::AsyncConsumer, - BasicProperties, Deliver, -}; +use crate::utils::Consumer; use anyhow::Result; use async_trait::async_trait; +use lapin::{message::Delivery, Channel, Connection}; use log::debug; use revolt_database::{events::rabbit::*, Database}; +#[derive(Clone)] +#[allow(unused)] pub struct FRReceivedConsumer { - #[allow(dead_code)] db: Database, authifier_db: authifier::Database, - conn: Option, - channel: Option, + connection: Arc, + channel: Arc, } -impl Channeled for FRReceivedConsumer { - fn get_connection(&self) -> Option<&Connection> { - if self.conn.is_none() { - None - } else { - Some(self.conn.as_ref().unwrap()) - } - } - - fn get_channel(&self) -> Option<&Channel> { - if self.channel.is_none() { - None - } else { - Some(self.channel.as_ref().unwrap()) - } - } - - fn set_connection(&mut self, conn: Connection) { - self.conn = Some(conn); - } - - fn set_channel(&mut self, channel: Channel) { - self.channel = Some(channel) - } -} - -impl FRReceivedConsumer { - pub fn new(db: Database, authifier_db: authifier::Database) -> FRReceivedConsumer { - FRReceivedConsumer { +#[async_trait] +impl Consumer for FRReceivedConsumer { + async fn create( + db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, + ) -> Self { + Self { db, authifier_db, - conn: None, - channel: None, + connection, + channel, } } - async fn consume_event( - &mut self, - _channel: &Channel, - _deliver: Deliver, - _basic_properties: BasicProperties, - content: Vec, - ) -> Result<()> { - let content = String::from_utf8(content)?; - let payload: FRReceivedPayload = serde_json::from_str(content.as_str())?; + fn channel(&self) -> &Arc { + &self.channel + } + + /// This consumer handles delegating messages into their respective platform queues. + async fn consume(&self, delivery: Delivery) -> Result<()> { + let payload: FRReceivedPayload = serde_json::from_slice(&delivery.data)?; debug!("Received FR received event"); @@ -80,36 +54,23 @@ impl FRReceivedConsumer { extras: HashMap::new(), }; - let args: BasicPublishArguments; + let routing_key = match sub.endpoint.as_str() { + "apn" => &config.pushd.apn.queue, + "fcm" => &config.pushd.fcm.queue, + endpoint => { + sendable.extras.insert("p256dh".to_string(), sub.p256dh); + sendable + .extras + .insert("endpoint".to_string(), endpoint.to_string()); - if sub.endpoint == "apn" { - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.apn.queue.as_str(), - ) - .finish(); - } else if sub.endpoint == "fcm" { - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.fcm.queue.as_str(), - ) - .finish(); - } else { - // web push (vapid) - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.vapid.queue.as_str(), - ) - .finish(); - sendable.extras.insert("p256dh".to_string(), sub.p256dh); - sendable - .extras - .insert("endpoint".to_string(), sub.endpoint.clone()); - } + &config.pushd.vapid.queue + } + }; let payload = serde_json::to_string(&sendable)?; - publish_message(self, payload.into(), args).await; + self.publish_message(payload.as_bytes(), &config.pushd.exchange, routing_key) + .await?; } } } @@ -117,24 +78,3 @@ impl FRReceivedConsumer { Ok(()) } } - -#[allow(unused_variables)] -#[async_trait] -impl AsyncConsumer for FRReceivedConsumer { - /// This consumer handles delegating messages into their respective platform queues. - async fn consume( - &mut self, - channel: &Channel, - deliver: Deliver, - basic_properties: BasicProperties, - content: Vec, - ) { - if let Err(err) = self - .consume_event(channel, deliver, basic_properties, content) - .await - { - revolt_config::capture_anyhow(&err); - eprintln!("Failed to process friend request received event: {err:?}"); - } - } -} diff --git a/crates/daemons/pushd/src/consumers/inbound/generic.rs b/crates/daemons/pushd/src/consumers/inbound/generic.rs index 302b1700..f413661d 100644 --- a/crates/daemons/pushd/src/consumers/inbound/generic.rs +++ b/crates/daemons/pushd/src/consumers/inbound/generic.rs @@ -1,70 +1,44 @@ -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; -use crate::consumers::inbound::internal::*; -use amqprs::{ - channel::{BasicPublishArguments, Channel}, - connection::Connection, - consumer::AsyncConsumer, - BasicProperties, Deliver, -}; +use crate::utils::Consumer; use anyhow::Result; use async_trait::async_trait; +use lapin::{message::Delivery, Channel, Connection}; use log::debug; use revolt_database::{events::rabbit::*, Database}; +#[derive(Clone)] +#[allow(unused)] pub struct GenericConsumer { - #[allow(dead_code)] db: Database, authifier_db: authifier::Database, - conn: Option, - channel: Option, + connection: Arc, + channel: Arc, } -impl Channeled for GenericConsumer { - fn get_connection(&self) -> Option<&Connection> { - if self.conn.is_none() { - None - } else { - Some(self.conn.as_ref().unwrap()) - } - } - - fn get_channel(&self) -> Option<&Channel> { - if self.channel.is_none() { - None - } else { - Some(self.channel.as_ref().unwrap()) - } - } - - fn set_connection(&mut self, conn: Connection) { - self.conn = Some(conn); - } - - fn set_channel(&mut self, channel: Channel) { - self.channel = Some(channel) - } -} - -impl GenericConsumer { - pub fn new(db: Database, authifier_db: authifier::Database) -> GenericConsumer { - GenericConsumer { +#[async_trait] +impl Consumer for GenericConsumer { + async fn create( + db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, + ) -> Self { + Self { db, authifier_db, - conn: None, - channel: None, + connection, + channel, } } - async fn consume_event( - &mut self, - _channel: &Channel, - _deliver: Deliver, - _basic_properties: BasicProperties, - content: Vec, - ) -> Result<()> { - let content = String::from_utf8(content)?; - let payload: MessageSentPayload = serde_json::from_str(content.as_str())?; + fn channel(&self) -> &Arc { + &self.channel + } + + /// This consumer handles delegating messages into their respective platform queues. + async fn consume(&self, delivery: Delivery) -> Result<()> { + let payload: MessageSentPayload = serde_json::from_slice(&delivery.data)?; debug!("Received message event on origin"); @@ -86,36 +60,23 @@ impl GenericConsumer { extras: HashMap::new(), }; - let args: BasicPublishArguments; + let routing_key = match sub.endpoint.as_str() { + "apn" => &config.pushd.apn.queue, + "fcm" => &config.pushd.fcm.queue, + endpoint => { + sendable.extras.insert("p256dh".to_string(), sub.p256dh); + sendable + .extras + .insert("endpoint".to_string(), endpoint.to_string()); - if sub.endpoint == "apn" { - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.apn.queue.as_str(), - ) - .finish(); - } else if sub.endpoint == "fcm" { - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.fcm.queue.as_str(), - ) - .finish(); - } else { - // web push (vapid) - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.vapid.queue.as_str(), - ) - .finish(); - sendable.extras.insert("p256dh".to_string(), sub.p256dh); - sendable - .extras - .insert("endpoint".to_string(), sub.endpoint.clone()); - } + &config.pushd.vapid.queue + } + }; let payload = serde_json::to_string(&sendable)?; - publish_message(self, payload.into(), args).await; + self.publish_message(payload.as_bytes(), &config.pushd.exchange, routing_key) + .await?; } } } @@ -123,24 +84,3 @@ impl GenericConsumer { Ok(()) } } - -#[allow(unused_variables)] -#[async_trait] -impl AsyncConsumer for GenericConsumer { - /// This consumer handles delegating messages into their respective platform queues. - async fn consume( - &mut self, - channel: &Channel, - deliver: Deliver, - basic_properties: BasicProperties, - content: Vec, - ) { - if let Err(err) = self - .consume_event(channel, deliver, basic_properties, content) - .await - { - revolt_config::capture_anyhow(&err); - eprintln!("Failed to process generic event: {err:?}"); - } - } -} diff --git a/crates/daemons/pushd/src/consumers/inbound/internal.rs b/crates/daemons/pushd/src/consumers/inbound/internal.rs deleted file mode 100644 index 387c08b5..00000000 --- a/crates/daemons/pushd/src/consumers/inbound/internal.rs +++ /dev/null @@ -1,53 +0,0 @@ -use amqprs::{ - channel::{BasicPublishArguments, Channel}, - connection::{Connection, OpenConnectionArguments}, - BasicProperties, -}; -use log::{debug, warn}; - -pub(crate) trait Channeled { - #[allow(unused)] - fn get_connection(&self) -> Option<&Connection>; - fn get_channel(&self) -> Option<&Channel>; - fn set_connection(&mut self, conn: Connection); - fn set_channel(&mut self, channel: Channel); -} - -pub(crate) async fn make_channel(consumer: &mut T) { - let config = revolt_config::config().await; - - let args = OpenConnectionArguments::new( - &config.rabbit.host, - config.rabbit.port, - &config.rabbit.username, - &config.rabbit.password, - ); - let conn = amqprs::connection::Connection::open(&args).await.unwrap(); - - let channel = conn.open_channel(None).await.unwrap(); - - consumer.set_connection(conn); - consumer.set_channel(channel); -} - -pub(crate) async fn publish_message( - consumer: &mut T, - payload: Vec, - args: BasicPublishArguments, -) { - let routing_key = &args.routing_key.clone(); - let mut channel = consumer.get_channel(); - if channel.is_none() { - make_channel(consumer).await; - channel = consumer.get_channel(); - } - - if let Some(chnl) = channel { - chnl.basic_publish(BasicProperties::default(), payload.clone(), args.clone()) - .await - .unwrap(); - debug!("Sent message to queue for target {}", routing_key); - } else { - warn!("Failed to unwrap channel (including attempt to make a channel)!") - } -} diff --git a/crates/daemons/pushd/src/consumers/inbound/mass_mention.rs b/crates/daemons/pushd/src/consumers/inbound/mass_mention.rs index 8d43cce8..55ab984f 100644 --- a/crates/daemons/pushd/src/consumers/inbound/mass_mention.rs +++ b/crates/daemons/pushd/src/consumers/inbound/mass_mention.rs @@ -1,17 +1,13 @@ use std::{ collections::{HashMap, HashSet}, hash::RandomState, + sync::Arc, }; -use crate::{consumers::inbound::internal::*, utils}; -use amqprs::{ - channel::{BasicPublishArguments, Channel}, - connection::Connection, - consumer::AsyncConsumer, - BasicProperties, Deliver, -}; +use crate::utils::{render_notification_content, Consumer}; use anyhow::Result; use async_trait::async_trait; +use lapin::{message::Delivery, Channel, Connection}; use revolt_database::{ events::rabbit::*, util::bulk_permissions::BulkDatabasePermissionQuery, Database, Member, MessageFlagsValue, @@ -19,52 +15,18 @@ use revolt_database::{ use revolt_models::v0::{MessageFlags, PushNotification}; use revolt_result::ToRevoltError; +#[derive(Clone)] +#[allow(unused)] pub struct MassMessageConsumer { - #[allow(dead_code)] db: Database, authifier_db: authifier::Database, - conn: Option, - channel: Option, -} - -impl Channeled for MassMessageConsumer { - fn get_connection(&self) -> Option<&Connection> { - if self.conn.is_none() { - None - } else { - Some(self.conn.as_ref().unwrap()) - } - } - - fn get_channel(&self) -> Option<&Channel> { - if self.channel.is_none() { - None - } else { - Some(self.channel.as_ref().unwrap()) - } - } - - fn set_connection(&mut self, conn: Connection) { - self.conn = Some(conn); - } - - fn set_channel(&mut self, channel: Channel) { - self.channel = Some(channel) - } + connection: Arc, + channel: Arc, } impl MassMessageConsumer { - pub fn new(db: Database, authifier_db: authifier::Database) -> MassMessageConsumer { - MassMessageConsumer { - db, - authifier_db, - conn: None, - channel: None, - } - } - async fn fire_notification_for_users( - &mut self, + &self, push: &PushNotification, users: &[String], ) -> Result<()> { @@ -84,56 +46,58 @@ impl MassMessageConsumer { extras: HashMap::new(), }; - let args: BasicPublishArguments; + let routing_key = match sub.endpoint.as_str() { + "apn" => &config.pushd.apn.queue, + "fcm" => &config.pushd.fcm.queue, + endpoint => { + sendable.extras.insert("p256dh".to_string(), sub.p256dh); + sendable + .extras + .insert("endpoint".to_string(), endpoint.to_string()); - if sub.endpoint == "apn" { - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.apn.queue.as_str(), - ) - .finish(); - } else if sub.endpoint == "fcm" { - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.fcm.queue.as_str(), - ) - .finish(); - } else { - // web push (vapid) - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.vapid.queue.as_str(), - ) - .finish(); - sendable.extras.insert("p256dh".to_string(), sub.p256dh); - sendable - .extras - .insert("endpoint".to_string(), sub.endpoint.clone()); - } + &config.pushd.vapid.queue + } + }; let payload = serde_json::to_string(&sendable)?; - publish_message(self, payload.into(), args).await; + self.publish_message(payload.as_bytes(), &config.pushd.exchange, routing_key) + .await?; } } } Ok(()) } +} - async fn consume_event( - &mut self, - _channel: &Channel, - _deliver: Deliver, - _basic_properties: BasicProperties, - content: Vec, - ) -> Result<()> { +#[async_trait] +impl Consumer for MassMessageConsumer { + async fn create( + db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, + ) -> Self { + Self { + db, + authifier_db, + connection, + channel, + } + } + + fn channel(&self) -> &Arc { + &self.channel + } + + /// This consumer handles adding mentions for all the users affected by a mass mention ping, and then sends out push notifications. + async fn consume(&self, delivery: Delivery) -> Result<()> { + let mut payload: MassMessageSentPayload = serde_json::from_slice(&delivery.data)?; let config = revolt_config::config().await; - let content = String::from_utf8(content)?; - let mut payload: MassMessageSentPayload = serde_json::from_str(content.as_str())?; for push in payload.notifications.iter_mut() { - if let Ok(body) = utils::render_notification_content(push, &self.db) + if let Ok(body) = render_notification_content(push, &self.db) .await .to_internal_error() { @@ -280,24 +244,3 @@ impl MassMessageConsumer { Ok(()) } } - -#[allow(unused_variables)] -#[async_trait] -impl AsyncConsumer for MassMessageConsumer { - /// This consumer handles adding mentions for all the users affected by a mass mention ping, and then sends out push notifications - async fn consume( - &mut self, - channel: &Channel, - deliver: Deliver, - basic_properties: BasicProperties, - content: Vec, - ) { - if let Err(err) = self - .consume_event(channel, deliver, basic_properties, content) - .await - { - revolt_config::capture_anyhow(&err); - eprintln!("Failed to process mass message event: {err:?}"); - } - } -} diff --git a/crates/daemons/pushd/src/consumers/inbound/message.rs b/crates/daemons/pushd/src/consumers/inbound/message.rs index 45de0b0b..e4aa04d5 100644 --- a/crates/daemons/pushd/src/consumers/inbound/message.rs +++ b/crates/daemons/pushd/src/consumers/inbound/message.rs @@ -1,76 +1,46 @@ -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; -use crate::{consumers::inbound::internal::*, utils}; -use amqprs::{ - channel::{BasicPublishArguments, Channel}, - connection::Connection, - consumer::AsyncConsumer, - BasicProperties, Deliver, -}; +use crate::utils::{render_notification_content, Consumer}; use anyhow::Result; use async_trait::async_trait; +use lapin::{message::Delivery, Channel, Connection}; use log::debug; use revolt_database::{events::rabbit::*, Database}; -use revolt_result::ToRevoltError; +#[derive(Clone)] +#[allow(unused)] pub struct MessageConsumer { - #[allow(dead_code)] db: Database, authifier_db: authifier::Database, - conn: Option, - channel: Option, + connection: Arc, + channel: Arc, } -impl Channeled for MessageConsumer { - fn get_connection(&self) -> Option<&Connection> { - if self.conn.is_none() { - None - } else { - Some(self.conn.as_ref().unwrap()) - } - } - - fn get_channel(&self) -> Option<&Channel> { - if self.channel.is_none() { - None - } else { - Some(self.channel.as_ref().unwrap()) - } - } - - fn set_connection(&mut self, conn: Connection) { - self.conn = Some(conn); - } - - fn set_channel(&mut self, channel: Channel) { - self.channel = Some(channel) - } -} - -impl MessageConsumer { - pub fn new(db: Database, authifier_db: authifier::Database) -> MessageConsumer { - MessageConsumer { +#[async_trait] +impl Consumer for MessageConsumer { + async fn create( + db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, + ) -> Self { + Self { db, authifier_db, - conn: None, - channel: None, + connection, + channel, } } - async fn consume_event( - &mut self, - _channel: &Channel, - _deliver: Deliver, - _basic_properties: BasicProperties, - content: Vec, - ) -> Result<()> { - let content = String::from_utf8(content)?; - let mut payload: MessageSentPayload = serde_json::from_str(content.as_str())?; + fn channel(&self) -> &Arc { + &self.channel + } - if let Ok(body) = utils::render_notification_content(&payload.notification, &self.db) - .await - .to_internal_error() - { + /// This consumer handles delegating messages into their respective platform queues. + async fn consume(&self, delivery: Delivery) -> Result<()> { + let mut payload: MessageSentPayload = serde_json::from_slice(&delivery.data)?; + + if let Ok(body) = render_notification_content(&payload.notification, &self.db).await { payload.notification.raw_body = Some(payload.notification.body); payload.notification.body = body; } @@ -95,36 +65,22 @@ impl MessageConsumer { extras: HashMap::new(), }; - let args: BasicPublishArguments; + let routing_key = match sub.endpoint.as_str() { + "apn" => &config.pushd.apn.queue, + "fcm" => &config.pushd.fcm.queue, + endpoint => { + sendable.extras.insert("p256dh".to_string(), sub.p256dh); + sendable + .extras + .insert("endpoint".to_string(), endpoint.to_string()); - if sub.endpoint == "apn" { - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.apn.queue.as_str(), - ) - .finish(); - } else if sub.endpoint == "fcm" { - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.fcm.queue.as_str(), - ) - .finish(); - } else { - // web push (vapid) - args = BasicPublishArguments::new( - config.pushd.exchange.as_str(), - config.pushd.vapid.queue.as_str(), - ) - .finish(); - sendable.extras.insert("p256dh".to_string(), sub.p256dh); - sendable - .extras - .insert("endpoint".to_string(), sub.endpoint.clone()); - } + &config.pushd.vapid.queue + } + }; let payload = serde_json::to_string(&sendable)?; - - publish_message(self, payload.into(), args).await; + self.publish_message(payload.as_bytes(), &config.pushd.exchange, routing_key) + .await?; } } } @@ -132,24 +88,3 @@ impl MessageConsumer { Ok(()) } } - -#[allow(unused_variables)] -#[async_trait] -impl AsyncConsumer for MessageConsumer { - /// This consumer handles delegating messages into their respective platform queues. - async fn consume( - &mut self, - channel: &Channel, - deliver: Deliver, - basic_properties: BasicProperties, - content: Vec, - ) { - if let Err(err) = self - .consume_event(channel, deliver, basic_properties, content) - .await - { - revolt_config::capture_anyhow(&err); - eprintln!("Failed to process message event: {err:?}"); - } - } -} diff --git a/crates/daemons/pushd/src/consumers/inbound/mod.rs b/crates/daemons/pushd/src/consumers/inbound/mod.rs index 4d6d36b2..7c93cdb6 100644 --- a/crates/daemons/pushd/src/consumers/inbound/mod.rs +++ b/crates/daemons/pushd/src/consumers/inbound/mod.rs @@ -3,6 +3,5 @@ pub mod dm_call; pub mod fr_accepted; pub mod fr_received; pub mod generic; -mod internal; pub mod mass_mention; pub mod message; diff --git a/crates/daemons/pushd/src/consumers/outbound/apn.rs b/crates/daemons/pushd/src/consumers/outbound/apn.rs index ec478c1f..feb6cb94 100644 --- a/crates/daemons/pushd/src/consumers/outbound/apn.rs +++ b/crates/daemons/pushd/src/consumers/outbound/apn.rs @@ -1,12 +1,13 @@ -use std::{borrow::Cow, collections::BTreeMap, io::Cursor}; +use std::{borrow::Cow, collections::BTreeMap, io::Cursor, sync::Arc}; -use amqprs::{channel::Channel as AmqpChannel, consumer::AsyncConsumer, BasicProperties, Deliver}; -use anyhow::{anyhow, Result}; +use crate::utils::Consumer; +use anyhow::Result; use async_trait::async_trait; use base64::{ engine::{self}, Engine as _, }; +use lapin::{message::Delivery, Channel as AMQPChannel, Connection}; use revolt_a2::{ request::{ notification::{DefaultAlert, NotificationOptions}, @@ -42,7 +43,7 @@ impl<'a> PayloadLike for MessagePayload<'a> { fn get_device_token(&self) -> &'a str { self.device_token } - fn get_options(&self) -> &NotificationOptions { + fn get_options(&self) -> &NotificationOptions<'a> { &self.options } } @@ -68,16 +69,20 @@ impl<'a> PayloadLike for CallStartStopPayload<'a> { fn get_device_token(&self) -> &'a str { self.device_token } - fn get_options(&self) -> &NotificationOptions { + fn get_options(&self) -> &NotificationOptions<'a> { &self.options } } // region: consumer +#[derive(Clone)] +#[allow(unused)] pub struct ApnsOutboundConsumer { - #[allow(dead_code)] db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, client: Client, } @@ -117,15 +122,21 @@ impl ApnsOutboundConsumer { } } -impl ApnsOutboundConsumer { - pub async fn new(db: Database) -> Result { +#[async_trait] +impl Consumer for ApnsOutboundConsumer { + async fn create( + db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, + ) -> Self { let config = revolt_config::config().await; if config.pushd.apn.pkcs8.is_empty() || config.pushd.apn.key_id.is_empty() || config.pushd.apn.team_id.is_empty() { - return Err("Missing APN keys."); + panic!("Missing APN keys."); } let endpoint = if config.pushd.apn.sandbox { @@ -148,18 +159,21 @@ impl ApnsOutboundConsumer { ) .expect("could not create APN client"); - Ok(ApnsOutboundConsumer { db, client }) + Self { + db, + authifier_db, + connection, + channel, + client, + } } - async fn consume_event( - &mut self, - _channel: &AmqpChannel, - _deliver: Deliver, - _basic_properties: BasicProperties, - content: Vec, - ) -> Result<()> { - let content = String::from_utf8(content)?; - let payload: PayloadToService = serde_json::from_str(content.as_str())?; + fn channel(&self) -> &Arc { + &self.channel + } + + async fn consume(&self, delivery: Delivery) -> Result<()> { + let payload: PayloadToService = serde_json::from_slice(&delivery.data)?; let payload_options = NotificationOptions { apns_id: None, @@ -170,20 +184,15 @@ impl ApnsOutboundConsumer { apns_collapse_id: None, }; - let resp: Result; - - match payload.notification { + let resp = match payload.notification { PayloadKind::FRReceived(alert) => { let loc_args = vec![Cow::from( - alert - .from_user - .display_name - .or(Some(format!( + alert.from_user.display_name.clone().unwrap_or_else(|| { + format!( "{}#{}", alert.from_user.username, alert.from_user.discriminator - ))) - .clone() - .ok_or_else(|| anyhow!("missing name"))?, + ) + }), )]; let apn_payload = Payload { @@ -216,20 +225,17 @@ impl ApnsOutboundConsumer { "Sending friend request received for user: {:}", &payload.user_id ); - resp = self.client.send(apn_payload).await; + self.client.send(apn_payload).await } PayloadKind::FRAccepted(alert) => { let loc_args = vec![Cow::from( - alert - .accepted_user - .display_name - .or(Some(format!( + alert.accepted_user.display_name.clone().unwrap_or_else(|| { + format!( "{}#{}", alert.accepted_user.username, alert.accepted_user.discriminator - ))) - .clone() - .ok_or_else(|| anyhow!("missing name"))?, + ) + }), )]; let apn_payload = Payload { @@ -262,7 +268,7 @@ impl ApnsOutboundConsumer { "Sending friend request accept for user: {:}", &payload.user_id ); - resp = self.client.send(apn_payload).await; + self.client.send(apn_payload).await } PayloadKind::Generic(alert) => { let apn_payload = Payload { @@ -295,7 +301,7 @@ impl ApnsOutboundConsumer { "Sending generic notification for user: {:}", &payload.user_id ); - resp = self.client.send(apn_payload).await; + self.client.send(apn_payload).await } PayloadKind::MessageNotification(alert) => { @@ -334,7 +340,7 @@ impl ApnsOutboundConsumer { "Sending message notification for user: {:}", &payload.user_id ); - resp = self.client.send(apn_payload).await; + self.client.send(apn_payload).await } PayloadKind::BadgeUpdate(badge) => { @@ -349,7 +355,7 @@ impl ApnsOutboundConsumer { }; debug!("Sending badge update for user: {:}", &payload.user_id); - resp = self.client.send(apn_payload).await; + self.client.send(apn_payload).await } PayloadKind::DmCallStartEnd(alert) => { @@ -378,58 +384,37 @@ impl ApnsOutboundConsumer { "Sending call start/stop notification for user: {:}", &payload.user_id ); - resp = self.client.send(apn_payload).await; + self.client.send(apn_payload).await } - } + }; - if let Err(err) = resp { - match err { - Error::ResponseError(Response { - error: - Some(ErrorBody { - reason: ErrorReason::BadDeviceToken | ErrorReason::Unregistered, - .. - }), - .. - }) => { - info!( - "Removing APNS subscription id {:} (user: {:}) due to invalid token", - &payload.session_id, &payload.user_id - ); - if let Err(err) = self - .db - .remove_push_subscription_by_session_id(&payload.session_id) - .await - { - revolt_config::capture_error(&err); - } - } - err => { + match resp { + Err(Error::ResponseError(Response { + error: + Some(ErrorBody { + reason: ErrorReason::BadDeviceToken | ErrorReason::Unregistered, + .. + }), + .. + })) => { + info!( + "Removing APNS subscription id {:} (user: {:}) due to invalid token", + &payload.session_id, &payload.user_id + ); + + if let Err(err) = self + .db + .remove_push_subscription_by_session_id(&payload.session_id) + .await + { revolt_config::capture_error(&err); } } - } + resp => { + resp?; + } + }; Ok(()) } } - -#[allow(unused_variables)] -#[async_trait] -impl AsyncConsumer for ApnsOutboundConsumer { - async fn consume( - &mut self, - channel: &AmqpChannel, - deliver: Deliver, - basic_properties: BasicProperties, - content: Vec, - ) { - if let Err(err) = self - .consume_event(channel, deliver, basic_properties, content) - .await - { - revolt_config::capture_anyhow(&err); - eprintln!("Failed to process APN event: {err:?}"); - } - } -} diff --git a/crates/daemons/pushd/src/consumers/outbound/fcm.rs b/crates/daemons/pushd/src/consumers/outbound/fcm.rs index 4ea21e0b..e3bd185c 100644 --- a/crates/daemons/pushd/src/consumers/outbound/fcm.rs +++ b/crates/daemons/pushd/src/consumers/outbound/fcm.rs @@ -1,14 +1,14 @@ -use std::{collections::HashMap, time::Duration}; +use std::{collections::HashMap, sync::Arc, time::Duration}; -use amqprs::{channel::Channel as AmqpChannel, consumer::AsyncConsumer, BasicProperties, Deliver}; - -use anyhow::{anyhow, bail, Result}; +use crate::utils::Consumer; +use anyhow::{bail, Result}; use async_trait::async_trait; use fcm_v1::{ auth::{Authenticator, ServiceAccountKey}, message::Message, Client, Error as FcmError, }; +use lapin::{message::Delivery, Channel as AMQPChannel, Connection}; use revolt_config::config; use revolt_database::{events::rabbit::*, Database}; use serde_json::Value; @@ -115,17 +115,31 @@ impl NotificationData { } } +#[derive(Clone)] +#[allow(unused)] pub struct FcmOutboundConsumer { db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, client: Client, } -impl FcmOutboundConsumer { - pub async fn new(db: Database) -> Result { +#[async_trait] +impl Consumer for FcmOutboundConsumer { + async fn create( + db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, + ) -> Self { let config = revolt_config::config().await; - Ok(FcmOutboundConsumer { + Self { db, + authifier_db, + connection, + channel, client: Client::new( Authenticator::service_account::<&str>(ServiceAccountKey { key_type: Some(config.pushd.fcm.key_type), @@ -145,33 +159,27 @@ impl FcmOutboundConsumer { false, Duration::from_secs(5), ), - }) + } } - async fn consume_event( - &mut self, - _channel: &AmqpChannel, - _deliver: Deliver, - _basic_properties: BasicProperties, - content: Vec, - ) -> Result<()> { - let content = String::from_utf8(content)?; - let payload: PayloadToService = serde_json::from_str(content.as_str())?; + fn channel(&self) -> &Arc { + &self.channel + } + + async fn consume(&self, delivery: Delivery) -> Result<()> { + let payload: PayloadToService = serde_json::from_slice(&delivery.data)?; #[allow(clippy::needless_late_init)] let resp: Result; match payload.notification { PayloadKind::FRReceived(alert) => { - let name = alert - .from_user - .display_name - .or(Some(format!( + let name = alert.from_user.display_name.clone().unwrap_or_else(|| { + format!( "{}#{}", alert.from_user.username, alert.from_user.discriminator - ))) - .clone() - .ok_or_else(|| anyhow!("missing name"))?; + ) + }); let data = NotificationData::FRReceived { id: alert.from_user.id, @@ -188,15 +196,12 @@ impl FcmOutboundConsumer { } PayloadKind::FRAccepted(alert) => { - let name = alert - .accepted_user - .display_name - .or(Some(format!( + let name = alert.accepted_user.display_name.clone().unwrap_or_else(|| { + format!( "{}#{}", alert.accepted_user.username, alert.accepted_user.discriminator - ))) - .clone() - .ok_or_else(|| anyhow!("missing name"))?; + ) + }); let data = NotificationData::FRAccepted { id: alert.accepted_user.id, @@ -269,43 +274,21 @@ impl FcmOutboundConsumer { } } - if let Err(err) = resp { - match err { - FcmError::Auth => { - if let Err(err) = self - .db - .remove_push_subscription_by_session_id(&payload.session_id) - .await - { - revolt_config::capture_error(&err); - } - } - err => { + match resp { + Err(FcmError::Auth) => { + if let Err(err) = self + .db + .remove_push_subscription_by_session_id(&payload.session_id) + .await + { revolt_config::capture_error(&err); } } - } + res => { + res?; + } + }; Ok(()) } } - -#[allow(unused_variables)] -#[async_trait] -impl AsyncConsumer for FcmOutboundConsumer { - async fn consume( - &mut self, - channel: &AmqpChannel, - deliver: Deliver, - basic_properties: BasicProperties, - content: Vec, - ) { - if let Err(err) = self - .consume_event(channel, deliver, basic_properties, content) - .await - { - revolt_config::capture_anyhow(&err); - eprintln!("Failed to process FCM event: {err:?}"); - } - } -} diff --git a/crates/daemons/pushd/src/consumers/outbound/vapid.rs b/crates/daemons/pushd/src/consumers/outbound/vapid.rs index 67ec31fb..0acd8be9 100644 --- a/crates/daemons/pushd/src/consumers/outbound/vapid.rs +++ b/crates/daemons/pushd/src/consumers/outbound/vapid.rs @@ -1,6 +1,6 @@ -use std::collections::HashMap; +use std::{collections::HashMap, sync::Arc}; -use amqprs::{channel::Channel as AmqpChannel, consumer::AsyncConsumer, BasicProperties, Deliver}; +use crate::utils::Consumer; use anyhow::{anyhow, bail, Result}; use async_trait::async_trait; @@ -8,46 +8,60 @@ use base64::{ engine::{self}, Engine as _, }; +use lapin::{message::Delivery, Channel as AMQPChannel, Connection}; use revolt_database::{events::rabbit::*, util::format_display_name, Database}; use web_push::{ ContentEncoding, IsahcWebPushClient, SubscriptionInfo, SubscriptionKeys, VapidSignatureBuilder, WebPushClient, WebPushError, WebPushMessageBuilder, }; +#[derive(Clone)] +#[allow(unused)] pub struct VapidOutboundConsumer { db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, client: IsahcWebPushClient, - pkey: Vec, + pkey: Arc>, } -impl VapidOutboundConsumer { - pub async fn new(db: Database) -> Result { +#[async_trait] +impl Consumer for VapidOutboundConsumer { + async fn create( + db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, + ) -> Self { let config = revolt_config::config().await; - if config.pushd.vapid.private_key.is_empty() | config.pushd.vapid.public_key.is_empty() { - bail!("no Vapid keys present"); + if config.pushd.vapid.private_key.is_empty() || config.pushd.vapid.public_key.is_empty() { + panic!("no Vapid keys present"); } - let web_push_private_key = engine::general_purpose::URL_SAFE_NO_PAD - .decode(config.pushd.vapid.private_key) - .expect("valid `VAPID_PRIVATE_KEY`"); + let web_push_private_key = Arc::new( + engine::general_purpose::URL_SAFE_NO_PAD + .decode(config.pushd.vapid.private_key) + .expect("valid `VAPID_PRIVATE_KEY`"), + ); - Ok(VapidOutboundConsumer { + Self { db, + authifier_db, + connection, + channel, client: IsahcWebPushClient::new().unwrap(), pkey: web_push_private_key, - }) + } } - async fn consume_event( - &mut self, - _channel: &AmqpChannel, - _deliver: Deliver, - _basic_properties: BasicProperties, - content: Vec, - ) -> Result<()> { - let content = String::from_utf8(content)?; - let payload: PayloadToService = serde_json::from_str(content.as_str())?; + fn channel(&self) -> &Arc { + &self.channel + } + + async fn consume(&self, delivery: Delivery) -> Result<()> { + let payload: PayloadToService = serde_json::from_slice(&delivery.data)?; let subscription = SubscriptionInfo { endpoint: payload @@ -65,10 +79,7 @@ impl VapidOutboundConsumer { }, }; - #[allow(clippy::needless_late_init)] - let payload_body: String; - - match payload.notification { + let payload_body = match payload.notification { PayloadKind::FRReceived(alert) => { let name = alert .from_user @@ -83,7 +94,7 @@ impl VapidOutboundConsumer { let mut body = HashMap::new(); body.insert("body", format!("{} sent you a friend request", name)); - payload_body = serde_json::to_string(&body)?; + serde_json::to_string(&body)? } PayloadKind::FRAccepted(alert) => { let name = alert @@ -99,14 +110,10 @@ impl VapidOutboundConsumer { let mut body = HashMap::new(); body.insert("body", format!("{} accepted your friend request", name)); - payload_body = serde_json::to_string(&body)?; - } - PayloadKind::Generic(alert) => { - payload_body = serde_json::to_string(&alert)?; - } - PayloadKind::MessageNotification(alert) => { - payload_body = serde_json::to_string(&alert)?; + serde_json::to_string(&body)? } + PayloadKind::Generic(alert) => serde_json::to_string(&alert)?, + PayloadKind::MessageNotification(alert) => serde_json::to_string(&alert)?, PayloadKind::DmCallStartEnd(alert) => { let initiator_name = if let Some(server_id) = self.db.fetch_channel(&alert.channel_id).await?.server() @@ -132,59 +139,41 @@ impl VapidOutboundConsumer { _ => bail!("Invalid DmCallStart/End channel type"), } - payload_body = serde_json::to_string(&body)?; + serde_json::to_string(&body)? } PayloadKind::BadgeUpdate(_) => { bail!("Vapid cannot handle badge updates and they should not be sent here."); } - } + }; - match VapidSignatureBuilder::from_pem(std::io::Cursor::new(&self.pkey), &subscription) { - Ok(sig_builder) => match sig_builder.build() { - Ok(signature) => { - let mut builder = WebPushMessageBuilder::new(&subscription); - builder.set_vapid_signature(signature); + let signature = VapidSignatureBuilder::from_pem( + std::io::Cursor::new(self.pkey.as_ref()), + &subscription, + )? + .build()?; - builder.set_payload(ContentEncoding::AesGcm, payload_body.as_bytes()); + let mut builder = WebPushMessageBuilder::new(&subscription); + builder.set_vapid_signature(signature); - match builder.build() { - Ok(msg) => { - if let Err(err) = self.client.send(msg).await { - if err == WebPushError::Unauthorized { - self.db - .remove_push_subscription_by_session_id(&payload.session_id) - .await?; - } - } + builder.set_payload(ContentEncoding::AesGcm, payload_body.as_bytes()); - Ok(()) - } - Err(err) => Err(err.into()), - } + let msg = builder.build()?; + + match self.client.send(msg).await { + Err(WebPushError::Unauthorized) => { + if let Err(err) = self + .db + .remove_push_subscription_by_session_id(&payload.session_id) + .await + { + revolt_config::capture_error(&err); } - Err(err) => Err(err.into()), - }, - Err(err) => Err(err.into()), - } - } -} + } + res => { + res?; + } + }; -#[allow(unused_variables)] -#[async_trait] -impl AsyncConsumer for VapidOutboundConsumer { - async fn consume( - &mut self, - channel: &AmqpChannel, - deliver: Deliver, - basic_properties: BasicProperties, - content: Vec, - ) { - if let Err(err) = self - .consume_event(channel, deliver, basic_properties, content) - .await - { - revolt_config::capture_anyhow(&err); - eprintln!("Failed to process Vapid event: {err:?}"); - } + Ok(()) } } diff --git a/crates/daemons/pushd/src/main.rs b/crates/daemons/pushd/src/main.rs index eaa2e1b8..b4824b65 100644 --- a/crates/daemons/pushd/src/main.rs +++ b/crates/daemons/pushd/src/main.rs @@ -1,17 +1,16 @@ #[macro_use] extern crate log; -use amqprs::{ - channel::{ - BasicConsumeArguments, Channel, ExchangeDeclareArguments, QueueBindArguments, - QueueDeclareArguments, - }, - connection::{Connection, OpenConnectionArguments}, - consumer::AsyncConsumer, - FieldTable, +use std::sync::Arc; + +use lapin::{ + options::{BasicConsumeOptions, ExchangeDeclareOptions, QueueBindOptions, QueueDeclareOptions}, + types::{AMQPValue, FieldTable}, + Channel, Connection, ConnectionProperties, }; use revolt_config::{config, Settings}; -use tokio::sync::Notify; +use revolt_database::Database; +use tokio::signal::ctrl_c; mod consumers; mod utils; @@ -24,6 +23,8 @@ use consumers::{ outbound::{apn::ApnsOutboundConsumer, fcm::FcmOutboundConsumer, vapid::VapidOutboundConsumer}, }; +use crate::utils::{Consumer, Delegate}; + #[tokio::main(flavor = "multi_thread", worker_threads = 2)] async fn main() { // Configure logging and environment @@ -43,7 +44,24 @@ async fn main() { panic!("Mongo is not in use, can't connect via authifier!") } - let mut connections: Vec<(Channel, Connection)> = Vec::new(); + let config = config().await; + + let connection = Arc::new( + Connection::connect( + &format!( + "amqp://{}:{}@{}:{}", + &config.rabbit.username, + &config.rabbit.password, + &config.rabbit.host, + &config.rabbit.port, + ), + ConnectionProperties::default(), + ) + .await + .expect("Failed to connect to RabbitMQ"), + ); + + let mut channels = Vec::new(); // An explainer of how this works: // The inbound connections are on separate routing keys, such that they only receive the proper payload @@ -54,171 +72,178 @@ async fn main() { // This'll require some interesting shimming if we need to add more events once this is in prod (different payloads between prod and test), // but that sounds like a problem for future us. - let config = config().await; - - // inbound: generic - connections.push( - make_queue_and_consume( + channels.push( + make_queue_and_consume::( + &db, + &authifier, + &connection, &config, &config.pushd.generic_queue, - config.pushd.get_generic_routing_key().as_str(), + &config.pushd.get_generic_routing_key(), None, - GenericConsumer::new(db.clone(), authifier.clone()), ) .await, ); - // inbound: messages - connections.push( - make_queue_and_consume( + channels.push( + make_queue_and_consume::( + &db, + &authifier, + &connection, &config, &config.pushd.message_queue, - config.pushd.get_message_routing_key().as_str(), + &config.pushd.get_message_routing_key(), None, - MessageConsumer::new(db.clone(), authifier.clone()), ) .await, ); - // inbound: FR received - connections.push( - make_queue_and_consume( + channels.push( + make_queue_and_consume::( + &db, + &authifier, + &connection, &config, &config.pushd.fr_received_queue, - config.pushd.get_fr_received_routing_key().as_str(), + &config.pushd.get_fr_received_routing_key(), None, - FRReceivedConsumer::new(db.clone(), authifier.clone()), ) .await, ); - // inbound: FR accepted - connections.push( - make_queue_and_consume( + channels.push( + make_queue_and_consume::( + &db, + &authifier, + &connection, &config, &config.pushd.fr_accepted_queue, - config.pushd.get_fr_accepted_routing_key().as_str(), + &config.pushd.get_fr_accepted_routing_key(), None, - FRAcceptedConsumer::new(db.clone(), authifier.clone()), ) .await, ); - // inbound: Mass Mentions - connections.push( - make_queue_and_consume( + channels.push( + make_queue_and_consume::( + &db, + &authifier, + &connection, &config, &config.pushd.mass_mention_queue, - config.pushd.get_mass_mention_routing_key().as_str(), + &config.pushd.get_mass_mention_routing_key(), None, - MassMessageConsumer::new(db.clone(), authifier.clone()), ) .await, ); - // inbound: Dm Calls - connections.push( - make_queue_and_consume( + channels.push( + make_queue_and_consume::( + &db, + &authifier, + &connection, &config, &config.pushd.dm_call_queue, - config.pushd.get_dm_call_routing_key().as_str(), + &config.pushd.get_dm_call_routing_key(), None, - DmCallConsumer::new(db.clone(), authifier.clone()), ) .await, ); if !config.pushd.apn.pkcs8.is_empty() { - connections.push( - make_queue_and_consume( + channels.push( + make_queue_and_consume::( + &db, + &authifier, + &connection, &config, &config.pushd.apn.queue, &config.pushd.apn.queue, None, - ApnsOutboundConsumer::new(db.clone()).await.unwrap(), ) .await, ); - let mut table = FieldTable::new(); - table.insert("x-message-deduplication".try_into().unwrap(), "true".into()); + let mut table = FieldTable::default(); + table.insert("x-message-deduplication".into(), AMQPValue::Boolean(true)); - connections.push( - make_queue_and_consume( + channels.push( + make_queue_and_consume::( + &db, + &authifier, + &connection, &config, &config.pushd.ack_queue, &config.pushd.ack_queue, Some(table), - AckConsumer::new(db.clone(), authifier.clone()), ) .await, ); } if !config.pushd.fcm.auth_uri.is_empty() { - connections.push( - make_queue_and_consume( + channels.push( + make_queue_and_consume::( + &db, + &authifier, + &connection, &config, &config.pushd.fcm.queue, &config.pushd.fcm.queue, None, - FcmOutboundConsumer::new(db.clone()).await.unwrap(), ) .await, - ) + ); } if !config.pushd.vapid.public_key.is_empty() { - connections.push( - make_queue_and_consume( + channels.push( + make_queue_and_consume::( + &db, + &authifier, + &connection, &config, &config.pushd.vapid.queue, &config.pushd.vapid.queue, None, - VapidOutboundConsumer::new(db.clone()).await.unwrap(), ) .await, - ) + ); } - let guard = Notify::new(); - guard.notified().await; + ctrl_c().await.unwrap(); - for (channel, conn) in connections { - channel.close().await.expect("Unable to close channel"); - conn.close().await.expect("Unable to close connection"); + for channel in channels { + let _ = channel.close(0, "close".into()).await; } } async fn make_queue_and_consume( + db: &Database, + authifier_db: &authifier::Database, + connection: &Arc, config: &Settings, queue_name: &str, routing_key: &str, queue_args: Option, - consumer: F, -) -> (Channel, Connection) +) -> Arc where - F: AsyncConsumer + Send + 'static, + F: Consumer, { - let connection = Connection::open(&OpenConnectionArguments::new( - &config.rabbit.host, - config.rabbit.port, - &config.rabbit.username, - &config.rabbit.password, - )) - .await - .unwrap(); - - let channel = connection.open_channel(None).await.unwrap(); + let channel = Arc::new(connection.create_channel().await.unwrap()); channel .exchange_declare( - ExchangeDeclareArguments::new(&config.pushd.exchange, "direct") - .durable(true) - .finish(), + config.pushd.exchange.clone().into(), + lapin::ExchangeKind::Direct, + ExchangeDeclareOptions { + durable: true, + ..Default::default() + }, + FieldTable::default(), ) .await - .expect("Failed to declare pushd exchange"); + .expect("Failed to declare exchange"); let mut queue_name = queue_name.to_string(); @@ -230,35 +255,59 @@ where let queue_name = queue_name.as_str(); - let mut args = QueueDeclareArguments::new(queue_name); - args.durable(true); - - if let Some(arg) = queue_args { - args.arguments(arg); - } - - let args = args.finish(); - _ = channel.queue_declare(args).await.unwrap().unwrap(); + let args = QueueDeclareOptions { + durable: true, + ..Default::default() + }; channel - .queue_bind(QueueBindArguments::new( - queue_name, - &config.pushd.exchange, - routing_key, - )) + .queue_declare(queue_name.into(), args, queue_args.unwrap_or_default()) + .await + .unwrap(); + + channel + .queue_bind( + queue_name.into(), + config.pushd.exchange.clone().into(), + routing_key.into(), + QueueBindOptions::default(), + FieldTable::default(), + ) .await .expect( "This probably means the revolt.notifications exchange does not exist in rabbitmq!", ); - let args = BasicConsumeArguments::new(queue_name, "") - .manual_ack(false) - .finish(); - - let routing_key = channel.basic_consume(consumer, args).await.unwrap(); + let consumer = channel + .basic_consume( + queue_name.into(), + "".into(), + BasicConsumeOptions { + no_ack: true, + ..Default::default() + }, + FieldTable::default(), + ) + .await + .unwrap(); info!( "Consuming routing key {} as queue {}, tag {}", - routing_key, queue_name, routing_key + routing_key, + queue_name, + consumer.tag() ); - (channel, connection) + + let delegate = Delegate( + F::create( + db.clone(), + authifier_db.clone(), + connection.clone(), + channel.clone(), + ) + .await, + ); + + consumer.set_delegate(delegate); + + channel } diff --git a/crates/daemons/pushd/src/utils/consumer.rs b/crates/daemons/pushd/src/utils/consumer.rs new file mode 100644 index 00000000..9006aa32 --- /dev/null +++ b/crates/daemons/pushd/src/utils/consumer.rs @@ -0,0 +1,91 @@ +use std::{ + future::{ready, Future}, + pin::Pin, + sync::Arc, +}; + +use anyhow::Result; +use async_trait::async_trait; +use lapin::{ + message::{Delivery, DeliveryResult}, + options::BasicPublishOptions, + BasicProperties, Channel, Connection, ConsumerDelegate, Error as AMQPError, +}; +use log::debug; +use revolt_database::Database; + +#[async_trait] +pub trait Consumer: Clone + Send + Sync + 'static { + async fn create( + db: Database, + authifier_db: authifier::Database, + connection: Arc, + channel: Arc, + ) -> Self; + fn channel(&self) -> &Arc; + async fn consume(&self, delivery: Delivery) -> Result<()>; + + async fn publish_message_with_options( + &self, + payload: &[u8], + exchange: &str, + routing_key: &str, + options: BasicPublishOptions, + properties: BasicProperties, + ) -> Result<(), AMQPError> { + let channel = self.channel(); + + channel + .basic_publish( + exchange.into(), + routing_key.into(), + options, + payload, + properties, + ) + .await?; + debug!("Sent message to queue for target {}", routing_key); + + Ok(()) + } + + async fn publish_message( + &self, + payload: &[u8], + exchange: &str, + routing_key: &str, + ) -> Result<(), AMQPError> { + self.publish_message_with_options( + payload, + exchange, + routing_key, + BasicPublishOptions::default(), + BasicProperties::default(), + ) + .await + } +} + +pub struct Delegate(pub C); + +impl ConsumerDelegate for Delegate { + fn on_new_delivery( + &self, + delivery: DeliveryResult, + ) -> Pin + Send>> { + match delivery { + Ok(Some(delivery)) => { + let consumer = self.0.clone(); + + Box::pin(async move { + if let Err(e) = consumer.consume(delivery).await { + revolt_config::capture_anyhow(&e); + log::error!("{e:?}"); + }; + }) + } + Ok(None) => Box::pin(ready(())), + Err(e) => Box::pin(async move { log::error!("Received bad delivery: {e:?}") }), + } + } +} diff --git a/crates/daemons/pushd/src/utils/mod.rs b/crates/daemons/pushd/src/utils/mod.rs index 62032e80..2a6866e0 100644 --- a/crates/daemons/pushd/src/utils/mod.rs +++ b/crates/daemons/pushd/src/utils/mod.rs @@ -1,2 +1,5 @@ mod renderer; +mod consumer; + pub use renderer::render_notification_content; +pub use consumer::{Consumer, Delegate}; \ No newline at end of file diff --git a/crates/daemons/voice-ingress/Cargo.toml b/crates/daemons/voice-ingress/Cargo.toml index 05ea4a60..e6996742 100644 --- a/crates/daemons/voice-ingress/Cargo.toml +++ b/crates/daemons/voice-ingress/Cargo.toml @@ -44,6 +44,3 @@ revolt-permissions = { workspace = true } livekit-api = { workspace = true } livekit-protocol = { workspace = true } livekit-runtime = { workspace = true, features = ["tokio"] } - -# RabbitMQ -amqprs = { workspace = true } diff --git a/crates/daemons/voice-ingress/src/main.rs b/crates/daemons/voice-ingress/src/main.rs index 727e1a25..45191c2c 100644 --- a/crates/daemons/voice-ingress/src/main.rs +++ b/crates/daemons/voice-ingress/src/main.rs @@ -13,7 +13,7 @@ mod guard; async fn main() -> Result<(), rocket::Error> { revolt_config::configure!(voice_ingress); - let amqp = AMQP::new_auto().await.unwrap(); + let amqp = AMQP::new_auto().await; let database = DatabaseInfo::Auto.connect().await.unwrap(); let voice_client = VoiceClient::from_revolt_config().await; diff --git a/crates/delta/Cargo.toml b/crates/delta/Cargo.toml index a74c5547..535da2af 100644 --- a/crates/delta/Cargo.toml +++ b/crates/delta/Cargo.toml @@ -64,7 +64,7 @@ schemars = { workspace = true } revolt_rocket_okapi = { workspace = true, features = ["swagger"] } # rabbit -amqprs = { workspace = true } +lapin = { workspace = true, features = ["tokio"] } # core authifier = { workspace = true } diff --git a/crates/delta/src/main.rs b/crates/delta/src/main.rs index 41f8fc02..c1dd1dac 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, events::client::EventV1}; use revolt_ratelimits::rocket as ratelimiter; use rocket::{Build, Rocket}; use rocket_cors::{AllowedOrigins, CorsOptions}; @@ -18,10 +17,6 @@ use rocket_prometheus::PrometheusMetrics; use std::net::Ipv4Addr; use std::str::FromStr; -use amqprs::{ - channel::ExchangeDeclareArguments, - connection::{Connection, OpenConnectionArguments}, -}; use async_std::channel::unbounded; use authifier::AuthifierEvent; use revolt_database::voice::VoiceClient; @@ -36,7 +31,6 @@ pub async fn web() -> Rocket { // Setup database let db = revolt_database::DatabaseInfo::Auto.connect().await.unwrap(); - log::info!("database_here {db:?}"); db.migrate_database().await.unwrap(); // Setup Authifier event channel @@ -96,33 +90,8 @@ pub async fn web() -> Rocket { // Voice handler let voice_client = VoiceClient::new(config.api.livekit.nodes.clone()); // Configure Rabbit - let connection = Connection::open(&OpenConnectionArguments::new( - &config.rabbit.host, - config.rabbit.port, - &config.rabbit.username, - &config.rabbit.password, - )) - .await - .expect("Failed to connect to RabbitMQ"); - let channel = connection - .open_channel(None) - .await - .expect("Failed to open RabbitMQ channel"); - - channel - .exchange_declare( - ExchangeDeclareArguments::new(&config.pushd.exchange, "direct") - .durable(true) - .finish(), - ) - .await - .expect("Failed to declare exchange"); - - let mut amqp = AMQP::new(connection, channel); - // amqp.configure_channels() - // .await - // .expect("Failed to configure channels"); + let amqp = AMQP::new_auto().await; // Launch background task workers revolt_database::tasks::start_workers(db.clone(), amqp.clone()); diff --git a/crates/delta/src/util/test.rs b/crates/delta/src/util/test.rs index a6251509..3c8d6835 100644 --- a/crates/delta/src/util/test.rs +++ b/crates/delta/src/util/test.rs @@ -4,7 +4,7 @@ use authifier::{ }; use futures::StreamExt; use rand::Rng; -use redis_kiss::redis::aio::PubSub; +use redis_kiss::{redis::aio::PubSub}; use revolt_database::{ events::client::EventV1, Channel, Database, Member, Message, PartialRole, Server, User, AMQP, }; @@ -25,8 +25,6 @@ pub struct TestHarness { impl TestHarness { pub async fn new() -> TestHarness { - let config = revolt_config::config().await; - let client = Client::tracked(crate::web().await) .await .expect("valid rocket instance"); @@ -49,19 +47,7 @@ impl TestHarness { .expect("`Authifier`") .clone(); - let connection = amqprs::connection::Connection::open( - &amqprs::connection::OpenConnectionArguments::new( - &config.rabbit.host, - config.rabbit.port, - &config.rabbit.username, - &config.rabbit.password, - ), - ) - .await - .unwrap(); - let channel = connection.open_channel(None).await.unwrap(); - - let amqp = AMQP::new(connection, channel); + let amqp = AMQP::new_auto().await; TestHarness { client,