From 853366a297f71a17d3045757c4eb04b6aa30462a Mon Sep 17 00:00:00 2001
From: Paul Makles <me@insrt.uk>
Date: Sun, 27 Aug 2023 15:18:19 +0100
Subject: [PATCH] feat: add support for Firebase Cloud Messaging

---
 .../database/src/models/messages/model.rs     |  4 +-
 crates/core/database/src/tasks/web_push.rs    | 63 +++++++++++++++----
 2 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/crates/core/database/src/models/messages/model.rs b/crates/core/database/src/models/messages/model.rs
index 92a1c6b6..a2d3aefd 100644
--- a/crates/core/database/src/models/messages/model.rs
+++ b/crates/core/database/src/models/messages/model.rs
@@ -2,7 +2,6 @@ use indexmap::{IndexMap, IndexSet};
 use iso8601_timestamp::Timestamp;
 use revolt_models::v0::{Embed, MessageAuthor, MessageSort, MessageWebhook, PushNotification};
 use revolt_result::Result;
-use serde_json::json;
 use ulid::Ulid;
 
 use crate::{
@@ -268,8 +267,7 @@ impl Message {
                     _ => vec![],
                 }
             },
-            json!(PushNotification::from(self.clone().into(), Some(author), &channel.id()).await)
-                .to_string(),
+            PushNotification::from(self.clone().into(), Some(author), &channel.id()).await,
         )
         .await;
 
diff --git a/crates/core/database/src/tasks/web_push.rs b/crates/core/database/src/tasks/web_push.rs
index fa74a3ca..8a055afe 100644
--- a/crates/core/database/src/tasks/web_push.rs
+++ b/crates/core/database/src/tasks/web_push.rs
@@ -8,7 +8,9 @@ use base64::{
 use deadqueue::limited::Queue;
 use once_cell::sync::Lazy;
 use revolt_config::config;
+use revolt_models::v0::PushNotification;
 use revolt_presence::filter_online;
+use serde_json::json;
 use web_push::{
     ContentEncoding, IsahcWebPushClient, SubscriptionInfo, SubscriptionKeys, VapidSignatureBuilder,
     WebPushClient, WebPushMessageBuilder,
@@ -19,14 +21,14 @@ use web_push::{
 struct PushTask {
     /// User IDs of the targets that are to receive this notification
     recipients: Vec<String>,
-    /// Raw JSON payload to send to clients
-    payload: String,
+    /// Push Notification
+    payload: PushNotification,
 }
 
 static Q: Lazy<Queue<PushTask>> = Lazy::new(|| Queue::new(10_000));
 
 /// Queue a new task for a worker
-pub async fn queue(recipients: Vec<String>, payload: String) {
+pub async fn queue(recipients: Vec<String>, payload: PushNotification) {
     if recipients.is_empty() {
         return;
     }
@@ -48,8 +50,15 @@ pub async fn queue(recipients: Vec<String>, payload: String) {
 /// Start a new worker
 pub async fn worker(db: Database) {
     let config = config().await;
-    let client = IsahcWebPushClient::new().unwrap();
-    let key = engine::general_purpose::URL_SAFE_NO_PAD
+
+    let web_push_client = IsahcWebPushClient::new().unwrap();
+    let fcm_client = if config.api.fcm.api_key.is_empty() {
+        None
+    } else {
+        Some(fcm::Client::new())
+    };
+
+    let web_push_private_key = engine::general_purpose::URL_SAFE_NO_PAD
         .decode(config.api.vapid.private_key)
         .expect("valid `VAPID_PRIVATE_KEY`");
 
@@ -61,7 +70,37 @@ pub async fn worker(db: Database) {
                 if let Some(sub) = session.subscription {
                     if sub.endpoint == "fcm" {
                         // Use Firebase Cloud Messaging
-                        info!("Skipping FCM client, unimplemented...");
+                        if let Some(client) = &fcm_client {
+                            let PushNotification {
+                                author,
+                                icon,
+                                image: _,
+                                body,
+                                tag,
+                                timestamp: _,
+                                url: _,
+                            } = &task.payload;
+
+                            let mut notification = fcm::NotificationBuilder::new();
+                            notification.title(author);
+                            notification.icon(icon);
+                            notification.body(body);
+                            notification.tag(tag);
+                            // TODO: expand support for fields
+                            let notification = notification.finalize();
+
+                            let mut message_builder =
+                                fcm::MessageBuilder::new(&config.api.fcm.api_key, &sub.auth);
+                            message_builder.notification(notification);
+
+                            if let Err(err) = client.send(message_builder.finalize()).await {
+                                error!("Failed to send FCM notification! {:?}", err);
+                            } else {
+                                info!("Sent FCM notification to {:?}.", session.id);
+                            }
+                        } else {
+                            info!("No FCM token was specified!");
+                        }
                     } else {
                         // Use Web Push Standard
                         let subscription = SubscriptionInfo {
@@ -73,20 +112,20 @@ pub async fn worker(db: Database) {
                         };
 
                         match VapidSignatureBuilder::from_pem(
-                            std::io::Cursor::new(&key),
+                            std::io::Cursor::new(&web_push_private_key),
                             &subscription,
                         ) {
                             Ok(sig_builder) => match sig_builder.build() {
                                 Ok(signature) => {
                                     let mut builder = WebPushMessageBuilder::new(&subscription);
                                     builder.set_vapid_signature(signature);
-                                    builder.set_payload(
-                                        ContentEncoding::AesGcm,
-                                        task.payload.as_bytes(),
-                                    );
+
+                                    let payload = json!(task.payload).to_string();
+                                    builder
+                                        .set_payload(ContentEncoding::AesGcm, payload.as_bytes());
 
                                     match builder.build() {
-                                        Ok(msg) => match client.send(msg).await {
+                                        Ok(msg) => match web_push_client.send(msg).await {
                                             Ok(_) => {
                                                 info!(
                                                     "Sent Web Push notification to {:?}.",