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

Feat: Push notification server (#387)

Browse Source 
* feat: create base of push daemon

Signed-off-by: IAmTomahawkx <iamtomahawkx@gmail.com>

* Add outbound senders

* Make web_push send to rabbit instead (temp stuff)

* feat: stability and friend requests

* make vapid fr stuff not suck

* swap naming of queue

* move pushd into daemons folder

* fix cargo file for move into daemons folder

* feat: probably working fcm push notifs

* comment out fcm webpush stuff since the config keys dont exist

* fix fcm, name queues according to their prod status and configure routing keys

* add pushd to docker

* mix: Remove old code, add stuff to pushd

* fix: lockfile

* feat: update rocket to 5.0.1

* fix: fix queues and ack bugs

* Move rabbit messsage processing into ack queue

* chore: update readme

* chore: optimizations for ack database hits

* pushd flowchart

* misc: update flowchart

* exit dependancy hell

* add rocket_impl flag to authifier

* make the tests file of delta actually compile

* fix: don't silence every push message

* fix: don't silence all messages

* add debug logging for sending data to rabbit from message events

* validate mentions at a server membership level

* put back that import that was actually important

* minor fix to lockfile

* update delta authifier

* feat: proper permissions for push notifications

* add unit test for mention sanitization

* remove local file dependancy on authifier

* update ports to proper defaults

* fixTM the node bindings

* Theoretically configure docker releases for pushd

Signed-off-by: IAmTomahawkx <iamtomahawkx@gmail.com>

* declare exchange in pushd and delta

* fix createbuckets script

Signed-off-by: IAmTomahawkx <iamtomahawkx@gmail.com>

* fix: reference db implementation

Signed-off-by: IAmTomahawkx <iamtomahawkx@gmail.com>

* fix: remove finally redundant code

Signed-off-by: IAmTomahawkx <iamtomahawkx@gmail.com>

* fix: changes

Signed-off-by: IAmTomahawkx <iamtomahawkx@gmail.com>

* fix: other changes

Signed-off-by: IAmTomahawkx <iamtomahawkx@gmail.com>

* fix: make channel name return result

Signed-off-by: IAmTomahawkx <iamtomahawkx@gmail.com>

---------

Signed-off-by: IAmTomahawkx <iamtomahawkx@gmail.com>
This commit is contained in:
Tom
2024-11-28 13:34:17 -08:00
committed by GitHub
parent ed5ded5e45
commit b55765d7c7
71 changed files with 3463 additions and 1059 deletions
Download Patch File Download Diff File Expand all files Collapse all files

138
crates/daemons/pushd/src/consumers/inbound/ack.rs Normal file
View File

@@ -0,0 +1,138 @@
use crate::consumers::inbound::internal::*;
use amqprs::{
channel::{BasicPublishArguments, Channel},
connection::Connection,
consumer::AsyncConsumer,
BasicProperties, Deliver,
};
use async_trait::async_trait;
use revolt_database::{events::rabbit::*, Database};
pub struct AckConsumer {
#[allow(dead_code)]
db: Database,
authifier_db: authifier::Database,
conn: Option<Connection>,
channel: Option<Channel>,
}
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 {
db,
authifier_db,
conn: None,
channel: None,
}
}
}
#[allow(unused_variables)]
#[async_trait]
impl AsyncConsumer for AckConsumer {
/// 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<u8>,
) {
let content = String::from_utf8(content).unwrap();
let payload: AckPayload = serde_json::from_str(content.as_str()).unwrap();
// 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;
if let Ok(u) = &unreads {
if u.is_empty() {
return;
}
} else {
return;
}
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()
.filter(|session| {
if let Some(sub) = &session.subscription {
sub.endpoint == "apn"
} else {
false
}
})
.collect();
if apple_sessions.is_empty() {
return;
}
// Step 3: calculate the actual mention count, since we have to send it out
let mut mention_count = 0;
for u in &unreads.unwrap() {
mention_count += u.mentions.as_ref().unwrap().len()
}
// Step 4: loop through each apple session and send the badge update
for session in apple_sessions {
let service_payload = PayloadToService {
notification: PayloadKind::BadgeUpdate(mention_count),
user_id: payload.user_id.clone(),
session_id: session.id.clone(),
token: session.subscription.as_ref().unwrap().auth.clone(),
extras: Default::default(),
};
let raw_service_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
);
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());
}
}
}
}
}

