From 54b6cf67ab11b65d4bfca3efd298186d9390ce48 Mon Sep 17 00:00:00 2001
From: Paul Makles <paulmakles@gmail.com>
Date: Sun, 17 Apr 2022 13:11:29 +0100
Subject: [PATCH] feat(tasks): implement web push

---
 src/routes/channels/message_send.rs | 43 +++++++++++++-
 src/tasks/mod.rs                    |  4 +-
 src/tasks/web_push.rs               | 91 +++++++++++++++++++++++++++++
 3 files changed, 135 insertions(+), 3 deletions(-)
 create mode 100644 src/tasks/web_push.rs

diff --git a/src/routes/channels/message_send.rs b/src/routes/channels/message_send.rs
index cf5f04f5..d4328512 100644
--- a/src/routes/channels/message_send.rs
+++ b/src/routes/channels/message_send.rs
@@ -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))
 }
diff --git a/src/tasks/mod.rs b/src/tasks/mod.rs
index a5106325..a71b25ee 100644
--- a/src/tasks/mod.rs
+++ b/src/tasks/mod.rs
@@ -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()));
     }
 }
diff --git a/src/tasks/web_push.rs b/src/tasks/web_push.rs
new file mode 100644
index 00000000..a166da38
--- /dev/null
+++ b/src/tasks/web_push.rs
@@ -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),
+                        }
+                    }
+                }
+            }
+        }
+    }
+}