Skip to main content

siderite_sessions/
memory.rs

1//! [`MemoryStore`]: a bounded in-process [`SessionStore`].
2
3use crate::store::{SessionRecord, SessionStore, StoreError};
4use async_trait::async_trait;
5use std::collections::HashMap;
6use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
7use std::time::{Duration, Instant};
8
9/// Capacity of [`MemoryStore::default`].
10const DEFAULT_CAPACITY: usize = 10_000;
11
12struct Entry {
13    record: SessionRecord,
14    expires: Instant,
15}
16
17#[derive(Default)]
18struct Inner {
19    entries: HashMap<String, Entry>,
20    /// Revocation epoch per user that was ever revoked.
21    epochs: HashMap<String, u64>,
22}
23
24impl Inner {
25    fn epoch(&self, user: &str) -> u64 {
26        self.epochs.get(user).copied().unwrap_or(0)
27    }
28
29    fn stale(&self, record: &SessionRecord) -> bool {
30        record
31            .user
32            .as_deref()
33            .is_some_and(|user| record.auth_epoch != self.epoch(user))
34    }
35
36    fn purge(&mut self, now: Instant) -> usize {
37        let before = self.entries.len();
38        self.entries.retain(|_, entry| entry.expires > now);
39        before - self.entries.len()
40    }
41}
42
43/// A session store in process memory, for tests and single-process apps.
44///
45/// It holds at most `capacity` sessions. When it is full, expired
46/// entries are dropped first; if none is expired a new session is
47/// refused with a [`StoreError`] (the request fails with `503`) rather
48/// than evicting somebody else's login. Run [`sweep`](Self::sweep) from a
49/// scheduled job to drop expired entries between writes.
50///
51/// Sessions do not survive a restart and are not shared between
52/// processes: use [`CacheStore`](crate::CacheStore) for that. Revoking a
53/// user remembers a small epoch counter for that user for the life of the
54/// process, so that a write racing the revocation cannot revive a session.
55#[derive(Clone)]
56pub struct MemoryStore {
57    inner: Arc<Mutex<Inner>>,
58    capacity: usize,
59}
60
61impl Default for MemoryStore {
62    fn default() -> Self {
63        Self::new(DEFAULT_CAPACITY)
64    }
65}
66
67impl std::fmt::Debug for MemoryStore {
68    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69        f.debug_struct("MemoryStore")
70            .field("capacity", &self.capacity)
71            .finish_non_exhaustive()
72    }
73}
74
75impl MemoryStore {
76    /// A store holding at most `capacity` sessions.
77    #[must_use]
78    pub fn new(capacity: usize) -> Self {
79        Self {
80            inner: Arc::default(),
81            capacity,
82        }
83    }
84
85    fn lock(&self) -> MutexGuard<'_, Inner> {
86        self.inner.lock().unwrap_or_else(PoisonError::into_inner)
87    }
88
89    /// Drop expired sessions and return how many were removed.
90    pub fn sweep(&self) -> usize {
91        self.lock().purge(Instant::now())
92    }
93
94    /// The number of stored sessions, expired ones included until swept.
95    #[must_use]
96    pub fn len(&self) -> usize {
97        self.lock().entries.len()
98    }
99
100    /// Whether no session is stored.
101    #[must_use]
102    pub fn is_empty(&self) -> bool {
103        self.len() == 0
104    }
105}
106
107#[async_trait]
108impl SessionStore for MemoryStore {
109    async fn load(
110        &self,
111        id_hash: &str,
112    ) -> Result<Option<SessionRecord>, StoreError> {
113        let now = Instant::now();
114        let inner = self.lock();
115        let live = inner.entries.get(id_hash).filter(|entry| {
116            entry.expires > now && !inner.stale(&entry.record)
117        });
118        Ok(live.map(|entry| entry.record.clone()))
119    }
120
121    async fn save(
122        &self,
123        id_hash: &str,
124        record: &SessionRecord,
125        ttl: Duration,
126    ) -> Result<(), StoreError> {
127        let now = Instant::now();
128        let mut inner = self.lock();
129        if inner.stale(record) {
130            return Err(StoreError::new("session was revoked"));
131        }
132        if !inner.entries.contains_key(id_hash)
133            && inner.entries.len() >= self.capacity
134        {
135            inner.purge(now);
136            if inner.entries.len() >= self.capacity {
137                return Err(StoreError::new("session store is full"));
138            }
139        }
140        let entry = Entry {
141            record: record.clone(),
142            expires: now.checked_add(ttl).unwrap_or(now),
143        };
144        inner.entries.insert(id_hash.to_owned(), entry);
145        Ok(())
146    }
147
148    async fn update(
149        &self,
150        id_hash: &str,
151        record: &SessionRecord,
152        ttl: Duration,
153    ) -> Result<bool, StoreError> {
154        let now = Instant::now();
155        let mut inner = self.lock();
156        if inner.stale(record) {
157            return Err(StoreError::new("session was revoked"));
158        }
159        let Some(entry) = inner.entries.get_mut(id_hash) else {
160            return Ok(false);
161        };
162        if entry.expires <= now {
163            return Ok(false);
164        }
165        entry.record = record.clone();
166        entry.expires = now.checked_add(ttl).unwrap_or(now);
167        Ok(true)
168    }
169
170    async fn delete(&self, id_hash: &str) -> Result<(), StoreError> {
171        self.lock().entries.remove(id_hash);
172        Ok(())
173    }
174
175    async fn revoke_user(&self, user: &str) -> Result<(), StoreError> {
176        let mut inner = self.lock();
177        let epoch = inner.epoch(user).saturating_add(1);
178        inner.epochs.insert(user.to_owned(), epoch);
179        inner
180            .entries
181            .retain(|_, entry| entry.record.user.as_deref() != Some(user));
182        Ok(())
183    }
184
185    async fn auth_epoch(&self, user: &str) -> Result<u64, StoreError> {
186        Ok(self.lock().epoch(user))
187    }
188}