121
crates/daemons/pushd/src/consumers/inbound/fr_accepted.rs Normal file
View File

@@ -0,0 +1,121 @@
use std::collections::HashMap;
use crate::consumers::inbound::internal::*;
use amqprs::{
channel::{BasicPublishArguments, Channel},
connection::Connection,
consumer::AsyncConsumer,
BasicProperties, Deliver,
};
use async_trait::async_trait;
use log::debug;
use revolt_database::{events::rabbit::*, Database};
pub struct FRAcceptedConsumer {
#[allow(dead_code)]
db: Database,
authifier_db: authifier::Database,
conn: Option<Connection>,
channel: Option<Channel>,
}
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 {
db,
authifier_db,
conn: None,
channel: None,
}
}
}
#[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<u8>,
) {
let content = String::from_utf8(content).unwrap();
let payload: FRAcceptedPayload = serde_json::from_str(content.as_str()).unwrap();
debug!("Received FR accept event");
if let Ok(sessions) = self.authifier_db.find_sessions(&payload.user).await {
let config = revolt_config::config().await;
for session in sessions {
if let Some(sub) = session.subscription {
let mut sendable = PayloadToService {
notification: PayloadKind::FRAccepted(payload.clone()),
token: sub.auth,
user_id: session.user_id,
session_id: session.id,
extras: HashMap::new(),
};
let args: BasicPublishArguments;
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("p265dh".to_string(), sub.p256dh);
sendable
.extras
.insert("endpoint".to_string(), sub.endpoint.clone());
}
let payload = serde_json::to_string(&sendable).unwrap();
publish_message(self, payload.into(), args).await;
}
}
}
}
}

121
crates/daemons/pushd/src/consumers/inbound/fr_received.rs Normal file
View File

@@ -0,0 +1,121 @@
use std::collections::HashMap;
use crate::consumers::inbound::internal::*;
use amqprs::{
channel::{BasicPublishArguments, Channel},
connection::Connection,
consumer::AsyncConsumer,
BasicProperties, Deliver,
};
use async_trait::async_trait;
use log::debug;
use revolt_database::{events::rabbit::*, Database};
pub struct FRReceivedConsumer {
#[allow(dead_code)]
db: Database,
authifier_db: authifier::Database,
conn: Option<Connection>,
channel: Option<Channel>,
}
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 {
db,
authifier_db,
conn: None,
channel: None,
}
}
}
#[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<u8>,
) {
let content = String::from_utf8(content).unwrap();
let payload: FRReceivedPayload = serde_json::from_str(content.as_str()).unwrap();
debug!("Received FR received event");
if let Ok(sessions) = self.authifier_db.find_sessions(&payload.user).await {
let config = revolt_config::config().await;
for session in sessions {
if let Some(sub) = session.subscription {
let mut sendable = PayloadToService {
notification: PayloadKind::FRReceived(payload.clone()),
token: sub.auth,
user_id: session.user_id,
session_id: session.id,
extras: HashMap::new(),
};
let args: BasicPublishArguments;
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("p265dh".to_string(), sub.p256dh);
sendable
.extras
.insert("endpoint".to_string(), sub.endpoint.clone());
}
let payload = serde_json::to_string(&sendable).unwrap();
publish_message(self, payload.into(), args).await;
}
}
}
}
}

127
crates/daemons/pushd/src/consumers/inbound/generic.rs Normal file
View File

