merge: pull request #333 from revoltchat/feat/apns

APNS fixes

Cargo.lock

@@ -3788,9 +3788,9 @@ dependencies = [
 
 [[package]]
 name = "revolt_a2"
-version = "0.10.0"
+version = "0.10.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "466eb5262fcbb26e6e10b8a0acf56eb7cc095008132ff4ebddfc44d8672f8066"
+checksum = "edbe1f79cb41271d3cd8f932d75dddeba963c19dc93d1ee6cbe0391b495ab2f5"
 dependencies = [
  "base64 0.21.3",
  "erased-serde",

crates/core/config/Revolt.toml

@@ -33,6 +33,7 @@ public_key = "BGcvgR-i2z4IQ5Mw841vJvkLjt8wY-FjmWrw83jOLCY52qcGZS0OF7nfLzuYbjsQIS
 api_key = ""
 
 [api.apn]
+sandbox = false
 pkcs8 = ""
 key_id = ""
 team_id = ""

crates/core/config/src/lib.rs

@@ -79,6 +79,7 @@ pub struct ApiFcm {
 
 #[derive(Deserialize, Debug, Clone)]
 pub struct ApiApn {
+    pub sandbox: bool,
     pub pkcs8: String,
     pub key_id: String,
     pub team_id: String,

crates/core/database/Cargo.toml

@@ -87,7 +87,7 @@ revolt_rocket_okapi = { version = "0.9.1", optional = true }
 # Notifications
 fcm = "0.9.2"
 web-push = "0.10.0"
-revolt_a2 = { version = "0.10.0", default-features = false, features = [
+revolt_a2 = { version = "0.10", default-features = false, features = [
     "ring",
 ] }
 

crates/core/database/src/models/channel_unreads/ops.rs

@@ -7,13 +7,13 @@ mod reference;
 
 #[async_trait]
 pub trait AbstractChannelUnreads: Sync + Send {
-    /// Acknowledge a message.
+    /// Acknowledge a message, and returns updated channel unread.
     async fn acknowledge_message(
         &self,
         channel_id: &str,
         user_id: &str,
         message_id: &str,
-    ) -> Result<()>;
+    ) -> Result<Option<ChannelUnread>>;
 
     /// Acknowledge many channels.
     async fn acknowledge_channels(&self, user_id: &str, channel_ids: &[String]) -> Result<()>;
@@ -28,4 +28,7 @@ pub trait AbstractChannelUnreads: Sync + Send {
 
     /// Fetch all channel unreads for a user.
     async fn fetch_unreads(&self, user_id: &str) -> Result<Vec<ChannelUnread>>;
+
+    /// Fetch unread for a specific user in a channel.
+    async fn fetch_unread(&self, user_id: &str, channel_id: &str) -> Result<Option<ChannelUnread>>;
 }

crates/core/database/src/models/channel_unreads/ops/mongodb.rs

@@ -1,4 +1,6 @@
 use bson::Document;
+use mongodb::options::FindOneAndUpdateOptions;
+use mongodb::options::ReturnDocument;
 use mongodb::options::UpdateOptions;
 use revolt_result::Result;
 use ulid::Ulid;
@@ -12,31 +14,35 @@ static COL: &str = "channel_unreads";
 
 #[async_trait]
 impl AbstractChannelUnreads for MongoDb {
-    /// Acknowledge a message.
+    /// Acknowledge a message, and returns updated channel unread.
     async fn acknowledge_message(
         &self,
         channel_id: &str,
         user_id: &str,
         message_id: &str,
-    ) -> Result<()> {
+    ) -> Result<Option<ChannelUnread>> {
-        self.col::<Document>(COL)
+        self.col::<ChannelUnread>(COL)
-            .update_one(
+            .find_one_and_update(
                 doc! {
                     "_id.channel": channel_id,
                     "_id.user": user_id,
                 },
                 doc! {
-                    "$unset": {
+                    "$pull": {
-                        "mentions": 1_i32
+                        "mentions": {
+                            "$lt": message_id
+                        }
                     },
                     "$set": {
                         "last_id": message_id
                     }
                 },
-                UpdateOptions::builder().upsert(true).build(),
+                FindOneAndUpdateOptions::builder()
+                    .upsert(true)
+                    .return_document(ReturnDocument::After)
+                    .build(),
             )
             .await
-            .map(|_| ())
             .map_err(|_| create_database_error!("update_one", COL))
     }
 
@@ -116,4 +122,18 @@ impl AbstractChannelUnreads for MongoDb {
             }
         )
     }
+
+    /// Fetch unread for a specific user in a channel.
+    async fn fetch_unread(&self, user_id: &str, channel_id: &str) -> Result<Option<ChannelUnread>> {
+        query!(
+            self,
+            find_one,
+            COL,
+            doc! {
+                "_id.user": user_id,
+                "_id.channel": channel_id
+            }
+        )
+    }
+
 }

crates/core/database/src/models/channel_unreads/ops/reference.rs

@@ -13,7 +13,7 @@ impl AbstractChannelUnreads for ReferenceDb {
         channel_id: &str,
         user_id: &str,
         message_id: &str,
-    ) -> Result<()> {
+    ) -> Result<Option<ChannelUnread>> {
         let mut unreads = self.channel_unreads.lock().await;
         let key = ChannelCompositeKey {
             channel: channel_id.to_string(),
@@ -27,14 +27,14 @@ impl AbstractChannelUnreads for ReferenceDb {
             unreads.insert(
                 key.clone(),
                 ChannelUnread {
-                    id: key,
+                    id: key.clone(),
                     last_id: Some(message_id.to_string()),
                     mentions: None,
                 },
             );
         }
 
-        Ok(())
+        Ok(unreads.get(&key).cloned())
     }
 
     /// Acknowledge many channels.
@@ -87,4 +87,14 @@ impl AbstractChannelUnreads for ReferenceDb {
             .cloned()
             .collect())
     }
+
+    /// Fetch unread for a specific user in a channel.
+    async fn fetch_unread(&self, user_id: &str, channel_id: &str) -> Result<Option<ChannelUnread>> {
+        let unreads = self.channel_unreads.lock().await;
+
+        Ok(unreads.get(&ChannelCompositeKey {
+            channel: channel_id.to_string(),
+            user: user_id.to_string()
+        }).cloned())
+    }
 }

crates/core/database/src/models/messages/model.rs

@@ -458,8 +458,8 @@ impl Message {
     ) -> Result<()> {
         self.send_without_notifications(
             db,
-            user,
+            user.clone(),
-            member,
+            member.clone(),
             matches!(channel, Channel::DirectMessage { .. }),
             generate_embeds,
         )
@@ -477,7 +477,7 @@ impl Message {
                     }
                 },
                 PushNotification::from(
-                    self.clone().into_model(None, None),
+                    self.clone().into_model(user, member),
                     Some(author),
                     channel.id(),
                 )

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

@@ -5,7 +5,9 @@ use deadqueue::limited::Queue;
 use once_cell::sync::Lazy;
 use std::{collections::HashMap, time::Duration};
 
-use super::DelayedTask;
+use revolt_result::Result;
+
+use super::{apple_notifications::{self, ApnJob}, DelayedTask};
 
 /// Enumeration of possible events
 #[derive(Debug, Eq, PartialEq)]
@@ -52,8 +54,43 @@ pub async fn queue(channel: String, user: String, event: AckEvent) {
     info!("Queue is using {} slots from {}.", Q.len(), Q.capacity());
 }
 
+pub async fn handle_ack_event(event: &AckEvent, db: &Database, authifier_db: &authifier::Database, user: &str, channel: &str) -> Result<()> {
+    match &event {
+        #[allow(clippy::disallowed_methods)] // event is sent by higher level function
+        AckEvent::AckMessage { id } => {
+            let unread = db.fetch_unread(user, channel).await?;
+            let updated = db.acknowledge_message(channel, user, id).await?;
+
+            if let (Some(before), Some(after)) = (unread, updated) {
+                let before_mentions = before.mentions.unwrap_or_default().len();
+                let after_mentions = after.mentions.unwrap_or_default().len();
+
+                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;
+                                }
+                            }
+                        }
+                    }
+                };
+
+            }
+        },
+        AckEvent::AddMention { ids } => {
+            db.add_mention_to_unread(channel, user, ids).await?;
+        }
+    };
+
+    Ok(())
+}
+
 /// Start a new worker
