feat: implement apple push notifications

crates/core/config/Revolt.toml

@@ -29,6 +29,11 @@ public_key = "BGcvgR-i2z4IQ5Mw841vJvkLjt8wY-FjmWrw83jOLCY52qcGZS0OF7nfLzuYbjsQIS
 [api.fcm]
 api_key = ""
 
+[api.apn]
+pkcs8 = ""
+key_id = ""
+team_id = ""
+
 [api.security]
 authifier_shield_key = ""
 voso_legacy_token = ""

crates/core/config/src/lib.rs

@@ -6,6 +6,8 @@ use futures_locks::RwLock;
 use once_cell::sync::Lazy;
 use serde::Deserialize;
 
+pub use sentry::capture_error;
+
 #[cfg(not(debug_assertions))]
 use std::env;
 
@@ -75,6 +77,13 @@ pub struct ApiFcm {
     pub api_key: String,
 }
 
+#[derive(Deserialize, Debug, Clone)]
+pub struct ApiApn {
+    pub pkcs8: String,
+    pub key_id: String,
+    pub team_id: String,
+}
+
 #[derive(Deserialize, Debug, Clone)]
 pub struct ApiSecurityCaptcha {
     pub hcaptcha_key: String,
@@ -100,6 +109,7 @@ pub struct Api {
     pub smtp: ApiSmtp,
     pub vapid: ApiVapid,
     pub fcm: ApiFcm,
+    pub apn: ApiApn,
     pub security: ApiSecurity,
     pub workers: ApiWorkers,
 }

crates/core/database/Cargo.toml

@@ -87,6 +87,9 @@ 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 = [
+    "ring",
+] }
 
 # Authifier
 authifier = { version = "1.0.8" }

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

@@ -58,4 +58,7 @@ pub trait AbstractUsers: Sync + Send {
 
     /// Delete a user by their id
     async fn delete_user(&self, id: &str) -> Result<()>;
+
+    /// Remove push subscription for a session by session id (TODO: remove)
+    async fn remove_push_subscription_by_session_id(&self, session_id: &str) -> Result<()>;
 }

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

@@ -316,6 +316,25 @@ impl AbstractUsers for MongoDb {
     async fn delete_user(&self, id: &str) -> Result<()> {
         query!(self, delete_one_by_id, COL, id).map(|_| ())
     }
+
+    /// Remove push subscription for a session by session id (TODO: remove)
+    async fn remove_push_subscription_by_session_id(&self, session_id: &str) -> Result<()> {
+        self.col::<User>("sessions")
+            .update_one(
+                doc! {
+                    "_id": session_id
+                },
+                doc! {
+                    "$unset": {
+                        "subscription": 1
+                    }
+                },
+                None,
+            )
+            .await
+            .map(|_| ())
+            .map_err(|_| create_database_error!("update_one", COL))
+    }
 }
 
 impl IntoDocumentPath for FieldsUser {

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

@@ -163,4 +163,9 @@ impl AbstractUsers for ReferenceDb {
             Err(create_error!(NotFound))
         }
     }
+
+    /// Remove push subscription for a session by session id (TODO: remove)
+    async fn remove_push_subscription_by_session_id(&self, _session_id: &str) -> Result<()> {
+        todo!()
+    }
 }

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

@@ -0,0 +1,113 @@
+use std::io::Cursor;
+
+use base64::{
+    engine::{self},
+    Engine as _,
+};
+use deadqueue::limited::Queue;
+use once_cell::sync::Lazy;
+use revolt_a2::{Client, ClientConfig, DefaultNotificationBuilder};
+use revolt_a2::{Error, ErrorBody, ErrorReason, NotificationBuilder, Response};
+use revolt_config::config;
+use revolt_models::v0::PushNotification;
+
+use crate::Database;
+
+/// Task information
+#[derive(Debug)]
+pub struct ApnTask {
+    /// Session Id
+    session_id: String,
+
+    /// Device token
+    device_token: String,
+
+    /// Title
+    title: String,
+
+    /// Body
+    body: String,
+
+    /// Thread Id
+    thread_id: String,
+}
+
+impl ApnTask {
+    pub fn from_notification(
+        session_id: String,
+        device_token: String,
+        notification: &PushNotification,
+    ) -> ApnTask {
+        ApnTask {
+            session_id,
+            device_token,
+            title: notification.author.to_string(),
+            body: notification.body.to_string(),
+            thread_id: notification.tag.to_string(),
+        }
+    }
+}
+
+static Q: Lazy<Queue<ApnTask>> = Lazy::new(|| Queue::new(10_000));
+
+/// Queue a new task for a worker
+pub async fn queue(task: ApnTask) {
+    Q.try_push(task).ok();
+    info!("Queue is using {} slots from {}.", Q.len(), Q.capacity());
+}
+
+/// 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 pkcs8 = engine::general_purpose::STANDARD
+        .decode(config.api.apn.pkcs8)
+        .expect("valid `pcks8`");
+
+    let client = Client::token(
+        &mut Cursor::new(pkcs8),
+        config.api.apn.key_id,
+        config.api.apn.team_id,
+        ClientConfig::default(),
+    )
+    .expect("could not create APN client");
+
+    loop {
+        let task = Q.pop().await;
+        let payload = DefaultNotificationBuilder::new()
+            .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 err {
+                Error::ResponseError(Response {
+                    error:
+                        Some(ErrorBody {
+                            reason: ErrorReason::BadDeviceToken | ErrorReason::Unregistered,
+                            ..
+                        }),
+                    ..
+                }) => {
+                    if let Err(err) = db
+                        .remove_push_subscription_by_session_id(&task.session_id)
+                        .await
+                    {
+                        revolt_config::capture_error(&err);
+                    }
+                }
+                err => {
+                    revolt_config::capture_error(&err);
+                }
+            }
+        }
+    }
+}

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

@@ -8,12 +8,15 @@ use std::time::Instant;
 const WORKER_COUNT: usize = 5;
 
 pub mod ack;
+pub mod apple_notifications;
 pub mod last_message_id;
 pub mod process_embeds;
 pub mod web_push;
 
 /// Spawn background workers
 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(last_message_id::worker(db.clone()));

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

@@ -6,6 +6,7 @@ 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;
@@ -16,6 +17,8 @@ use web_push::{
     WebPushClient, WebPushMessageBuilder,
 };
 
+use super::apple_notifications;
+
 /// Task information
 #[derive(Debug)]
 struct PushTask {
@@ -101,6 +104,15 @@ pub async fn worker(db: Database) {
                         } else {
                             info!("No FCM token was specified!");
                         }
+                    } else if sub.endpoint == "apn" {
+                        apple_notifications::queue(
+                            apple_notifications::ApnTask::from_notification(
+                                session.id,
+                                sub.auth,
+                                &task.payload,
+                            ),
+                        )
+                        .await;
                     } else {
                         // Use Web Push Standard
                         let subscription = SubscriptionInfo {