@@ -0,0 +1,127 @@
use std::collections::HashMap;
use crate::consumers::inbound::internal::*;
use amqprs::{
channel::{BasicPublishArguments, Channel},
connection::Connection,
consumer::AsyncConsumer,
BasicProperties, Deliver,
};
use async_trait::async_trait;
use log::debug;
use revolt_database::{events::rabbit::*, Database};
pub struct GenericConsumer {
#[allow(dead_code)]
db: Database,
authifier_db: authifier::Database,
conn: Option<Connection>,
channel: Option<Channel>,
}
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 {
db,
authifier_db,
conn: None,
channel: None,
}
}
}
#[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<u8>,
) {
let content = String::from_utf8(content).unwrap();
let payload: MessageSentPayload = serde_json::from_str(content.as_str()).unwrap();
debug!("Received message event on origin");
if let Ok(sessions) = self
.authifier_db
.find_sessions_with_subscription(&payload.users)
.await
{
let config = revolt_config::config().await;
for session in sessions {
if let Some(sub) = session.subscription {
let mut sendable = PayloadToService {
notification: PayloadKind::MessageNotification(
payload.notification.clone(),
),
token: sub.auth,
user_id: session.user_id,
session_id: session.id,
extras: HashMap::new(),
};
let args: BasicPublishArguments;
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("p265dh".to_string(), sub.p256dh);
sendable
.extras
.insert("endpoint".to_string(), sub.endpoint.clone());
}
let payload = serde_json::to_string(&sendable).unwrap();
publish_message(self, payload.into(), args).await;
}
}
}
}
}

53
crates/daemons/pushd/src/consumers/inbound/internal.rs Normal file
View File

@@ -0,0 +1,53 @@
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<T: Channeled>(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<T: Channeled>(
consumer: &mut T,
payload: Vec<u8>,
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)!")
}
}

127
crates/daemons/pushd/src/consumers/inbound/message.rs Normal file
View File

@@ -0,0 +1,127 @@
use std::collections::HashMap;
use crate::consumers::inbound::internal::*;
use amqprs::{
channel::{BasicPublishArguments, Channel},
connection::Connection,
consumer::AsyncConsumer,
BasicProperties, Deliver,
};
use async_trait::async_trait;
use log::debug;
use revolt_database::{events::rabbit::*, Database};
pub struct MessageConsumer {
#[allow(dead_code)]
db: Database,
authifier_db: authifier::Database,
conn: Option<Connection>,
channel: Option<Channel>,
}
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 {
db,
authifier_db,
conn: None,
channel: None,
}
}
}
#[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<u8>,
) {
let content = String::from_utf8(content).unwrap();
let payload: MessageSentPayload = serde_json::from_str(content.as_str()).unwrap();
debug!("Received message event on origin");
if let Ok(sessions) = self
.authifier_db
.find_sessions_with_subscription(&payload.users)
.await
{
let config = revolt_config::config().await;
for session in sessions {
if let Some(sub) = session.subscription {
let mut sendable = PayloadToService {
notification: PayloadKind::MessageNotification(
payload.notification.clone(),
),
token: sub.auth,
user_id: session.user_id,
session_id: session.id,
extras: HashMap::new(),
};
let args: BasicPublishArguments;
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("p265dh".to_string(), sub.p256dh);
sendable
.extras
.insert("endpoint".to_string(), sub.endpoint.clone());
}
let payload = serde_json::to_string(&sendable).unwrap();
publish_message(self, payload.into(), args).await;
}
}
}
}
}

6
crates/daemons/pushd/src/consumers/inbound/mod.rs Normal file
View File

@@ -0,0 +1,6 @@
pub mod ack;
pub mod fr_accepted;
pub mod fr_received;
pub mod generic;
mod internal;
pub mod message;

2
crates/daemons/pushd/src/consumers/mod.rs Normal file
View File

@@ -0,0 +1,2 @@
pub mod inbound;
pub mod outbound;

338
crates/daemons/pushd/src/consumers/outbound/apn.rs Normal file
View File

