add Badge update capabilities

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

crates/core/database/src/tasks/apple_notifications.rs

@@ -9,10 +9,9 @@ use once_cell::sync::Lazy;
 use revolt_a2::{
     request::{
         notification::{DefaultAlert, NotificationOptions},
-        payload::{APSAlert, PayloadLike, APS},
+        payload::{APSAlert, APSSound, PayloadLike, APS},
     },
-    Client, ClientConfig, DefaultNotificationBuilder, Endpoint, Error, ErrorBody, ErrorReason,
+    Client, ClientConfig, Endpoint, Error, ErrorBody, ErrorReason, Priority, PushType, Response,
-    NotificationBuilder, Priority, PushType, Response,
 };
 use revolt_config::config;
 use revolt_models::v0::{Message, PushNotification};
@@ -37,11 +36,11 @@ struct Payload<'a> {
     #[serde(skip_serializing)]
     device_token: &'a str,
 
-    message: Message,
+    message: &'a Message,
-    url: String,
+    url: &'a str,
-    authorAvatar: String,
+    authorAvatar: &'a str,
-    authorDisplayName: String,
+    authorDisplayName: &'a str,
-    channelName: String,
+    channelName: &'a str,
 }
 
 impl<'a> PayloadLike for Payload<'a> {
@@ -55,13 +54,16 @@ impl<'a> PayloadLike for Payload<'a> {
 
 /// Task information
 #[derive(Debug)]
-pub struct ApnTask {
+pub struct AlertJob {
     /// Session Id
     session_id: String,
 
     /// Device token
     device_token: String,
 
+    /// User Id
+    user_id: String,
+
     /// Title
     title: String,
 
@@ -78,7 +80,7 @@ pub struct ApnTask {
     custom_payload: ApnPayload,
 }
 
-impl ApnTask {
+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"
@@ -90,38 +92,96 @@ impl ApnTask {
             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,
-    ) -> ApnTask {
+    ) -> ApnJob {
-        ApnTask {
+        ApnJob {
-            session_id,
+            job_type: JobType::Alert(AlertJob {
-            device_token,
+                session_id,
-            title: ApnTask::format_title(notification),
+                device_token,
-            body: notification.body.to_string(),
+                user_id,
-            thread_id: notification.tag.to_string(),
+                title: AlertJob::format_title(notification),
-            category: "ALERT_MESSAGE".to_string(),
+                body: notification.body.to_string(),
-            custom_payload: ApnPayload {
+                thread_id: notification.tag.to_string(),
-                message: notification.message.clone(),
+                category: "ALERT_MESSAGE".to_string(),
-                url: notification.url.clone(),
+                custom_payload: ApnPayload {
-                authorAvatar: notification.icon.clone(),
+                    message: notification.message.clone(),
-                authorDisplayName: notification.author.clone(),
+                    url: notification.url.clone(),
-                channelName: "#fetchchannelnamehere".to_string(), // TODO: get actual channel name
+                    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,
+            }),
         }
     }
 }
 
-static Q: Lazy<Queue<ApnTask>> = Lazy::new(|| Queue::new(10_000));
+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: ApnTask) {
+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;
@@ -164,38 +224,63 @@ pub async fn worker(db: Database) {
 
     loop {
         let task = Q.pop().await;
+        let payload: AssembledPayload;
 
-        let payload: Payload = Payload {
+        match task.job_type {
-            aps: APS {
+            JobType::Alert(ref alert) => {
-                alert: Some(APSAlert::Default(DefaultAlert {
+                payload = AssembledPayload::Alert(Payload {
-                    title: Some(&task.title),
+                    aps: APS {
-                    subtitle: None,
+                        alert: Some(APSAlert::Default(DefaultAlert {
-                    body: Some(&task.body),
+                            title: Some(&alert.title),
-                    title_loc_key: None,
+                            subtitle: None,
-                    title_loc_args: None,
+                            body: Some(&alert.body),
-                    action_loc_key: None,
+                            title_loc_key: None,
-                    loc_key: None,
+                            title_loc_args: None,
-                    loc_args: None,
+                            action_loc_key: None,
-                    launch_image: None,
+                            loc_key: None,
-                })),
+                            loc_args: None,
-                badge: Some(1),
+                            launch_image: None,
-                sound: None,
+                        })),
-                thread_id: Some(&task.thread_id),
+                        badge: get_badge_count(&db, &alert.user_id).await,
-                content_available: Some(1),
+                        sound: Some(APSSound::Sound("default")),
-                category: Some(&task.category),
+                        thread_id: Some(&alert.thread_id),
-                mutable_content: Some(1),
+                        content_available: None,
-                url_args: None,
+                        category: Some(&alert.category),
-            },
+                        mutable_content: Some(1),
-            device_token: &task.device_token,
+                        url_args: None,
-            options: payload_options.clone(),
+                    },
-            message: task.custom_payload.message,
+                    device_token: &alert.device_token,
-            url: task.custom_payload.url,
+                    options: payload_options.clone(),
-            authorAvatar: task.custom_payload.authorAvatar,
+                    message: &alert.custom_payload.message,
-            authorDisplayName: task.custom_payload.authorDisplayName,
+                    url: &alert.custom_payload.url,
-            channelName: task.custom_payload.channelName,
+                    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,
         };
-
-        let resp = client.send(payload).await;
         //println!("response from APNS: {:?}", resp);
 
         if let Err(err) = resp {
@@ -209,7 +294,10 @@ pub async fn worker(db: Database) {
                     ..
                 }) => {
                     if let Err(err) = db
-                        .remove_push_subscription_by_session_id(&task.session_id)
+                        .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);

crates/core/database/src/tasks/web_push.rs

@@ -105,13 +105,12 @@ pub async fn worker(db: Database) {
                             info!("No FCM token was specified!");
                         }
                     } else if sub.endpoint == "apn" {
-                        apple_notifications::queue(
+                        apple_notifications::queue(apple_notifications::ApnJob::from_notification(
-                            apple_notifications::ApnTask::from_notification(
+                            session.id,
-                                session.id,
+                            session.user_id,
-                                sub.auth,
+                            sub.auth,
-                                &task.payload,
+                            &task.payload,
-                            ),
+                        ))
-                        )
                         .await;
                     } else {
                         // Use Web Push Standard