forked from jmug/stoatchat
Feat: Push notification server (#387)
* 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
GitHub
@@ -1,21 +1,27 @@
|
||||
// Queue Type: Debounced
|
||||
use crate::Database;
|
||||
use crate::{Database, Message, AMQP};
|
||||
|
||||
use deadqueue::limited::Queue;
|
||||
use once_cell::sync::Lazy;
|
||||
use std::{collections::HashMap, time::Duration};
|
||||
use revolt_models::v0::PushNotification;
|
||||
use rocket::form::validate::Contains;
|
||||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
time::Duration,
|
||||
};
|
||||
use validator::HasLen;
|
||||
|
||||
use revolt_result::Result;
|
||||
|
||||
use super::{apple_notifications::{self, ApnJob}, DelayedTask};
|
||||
use super::DelayedTask;
|
||||
|
||||
/// Enumeration of possible events
|
||||
#[derive(Debug, Eq, PartialEq)]
|
||||
pub enum AckEvent {
|
||||
/// Add mentions for a user in a channel
|
||||
AddMention {
|
||||
/// Message IDs
|
||||
ids: Vec<String>,
|
||||
/// Add mentions for a channel
|
||||
ProcessMessage {
|
||||
/// push notification, message, recipients, push silenced
|
||||
messages: Vec<(Option<PushNotification>, Message, Vec<String>, bool)>,
|
||||
},
|
||||
|
||||
/// Acknowledge message in a channel for a user
|
||||
@@ -30,7 +36,7 @@ struct Data {
|
||||
/// Channel to ack
|
||||
channel: String,
|
||||
/// User to ack for
|
||||
user: String,
|
||||
user: Option<String>,
|
||||
/// Event
|
||||
event: AckEvent,
|
||||
}
|
||||
@@ -43,21 +49,49 @@ struct Task {
|
||||
static Q: Lazy<Queue<Data>> = Lazy::new(|| Queue::new(10_000));
|
||||
|
||||
/// Queue a new task for a worker
|
||||
pub async fn queue(channel: String, user: String, event: AckEvent) {
|
||||
pub async fn queue_ack(channel: String, user: String, event: AckEvent) {
|
||||
Q.try_push(Data {
|
||||
channel,
|
||||
user,
|
||||
user: Some(user),
|
||||
event,
|
||||
})
|
||||
.ok();
|
||||
|
||||
info!("Queue is using {} slots from {}.", Q.len(), Q.capacity());
|
||||
info!(
|
||||
"Queue is using {} slots from {}. Queued type: ACK",
|
||||
Q.len(),
|
||||
Q.capacity()
|
||||
);
|
||||
}
|
||||
|
||||
pub async fn handle_ack_event(event: &AckEvent, db: &Database, authifier_db: &authifier::Database, user: &str, channel: &str) -> Result<()> {
|
||||
pub async fn queue_message(channel: String, event: AckEvent) {
|
||||
Q.try_push(Data {
|
||||
channel,
|
||||
user: None,
|
||||
event,
|
||||
})
|
||||
.ok();
|
||||
|
||||
info!(
|
||||
"Queue is using {} slots from {}. Queued type: MENTION",
|
||||
Q.len(),
|
||||
Q.capacity()
|
||||
);
|
||||
}
|
||||
|
||||
pub async fn handle_ack_event(
|
||||
event: &AckEvent,
|
||||
db: &Database,
|
||||
amqp: &AMQP,
|
||||
user: &Option<String>,
|
||||
channel: &str,
|
||||
) -> Result<()> {
|
||||
match &event {
|
||||
#[allow(clippy::disallowed_methods)] // event is sent by higher level function
|
||||
AckEvent::AckMessage { id } => {
|
||||
let user = user.as_ref().unwrap();
|
||||
let user: &str = user.as_str();
|
||||
|
||||
let unread = db.fetch_unread(user, channel).await?;
|
||||
let updated = db.acknowledge_message(channel, user, id).await?;
|
||||
|
||||
@@ -68,21 +102,67 @@ pub async fn handle_ack_event(event: &AckEvent, db: &Database, authifier_db: &au
|
||||
let mentions_acked = before_mentions - after_mentions;
|
||||
|
||||
if mentions_acked > 0 {
|
||||
if let Ok(sessions) = authifier_db.find_sessions(user).await {
|
||||
for session in sessions {
|
||||
if let Some(sub) = session.subscription {
|
||||
if sub.endpoint == "apn" {
|
||||
apple_notifications::queue(ApnJob::from_ack(session.id, user.to_string(), sub.auth)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Err(err) = amqp
|
||||
.ack_message(user.to_string(), channel.to_string(), id.to_owned())
|
||||
.await
|
||||
{
|
||||
revolt_config::capture_error(&err);
|
||||
}
|
||||
};
|
||||
|
||||
}
|
||||
},
|
||||
AckEvent::AddMention { ids } => {
|
||||
db.add_mention_to_unread(channel, user, ids).await?;
|
||||
}
|
||||
AckEvent::ProcessMessage { messages } => {
|
||||
let mut users: HashSet<&String> = HashSet::new();
|
||||
debug!(
|
||||
"Processing {} messages from channel {}",
|
||||
messages.len(),
|
||||
messages[0].1.channel
|
||||
);
|
||||
|
||||
// find all the users we'll be notifying
|
||||
messages
|
||||
.iter()
|
||||
.for_each(|(_, _, recipents, _)| users.extend(recipents.iter()));
|
||||
|
||||
debug!("Found {} users to notify.", users.len());
|
||||
|
||||
for user in users {
|
||||
let message_ids: Vec<String> = messages
|
||||
.iter()
|
||||
.filter(|(_, _, recipients, _)| recipients.contains(user))
|
||||
.map(|(_, message, _, _)| message.id.clone())
|
||||
.collect();
|
||||
|
||||
if !message_ids.is_empty() {
|
||||
db.add_mention_to_unread(channel, user, &message_ids)
|
||||
.await?;
|
||||
}
|
||||
debug!("Added {} mentions for user {}", message_ids.len(), &user);
|
||||
}
|
||||
|
||||
for (push, _, recipients, silenced) in messages {
|
||||
if *silenced || recipients.is_empty() || push.is_none() {
|
||||
debug!(
|
||||
"Rejecting push: silenced: {}, recipient count: {}, push exists: {:?}",
|
||||
*silenced,
|
||||
recipients.length(),
|
||||
push
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
debug!(
|
||||
"Sending push event to AMQP; message {} for {} users",
|
||||
push.as_ref().unwrap().message.id,
|
||||
recipients.len()
|
||||
);
|
||||
if let Err(err) = amqp
|
||||
.message_sent(recipients.clone(), push.clone().unwrap())
|
||||
.await
|
||||
{
|
||||
revolt_config::capture_error(&err);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
@@ -90,9 +170,9 @@ pub async fn handle_ack_event(event: &AckEvent, db: &Database, authifier_db: &au
|
||||
}
|
||||
|
||||
/// Start a new worker
|
||||
pub async fn worker(db: Database, authifier_db: authifier::Database) {
|
||||
let mut tasks = HashMap::<(String, String), DelayedTask<Task>>::new();
|
||||
let mut keys = vec![];
|
||||
pub async fn worker(db: Database, amqp: AMQP) {
|
||||
let mut tasks = HashMap::<(Option<String>, String, u8), DelayedTask<Task>>::new();
|
||||
let mut keys: Vec<(Option<String>, String, u8)> = vec![];
|
||||
|
||||
loop {
|
||||
// Find due tasks.
|
||||
@@ -106,12 +186,13 @@ pub async fn worker(db: Database, authifier_db: authifier::Database) {
|
||||
for key in &keys {
|
||||
if let Some(task) = tasks.remove(key) {
|
||||
let Task { event } = task.data;
|
||||
let (user, channel) = key;
|
||||
let (user, channel, _) = key;
|
||||
|
||||
if let Err(err) = handle_ack_event(&event, &db, &authifier_db, user, channel).await {
|
||||
error!("{err:?} for {event:?}. ({user}, {channel})");
|
||||
if let Err(err) = handle_ack_event(&event, &db, &amqp, user, channel).await {
|
||||
revolt_config::capture_error(&err);
|
||||
error!("{err:?} for {event:?}. ({user:?}, {channel})");
|
||||
} else {
|
||||
info!("User {user} ack in {channel} with {event:?}");
|
||||
info!("User {user:?} ack in {channel} with {event:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -126,20 +207,41 @@ pub async fn worker(db: Database, authifier_db: authifier::Database) {
|
||||
mut event,
|
||||
}) = Q.try_pop()
|
||||
{
|
||||
let key = (user, channel);
|
||||
let key: (Option<String>, String, u8) = (
|
||||
user,
|
||||
channel,
|
||||
match &event {
|
||||
AckEvent::AckMessage { .. } => 0,
|
||||
AckEvent::ProcessMessage { .. } => 1,
|
||||
},
|
||||
);
|
||||
if let Some(task) = tasks.get_mut(&key) {
|
||||
task.delay();
|
||||
|
||||
match &mut event {
|
||||
AckEvent::AddMention { ids } => {
|
||||
if let AckEvent::AddMention { ids: existing } = &mut task.data.event {
|
||||
existing.append(ids);
|
||||
AckEvent::ProcessMessage { messages: new_data } => {
|
||||
if let AckEvent::ProcessMessage { messages: existing } =
|
||||
&mut task.data.event
|
||||
{
|
||||
// add the new message to the list of messages to be processed.
|
||||
existing.append(new_data);
|
||||
|
||||
// put a cap on the amount of messages that can be queued, for particularly active channels
|
||||
if (existing.length() as u16)
|
||||
< revolt_config::config()
|
||||
.await
|
||||
.features
|
||||
.advanced
|
||||
.process_message_delay_limit
|
||||
{
|
||||
task.delay();
|
||||
}
|
||||
} else {
|
||||
task.data.event = event;
|
||||
panic!("Somehow got an ack message in the add mention arm");
|
||||
}
|
||||
}
|
||||
AckEvent::AckMessage { .. } => {
|
||||
// replace the last acked message with the new acked message
|
||||
task.data.event = event;
|
||||
task.delay();
|
||||
}
|
||||
}
|
||||
} else {
|
||||
|
||||
@@ -1,314 +0,0 @@
|
||||
use std::io::Cursor;
|
||||
|
||||
use base64::{
|
||||
engine::{self},
|
||||
Engine as _,
|
||||
};
|
||||
use deadqueue::limited::Queue;
|
||||
use once_cell::sync::Lazy;
|
||||
use revolt_a2::{
|
||||
request::{
|
||||
notification::{DefaultAlert, NotificationOptions},
|
||||
payload::{APSAlert, APSSound, PayloadLike, APS},
|
||||
},
|
||||
Client, ClientConfig, Endpoint, Error, ErrorBody, ErrorReason, Priority, PushType, Response,
|
||||
};
|
||||
use revolt_config::config;
|
||||
use revolt_models::v0::{Message, PushNotification};
|
||||
|
||||
use crate::Database;
|
||||
|
||||
/// Payload information, before assembly
|
||||
#[derive(Debug)]
|
||||
#[allow(non_snake_case)]
|
||||
pub struct ApnPayload {
|
||||
message: Message,
|
||||
url: String,
|
||||
authorAvatar: String,
|
||||
authorDisplayName: String,
|
||||
channelName: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Debug)]
|
||||
#[allow(non_snake_case)]
|
||||
struct Payload<'a> {
|
||||
aps: APS<'a>,
|
||||
#[serde(skip_serializing)]
|
||||
options: NotificationOptions<'a>,
|
||||
#[serde(skip_serializing)]
|
||||
device_token: &'a str,
|
||||
|
||||
message: &'a Message,
|
||||
url: &'a str,
|
||||
authorAvatar: &'a str,
|
||||
authorDisplayName: &'a str,
|
||||
channelName: &'a str,
|
||||
}
|
||||
|
||||
impl<'a> PayloadLike for Payload<'a> {
|
||||
fn get_device_token(&self) -> &'a str {
|
||||
self.device_token
|
||||
}
|
||||
fn get_options(&self) -> &NotificationOptions {
|
||||
&self.options
|
||||
}
|
||||
}
|
||||
|
||||
/// Task information
|
||||
#[derive(Debug)]
|
||||
pub struct AlertJob {
|
||||
/// Session Id
|
||||
session_id: String,
|
||||
|
||||
/// Device token
|
||||
device_token: String,
|
||||
|
||||
/// User Id
|
||||
user_id: String,
|
||||
|
||||
/// Title
|
||||
title: String,
|
||||
|
||||
/// Body
|
||||
body: String,
|
||||
|
||||
/// Thread Id
|
||||
thread_id: String,
|
||||
|
||||
/// Category (informs the client what kind of notification is being sent.)
|
||||
category: String,
|
||||
|
||||
/// Payload used by the iOS client to modify the notification
|
||||
custom_payload: ApnPayload,
|
||||
}
|
||||
|
||||
impl AlertJob {
|
||||
fn format_title(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.
|
||||
format!(
|
||||
"{} in {}",
|
||||
notification.author, notification.message.channel
|
||||
) // TODO: this absolutely needs a channel name
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct BadgeJob {
|
||||
/// Session Id
|
||||
session_id: String,
|
||||
|
||||
/// Device token
|
||||
device_token: String,
|
||||
|
||||
/// User Id
|
||||
user_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum JobType {
|
||||
Alert(AlertJob),
|
||||
Badge(BadgeJob),
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ApnJob {
|
||||
job_type: JobType,
|
||||
}
|
||||
|
||||
impl ApnJob {
|
||||
pub fn from_notification(
|
||||
session_id: String,
|
||||
user_id: String,
|
||||
device_token: String,
|
||||
notification: &PushNotification,
|
||||
) -> ApnJob {
|
||||
ApnJob {
|
||||
job_type: JobType::Alert(AlertJob {
|
||||
session_id,
|
||||
device_token,
|
||||
user_id,
|
||||
title: AlertJob::format_title(notification),
|
||||
body: notification.body.to_string(),
|
||||
thread_id: notification.tag.to_string(),
|
||||
category: "ALERT_MESSAGE".to_string(),
|
||||
custom_payload: ApnPayload {
|
||||
message: notification.message.clone(),
|
||||
url: notification.url.clone(),
|
||||
authorAvatar: notification.icon.clone(),
|
||||
authorDisplayName: notification.author.clone(),
|
||||
channelName: "#fetchchannelnamehere".to_string(), // TODO: get actual channel name
|
||||
},
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_ack(session_id: String, user_id: String, device_token: String) -> ApnJob {
|
||||
ApnJob {
|
||||
job_type: JobType::Badge(BadgeJob {
|
||||
session_id,
|
||||
device_token,
|
||||
user_id,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum AssembledPayload<'a> {
|
||||
Alert(Payload<'a>),
|
||||
Default(revolt_a2::request::payload::Payload<'a>),
|
||||
}
|
||||
|
||||
static Q: Lazy<Queue<ApnJob>> = Lazy::new(|| Queue::new(10_000));
|
||||
|
||||
/// Queue a new task for a worker
|
||||
pub async fn queue(task: ApnJob) {
|
||||
Q.try_push(task).ok();
|
||||
info!("Queue is using {} slots from {}.", Q.len(), Q.capacity());
|
||||
}
|
||||
|
||||
async fn get_badge_count(db: &Database, user: &str) -> Option<u32> {
|
||||
if let Ok(unreads) = db.fetch_unreads(user).await {
|
||||
let mut mention_count = 0;
|
||||
for channel in unreads {
|
||||
if let Some(mentions) = channel.mentions {
|
||||
mention_count += mentions.len() as u32
|
||||
}
|
||||
}
|
||||
|
||||
return Some(mention_count);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Start a new worker
|
||||
pub async fn worker(db: Database) {
|
||||
let config = config().await;
|
||||
if config.api.apn.pkcs8.is_empty()
|
||||
|| config.api.apn.key_id.is_empty()
|
||||
|| config.api.apn.team_id.is_empty()
|
||||
{
|
||||
eprintln!("Missing APN keys.");
|
||||
return;
|
||||
}
|
||||
|
||||
let endpoint = if config.api.apn.sandbox {
|
||||
Endpoint::Sandbox
|
||||
} else {
|
||||
Endpoint::Production
|
||||
};
|
||||
|
||||
let pkcs8 = engine::general_purpose::STANDARD
|
||||
.decode(config.api.apn.pkcs8)
|
||||
.expect("valid `pcks8`");
|
||||
|
||||
let client_config = ClientConfig::new(endpoint);
|
||||
|
||||
let client = Client::token(
|
||||
&mut Cursor::new(pkcs8),
|
||||
config.api.apn.key_id,
|
||||
config.api.apn.team_id,
|
||||
client_config,
|
||||
)
|
||||
.expect("could not create APN client");
|
||||
|
||||
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,
|
||||
};
|
||||
|
||||
loop {
|
||||
let task = Q.pop().await;
|
||||
let payload: AssembledPayload;
|
||||
|
||||
match task.job_type {
|
||||
JobType::Alert(ref alert) => {
|
||||
payload = AssembledPayload::Alert(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: get_badge_count(&db, &alert.user_id).await,
|
||||
sound: Some(APSSound::Sound("default")),
|
||||
thread_id: Some(&alert.thread_id),
|
||||
content_available: None,
|
||||
category: Some(&alert.category),
|
||||
mutable_content: Some(1),
|
||||
url_args: None,
|
||||
},
|
||||
device_token: &alert.device_token,
|
||||
options: payload_options.clone(),
|
||||
message: &alert.custom_payload.message,
|
||||
url: &alert.custom_payload.url,
|
||||
authorAvatar: &alert.custom_payload.authorAvatar,
|
||||
authorDisplayName: &alert.custom_payload.authorDisplayName,
|
||||
channelName: &alert.custom_payload.channelName,
|
||||
});
|
||||
}
|
||||
JobType::Badge(ref alert) => {
|
||||
payload = AssembledPayload::Default(revolt_a2::request::payload::Payload {
|
||||
aps: APS {
|
||||
alert: None,
|
||||
badge: get_badge_count(&db, &alert.user_id).await,
|
||||
sound: None,
|
||||
thread_id: None,
|
||||
content_available: None,
|
||||
category: None,
|
||||
mutable_content: None,
|
||||
url_args: None,
|
||||
},
|
||||
device_token: &alert.device_token,
|
||||
options: payload_options.clone(),
|
||||
data: std::collections::BTreeMap::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
let resp = match payload {
|
||||
AssembledPayload::Alert(p) => client.send(p).await,
|
||||
AssembledPayload::Default(p) => client.send(p).await,
|
||||
};
|
||||
//println!("response from APNS: {:?}", resp);
|
||||
|
||||
if let Err(err) = resp {
|
||||
match err {
|
||||
Error::ResponseError(Response {
|
||||
error:
|
||||
Some(ErrorBody {
|
||||
reason: ErrorReason::BadDeviceToken | ErrorReason::Unregistered,
|
||||
..
|
||||
}),
|
||||
..
|
||||
}) => {
|
||||
if let Err(err) = db
|
||||
.remove_push_subscription_by_session_id(match task.job_type {
|
||||
JobType::Alert(ref a) => &a.session_id.as_str(),
|
||||
JobType::Badge(ref a) => &a.session_id.as_str(),
|
||||
})
|
||||
.await
|
||||
{
|
||||
revolt_config::capture_error(&err);
|
||||
}
|
||||
}
|
||||
err => {
|
||||
revolt_config::capture_error(&err);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,6 @@
|
||||
//! Semi-important background task management
|
||||
|
||||
use crate::Database;
|
||||
use crate::{Database, AMQP};
|
||||
|
||||
use async_std::task;
|
||||
use std::time::Instant;
|
||||
@@ -8,22 +8,18 @@ use std::time::Instant;
|
||||
const WORKER_COUNT: usize = 5;
|
||||
|
||||
pub mod ack;
|
||||
pub mod apple_notifications;
|
||||
pub mod authifier_relay;
|
||||
pub mod last_message_id;
|
||||
pub mod process_embeds;
|
||||
pub mod web_push;
|
||||
|
||||
/// Spawn background workers
|
||||
pub fn start_workers(db: Database, authifier_db: authifier::Database) {
|
||||
pub fn start_workers(db: Database, amqp: AMQP) {
|
||||
task::spawn(authifier_relay::worker());
|
||||
task::spawn(apple_notifications::worker(db.clone()));
|
||||
|
||||
for _ in 0..WORKER_COUNT {
|
||||
task::spawn(ack::worker(db.clone(), authifier_db.clone()));
|
||||
task::spawn(ack::worker(db.clone(), amqp.clone()));
|
||||
task::spawn(last_message_id::worker(db.clone()));
|
||||
task::spawn(process_embeds::worker(db.clone()));
|
||||
task::spawn(web_push::worker(db.clone(), authifier_db.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,198 +0,0 @@
|
||||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use authifier::Database as AuthifierDatabase;
|
||||
use base64::{
|
||||
engine::{self},
|
||||
Engine as _,
|
||||
};
|
||||
use deadqueue::limited::Queue;
|
||||
use fcm_v1::auth::{Authenticator, ServiceAccountKey};
|
||||
use once_cell::sync::Lazy;
|
||||
use revolt_config::{config, report_internal_error};
|
||||
use revolt_models::v0::PushNotification;
|
||||
use revolt_presence::filter_online;
|
||||
use serde_json::json;
|
||||
use web_push::{
|
||||
ContentEncoding, IsahcWebPushClient, SubscriptionInfo, SubscriptionKeys, VapidSignatureBuilder,
|
||||
WebPushClient, WebPushMessageBuilder,
|
||||
};
|
||||
|
||||
use crate::Database;
|
||||
|
||||
use super::apple_notifications;
|
||||
|
||||
/// Task information
|
||||
#[derive(Debug)]
|
||||
struct PushTask {
|
||||
/// User IDs of the targets that are to receive this notification
|
||||
recipients: Vec<String>,
|
||||
/// Push Notification
|
||||
payload: PushNotification,
|
||||
}
|
||||
|
||||
static Q: Lazy<Queue<PushTask>> = Lazy::new(|| Queue::new(10_000));
|
||||
|
||||
/// Queue a new task for a worker
|
||||
pub async fn queue(recipients: Vec<String>, payload: PushNotification) {
|
||||
if recipients.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let online_ids = filter_online(&recipients).await;
|
||||
let recipients = (&recipients.into_iter().collect::<HashSet<String>>() - &online_ids)
|
||||
.into_iter()
|
||||
.collect::<Vec<String>>();
|
||||
|
||||
Q.try_push(PushTask {
|
||||
recipients,
|
||||
payload,
|
||||
})
|
||||
.ok();
|
||||
|
||||
info!("Queue is using {} slots from {}.", Q.len(), Q.capacity());
|
||||
}
|
||||
|
||||
/// Start a new worker
|
||||
pub async fn worker(db: Database, authifier_db: AuthifierDatabase) {
|
||||
let config = config().await;
|
||||
|
||||
let web_push_client = IsahcWebPushClient::new().unwrap();
|
||||
let fcm_client = if config.api.fcm.key_type.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(fcm_v1::Client::new(
|
||||
Authenticator::service_account::<&str>(ServiceAccountKey {
|
||||
key_type: Some(config.api.fcm.key_type),
|
||||
project_id: Some(config.api.fcm.project_id.clone()),
|
||||
private_key_id: Some(config.api.fcm.private_key_id),
|
||||
private_key: config.api.fcm.private_key,
|
||||
client_email: config.api.fcm.client_email,
|
||||
client_id: Some(config.api.fcm.client_id),
|
||||
auth_uri: Some(config.api.fcm.auth_uri),
|
||||
token_uri: config.api.fcm.token_uri,
|
||||
auth_provider_x509_cert_url: Some(config.api.fcm.auth_provider_x509_cert_url),
|
||||
client_x509_cert_url: Some(config.api.fcm.client_x509_cert_url),
|
||||
})
|
||||
.await
|
||||
.unwrap(),
|
||||
config.api.fcm.project_id,
|
||||
false,
|
||||
Duration::from_secs(5),
|
||||
))
|
||||
};
|
||||
|
||||
let web_push_private_key = engine::general_purpose::URL_SAFE_NO_PAD
|
||||
.decode(config.api.vapid.private_key)
|
||||
.expect("valid `VAPID_PRIVATE_KEY`");
|
||||
|
||||
loop {
|
||||
let task = Q.pop().await;
|
||||
|
||||
if let Ok(sessions) = authifier_db
|
||||
.find_sessions_with_subscription(&task.recipients)
|
||||
.await
|
||||
{
|
||||
for session in sessions {
|
||||
if let Some(sub) = session.subscription {
|
||||
if sub.endpoint == "fcm" {
|
||||
// Use Firebase Cloud Messaging
|
||||
if let Some(client) = &fcm_client {
|
||||
let message = fcm_v1::message::Message {
|
||||
token: Some(sub.auth),
|
||||
data: Some(HashMap::from([(
|
||||
"payload".to_owned(),
|
||||
serde_json::Value::String(json!(&task.payload).to_string()),
|
||||
)])),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
if let Err(err) = client.send(&message).await {
|
||||
error!("Failed to send FCM notification! {:?}", err);
|
||||
|
||||
if let fcm_v1::Error::FCM(fcm_error) = err {
|
||||
if fcm_error.contains("404 (Not Found)") {
|
||||
println!("Unregistering {:?}", session.id);
|
||||
|
||||
report_internal_error!(
|
||||
db.remove_push_subscription_by_session_id(&session.id)
|
||||
.await
|
||||
)
|
||||
.ok();
|
||||
}
|
||||
}
|
||||
} else {
|
||||
info!("Sent FCM notification to {:?}.", session.id);
|
||||
}
|
||||
} else {
|
||||
info!("No FCM token was specified!");
|
||||
}
|
||||
} else if sub.endpoint == "apn" {
|
||||
apple_notifications::queue(apple_notifications::ApnJob::from_notification(
|
||||
session.id,
|
||||
session.user_id,
|
||||
sub.auth,
|
||||
&task.payload,
|
||||
))
|
||||
.await;
|
||||
} else {
|
||||
// Use Web Push Standard
|
||||
let subscription = SubscriptionInfo {
|
||||
endpoint: sub.endpoint,
|
||||
keys: SubscriptionKeys {
|
||||
auth: sub.auth,
|
||||
p256dh: sub.p256dh,
|
||||
},
|
||||
};
|
||||
|
||||
match VapidSignatureBuilder::from_pem(
|
||||
std::io::Cursor::new(&web_push_private_key),
|
||||
&subscription,
|
||||
) {
|
||||
Ok(sig_builder) => match sig_builder.build() {
|
||||
Ok(signature) => {
|
||||
let mut builder = WebPushMessageBuilder::new(&subscription);
|
||||
builder.set_vapid_signature(signature);
|
||||
|
||||
let payload = json!(task.payload).to_string();
|
||||
builder
|
||||
.set_payload(ContentEncoding::AesGcm, payload.as_bytes());
|
||||
|
||||
match builder.build() {
|
||||
Ok(msg) => match web_push_client.send(msg).await {
|
||||
Ok(_) => {
|
||||
info!(
|
||||
"Sent Web Push notification to {:?}.",
|
||||
session.id
|
||||
)
|
||||
}
|
||||
Err(err) => {
|
||||
error!("Hit error sending Web Push! {:?}", err)
|
||||
}
|
||||
},
|
||||
Err(err) => {
|
||||
error!(
|
||||
"Failed to build message for {}! {:?}",
|
||||
session.user_id, err
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => error!(
|
||||
"Failed to build signature for {}! {:?}",
|
||||
session.user_id, err
|
||||
),
|
||||
},
|
||||
Err(err) => error!(
|
||||
"Failed to create signature builder for {}! {:?}",
|
||||
session.user_id, err
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user