@@ -0,0 +1,338 @@
use std::{borrow::Cow, collections::BTreeMap, io::Cursor};
use amqprs::{channel::Channel as AmqpChannel, consumer::AsyncConsumer, BasicProperties, Deliver};
use async_trait::async_trait;
use base64::{
engine::{self},
Engine as _,
};
use revolt_a2::{
request::{
notification::{DefaultAlert, NotificationOptions},
payload::{APSAlert, APSSound, Payload, PayloadLike, APS},
},
Client, ClientConfig, Endpoint, Error, ErrorBody, ErrorReason, Priority, PushType, Response,
};
use revolt_database::{events::rabbit::*, Database};
use revolt_models::v0::{Channel, Message, PushNotification};
use serde::Serialize;
// region: payload
#[derive(Serialize, Debug)]
struct MessagePayload<'a> {
aps: APS<'a>,
#[serde(skip_serializing)]
options: NotificationOptions<'a>,
#[serde(skip_serializing)]
device_token: &'a str,
message: &'a Message,
url: &'a str,
#[serde(rename = "camelCase")]
author_avatar: &'a str,
#[serde(rename = "camelCase")]
author_display_name: &'a str,
#[serde(rename = "camelCase")]
channel_name: &'a str,
}
impl<'a> PayloadLike for MessagePayload<'a> {
fn get_device_token(&self) -> &'a str {
self.device_token
}
fn get_options(&self) -> &NotificationOptions {
&self.options
}
}
// region: consumer
pub struct ApnsOutboundConsumer {
#[allow(dead_code)]
db: Database,
client: Client,
}
impl ApnsOutboundConsumer {
fn format_title(&self, notification: &PushNotification) -> String {
// ideally this changes depending on context
// in a server, it would look like "Sendername, #channelname in servername"
// in a group, it would look like "Sendername in groupname"
// in a dm it should just be "Sendername".
// not sure how feasible all those are given the PushNotification object as it currently stands.
match &notification.channel {
Channel::DirectMessage { .. } => notification.author.clone(),
Channel::Group { name, .. } => format!("{}, #{}", notification.author, name),
Channel::TextChannel { name, .. } | Channel::VoiceChannel { name, .. } => {
format!("{} in #{}", notification.author, name)
}
_ => "Unknown".to_string(),
}
}
async fn get_badge_count(&self, user: &str) -> Option<u32> {
if let Ok(unreads) = self.db.fetch_unread_mentions(user).await {
let mut mention_count = 0;
for channel in unreads {
if let Some(mentions) = channel.mentions {
mention_count += mentions.len() as u32
}
}
println!("Got badge count for APN: {}", mention_count);
return Some(mention_count);
}
None
}
}
impl ApnsOutboundConsumer {
pub async fn new(db: Database) -> Result<ApnsOutboundConsumer, &'static str> {
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.");
}
let endpoint = if config.pushd.apn.sandbox {
Endpoint::Sandbox
} else {
Endpoint::Production
};
let pkcs8 = engine::general_purpose::STANDARD
.decode(config.pushd.apn.pkcs8.clone())
.expect("valid `pcks8`");
let client_config = ClientConfig::new(endpoint);
let client = Client::token(
&mut Cursor::new(pkcs8),
config.pushd.apn.key_id.clone(),
config.pushd.apn.team_id.clone(),
client_config,
)
.expect("could not create APN client");
Ok(ApnsOutboundConsumer { db, client })
}
}
#[allow(unused_variables)]
#[async_trait]
impl AsyncConsumer for ApnsOutboundConsumer {
async fn consume(
&mut self,
channel: &AmqpChannel,
deliver: Deliver,
basic_properties: BasicProperties,
content: Vec<u8>,
) {
let content = String::from_utf8(content).unwrap();
let payload: PayloadToService = serde_json::from_str(content.as_str()).unwrap();
let payload_options = NotificationOptions {
apns_id: None,
apns_push_type: Some(PushType::Alert),
apns_expiration: None,
apns_priority: Some(Priority::High),
apns_topic: Some("chat.revolt.app"),
apns_collapse_id: None,
};
let resp: Result<Response, Error>;
match payload.notification {
PayloadKind::FRReceived(alert) => {
let loc_args = vec![Cow::from(
alert
.from_user
.display_name
.or(Some(format!(
"{}#{}",
alert.from_user.username, alert.from_user.discriminator
)))
.clone()
.unwrap(),
)];
let apn_payload = Payload {
aps: APS {
alert: Some(APSAlert::Default(DefaultAlert {
title: None,
subtitle: None,
body: None,
title_loc_key: None,
title_loc_args: None,
action_loc_key: None,
loc_key: Some("push.fr.received"),
loc_args: Some(loc_args),
launch_image: None,
})),
badge: self.get_badge_count(&payload.user_id).await,
sound: Some(APSSound::Sound("default")),
thread_id: None,
content_available: None,
category: None,
mutable_content: Some(1),
url_args: None,
},
device_token: &payload.token,
options: payload_options.clone(),
data: BTreeMap::new(),
};
resp = 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.username, alert.accepted_user.discriminator
)))
.clone()
.unwrap(),
)];
let apn_payload = Payload {
aps: APS {
alert: Some(APSAlert::Default(DefaultAlert {
title: None,
subtitle: None,
body: None,
title_loc_key: None,
title_loc_args: None,
action_loc_key: None,
loc_key: Some("push.fr.accepted"),
loc_args: Some(loc_args),
launch_image: None,
})),
badge: self.get_badge_count(&payload.user_id).await,
sound: Some(APSSound::Sound("default")),
thread_id: None,
content_available: None,
category: None,
mutable_content: Some(1),
url_args: None,
},
device_token: &payload.token,
options: payload_options.clone(),
data: BTreeMap::new(),
};
resp = self.client.send(apn_payload).await;
}
PayloadKind::Generic(alert) => {
let apn_payload = Payload {
aps: APS {
alert: Some(APSAlert::Default(DefaultAlert {
title: Some(&alert.title),
subtitle: None,
body: Some(&alert.body),
title_loc_key: None,
title_loc_args: None,
action_loc_key: None,
loc_key: None,
loc_args: None,
launch_image: None,
})),
badge: self.get_badge_count(&payload.user_id).await,
sound: Some(APSSound::Sound("default")),
thread_id: None,
content_available: None,
category: None,
mutable_content: Some(1),
url_args: None,
},
device_token: &payload.token,
options: payload_options.clone(),
data: BTreeMap::new(),
};
resp = self.client.send(apn_payload).await;
}
PayloadKind::MessageNotification(alert) => {
let title = self.format_title(&alert);
let apn_payload = MessagePayload {
aps: APS {
alert: Some(APSAlert::Default(DefaultAlert {
title: Some(&title),
subtitle: None,
body: Some(&alert.body),
title_loc_key: None,
title_loc_args: None,
action_loc_key: None,
loc_key: None,
loc_args: None,
launch_image: None,
})),
badge: self.get_badge_count(&payload.user_id).await,
sound: Some(APSSound::Sound("default")),
thread_id: Some(alert.channel.id()),
content_available: None,
category: None,
mutable_content: Some(1),
url_args: None,
},
device_token: &payload.token,
options: payload_options.clone(),
message: &alert.message,
url: &alert.url,
author_avatar: &alert.icon,
author_display_name: &alert.author,
channel_name: alert.channel.name().unwrap_or(&title),
};
resp = self.client.send(apn_payload).await;
}
PayloadKind::BadgeUpdate(badge) => {
let apn_payload = Payload {
aps: APS {
badge: Some(badge as u32),
..Default::default()
},
device_token: &payload.token,
options: payload_options.clone(),
data: BTreeMap::new(),
};
resp = self.client.send(apn_payload).await;
}
}
if let Err(err) = resp {
match err {
Error::ResponseError(Response {
error:
Some(ErrorBody {
reason: ErrorReason::BadDeviceToken | ErrorReason::Unregistered,
..
}),
..
}) => {
if let Err(err) = self
.db
.remove_push_subscription_by_session_id(&payload.session_id)
.await
{
revolt_config::capture_error(&err);
}
}
err => {
revolt_config::capture_error(&err);
}
}
}
}
}

