Skip to main content

this/events/sinks/
device_tokens.rs

1//! Device token store for push notifications
2//!
3//! Stores device push tokens (Expo, APNs, FCM) per user, enabling
4//! the `PushNotificationSink` to look up where to send push notifications.
5//!
6//! # Token lifecycle
7//!
8//! 1. Client registers a token: `register(user_id, token, platform)`
9//! 2. Push sink delivers to all user tokens: `get_tokens(user_id)`
10//! 3. Client unregisters on logout: `unregister(user_id, token)`
11
12use chrono::{DateTime, Utc};
13use serde::{Deserialize, Serialize};
14use std::collections::HashMap;
15use tokio::sync::RwLock;
16
17/// Supported push notification platforms
18#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
19#[serde(rename_all = "snake_case")]
20pub enum Platform {
21    /// iOS (APNs or Expo)
22    Ios,
23    /// Android (FCM or Expo)
24    Android,
25    /// Web (Web Push or Expo)
26    Web,
27}
28
29impl std::fmt::Display for Platform {
30    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31        match self {
32            Platform::Ios => write!(f, "ios"),
33            Platform::Android => write!(f, "android"),
34            Platform::Web => write!(f, "web"),
35        }
36    }
37}
38
39/// A registered device push token
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct DeviceToken {
42    /// The push token string (e.g., Expo push token "ExponentPushToken\[xxx\]")
43    pub token: String,
44
45    /// Platform this token belongs to
46    pub platform: Platform,
47
48    /// When this token was registered
49    pub registered_at: DateTime<Utc>,
50}
51
52/// In-memory device token store
53///
54/// Thread-safe store for device tokens, keyed by user ID.
55/// Each user can have multiple tokens (multiple devices).
56#[derive(Debug)]
57pub struct DeviceTokenStore {
58    tokens: RwLock<HashMap<String, Vec<DeviceToken>>>,
59}
60
61impl DeviceTokenStore {
62    /// Create an empty device token store
63    pub fn new() -> Self {
64        Self {
65            tokens: RwLock::new(HashMap::new()),
66        }
67    }
68
69    /// Register a device token for a user
70    ///
71    /// If the same token already exists for this user, it is updated
72    /// (platform and registered_at are refreshed).
73    pub async fn register(&self, user_id: &str, token: String, platform: Platform) {
74        let mut store = self.tokens.write().await;
75        let user_tokens = store.entry(user_id.to_string()).or_default();
76
77        // Update existing token or add new
78        if let Some(existing) = user_tokens.iter_mut().find(|t| t.token == token) {
79            existing.platform = platform;
80            existing.registered_at = Utc::now();
81        } else {
82            user_tokens.push(DeviceToken {
83                token,
84                platform,
85                registered_at: Utc::now(),
86            });
87        }
88    }
89
90    /// Unregister a device token for a user
91    ///
92    /// Returns true if the token was found and removed.
93    pub async fn unregister(&self, user_id: &str, token: &str) -> bool {
94        let mut store = self.tokens.write().await;
95        if let Some(user_tokens) = store.get_mut(user_id) {
96            let len_before = user_tokens.len();
97            user_tokens.retain(|t| t.token != token);
98            return user_tokens.len() < len_before;
99        }
100        false
101    }
102
103    /// Get all device tokens for a user
104    pub async fn get_tokens(&self, user_id: &str) -> Vec<DeviceToken> {
105        let store = self.tokens.read().await;
106        store.get(user_id).cloned().unwrap_or_default()
107    }
108
109    /// Get token count for a user
110    pub async fn token_count(&self, user_id: &str) -> usize {
111        let store = self.tokens.read().await;
112        store.get(user_id).map(|t| t.len()).unwrap_or(0)
113    }
114
115    /// Remove all tokens for a user (e.g., on account deletion)
116    pub async fn remove_all(&self, user_id: &str) -> usize {
117        let mut store = self.tokens.write().await;
118        store.remove(user_id).map(|t| t.len()).unwrap_or(0)
119    }
120}
121
122impl Default for DeviceTokenStore {
123    fn default() -> Self {
124        Self::new()
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131
132    #[tokio::test]
133    async fn test_register_and_get_tokens() {
134        let store = DeviceTokenStore::new();
135
136        store
137            .register("user-A", "token-1".to_string(), Platform::Ios)
138            .await;
139        store
140            .register("user-A", "token-2".to_string(), Platform::Android)
141            .await;
142
143        let tokens = store.get_tokens("user-A").await;
144        assert_eq!(tokens.len(), 2);
145        assert_eq!(tokens[0].token, "token-1");
146        assert_eq!(tokens[0].platform, Platform::Ios);
147        assert_eq!(tokens[1].token, "token-2");
148        assert_eq!(tokens[1].platform, Platform::Android);
149    }
150
151    #[tokio::test]
152    async fn test_unregister() {
153        let store = DeviceTokenStore::new();
154
155        store
156            .register("user-A", "token-1".to_string(), Platform::Ios)
157            .await;
158        store
159            .register("user-A", "token-2".to_string(), Platform::Android)
160            .await;
161
162        assert_eq!(store.token_count("user-A").await, 2);
163
164        let removed = store.unregister("user-A", "token-1").await;
165        assert!(removed);
166        assert_eq!(store.token_count("user-A").await, 1);
167
168        let tokens = store.get_tokens("user-A").await;
169        assert_eq!(tokens[0].token, "token-2");
170    }
171
172    #[tokio::test]
173    async fn test_unregister_nonexistent() {
174        let store = DeviceTokenStore::new();
175        assert!(!store.unregister("user-A", "nonexistent").await);
176    }
177
178    #[tokio::test]
179    async fn test_register_duplicate_updates() {
180        let store = DeviceTokenStore::new();
181
182        store
183            .register("user-A", "token-1".to_string(), Platform::Ios)
184            .await;
185        // Re-register same token but different platform
186        store
187            .register("user-A", "token-1".to_string(), Platform::Android)
188            .await;
189
190        let tokens = store.get_tokens("user-A").await;
191        assert_eq!(tokens.len(), 1); // No duplicate
192        assert_eq!(tokens[0].platform, Platform::Android); // Updated
193    }
194
195    #[tokio::test]
196    async fn test_get_tokens_empty() {
197        let store = DeviceTokenStore::new();
198        assert!(store.get_tokens("nonexistent").await.is_empty());
199        assert_eq!(store.token_count("nonexistent").await, 0);
200    }
201
202    #[tokio::test]
203    async fn test_remove_all() {
204        let store = DeviceTokenStore::new();
205
206        store
207            .register("user-A", "token-1".to_string(), Platform::Ios)
208            .await;
209        store
210            .register("user-A", "token-2".to_string(), Platform::Android)
211            .await;
212
213        let removed = store.remove_all("user-A").await;
214        assert_eq!(removed, 2);
215        assert!(store.get_tokens("user-A").await.is_empty());
216    }
217
218    #[tokio::test]
219    async fn test_separate_users() {
220        let store = DeviceTokenStore::new();
221
222        store
223            .register("user-A", "token-a".to_string(), Platform::Ios)
224            .await;
225        store
226            .register("user-B", "token-b".to_string(), Platform::Android)
227            .await;
228
229        assert_eq!(store.token_count("user-A").await, 1);
230        assert_eq!(store.token_count("user-B").await, 1);
231
232        // Unregistering from one user doesn't affect another
233        store.unregister("user-A", "token-a").await;
234        assert_eq!(store.token_count("user-A").await, 0);
235        assert_eq!(store.token_count("user-B").await, 1);
236    }
237}