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