199
crates/daemons/pushd/src/consumers/outbound/fcm.rs Normal file
View File

@@ -0,0 +1,199 @@
use std::{collections::HashMap, time::Duration};
use amqprs::{channel::Channel as AmqpChannel, consumer::AsyncConsumer, BasicProperties, Deliver};
use async_trait::async_trait;
use fcm_v1::{
android::AndroidConfig,
auth::{Authenticator, ServiceAccountKey},
message::{Message, Notification},
Client, Error as FcmError,
};
use revolt_database::{events::rabbit::*, Database};
use revolt_models::v0::{Channel, PushNotification};
use serde_json::Value;
pub struct FcmOutboundConsumer {
db: Database,
client: Client,
}
impl FcmOutboundConsumer {
fn format_title(&self, notification: &PushNotification) -> String {
// ideally this changes depending on context
// in a server, it would look like "Sendername, #channelname in servername"
// in a group, it would look like "Sendername in groupname"
// in a dm it should just be "Sendername".
// not sure how feasible all those are given the PushNotification object as it currently stands.
match &notification.channel {
Channel::DirectMessage { .. } => notification.author.clone(),
Channel::Group { name, .. } => format!("{}, #{}", notification.author, name),
Channel::TextChannel { name, .. } | Channel::VoiceChannel { name, .. } => {
format!("{} in #{}", notification.author, name)
}
_ => "Unknown".to_string(),
}
}
}
impl FcmOutboundConsumer {
pub async fn new(db: Database) -> Result<FcmOutboundConsumer, &'static str> {
let config = revolt_config::config().await;
Ok(FcmOutboundConsumer {
db,
client: Client::new(
Authenticator::service_account::<&str>(ServiceAccountKey {
key_type: Some(config.pushd.fcm.key_type),
project_id: Some(config.pushd.fcm.project_id.clone()),
private_key_id: Some(config.pushd.fcm.private_key_id),
private_key: config.pushd.fcm.private_key,
client_email: config.pushd.fcm.client_email,
client_id: Some(config.pushd.fcm.client_id),
auth_uri: Some(config.pushd.fcm.auth_uri),
token_uri: config.pushd.fcm.token_uri,
auth_provider_x509_cert_url: Some(config.pushd.fcm.auth_provider_x509_cert_url),
client_x509_cert_url: Some(config.pushd.fcm.client_x509_cert_url),
})
.await
.unwrap(),
config.pushd.fcm.project_id,
false,
Duration::from_secs(5),
),
})
}
}
#[allow(unused_variables)]
#[async_trait]
impl AsyncConsumer for FcmOutboundConsumer {
async fn consume(
&mut self,
channel: &AmqpChannel,
deliver: Deliver,
basic_properties: BasicProperties,
content: Vec<u8>,
) {
let content = String::from_utf8(content).unwrap();
let payload: PayloadToService = serde_json::from_str(content.as_str()).unwrap();
let config = revolt_config::config().await;
#[allow(clippy::needless_late_init)]
let resp: Result<Message, FcmError>;
match payload.notification {
PayloadKind::FRReceived(alert) => {
let name = alert
.from_user
.display_name
.or(Some(format!(
"{}#{}",
alert.from_user.username, alert.from_user.discriminator
)))
.clone()
.unwrap();
let mut data = HashMap::new();
data.insert(
"type".to_string(),
Value::String("push.fr.receive".to_string()),
);
data.insert("id".to_string(), Value::String(alert.from_user.id));
data.insert("username".to_string(), Value::String(name));
let msg = Message {
token: Some(payload.token),
data: Some(data),
..Default::default()
};
resp = self.client.send(&msg).await;
}
PayloadKind::FRAccepted(alert) => {
let name = alert
.accepted_user
.display_name
.or(Some(format!(
"{}#{}",
alert.accepted_user.username, alert.accepted_user.discriminator
)))
.clone()
.unwrap();
let mut data: HashMap<String, Value> = HashMap::new();
data.insert(
"type".to_string(),
Value::String("push.fr.accept".to_string()),
);
data.insert("id".to_string(), Value::String(alert.accepted_user.id));
data.insert("username".to_string(), Value::String(name));
let msg = Message {
token: Some(payload.token),
data: Some(data),
..Default::default()
};
resp = self.client.send(&msg).await;
}
PayloadKind::Generic(alert) => {
let msg = Message {
token: Some(payload.token),
notification: Some(Notification {
title: Some(alert.title),
body: Some(alert.body),
image: alert.icon,
}),
..Default::default()
};
resp = self.client.send(&msg).await;
}
PayloadKind::MessageNotification(alert) => {
let title = self.format_title(&alert);
let msg = Message {
token: Some(payload.token),
notification: Some(Notification {
title: Some(title),
body: Some(alert.body),
image: Some(alert.icon),
}),
android: Some(AndroidConfig {
collapse_key: Some(alert.tag),
..Default::default()
}),
..Default::default()
};
resp = self.client.send(&msg).await;
}
PayloadKind::BadgeUpdate(_) => {
panic!("FCM cannot handle badge updates, and they should not be sent here.")
}
}
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 => {
revolt_config::capture_error(&err);
}
}
}
}
}

