1use super::StoreError;
4use std::collections::HashMap;
5use std::sync::Mutex;
6
7#[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 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#[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 {
39 fn put(&self, rec: SessionRecord) -> Result<(), StoreError>;
40 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 fn revoke_subject(&self, subject: &str) -> Result<usize, StoreError>;
47 fn recent_logins(&self, subject: &str) -> Result<Vec<LoginPoint>, StoreError>;
49 fn record_login(&self, subject: &str, point: LoginPoint) -> Result<(), StoreError>;
50 fn purge_expired(&self, now: u64) -> Result<usize, StoreError>;
53}
54
55pub const MAX_LOGINS_PER_SUBJECT: usize = 10;
57
58pub const LOGIN_HISTORY_KEEP_SECS: u64 = 604_800;
71
72#[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 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 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 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 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 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 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 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 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 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 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 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 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 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}