r_token/
memory.rs

1use crate::RTokenError;
2use crate::models::RTokenInfo;
3use chrono::Utc;
4use std::{
5    collections::HashMap,
6    sync::{Arc, Mutex},
7};
8
9/// Issues, stores, and revokes authentication tokens.
10///
11/// This type is designed to be stored in actix-web application state
12/// (e.g. `web::Data<RTokenManager>`). Internally it uses an `Arc<Mutex<...>>`,
13/// so `Clone` creates another handle to the same shared store.
14///
15/// Tokens are generated as UUID v4 strings. Each token is associated with:
16/// - a user id (`String`)
17/// - an expiration timestamp (Unix epoch milliseconds)
18///
19/// ## 繁體中文
20///
21/// 負責簽發、儲存與註銷 token 的管理器。
22///
23/// 一般會放在 actix-web 的 application state 中(例如 `web::Data<RTokenManager>`)。
24/// 內部以 `Arc<Mutex<...>>` 共享狀態,因此 `Clone` 只是在同一份映射表上增加一個引用。
25///
26/// token 以 UUID v4 字串產生,並會綁定:
27/// - 使用者 id(`String`)
28/// - 到期時間(Unix epoch 毫秒)
29#[derive(Clone, Default)]
30pub struct RTokenManager {
31    /// In-memory token store.
32    ///
33    /// ## 繁體中文
34    ///
35    /// 記憶體中的 token 儲存表。
36    // store: Arc<Mutex<HashMap<String, String>>>,
37    store: Arc<Mutex<HashMap<String, RTokenInfo>>>,
38}
39
40impl RTokenManager {
41    /// Creates an empty manager.
42    ///
43    /// ## 繁體中文
44    ///
45    /// 建立一個空的管理器。
46    pub fn new() -> Self {
47        Self {
48            store: Arc::new(Mutex::new(HashMap::new())),
49        }
50    }
51
52    /// Issues a new token for the given user id.
53    ///
54    /// `expire_time` is treated as TTL in seconds. The token will be considered invalid
55    /// once the stored expiration timestamp is earlier than the current time.
56    ///
57    /// Returns [`RTokenError::MutexPoisoned`] if the internal mutex is poisoned.
58    ///
59    /// ## 繁體中文
60    ///
61    /// 為指定使用者 id 簽發新 token。
62    ///
63    /// `expire_time` 會被視為 TTL(秒)。當儲存的到期時間早於目前時間時,token 會被視為無效。
64    ///
65    /// 若內部 mutex 發生 poisoned,會回傳 [`RTokenError::MutexPoisoned`]。
66    // pub fn login(&self, id: &str, expire_time: u64,role:impl Into<Vec<String>>) -> Result<String, RTokenError> {
67    pub fn login(&self, id: &str, expire_time: u64) -> Result<String, RTokenError> {
68        let token = uuid::Uuid::new_v4().to_string();
69        // Acquire the write lock and insert the token-user mapping into the store
70        // 获取写锁并将 Token-用户映射关系插入到存储中
71        // #[allow(clippy::unwrap_used)]
72        // self.store.lock().unwrap().insert(token.clone(), id.to_string());
73        let now = Utc::now();
74        let ttl = chrono::Duration::seconds(expire_time as i64);
75        let deadline = now + ttl;
76        let expire_time = deadline.timestamp_millis() as u64;
77        let info = RTokenInfo {
78            user_id: id.to_string(),
79            expire_at: expire_time,
80            roles: Vec::new(),
81        };
82        self.store
83            .lock()
84            .map_err(|_| RTokenError::MutexPoisoned)?
85            .insert(token.clone(), info);
86        Ok(token)
87    }
88
89    #[cfg(feature = "rbac")]
90    pub fn login_with_roles(
91        &self,
92        id: &str,
93        expire_time: u64,
94        role: impl Into<Vec<String>>,
95    ) -> Result<String, RTokenError> {
96        let token = uuid::Uuid::new_v4().to_string();
97        let now = Utc::now();
98        let ttl = chrono::Duration::seconds(expire_time as i64);
99        let deadline = now + ttl;
100        let expire_time = deadline.timestamp_millis() as u64;
101        let info = RTokenInfo {
102            user_id: id.to_string(),
103            expire_at: expire_time,
104            roles: role.into(),
105        };
106        self.store
107            .lock()
108            .map_err(|_| RTokenError::MutexPoisoned)?
109            .insert(token.clone(), info);
110        Ok(token)
111    }
112
113    // pub fn set_role(&self, token: &str, role: impl Into<Vec<String>>) -> Result<(), RTokenError> {
114    #[cfg(feature = "rbac")]
115    pub fn set_roles(&self, token: &str, roles: impl Into<Vec<String>>) -> Result<(), RTokenError> {
116        let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
117        if let Some(info) = store.get_mut(token) {
118            info.roles = roles.into();
119        }
120        Ok(())
121    }
122
123    #[cfg(feature = "rbac")]
124    pub fn get_roles(&self, token: &str) -> Result<Option<Vec<String>>, RTokenError> {
125        let store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
126        Ok(store.get(token).map(|info| info.roles.clone()))
127    }
128
129    /// Revokes a token by removing it from the in-memory store.
130    ///
131    /// This operation is idempotent: removing a non-existing token is treated as success.
132    /// Returns [`RTokenError::MutexPoisoned`] if the internal mutex is poisoned.
133    ///
134    /// ## 繁體中文
135    ///
136    /// 從記憶體儲存表中移除 token,以達到註銷效果。
137    ///
138    /// 此操作具冪等性:移除不存在的 token 也視為成功。
139    /// 若內部 mutex 發生 poisoned,會回傳 [`RTokenError::MutexPoisoned`]。
140    pub fn logout(&self, token: &str) -> Result<(), RTokenError> {
141        // self.store.lock().unwrap().remove(token);
142        self.store
143            .lock()
144            .map_err(|_| RTokenError::MutexPoisoned)?
145            .remove(token);
146        Ok(())
147    }
148
149    pub fn expires_at(&self, token: &str) -> Result<Option<u64>, RTokenError> {
150        let store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
151        Ok(store.get(token).map(|info| info.expire_at))
152    }
153
154    pub fn ttl_seconds(&self, token: &str) -> Result<Option<i64>, RTokenError> {
155        let now_ms = Utc::now().timestamp_millis() as u64;
156        let store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
157        let Some(expire_at) = store.get(token).map(|info| info.expire_at) else {
158            return Ok(None);
159        };
160
161        if expire_at <= now_ms {
162            return Ok(Some(0));
163        }
164
165        let remaining_ms = expire_at - now_ms;
166        let remaining_seconds = remaining_ms.div_ceil(1000) as i64;
167        Ok(Some(remaining_seconds))
168    }
169
170    pub fn renew(&self, token: &str, ttl_seconds: u64) -> Result<bool, RTokenError> {
171        let now = Utc::now();
172        let ttl = chrono::Duration::seconds(ttl_seconds as i64);
173        let expire_at = (now + ttl).timestamp_millis() as u64;
174
175        let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
176        let Some(info) = store.get_mut(token) else {
177            return Ok(false);
178        };
179
180        if info.expire_at < Utc::now().timestamp_millis() as u64 {
181            store.remove(token);
182            return Ok(false);
183        }
184
185        info.expire_at = expire_at;
186        Ok(true)
187    }
188
189    pub fn rotate(&self, token: &str, ttl_seconds: u64) -> Result<Option<String>, RTokenError> {
190        let now = Utc::now();
191        let ttl = chrono::Duration::seconds(ttl_seconds as i64);
192        let expire_at = (now + ttl).timestamp_millis() as u64;
193
194        let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
195        let Some(info) = store.get(token).cloned() else {
196            return Ok(None);
197        };
198
199        if info.expire_at < Utc::now().timestamp_millis() as u64 {
200            store.remove(token);
201            return Ok(None);
202        }
203
204        let new_token = uuid::Uuid::new_v4().to_string();
205        let new_info = RTokenInfo {
206            user_id: info.user_id,
207            expire_at,
208            roles: info.roles,
209        };
210
211        store.remove(token);
212        store.insert(new_token.clone(), new_info);
213        Ok(Some(new_token))
214    }
215
216    pub fn prune_expired(&self) -> Result<usize, RTokenError> {
217        let now = Utc::now().timestamp_millis() as u64;
218        let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
219
220        let original_len = store.len();
221        store.retain(|_token, info| info.expire_at >= now);
222        Ok(original_len - store.len())
223    }
224
225    /// Validates a token and returns the associated user id if present.
226    ///
227    /// Behavior:
228    /// - Returns `Ok(Some(user_id))` when the token exists and is not expired.
229    /// - Returns `Ok(None)` when the token does not exist or is expired.
230    /// - Expired tokens are removed from the in-memory store during validation.
231    ///
232    /// ## 繁體中文
233    ///
234    /// 驗證 token,若有效則回傳對應的使用者 id。
235    ///
236    /// 行為:
237    /// - token 存在且未過期:回傳 `Ok(Some(user_id))`
238    /// - token 不存在或已過期:回傳 `Ok(None)`
239    /// - 若 token 已過期,會在驗證時從記憶體儲存表中移除
240    pub fn validate(&self, token: &str) -> Result<Option<String>, RTokenError> {
241        #[cfg(feature = "rbac")]
242        {
243            Ok(self
244                .validate_with_roles(token)?
245                .map(|(user_id, _roles)| user_id))
246        }
247
248        #[cfg(not(feature = "rbac"))]
249        {
250            let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
251            let Some(info) = store.get(token) else {
252                return Ok(None);
253            };
254
255            if info.expire_at < Utc::now().timestamp_millis() as u64 {
256                store.remove(token);
257                return Ok(None);
258            }
259
260            Ok(Some(info.user_id.clone()))
261        }
262    }
263
264    #[cfg(feature = "rbac")]
265    /// Validates a token and returns both user id and roles (RBAC enabled).
266    ///
267    /// This has the same expiration behavior as [`RTokenManager::validate`].
268    ///
269    /// ## 繁體中文
270    ///
271    /// 驗證 token,並在 RBAC 啟用時同時回傳使用者 id 與角色列表。
272    ///
273    /// 到期行為與 [`RTokenManager::validate`] 相同。
274    pub fn validate_with_roles(
275        &self,
276        token: &str,
277    ) -> Result<Option<(String, Vec<String>)>, RTokenError> {
278        let mut store = self.store.lock().map_err(|_| RTokenError::MutexPoisoned)?;
279        let Some(info) = store.get(token) else {
280            return Ok(None);
281        };
282
283        if info.expire_at < Utc::now().timestamp_millis() as u64 {
284            store.remove(token);
285            return Ok(None);
286        }
287
288        Ok(Some((info.user_id.clone(), info.roles.clone())))
289    }
290}
291
292/// An authenticated request context extracted from actix-web.
293///
294/// If extraction succeeds, `id` is the user id previously passed to
295/// [`RTokenManager::login`], and `token` is the original token from the request.
296///
297/// The token is read from `Authorization` header. Both of the following formats
298/// are accepted:
299/// - `Authorization: <token>`
300/// - `Authorization: Bearer <token>`
301///
302/// ## 繁體中文
303///
304/// 由 actix-web 自動抽取的已驗證使用者上下文。
305///
306/// Extractor 成功時:
307/// - `id` 會是先前傳給 [`RTokenManager::login`] 的使用者 id
308/// - `token` 會是請求中帶來的 token 原文
309///
310/// token 會從 `Authorization` header 讀取,支援以下格式:
311/// - `Authorization: <token>`
312/// - `Authorization: Bearer <token>`
313#[cfg(feature = "actix")]
314#[derive(Debug)]
315pub struct RUser {
316    /// The user id associated with the token.
317    ///
318    /// ## 繁體中文
319    ///
320    /// 與 token 綁定的使用者 id。
321    pub id: String,
322
323    /// The raw token string from the request.
324    ///
325    /// ## 繁體中文
326    ///
327    /// 來自請求的 token 字串原文。
328    pub token: String,
329    #[cfg(feature = "rbac")]
330    pub roles: Vec<String>,
331}
332
333#[cfg(feature = "rbac")]
334impl RUser {
335    pub fn has_role(&self, role: &str) -> bool {
336        self.roles.iter().any(|r| r == role)
337    }
338}
339
340/// Extracts [`RUser`] from an actix-web request.
341///
342/// Failure modes:
343/// - 500: manager is missing from `app_data`, or mutex is poisoned
344/// - 401: token is missing, invalid, or expired
345///
346/// ## 繁體中文
347///
348/// 從 actix-web 請求中抽取 [`RUser`]。
349///
350/// 失敗情況:
351/// - 500:`app_data` 中找不到管理器,或 mutex poisoned
352/// - 401:token 缺失、無效、或已過期
353#[cfg(feature = "actix")]
354impl actix_web::FromRequest for RUser {
355    type Error = actix_web::Error;
356    type Future = std::future::Ready<Result<Self, Self::Error>>;
357
358    fn from_request(
359        req: &actix_web::HttpRequest,
360        _payload: &mut actix_web::dev::Payload,
361    ) -> Self::Future {
362        use actix_web::web;
363
364        // 獲取管理器
365        let manager = match req.app_data::<web::Data<RTokenManager>>() {
366            Some(m) => m,
367            None => {
368                return std::future::ready(Err(actix_web::error::ErrorInternalServerError(
369                    "Token manager not found",
370                )));
371            }
372        };
373        let token = match crate::extract_token_from_request(req) {
374            Some(token) => token,
375            None => {
376                return std::future::ready(Err(actix_web::error::ErrorUnauthorized(
377                    "Unauthorized",
378                )));
379            }
380        };
381
382        #[cfg(feature = "rbac")]
383        {
384            let user_info = match manager.validate_with_roles(&token) {
385                Ok(user_info) => user_info,
386                Err(_) => {
387                    return std::future::ready(Err(actix_web::error::ErrorInternalServerError(
388                        "Mutex poisoned",
389                    )));
390                }
391            };
392
393            if let Some((user_id, roles)) = user_info {
394                return std::future::ready(Ok(RUser {
395                    id: user_id,
396                    token,
397                    roles,
398                }));
399            }
400
401            std::future::ready(Err(actix_web::error::ErrorUnauthorized("Invalid token")))
402        }
403
404        #[cfg(not(feature = "rbac"))]
405        {
406            let user_id = match manager.validate(&token) {
407                Ok(user_id) => user_id,
408                Err(_) => {
409                    return std::future::ready(Err(actix_web::error::ErrorInternalServerError(
410                        "Mutex poisoned",
411                    )));
412                }
413            };
414
415            if let Some(user_id) = user_id {
416                return std::future::ready(Ok(RUser { id: user_id, token }));
417            }
418
419            std::future::ready(Err(actix_web::error::ErrorUnauthorized("Invalid token")))
420        }
421    }
422}