forked from jmug/stoatchat
feat(tasks): implement web push
This commit is contained in:
@@ -3,9 +3,12 @@ use std::collections::HashSet;
|
||||
use revolt_quark::{
|
||||
models::{
|
||||
message::{Content, Masquerade, Reply, SendableEmbed},
|
||||
Message, User,
|
||||
Channel, Message, User,
|
||||
},
|
||||
perms, Db, Error, Permission, Ref, Result,
|
||||
perms,
|
||||
presence::presence_filter_online,
|
||||
types::push::PushNotification,
|
||||
Db, Error, Permission, Ref, Result,
|
||||
};
|
||||
|
||||
use regex::Regex;
|
||||
@@ -15,6 +18,7 @@ use ulid::Ulid;
|
||||
use validator::Validate;
|
||||
|
||||
use crate::util::idempotency::IdempotencyKey;
|
||||
use crate::util::variables::{APP_URL, AUTUMN_URL, PUBLIC_URL};
|
||||
|
||||
#[derive(Validate, Serialize, Deserialize, JsonSchema)]
|
||||
pub struct DataMessageSend {
|
||||
@@ -181,5 +185,40 @@ pub async fn message_send(
|
||||
)
|
||||
.await;
|
||||
|
||||
// ! FIXME: this should be part of quark.
|
||||
// ! Actually, so should be the thing above
|
||||
// ! probably the entire task queue system
|
||||
// ! should be moved
|
||||
crate::tasks::web_push::queue(
|
||||
{
|
||||
let mut target_ids = vec![];
|
||||
match &channel {
|
||||
Channel::DirectMessage { recipients, .. } | Channel::Group { recipients, .. } => {
|
||||
target_ids = (&recipients.iter().cloned().collect::<HashSet<String>>()
|
||||
- &presence_filter_online(recipients).await)
|
||||
.into_iter()
|
||||
.collect::<Vec<String>>();
|
||||
}
|
||||
Channel::TextChannel { .. } => {
|
||||
if let Some(mentions) = &message.mentions {
|
||||
target_ids.append(&mut mentions.clone());
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
};
|
||||
target_ids
|
||||
},
|
||||
json!(PushNotification::new(
|
||||
message.clone(),
|
||||
user,
|
||||
channel.id(),
|
||||
&*AUTUMN_URL,
|
||||
&*PUBLIC_URL,
|
||||
&*APP_URL,
|
||||
))
|
||||
.to_string(),
|
||||
)
|
||||
.await;
|
||||
|
||||
Ok(Json(message))
|
||||
}
|
||||
|
||||
@@ -3,13 +3,15 @@
|
||||
use async_std::task;
|
||||
use revolt_quark::Database;
|
||||
|
||||
const WORKER_COUNT: usize = 10;
|
||||
const WORKER_COUNT: usize = 5;
|
||||
|
||||
pub mod process_embeds;
|
||||
pub mod web_push;
|
||||
|
||||
/// Spawn background workers
|
||||
pub async fn start_workers(db: Database) {
|
||||
for _ in 0..WORKER_COUNT {
|
||||
task::spawn(process_embeds::worker(db.clone()));
|
||||
task::spawn(web_push::worker(db.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
91
src/tasks/web_push.rs
Normal file
91
src/tasks/web_push.rs
Normal file
@@ -0,0 +1,91 @@
|
||||
use crate::util::variables::VAPID_PRIVATE_KEY;
|
||||
|
||||
use deadqueue::limited::Queue;
|
||||
use futures::StreamExt;
|
||||
use rauth::entities::Session;
|
||||
use revolt_quark::{bson::doc, r#impl::mongo::MongoDb, Database};
|
||||
use web_push::{
|
||||
ContentEncoding, SubscriptionInfo, SubscriptionKeys, VapidSignatureBuilder, WebPushClient,
|
||||
WebPushMessageBuilder,
|
||||
};
|
||||
|
||||
#[derive(Debug)]
|
||||
struct PushTask {
|
||||
recipients: Vec<String>,
|
||||
payload: String,
|
||||
}
|
||||
|
||||
lazy_static! {
|
||||
static ref Q: Queue<PushTask> = Queue::new(10_000);
|
||||
}
|
||||
|
||||
pub async fn queue(recipients: Vec<String>, payload: String) {
|
||||
if recipients.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
Q.push(PushTask {
|
||||
recipients,
|
||||
payload,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
pub async fn worker(db: Database) {
|
||||
let client = WebPushClient::new();
|
||||
let key = base64::decode_config(VAPID_PRIVATE_KEY.clone(), base64::URL_SAFE)
|
||||
.expect("valid `VAPID_PRIVATE_KEY`");
|
||||
|
||||
if let Database::MongoDb(MongoDb(db)) = db {
|
||||
loop {
|
||||
let task = Q.pop().await;
|
||||
|
||||
// ! FIXME: this is hard-coded until rauth is merged into quark
|
||||
if let Ok(mut cursor) = db
|
||||
.database("revolt")
|
||||
.collection::<Session>("sessions")
|
||||
.find(
|
||||
doc! {
|
||||
"user_id": {
|
||||
"$in": task.recipients
|
||||
},
|
||||
"subscription": {
|
||||
"$exists": true
|
||||
}
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
while let Some(Ok(session)) = cursor.next().await {
|
||||
if let Some(sub) = session.subscription {
|
||||
let subscription = SubscriptionInfo {
|
||||
endpoint: sub.endpoint,
|
||||
keys: SubscriptionKeys {
|
||||
auth: sub.auth,
|
||||
p256dh: sub.p256dh,
|
||||
},
|
||||
};
|
||||
|
||||
let mut builder = WebPushMessageBuilder::new(&subscription).unwrap();
|
||||
let sig_builder = VapidSignatureBuilder::from_pem(
|
||||
std::io::Cursor::new(&key),
|
||||
&subscription,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let signature = sig_builder.build().unwrap();
|
||||
builder.set_vapid_signature(signature);
|
||||
builder.set_payload(ContentEncoding::AesGcm, task.payload.as_bytes());
|
||||
|
||||
let msg = builder.build().unwrap();
|
||||
match client.send(msg).await {
|
||||
Ok(_) => info!("Sent Web Push notification to {:?}.", session.id),
|
||||
Err(err) => error!("Hit error sending Web Push! {:?}", err),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user