-pub async fn worker(db: Database) {
+pub async fn worker(db: Database, authifier_db: authifier::Database) {
     let mut tasks = HashMap::<(String, String), DelayedTask<Task>>::new();
     let mut keys = vec![];
 
@@ -71,13 +108,7 @@ pub async fn worker(db: Database) {
                 let Task { event } = task.data;
                 let (user, channel) = key;
 
-                if let Err(err) = match &event {
+                if let Err(err) = handle_ack_event(&event, &db, &authifier_db, user, channel).await {
-                    #[allow(clippy::disallowed_methods)] // event is sent by higher level function
-                    AckEvent::AckMessage { id } => db.acknowledge_message(channel, user, id).await,
-                    AckEvent::AddMention { ids } => {
-                        db.add_mention_to_unread(channel, user, ids).await
-                    }
-                } {
                     error!("{err:?} for {event:?}. ({user}, {channel})");
                 } else {
                     info!("User {user} ack in {channel} with {event:?}");

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

@@ -6,22 +6,64 @@ use base64::{
 };
 use deadqueue::limited::Queue;
 use once_cell::sync::Lazy;
-use revolt_a2::{Client, ClientConfig, DefaultNotificationBuilder};
+use revolt_a2::{
-use revolt_a2::{Error, ErrorBody, ErrorReason, NotificationBuilder, Response};
+    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::PushNotification;
+use revolt_models::v0::{Message, PushNotification};
 
 use crate::Database;
 
+/// Payload information, before assembly
+#[derive(Debug)]
+pub struct ApnPayload {
+    message: Message,
+    url: String,
+    authorAvatar: String,
+    authorDisplayName: String,
+    channelName: String,
+}
+
+#[derive(Serialize, Debug)]
+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 ApnTask {
+pub struct AlertJob {
     /// Session Id
     session_id: String,
 
     /// Device token
     device_token: String,
 
+    /// User Id
+    user_id: String,
+
     /// Title
     title: String,
 
@@ -30,32 +72,116 @@ pub struct ApnTask {
 
     /// 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 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"
+        // 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,
-    ) -> ApnTask {
+    ) -> ApnJob {
-        ApnTask {
+        ApnJob {
-            session_id,
+            job_type: JobType::Alert(AlertJob {
-            device_token,
+                session_id,
-            title: notification.author.to_string(),
+                device_token,
-            body: notification.body.to_string(),
+                user_id,
-            thread_id: notification.tag.to_string(),
+                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,
+            }),
         }
     }
 }
 
-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;
@@ -67,27 +193,97 @@ pub async fn worker(db: Database) {
         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,
-        ClientConfig::default(),
+        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 = DefaultNotificationBuilder::new()
+        let payload: AssembledPayload;
-            .set_title(&task.title)
-            .set_body(&task.body)
-            .set_thread_id(&task.thread_id)
-            .build(&task.device_token, Default::default());
 
-        if let Err(err) = client.send(payload).await {
+        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:
@@ -98,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/mod.rs

@@ -18,7 +18,7 @@ pub async fn start_workers(db: Database, authifier_db: authifier::Database) {
     task::spawn(apple_notifications::worker(db.clone()));
 
     for _ in 0..WORKER_COUNT {
-        task::spawn(ack::worker(db.clone()));
+        task::spawn(ack::worker(db.clone(), authifier_db.clone()));
         task::spawn(last_message_id::worker(db.clone()));
         task::spawn(process_embeds::worker(db.clone()));
         task::spawn(web_push::worker(authifier_db.clone()));

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

@@ -6,7 +6,6 @@ use base64::{
     Engine as _,
 };
 use deadqueue::limited::Queue;
-use fcm::FcmError;
 use once_cell::sync::Lazy;
 use revolt_config::config;
 use revolt_models::v0::PushNotification;
@@ -82,6 +81,7 @@ pub async fn worker(db: Database) {
                                 tag,
                                 timestamp: _,
                                 url: _,
+                                message: _,
                             } = &task.payload;
 
                             let mut notification = fcm::NotificationBuilder::new();
@@ -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

crates/core/models/src/v0/embeds.rs

@@ -98,38 +98,38 @@ auto_derived!(
     pub struct WebsiteMetadata {
         /// Direct URL to web page
         #[serde(skip_serializing_if = "Option::is_none")]
-        url: Option<String>,
+        pub url: Option<String>,
         /// Original direct URL
         #[serde(skip_serializing_if = "Option::is_none")]
-        original_url: Option<String>,
+        pub original_url: Option<String>,
         /// Remote content
         #[serde(skip_serializing_if = "Option::is_none")]
-        special: Option<Special>,
+        pub special: Option<Special>,
 
         /// Title of website
         #[serde(skip_serializing_if = "Option::is_none")]
-        title: Option<String>,
+        pub title: Option<String>,
         /// Description of website
         #[serde(skip_serializing_if = "Option::is_none")]
-        description: Option<String>,
+        pub description: Option<String>,
         /// Embedded image
         #[serde(skip_serializing_if = "Option::is_none")]
-        image: Option<Image>,
+        pub image: Option<Image>,
         /// Embedded video
         #[serde(skip_serializing_if = "Option::is_none")]
-        video: Option<Video>,
+        pub video: Option<Video>,
 
         // #[serde(skip_serializing_if = "Option::is_none")]
         // opengraph_type: Option<String>,
         /// Site name
         #[serde(skip_serializing_if = "Option::is_none")]
-        site_name: Option<String>,
+        pub site_name: Option<String>,
         /// URL to site icon
         #[serde(skip_serializing_if = "Option::is_none")]
-        icon_url: Option<String>,
+        pub icon_url: Option<String>,
         /// CSS Colour
         #[serde(skip_serializing_if = "Option::is_none")]
-        colour: Option<String>,
+        pub colour: Option<String>,
     }
 
     /// Text Embed

crates/core/models/src/v0/messages.rs

@@ -207,6 +207,8 @@ auto_derived!(
         pub timestamp: u64,
         /// URL to open when clicking notification
         pub url: String,
+        /// The message object itself, to send to clients for processing
+        pub message: Message,
     }
 
     /// Representation of a text embed before it is sent.
@@ -453,15 +455,30 @@ impl PushNotification {
             format!("{}/assets/logo.png", config.hosts.app)
         };
 
-        let image = msg.attachments.and_then(|attachments| {
+        let image = msg.attachments.as_ref().and_then(|attachments| {
             attachments
                 .first()
                 .map(|v| format!("{}/attachments/{}", config.hosts.autumn, v.id))
         });
 
-        let body = if let Some(sys) = msg.system {
+        let body = if let Some(ref sys) = msg.system {
-            sys.into()
+            sys.clone().into()
-        } else if let Some(text) = msg.content {
+        } else if let Some(ref text) = msg.content {
+            text.clone()
+        } else if let Some(text) = msg.embeds.as_ref().and_then(|embeds| match embeds.first() {
+            Some(Embed::Image(_)) => Some("Sent an image".to_string()),
+            Some(Embed::Video(_)) => Some("Sent a video".to_string()),
+            Some(Embed::Text(e)) => e
+                .description
+                .clone()
+                .or(e.title.clone().or(Some("Empty Embed".to_string()))),
+            Some(Embed::Website(e)) => e.title.clone().or(e
+                .description
+                .clone()
+                .or(e.site_name.clone().or(Some("Empty Embed".to_string())))),
+            Some(Embed::None) => Some("Empty Message".to_string()), // ???
+            None => Some("Empty Message".to_string()),              // ??
+        }) {
             text
         } else {
             "Empty Message".to_string()
@@ -482,6 +499,7 @@ impl PushNotification {
             tag: channel_id.to_string(),
             timestamp,
             url: format!("{}/channel/{}/{}", config.hosts.app, channel_id, msg.id),
+            message: msg,
         }
     }
 }