1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

use std::collections::HashSet;

use authifier::Database;
use base64::{
    engine::{self},
    Engine as _,
};
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,
};

/// Task information
#[derive(Debug)]
struct PushTask {
    /// User IDs of the targets that are to receive this notification
    recipients: Vec<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: PushNotification) {
    if recipients.is_empty() {
        return;
    }

    let online_ids = filter_online(&recipients).await;
    let recipients = (&recipients.into_iter().collect::<HashSet<String>>() - &online_ids)
        .into_iter()
        .collect::<Vec<String>>();

    Q.try_push(PushTask {
        recipients,
        payload,
    })
    .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;

    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`");

    loop {
        let task = Q.pop().await;

        if let Ok(sessions) = db.find_sessions_with_subscription(&task.recipients).await {
            for session in sessions {
                if let Some(sub) = session.subscription {
                    if sub.endpoint == "fcm" {
                        // Use Firebase Cloud Messaging
                        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 {
                            endpoint: sub.endpoint,
                            keys: SubscriptionKeys {
                                auth: sub.auth,
                                p256dh: sub.p256dh,
                            },
                        };

                        match VapidSignatureBuilder::from_pem(
                            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);

                                    let payload = json!(task.payload).to_string();
                                    builder
                                        .set_payload(ContentEncoding::AesGcm, payload.as_bytes());

                                    match builder.build() {
                                        Ok(msg) => match web_push_client.send(msg).await {
                                            Ok(_) => {
                                                info!(
                                                    "Sent Web Push notification to {:?}.",
                                                    session.id
                                                )
                                            }
                                            Err(err) => {
                                                error!("Hit error sending Web Push! {:?}", err)
                                            }
                                        },
                                        Err(err) => {
                                            error!(
                                                "Failed to build message for {}! {:?}",
                                                session.user_id, err
                                            )
                                        }
                                    }
                                }
                                Err(err) => error!(
                                    "Failed to build signature for {}! {:?}",
                                    session.user_id, err
                                ),
                            },
                            Err(err) => error!(
                                "Failed to create signature builder for {}! {:?}",
                                session.user_id, err
                            ),
                        }
                    }
                }
            }
        }
    }
}