1use crate::random;
8use rustlavel_core::Json;
9use std::collections::{BTreeMap, BTreeSet};
10
11pub const ID_BYTES: usize = 32;
18
19pub fn is_valid_id(id: &str) -> bool {
26 (ID_BYTES * 2..=128).contains(&id.len())
27 && id.bytes().all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
28}
29
30#[derive(Debug, Clone, PartialEq)]
32pub struct Session {
33 id: String,
34 data: BTreeMap<String, Json>,
35 fresh_flash: BTreeSet<String>,
37 stale_flash: BTreeSet<String>,
39 updated_at: u64,
41 dirty: bool,
42}
43
44impl Session {
45 pub const TOKEN_KEY: &'static str = "_token";
48
49 pub fn new() -> Self {
51 Session::with_id(Session::new_id())
52 }
53
54 pub fn with_id(id: impl Into<String>) -> Self {
56 Session {
57 id: id.into(),
58 data: BTreeMap::new(),
59 fresh_flash: BTreeSet::new(),
60 stale_flash: BTreeSet::new(),
61 updated_at: crate::unix_now(),
62 dirty: false,
63 }
64 }
65
66 pub fn new_id() -> String {
68 random::hex(ID_BYTES)
69 }
70
71 pub fn id(&self) -> &str {
72 &self.id
73 }
74
75 pub fn updated_at(&self) -> u64 {
77 self.updated_at
78 }
79
80 pub fn is_dirty(&self) -> bool {
85 self.dirty
86 }
87
88 pub fn get(&self, key: &str) -> Option<&Json> {
89 self.data.get(key)
90 }
91
92 pub fn get_string(&self, key: &str) -> Option<String> {
94 match self.data.get(key)? {
95 Json::String(value) => Some(value.clone()),
96 Json::Null => None,
97 other => Some(other.to_string()),
98 }
99 }
100
101 pub fn has(&self, key: &str) -> bool {
102 self.data.contains_key(key)
103 }
104
105 pub fn put(&mut self, key: impl Into<String>, value: impl Into<Json>) {
106 self.data.insert(key.into(), value.into());
107 self.dirty = true;
108 }
109
110 pub fn forget(&mut self, key: &str) -> Option<Json> {
112 let removed = self.data.remove(key);
113 self.fresh_flash.remove(key);
114 self.stale_flash.remove(key);
115 if removed.is_some() {
116 self.dirty = true;
117 }
118 removed
119 }
120
121 pub fn flush(&mut self) {
125 self.data.clear();
126 self.fresh_flash.clear();
127 self.stale_flash.clear();
128 self.dirty = true;
129 }
130
131 pub fn all(&self) -> &BTreeMap<String, Json> {
132 &self.data
133 }
134
135 pub fn flash(&mut self, key: impl Into<String>, value: impl Into<Json>) {
141 let key = key.into();
142 self.fresh_flash.insert(key.clone());
143 self.put(key, value);
144 }
145
146 pub fn keep(&mut self, key: &str) {
148 if self.stale_flash.remove(key) {
149 self.fresh_flash.insert(key.to_string());
150 self.dirty = true;
151 }
152 }
153
154 pub fn reflash(&mut self) {
156 let carried = std::mem::take(&mut self.stale_flash);
157 if !carried.is_empty() {
158 self.dirty = true;
159 }
160 self.fresh_flash.extend(carried);
161 }
162
163 pub fn age_flash(&mut self) {
169 let expiring: Vec<String> =
170 self.stale_flash.difference(&self.fresh_flash).cloned().collect();
171 for key in expiring {
172 self.data.remove(&key);
173 self.dirty = true;
174 }
175 self.stale_flash = std::mem::take(&mut self.fresh_flash);
176 }
177
178 pub fn regenerate(&mut self) -> &str {
189 self.id = Session::new_id();
190 self.forget(Session::TOKEN_KEY);
204 self.dirty = true;
205 &self.id
206 }
207
208 pub fn invalidate(&mut self) -> &str {
210 self.flush();
211 self.regenerate()
212 }
213
214 pub fn token(&mut self) -> String {
220 if let Some(token) = self.get_string(Session::TOKEN_KEY) {
221 return token;
222 }
223 let token = random::hex(ID_BYTES);
224 self.put(Session::TOKEN_KEY, token.clone());
225 token
226 }
227
228 pub fn to_json(&self) -> Json {
230 Json::object([
231 ("data", Json::Object(self.data.clone())),
232 (
233 "flash",
234 Json::Array(self.stale_flash.iter().cloned().map(Json::String).collect()),
235 ),
236 ("updated_at", Json::Number(self.updated_at as f64)),
237 ])
238 }
239
240 pub fn from_json(id: impl Into<String>, payload: &Json) -> Self {
246 let mut session = Session::with_id(id);
247
248 if let Some(Json::Object(data)) = payload.get("data") {
249 session.data = data.clone();
250 }
251 if let Some(Json::Array(flash)) = payload.get("flash") {
252 session.stale_flash =
253 flash.iter().filter_map(|key| key.as_str().map(str::to_string)).collect();
254 }
255 if let Some(updated_at) = payload.get("updated_at").and_then(Json::as_f64) {
256 session.updated_at = updated_at as u64;
257 }
258
259 session.dirty = false;
260 session
261 }
262
263 pub fn touch(&mut self) {
267 self.updated_at = crate::unix_now();
268 }
269}
270
271impl Default for Session {
272 fn default() -> Self {
273 Session::new()
274 }
275}
276
277#[cfg(test)]
278mod tests {
279
280 #[test]
289 fn rotating_the_id_rotates_the_csrf_token() {
290 let mut session = Session::new();
291 let planted = session.token();
292 let before = session.id().to_string();
293
294 session.regenerate();
295
296 assert_ne!(session.id(), before, "the id did not rotate");
297 assert_ne!(
298 session.token(),
299 planted,
300 "the token survived the rotation, so a planted one still works after sign-in"
301 );
302 }
303
304 #[test]
307 fn rotating_does_not_mint_a_token_nobody_asked_for() {
308 let mut session = Session::new();
309 session.regenerate();
310 assert!(session.get_string(Session::TOKEN_KEY).is_none());
311 }
312 use super::*;
313
314 #[test]
315 fn new_sessions_get_a_long_random_hexadecimal_id() {
316 let first = Session::new();
317 let second = Session::new();
318
319 assert_eq!(first.id().len(), ID_BYTES * 2);
320 assert!(is_valid_id(first.id()));
321 assert_ne!(first.id(), second.id());
322 }
323
324 #[test]
325 fn id_validation_rejects_anything_that_is_not_our_own_shape() {
326 assert!(is_valid_id(&Session::new_id()));
327 assert!(!is_valid_id(""));
328 assert!(!is_valid_id("short"));
329 assert!(!is_valid_id("../../../etc/passwd"));
330 assert!(!is_valid_id(&"g".repeat(64)));
332 assert!(!is_valid_id(&"A".repeat(64)));
333 assert!(!is_valid_id(&"a".repeat(129)));
334 }
335
336 #[test]
337 fn values_are_put_read_and_forgotten() {
338 let mut session = Session::new();
339 assert!(!session.is_dirty());
340
341 session.put("name", "Ada");
342 session.put("visits", 3);
343
344 assert_eq!(session.get_string("name").as_deref(), Some("Ada"));
345 assert_eq!(session.get("visits").and_then(Json::as_i64), Some(3));
346 assert!(session.has("name"));
347 assert!(session.is_dirty());
348
349 assert_eq!(session.forget("name"), Some(Json::from("Ada")));
350 assert!(!session.has("name"));
351 assert_eq!(session.forget("name"), None);
352 }
353
354 #[test]
355 fn flush_empties_the_session_but_keeps_the_id() {
356 let mut session = Session::new();
357 let id = session.id().to_string();
358 session.put("a", 1);
359 session.flash("b", 2);
360
361 session.flush();
362
363 assert!(session.all().is_empty());
364 assert_eq!(session.id(), id);
365 }
366
367 #[test]
368 fn flash_data_survives_exactly_one_request() {
369 let mut session = Session::new();
371 session.flash("status", "Profile saved");
372 assert_eq!(session.get_string("status").as_deref(), Some("Profile saved"));
373 session.age_flash();
374
375 assert_eq!(session.get_string("status").as_deref(), Some("Profile saved"));
377 session.age_flash();
378
379 assert_eq!(session.get_string("status"), None);
381 }
382
383 #[test]
384 fn ordinary_values_are_not_swept_away_with_the_flash() {
385 let mut session = Session::new();
386 session.put("user_id", 7);
387 session.flash("status", "hi");
388
389 session.age_flash();
390 session.age_flash();
391
392 assert_eq!(session.get("user_id").and_then(Json::as_i64), Some(7));
393 assert_eq!(session.get_string("status"), None);
394 }
395
396 #[test]
397 fn re_flashing_a_key_gives_it_another_request() {
398 let mut session = Session::new();
399 session.flash("status", "first");
400 session.age_flash();
401
402 session.flash("status", "second");
405 session.age_flash();
406
407 assert_eq!(session.get_string("status").as_deref(), Some("second"));
408 }
409
410 #[test]
411 fn keep_and_reflash_extend_flashed_values() {
412 let mut session = Session::new();
413 session.flash("a", 1);
414 session.flash("b", 2);
415 session.age_flash();
416
417 session.keep("a");
418 session.age_flash();
419
420 assert!(session.has("a"));
421 assert!(!session.has("b"));
422
423 session.flash("c", 3);
424 session.age_flash();
425 session.reflash();
426 session.age_flash();
427
428 assert!(session.has("c"));
429 }
430
431 #[test]
432 fn regenerate_rotates_the_id_and_keeps_the_data() {
433 let mut session = Session::new();
434 session.put("user_id", 7);
435 let before = session.id().to_string();
436
437 let after = session.regenerate().to_string();
438
439 assert_ne!(before, after);
440 assert!(is_valid_id(&after));
441 assert_eq!(session.get("user_id").and_then(Json::as_i64), Some(7));
442 }
443
444 #[test]
445 fn invalidate_rotates_the_id_and_drops_the_data() {
446 let mut session = Session::new();
447 session.put("user_id", 7);
448 let before = session.id().to_string();
449
450 session.invalidate();
451
452 assert_ne!(session.id(), before);
453 assert!(session.all().is_empty());
454 }
455
456 #[test]
457 fn the_csrf_token_is_generated_once_and_then_reused() {
458 let mut session = Session::new();
459 let first = session.token();
460
461 assert_eq!(first.len(), ID_BYTES * 2);
462 assert_eq!(session.token(), first);
463 assert_eq!(session.get_string(Session::TOKEN_KEY), Some(first));
464 }
465
466 #[test]
467 fn serialization_round_trips_data_and_pending_flash() {
468 let mut session = Session::new();
469 session.put("user_id", 7);
470 session.flash("status", "saved");
471 session.age_flash();
472
473 let restored = Session::from_json(session.id(), &session.to_json());
474
475 assert_eq!(restored.get("user_id").and_then(Json::as_i64), Some(7));
476 assert_eq!(restored.get_string("status").as_deref(), Some("saved"));
477 assert!(!restored.is_dirty(), "a freshly loaded session has not been changed yet");
478
479 let mut restored = restored;
481 restored.age_flash();
482 assert_eq!(restored.get_string("status"), None);
483 }
484
485 #[test]
486 fn a_corrupt_payload_yields_an_empty_session_rather_than_an_error() {
487 let session = Session::from_json(Session::new_id(), &Json::String("garbage".into()));
488
489 assert!(session.all().is_empty());
490 assert!(is_valid_id(session.id()));
491 }
492}