Skip to main content

security_rust/session/
store.rs

1// Copyright (c) 2026 erik <erik@erik.xyz> — https://erik.xyz
2
3use super::StoreError;
4use std::collections::HashMap;
5use std::sync::Mutex;
6
7/// 会话记录 —— 这张表就是「token 会话」。
8#[derive(Debug, Clone, PartialEq)]
9pub struct SessionRecord {
10    pub token: String,
11    pub subject: String,
12    pub fingerprint: String,
13    pub location: Option<String>,
14    pub coords: Option<(f64, f64)>,
15    /// 登录时的签名基线。
16    pub signature: Option<String>,
17    pub issued_at: u64,
18    pub last_seen: u64,
19    pub expires_at: u64,
20    pub revoked: bool,
21}
22
23/// 一次登录的位置快照,用于异地检测与不可能旅行。
24#[derive(Debug, Clone, PartialEq)]
25pub struct LoginPoint {
26    pub location: Option<String>,
27    pub coords: Option<(f64, f64)>,
28    pub at: u64,
29}
30
31/// 会话存储后端。
32///
33/// guard **按值**持有 store(`SessionGuard<S>` 的字段就是 `S`),而本 trait
34/// 没有为 `Arc<T>` 提供 blanket impl —— 所以 `Arc<MemoryStore>` 不是
35/// `SessionStore`,两个 guard 想共享同一份存储得包一层 newtype 并转发全部方法。
36/// 多线程场景不必如此:`SessionGuard<S>` 本身是 `Send + Sync`,共享一个
37/// `Arc<SessionGuard<MemoryStore>>`(见 `examples/axum_middleware.rs`)即可。
38pub trait SessionStore: Send + Sync {
39    fn put(&self, rec: SessionRecord) -> Result<(), StoreError>;
40    /// 原样返回记录,**不过滤过期**。过期判定归 `guard` ——
41    /// 若此处对过期记录返回 `None`,`verify` 将无法区分 `TokenExpired` 与 `TokenUnknown`。
42    fn get(&self, token: &str) -> Result<Option<SessionRecord>, StoreError>;
43    fn touch(&self, token: &str, now: u64) -> Result<(), StoreError>;
44    fn revoke(&self, token: &str) -> Result<(), StoreError>;
45    /// 吊销某 subject 的全部会话,返回受影响条数。
46    fn revoke_subject(&self, subject: &str) -> Result<usize, StoreError>;
47    /// 该 subject 的登录历史,按时间升序(最旧在前)。
48    fn recent_logins(&self, subject: &str) -> Result<Vec<LoginPoint>, StoreError>;
49    fn record_login(&self, subject: &str, point: LoginPoint) -> Result<(), StoreError>;
50    /// 清除已过期记录,并回收超过 [`LOGIN_HISTORY_KEEP_SECS`] 未再登录的
51    /// subject 的登录历史;返回清除的会话条数。
52    fn purge_expired(&self, now: u64) -> Result<usize, StoreError>;
53}
54
55/// 每个 subject 保留的登录历史条数上限。
56pub const MAX_LOGINS_PER_SUBJECT: usize = 10;
57
58/// 登录历史的保留时长(秒):超过这段时间没有新登录点的 subject,
59/// 其登录历史会在 [`SessionStore::purge_expired`] 时整条删除。
60///
61/// **取舍**:历史只在 `bind` / `verify` 里被读其中的 `last()` 一条,因此回收一个
62/// 休眠 subject 的历史,代价是他的**下一次登录少做一次异地 / 不可能旅行判定**
63/// (之后历史立即重建)。这是漏报而非误报——`location_changed` 缺输入返回
64/// `false`、`impossible_travel` 缺输入返回 `None`,删掉输入只会让结论变成「不报」。
65/// 这个常量就是旋钮:调大 = 更不容易漏报、内存占用更高。
66///
67/// 取 7 天是宽裕余量:地表最远两点约 20000 km,默认 900 km/h 阈值下不可能旅行
68/// 判定在约 22 小时后本就不可能成立。单 subject 的历史条数由
69/// [`MAX_LOGINS_PER_SUBJECT`] 限死,**无上限的是 subject 数量**,只能靠这个时间窗口回收。
70pub const LOGIN_HISTORY_KEEP_SECS: u64 = 604_800;
71
72/// 内存后端。无后台线程 —— 过期判定归 guard,内存回收靠 `purge_expired`。
73#[derive(Debug)]
74pub struct MemoryStore {
75    sessions: Mutex<HashMap<String, SessionRecord>>,
76    logins: Mutex<HashMap<String, Vec<LoginPoint>>>,
77}
78
79impl Default for MemoryStore {
80    fn default() -> Self {
81        Self::new()
82    }
83}
84
85impl MemoryStore {
86    pub fn new() -> Self {
87        Self {
88            sessions: Mutex::new(HashMap::new()),
89            logins: Mutex::new(HashMap::new()),
90        }
91    }
92
93    /// 互斥锁获取。线程 panic 导致锁中毒时恢复内部数据而非永久 Err:
94    /// 本 store 的每个操作都是单次 HashMap 读/写,不存在「改到一半」的不变量,
95    /// 恢复是安全的;而让一次 panic 永久锁死整个会话存储是一种自我 DoS。
96    fn lock<T>(m: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
97        m.lock().unwrap_or_else(|e| e.into_inner())
98    }
99}
100
101impl SessionStore for MemoryStore {
102    fn put(&self, rec: SessionRecord) -> Result<(), StoreError> {
103        Self::lock(&self.sessions).insert(rec.token.clone(), rec);
104        Ok(())
105    }
106
107    fn get(&self, token: &str) -> Result<Option<SessionRecord>, StoreError> {
108        Ok(Self::lock(&self.sessions).get(token).cloned())
109    }
110
111    fn touch(&self, token: &str, now: u64) -> Result<(), StoreError> {
112        if let Some(r) = Self::lock(&self.sessions).get_mut(token) {
113            r.last_seen = now;
114        }
115        Ok(())
116    }
117
118    fn revoke(&self, token: &str) -> Result<(), StoreError> {
119        if let Some(r) = Self::lock(&self.sessions).get_mut(token) {
120            r.revoked = true;
121        }
122        Ok(())
123    }
124
125    // ponytail: 持全局锁的 O(n) 全表扫描。内存后端规模下可接受;
126    // 若单 subject 会话数上到万级或需跨实例,加 subject -> tokens 索引。
127    fn revoke_subject(&self, subject: &str) -> Result<usize, StoreError> {
128        let mut n = 0;
129        for r in Self::lock(&self.sessions).values_mut() {
130            if r.subject == subject && !r.revoked {
131                r.revoked = true;
132                n += 1;
133            }
134        }
135        Ok(n)
136    }
137
138    fn recent_logins(&self, subject: &str) -> Result<Vec<LoginPoint>, StoreError> {
139        Ok(Self::lock(&self.logins)
140            .get(subject)
141            .cloned()
142            .unwrap_or_default())
143    }
144
145    fn record_login(&self, subject: &str, point: LoginPoint) -> Result<(), StoreError> {
146        let mut g = Self::lock(&self.logins);
147        let v = g.entry(subject.to_string()).or_default();
148        v.push(point);
149        // 有界:只保留最近 MAX_LOGINS_PER_SUBJECT 条
150        if v.len() > MAX_LOGINS_PER_SUBJECT {
151            v.drain(..v.len() - MAX_LOGINS_PER_SUBJECT);
152        }
153        Ok(())
154    }
155
156    fn purge_expired(&self, now: u64) -> Result<usize, StoreError> {
157        let mut g = Self::lock(&self.sessions);
158        let before = g.len();
159        g.retain(|_, r| r.expires_at > now);
160        let removed = before - g.len();
161        drop(g);
162
163        // 登录历史没有 expires_at 可依,按最后一个登录点的年龄回收。
164        // 扫全量而非取 `v.last()`:`at` 只在时钟单调时才随插入递增,回拨会让
165        // `last()` 指向更旧的点,把窗口内仍有效的历史整条丢掉。
166        let horizon = now.saturating_sub(LOGIN_HISTORY_KEEP_SECS);
167        Self::lock(&self.logins).retain(|_, v| v.iter().any(|p| p.at > horizon));
168        Ok(removed)
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175
176    fn rec(token: &str, subject: &str, expires_at: u64) -> SessionRecord {
177        SessionRecord {
178            token: token.into(),
179            subject: subject.into(),
180            fingerprint: "ip=1.2.3.4|ua=curl".into(),
181            location: Some("CN-BJ".into()),
182            coords: Some((39.9042, 116.4074)),
183            signature: None,
184            issued_at: 1_000,
185            last_seen: 1_000,
186            expires_at,
187            revoked: false,
188        }
189    }
190
191    #[test]
192    fn put_then_get_roundtrip() {
193        let s = MemoryStore::new();
194        s.put(rec("t1", "u1", 2_000)).unwrap();
195        let got = s.get("t1").unwrap().expect("present");
196        assert_eq!(got.token, "t1");
197        assert_eq!(got.subject, "u1");
198    }
199
200    #[test]
201    fn get_unknown_token_is_none() {
202        let s = MemoryStore::new();
203        assert!(s.get("nope").unwrap().is_none());
204    }
205
206    #[test]
207    fn get_returns_expired_record_unfiltered() {
208        // 关键契约:store 不做过期过滤,过期判定归 guard。
209        // 否则 verify 无法区分 TokenExpired 与 TokenUnknown。
210        let s = MemoryStore::new();
211        s.put(rec("t1", "u1", 500)).unwrap();
212        let got = s.get("t1").unwrap().expect("must still be returned");
213        assert_eq!(got.expires_at, 500);
214    }
215
216    #[test]
217    fn touch_updates_last_seen_only() {
218        let s = MemoryStore::new();
219        s.put(rec("t1", "u1", 9_999)).unwrap();
220        s.touch("t1", 1_500).unwrap();
221        let got = s.get("t1").unwrap().unwrap();
222        assert_eq!(got.last_seen, 1_500);
223        assert_eq!(got.issued_at, 1_000);
224        assert_eq!(got.expires_at, 9_999);
225    }
226
227    #[test]
228    fn touch_unknown_token_is_ok() {
229        let s = MemoryStore::new();
230        assert!(s.touch("nope", 1_500).is_ok());
231    }
232
233    #[test]
234    fn revoke_marks_record_not_deletes_it() {
235        let s = MemoryStore::new();
236        s.put(rec("t1", "u1", 9_999)).unwrap();
237        s.revoke("t1").unwrap();
238        let got = s.get("t1").unwrap().expect("kept for diagnosis");
239        assert!(got.revoked);
240    }
241
242    #[test]
243    fn revoke_unknown_token_is_ok() {
244        let s = MemoryStore::new();
245        assert!(s.revoke("nope").is_ok());
246    }
247
248    #[test]
249    fn revoke_subject_counts_only_newly_revoked() {
250        let s = MemoryStore::new();
251        s.put(rec("t1", "u1", 9_999)).unwrap();
252        s.put(rec("t2", "u1", 9_999)).unwrap();
253        s.put(rec("t3", "u2", 9_999)).unwrap();
254        assert_eq!(s.revoke_subject("u1").unwrap(), 2);
255        // 再调一次不该重复计数
256        assert_eq!(s.revoke_subject("u1").unwrap(), 0);
257        assert!(!s.get("t3").unwrap().unwrap().revoked);
258    }
259
260    #[test]
261    fn revoke_subject_unknown_is_zero() {
262        let s = MemoryStore::new();
263        assert_eq!(s.revoke_subject("ghost").unwrap(), 0);
264    }
265
266    fn point(loc: &str, at: u64) -> LoginPoint {
267        LoginPoint {
268            location: Some(loc.into()),
269            coords: None,
270            at,
271        }
272    }
273
274    #[test]
275    fn login_history_is_ascending() {
276        let s = MemoryStore::new();
277        s.record_login("u1", point("CN-BJ", 100)).unwrap();
278        s.record_login("u1", point("CN-SH", 200)).unwrap();
279        let h = s.recent_logins("u1").unwrap();
280        assert_eq!(h.len(), 2);
281        assert_eq!(h[0].at, 100);
282        assert_eq!(h[1].at, 200);
283    }
284
285    #[test]
286    fn login_history_is_bounded() {
287        let s = MemoryStore::new();
288        for i in 0..(MAX_LOGINS_PER_SUBJECT as u64 + 5) {
289            s.record_login("u1", point("CN-BJ", 100 + i)).unwrap();
290        }
291        let h = s.recent_logins("u1").unwrap();
292        assert_eq!(h.len(), MAX_LOGINS_PER_SUBJECT);
293        // 保留的是最近的:最旧的一条应为 at = 105
294        assert_eq!(h[0].at, 105);
295    }
296
297    #[test]
298    fn recent_logins_unknown_subject_is_empty() {
299        let s = MemoryStore::new();
300        assert!(s.recent_logins("ghost").unwrap().is_empty());
301    }
302
303    #[test]
304    fn purge_expired_removes_only_expired() {
305        let s = MemoryStore::new();
306        s.put(rec("old", "u1", 500)).unwrap();
307        s.put(rec("new", "u1", 5_000)).unwrap();
308        assert_eq!(s.purge_expired(1_000).unwrap(), 1);
309        assert!(s.get("old").unwrap().is_none());
310        assert!(s.get("new").unwrap().is_some());
311    }
312
313    #[test]
314    fn purge_expired_boundary_is_exclusive() {
315        // expires_at == now 视为已过期
316        let s = MemoryStore::new();
317        s.put(rec("t", "u1", 1_000)).unwrap();
318        assert_eq!(s.purge_expired(1_000).unwrap(), 1);
319    }
320
321    #[test]
322    fn purge_expired_on_empty_is_zero() {
323        let s = MemoryStore::new();
324        assert_eq!(s.purge_expired(1_000).unwrap(), 0);
325    }
326
327    #[test]
328    fn purge_expired_reclaims_dormant_login_history() {
329        // 登录历史此前只增不减:凭据填充 / 换用户名爆破每个 subject 永久占一份条目。
330        // 这里钉住回收路径 —— 没有它,purge_expired 对 logins 完全无效。
331        let s = MemoryStore::new();
332        let now = 10_000_000;
333        let stale = now - LOGIN_HISTORY_KEEP_SECS;
334        s.record_login("dormant", point("CN-BJ", stale - 1))
335            .unwrap();
336        // 窗口边界(恰好在 horizon 上)同样算休眠,与 sessions 的 `expires_at > now` 一致
337        s.record_login("edge", point("CN-BJ", stale)).unwrap();
338        s.record_login("active", point("CN-SH", now - 10)).unwrap();
339
340        assert_eq!(s.purge_expired(now).unwrap(), 0, "无过期会话");
341
342        assert!(s.recent_logins("dormant").unwrap().is_empty());
343        assert!(s.recent_logins("edge").unwrap().is_empty());
344        assert_eq!(s.recent_logins("active").unwrap().len(), 1);
345        // 「删掉键」才算真的回收,空 Vec 仍会占住 map 条目
346        assert_eq!(MemoryStore::lock(&s.logins).len(), 1);
347    }
348
349    #[test]
350    fn purge_expired_keeps_history_after_clock_rollback() {
351        let s = MemoryStore::new();
352        let now = 10_000_000;
353        // 时钟回拨:后写入的点 at 反而更小,`last()` 会指向窗口外的那条
354        s.record_login("u1", point("CN-BJ", now - 10)).unwrap();
355        s.record_login("u1", point("CN-BJ", now - 2 * LOGIN_HISTORY_KEEP_SECS))
356            .unwrap();
357
358        assert_eq!(s.purge_expired(now).unwrap(), 0);
359        assert_eq!(
360            s.recent_logins("u1").unwrap().len(),
361            2,
362            "取 last() 会把窗口内仍有效的历史整条丢掉"
363        );
364    }
365
366    #[test]
367    fn lock_recovers_from_poisoned_mutex() {
368        // 钉住 lock() 的恢复不变量:一次 panic 不能永久锁死会话存储。
369        let m = Mutex::new(rec("t1", "u1", 9_999));
370        std::panic::catch_unwind(|| {
371            let _guard = m.lock().unwrap();
372            panic!("poison");
373        })
374        .unwrap_err();
375        assert!(m.is_poisoned());
376
377        let g = MemoryStore::lock(&m);
378        assert_eq!(g.token, "t1");
379        assert_eq!(g.expires_at, 9_999);
380    }
381}