3
crates/daemons/pushd/src/consumers/outbound/mod.rs Normal file
View File

@@ -0,0 +1,3 @@
pub mod apn;
pub mod fcm;
pub mod vapid;

149
crates/daemons/pushd/src/consumers/outbound/vapid.rs Normal file
View File

@@ -0,0 +1,149 @@
use std::collections::HashMap;
use amqprs::{channel::Channel as AmqpChannel, consumer::AsyncConsumer, BasicProperties, Deliver};
use async_trait::async_trait;
use base64::{
engine::{self},
Engine as _,
};
use revolt_database::{events::rabbit::*, Database};
// use revolt_models::v0::{Channel, PushNotification};
use web_push::{
ContentEncoding, IsahcWebPushClient, SubscriptionInfo, SubscriptionKeys, VapidSignatureBuilder,
WebPushClient, WebPushError, WebPushMessageBuilder,
};
pub struct VapidOutboundConsumer {
db: Database,
client: IsahcWebPushClient,
pkey: Vec<u8>,
}
impl VapidOutboundConsumer {
pub async fn new(db: Database) -> Result<VapidOutboundConsumer, &'static str> {
let config = revolt_config::config().await;
if config.pushd.vapid.private_key.is_empty() | config.pushd.vapid.public_key.is_empty() {
return Err("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`");
Ok(VapidOutboundConsumer {
db,
client: IsahcWebPushClient::new().unwrap(),
pkey: web_push_private_key,
})
}
}
#[allow(unused_variables)]
#[async_trait]
impl AsyncConsumer for VapidOutboundConsumer {
async fn consume(
&mut self,
channel: &AmqpChannel,
deliver: Deliver,
basic_properties: BasicProperties,
content: Vec<u8>,
) {
let content = String::from_utf8(content).unwrap();
let payload: PayloadToService = serde_json::from_str(content.as_str()).unwrap();
let config = revolt_config::config().await;
let subscription = SubscriptionInfo {
endpoint: payload.extras.get("endpoint").unwrap().clone(),
keys: SubscriptionKeys {
auth: payload.token,
p256dh: payload.extras.get("p256dh").unwrap().clone(),
},
};
#[allow(clippy::needless_late_init)]
let payload_body: String;
match payload.notification {
PayloadKind::FRReceived(alert) => {
let name = alert
.from_user
.display_name
.or(Some(format!(
"{}#{}",
alert.from_user.username, alert.from_user.discriminator
)))
.clone()
.unwrap();
let mut body = HashMap::new();
body.insert("body", format!("{} sent you a friend request", name));
payload_body = serde_json::to_string(&body).unwrap();
}
PayloadKind::FRAccepted(alert) => {
let name = alert
.accepted_user
.display_name
.or(Some(format!(
"{}#{}",
alert.accepted_user.username, alert.accepted_user.discriminator
)))
.clone()
.unwrap();
let mut body = HashMap::new();
body.insert("body", format!("{} accepted your friend request", name));
payload_body = serde_json::to_string(&body).unwrap();
}
PayloadKind::Generic(alert) => {
payload_body = serde_json::to_string(&alert).unwrap();
}
PayloadKind::MessageNotification(alert) => {
payload_body = serde_json::to_string(&alert).unwrap();
}
PayloadKind::BadgeUpdate(_) => {
panic!("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);
builder.set_payload(ContentEncoding::AesGcm, payload_body.as_bytes());
match builder.build() {
Ok(msg) => {
if let Err(err) = self.client.send(msg).await {
if 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) => {
revolt_config::capture_error(&err);
}
}
}
Err(err) => {
revolt_config::capture_error(&err);
}
},
Err(err) => {
revolt_config::capture_error(&err);
}
}
}
}

233
crates/daemons/pushd/src/main.rs Normal file
View File

@@ -0,0 +1,233 @@
use amqprs::{
channel::{
BasicConsumeArguments, Channel, ExchangeDeclareArguments, QueueBindArguments,
QueueDeclareArguments,
},
connection::{Connection, OpenConnectionArguments},
consumer::AsyncConsumer,
FieldTable,
};
use revolt_config::{config, Settings};
use tokio::sync::Notify;
mod consumers;
use consumers::{
inbound::{
ack::AckConsumer, fr_accepted::FRAcceptedConsumer, fr_received::FRReceivedConsumer,
generic::GenericConsumer, message::MessageConsumer,
},
outbound::{apn::ApnsOutboundConsumer, fcm::FcmOutboundConsumer, vapid::VapidOutboundConsumer},
};
#[tokio::main(flavor = "multi_thread", worker_threads = 2)]
async fn main() {
let config = config().await;
// Setup database
let db = revolt_database::DatabaseInfo::Auto.connect().await.unwrap();
let authifier: authifier::Database;
if let Some(client) = match &db {
revolt_database::Database::Reference(_) => None,
revolt_database::Database::MongoDb(mongo) => Some(mongo),
} {
authifier =
authifier::Database::MongoDb(authifier::database::MongoDb(client.database("revolt")));
} else {
panic!("Mongo is not in use, can't connect via authifier!")
}
let mut connections: Vec<(Channel, Connection)> = Vec::new();
// An explainer of how this works:
// The inbound connections are on separate routing keys, such that they only receive the proper payload
// from their respective api (prod or test).
// However, the outbound queues that go to the services are routed to receive from both, so that messages
// sent from beta are still notified on prod, and vice versa.
// 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.
// inbound: generic
connections.push(
make_queue_and_consume(
&config,
&config.pushd.generic_queue,
config.pushd.get_generic_routing_key().as_str(),
None,
GenericConsumer::new(db.clone(), authifier.clone()),
)
.await,
);
// inbound: messages
connections.push(
make_queue_and_consume(
&config,
&config.pushd.message_queue,
config.pushd.get_message_routing_key().as_str(),
None,
MessageConsumer::new(db.clone(), authifier.clone()),
)
.await,
);
// inbound: FR received
connections.push(
make_queue_and_consume(
&config,
&config.pushd.fr_received_queue,
config.pushd.get_fr_received_routing_key().as_str(),
None,
FRReceivedConsumer::new(db.clone(), authifier.clone()),
)
.await,
);
// inbound: FR accepted
connections.push(
make_queue_and_consume(
&config,
&config.pushd.fr_accepted_queue,
config.pushd.get_fr_accepted_routing_key().as_str(),
None,
FRAcceptedConsumer::new(db.clone(), authifier.clone()),
)
.await,
);
if !config.pushd.apn.pkcs8.is_empty() {
connections.push(
make_queue_and_consume(
&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());
connections.push(
make_queue_and_consume(
&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(
&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(
&config,
&config.pushd.vapid.queue,
&config.pushd.vapid.queue,
None,
VapidOutboundConsumer::new(db.clone()).await.unwrap(),
)
.await,
)
}
let guard = Notify::new();
guard.notified().await;
for (channel, conn) in connections {
channel.close().await.expect("Unable to close channel");
conn.close().await.expect("Unable to close connection");
}
}
async fn make_queue_and_consume<F>(
config: &Settings,
queue_name: &str,
routing_key: &str,
queue_args: Option<FieldTable>,
consumer: F,
) -> (Channel, Connection)
where
F: AsyncConsumer + Send + 'static,
{
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();
channel
.exchange_declare(
ExchangeDeclareArguments::new(&config.pushd.exchange, "direct")
.durable(true)
.finish(),
)
.await
.expect("Failed to declare pushd exchange");
let mut queue_name = queue_name.to_string();
if config.pushd.production {
queue_name += "-prd";
} else {
queue_name += "-tst";
}
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();
channel
.queue_bind(QueueBindArguments::new(
queue_name,
&config.pushd.exchange,
routing_key,
))
.await
.expect(
"This probably means the revolt.notifications exchange does not exist in rabbitmq!",
);
let args = BasicConsumeArguments::new(queue_name, "")
.manual_ack(false)
.finish();
channel.basic_consume(consumer, args).await.unwrap();
log::info!(
"Consuming routing key {} as queue {}",
routing_key,
queue_name
);
(channel, connection)
}