rustlavel_auth/
session.rs1use 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.dirty = true;
191 &self.id
192 }
193
194 pub fn invalidate(&mut self) -> &str {
196 self.flush();
197 self.regenerate()
198 }
199
200 pub fn token(&mut self) -> String {
206 if let Some(token) = self.get_string(Session::TOKEN_KEY) {
207 return token;
208 }
209 let token = random::hex(ID_BYTES);
210 self.put(Session::TOKEN_KEY, token.clone());
211 token
212 }
213
214 pub fn to_json(&self) -> Json {
216 Json::object([
217 ("data", Json::Object(self.data.clone())),
218 (
219 "flash",
220 Json::Array(self.stale_flash.iter().cloned().map(Json::String).collect()),
221 ),
222 ("updated_at", Json::Number(self.updated_at as f64)),
223 ])
224 }
225
226 pub fn from_json(id: impl Into<String>, payload: &Json) -> Self {
232 let mut session = Session::with_id(id);
233
234 if let Some(Json::Object(data)) = payload.get("data") {
235 session.data = data.clone();
236 }
237 if let Some(Json::Array(flash)) = payload.get("flash") {
238 session.stale_flash =
239 flash.iter().filter_map(|key| key.as_str().map(str::to_string)).collect();
240 }
241 if let Some(updated_at) = payload.get("updated_at").and_then(Json::as_f64) {
242 session.updated_at = updated_at as u64;
243 }
244
245 session.dirty = false;
246 session
247 }
248
249 pub fn touch(&mut self) {
253 self.updated_at = crate::unix_now();
254 }
255}
256
257impl Default for Session {
258 fn default() -> Self {
259 Session::new()
260 }
261}
262
263#[cfg(test)]
264mod tests {
265 use super::*;
266
267 #[test]
268 fn new_sessions_get_a_long_random_hexadecimal_id() {
269 let first = Session::new();
270 let second = Session::new();
271
272 assert_eq!(first.id().len(), ID_BYTES * 2);
273 assert!(is_valid_id(first.id()));
274 assert_ne!(first.id(), second.id());
275 }
276
277 #[test]
278 fn id_validation_rejects_anything_that_is_not_our_own_shape() {
279 assert!(is_valid_id(&Session::new_id()));
280 assert!(!is_valid_id(""));
281 assert!(!is_valid_id("short"));
282 assert!(!is_valid_id("../../../etc/passwd"));
283 assert!(!is_valid_id(&"g".repeat(64)));
285 assert!(!is_valid_id(&"A".repeat(64)));
286 assert!(!is_valid_id(&"a".repeat(129)));
287 }
288
289 #[test]
290 fn values_are_put_read_and_forgotten() {
291 let mut session = Session::new();
292 assert!(!session.is_dirty());
293
294 session.put("name", "Ada");
295 session.put("visits", 3);
296
297 assert_eq!(session.get_string("name").as_deref(), Some("Ada"));
298 assert_eq!(session.get("visits").and_then(Json::as_i64), Some(3));
299 assert!(session.has("name"));
300 assert!(session.is_dirty());
301
302 assert_eq!(session.forget("name"), Some(Json::from("Ada")));
303 assert!(!session.has("name"));
304 assert_eq!(session.forget("name"), None);
305 }
306
307 #[test]
308 fn flush_empties_the_session_but_keeps_the_id() {
309 let mut session = Session::new();
310 let id = session.id().to_string();
311 session.put("a", 1);
312 session.flash("b", 2);
313
314 session.flush();
315
316 assert!(session.all().is_empty());
317 assert_eq!(session.id(), id);
318 }
319
320 #[test]
321 fn flash_data_survives_exactly_one_request() {
322 let mut session = Session::new();
324 session.flash("status", "Profile saved");
325 assert_eq!(session.get_string("status").as_deref(), Some("Profile saved"));
326 session.age_flash();
327
328 assert_eq!(session.get_string("status").as_deref(), Some("Profile saved"));
330 session.age_flash();
331
332 assert_eq!(session.get_string("status"), None);
334 }
335
336 #[test]
337 fn ordinary_values_are_not_swept_away_with_the_flash() {
338 let mut session = Session::new();
339 session.put("user_id", 7);
340 session.flash("status", "hi");
341
342 session.age_flash();
343 session.age_flash();
344
345 assert_eq!(session.get("user_id").and_then(Json::as_i64), Some(7));
346 assert_eq!(session.get_string("status"), None);
347 }
348
349 #[test]
350 fn re_flashing_a_key_gives_it_another_request() {
351 let mut session = Session::new();
352 session.flash("status", "first");
353 session.age_flash();
354
355 session.flash("status", "second");
358 session.age_flash();
359
360 assert_eq!(session.get_string("status").as_deref(), Some("second"));
361 }
362
363 #[test]
364 fn keep_and_reflash_extend_flashed_values() {
365 let mut session = Session::new();
366 session.flash("a", 1);
367 session.flash("b", 2);
368 session.age_flash();
369
370 session.keep("a");
371 session.age_flash();
372
373 assert!(session.has("a"));
374 assert!(!session.has("b"));
375
376 session.flash("c", 3);
377 session.age_flash();
378 session.reflash();
379 session.age_flash();
380
381 assert!(session.has("c"));
382 }
383
384 #[test]
385 fn regenerate_rotates_the_id_and_keeps_the_data() {
386 let mut session = Session::new();
387 session.put("user_id", 7);
388 let before = session.id().to_string();
389
390 let after = session.regenerate().to_string();
391
392 assert_ne!(before, after);
393 assert!(is_valid_id(&after));
394 assert_eq!(session.get("user_id").and_then(Json::as_i64), Some(7));
395 }
396
397 #[test]
398 fn invalidate_rotates_the_id_and_drops_the_data() {
399 let mut session = Session::new();
400 session.put("user_id", 7);
401 let before = session.id().to_string();
402
403 session.invalidate();
404
405 assert_ne!(session.id(), before);
406 assert!(session.all().is_empty());
407 }
408
409 #[test]
410 fn the_csrf_token_is_generated_once_and_then_reused() {
411 let mut session = Session::new();
412 let first = session.token();
413
414 assert_eq!(first.len(), ID_BYTES * 2);
415 assert_eq!(session.token(), first);
416 assert_eq!(session.get_string(Session::TOKEN_KEY), Some(first));
417 }
418
419 #[test]
420 fn serialization_round_trips_data_and_pending_flash() {
421 let mut session = Session::new();
422 session.put("user_id", 7);
423 session.flash("status", "saved");
424 session.age_flash();
425
426 let restored = Session::from_json(session.id(), &session.to_json());
427
428 assert_eq!(restored.get("user_id").and_then(Json::as_i64), Some(7));
429 assert_eq!(restored.get_string("status").as_deref(), Some("saved"));
430 assert!(!restored.is_dirty(), "a freshly loaded session has not been changed yet");
431
432 let mut restored = restored;
434 restored.age_flash();
435 assert_eq!(restored.get_string("status"), None);
436 }
437
438 #[test]
439 fn a_corrupt_payload_yields_an_empty_session_rather_than_an_error() {
440 let session = Session::from_json(Session::new_id(), &Json::String("garbage".into()));
441
442 assert!(session.all().is_empty());
443 assert!(is_valid_id(session.id()));
444 }
445}