Skip to main content

arc_es_sqlite/
session.rs

1//! SQLite-backed [`SessionStore`] (HIPAA-4).
2//!
3//! Backend-agnostic per the workspace memory: the trait lives in
4//! `arc-core`; this is one of several implementations. Postgres and Redis
5//! variants are slot-in replacements.
6
7use arc_core::session::{SessionRecord, SessionStore, SessionStoreError};
8use async_trait::async_trait;
9use diesel::prelude::*;
10use diesel::r2d2::{self, ConnectionManager};
11use diesel::sqlite::SqliteConnection;
12use std::sync::Arc;
13use uuid::Uuid;
14
15mod schema {
16    diesel::table! {
17        jwt_sessions (jti) {
18            jti -> Text,
19            actor_id -> Text,
20            created_at_us -> BigInt,
21            expires_at_us -> BigInt,
22            revoked_at_us -> Nullable<BigInt>,
23        }
24    }
25}
26
27use schema::jwt_sessions;
28
29#[derive(Debug, Insertable, Clone)]
30#[diesel(table_name = jwt_sessions)]
31struct NewSessionRow {
32    jti: String,
33    actor_id: String,
34    created_at_us: i64,
35    expires_at_us: i64,
36    revoked_at_us: Option<i64>,
37}
38
39#[derive(Debug, Queryable, Clone)]
40struct SessionRow {
41    jti: String,
42    actor_id: String,
43    created_at_us: i64,
44    expires_at_us: i64,
45    revoked_at_us: Option<i64>,
46}
47
48impl SessionRow {
49    fn into_record(self) -> Result<SessionRecord, SessionStoreError> {
50        let jti = Uuid::parse_str(&self.jti)
51            .map_err(|e| SessionStoreError::Sink(format!("malformed jti UUID in DB row: {e}")))?;
52        Ok(SessionRecord {
53            jti,
54            actor_id: self.actor_id,
55            created_at_us: self.created_at_us,
56            expires_at_us: self.expires_at_us,
57            revoked_at_us: self.revoked_at_us,
58        })
59    }
60}
61
62type Pool = r2d2::Pool<ConnectionManager<SqliteConnection>>;
63
64/// Durable JWT session store backed by SQLite.
65#[derive(Clone)]
66pub struct SqliteSessionStore {
67    pool: Arc<Pool>,
68}
69
70impl SqliteSessionStore {
71    pub async fn new(database_url: &str) -> Result<Self, SessionStoreError> {
72        let manager = ConnectionManager::<SqliteConnection>::new(database_url);
73        let pool = Pool::builder()
74            .max_size(10)
75            .build(manager)
76            .map_err(|e| SessionStoreError::Sink(format!("failed to create pool: {e}")))?;
77        Ok(Self {
78            pool: Arc::new(pool),
79        })
80    }
81
82    pub fn with_pool(pool: Pool) -> Self {
83        Self {
84            pool: Arc::new(pool),
85        }
86    }
87}
88
89async fn run_blocking<F, T>(f: F) -> Result<T, SessionStoreError>
90where
91    F: FnOnce() -> Result<T, SessionStoreError> + Send + 'static,
92    T: Send + 'static,
93{
94    tokio::task::spawn_blocking(f)
95        .await
96        .map_err(|e| SessionStoreError::Sink(format!("join error: {e}")))?
97}
98
99#[async_trait]
100impl SessionStore for SqliteSessionStore {
101    async fn record_session(&self, record: SessionRecord) -> Result<(), SessionStoreError> {
102        if record.actor_id.trim().is_empty() {
103            return Err(SessionStoreError::Validation("actor_id empty".into()));
104        }
105        if record.expires_at_us <= record.created_at_us {
106            return Err(SessionStoreError::Validation(
107                "expires_at_us must be > created_at_us".into(),
108            ));
109        }
110
111        let row = NewSessionRow {
112            jti: record.jti.to_string(),
113            actor_id: record.actor_id,
114            created_at_us: record.created_at_us,
115            expires_at_us: record.expires_at_us,
116            revoked_at_us: record.revoked_at_us,
117        };
118        let pool = self.pool.clone();
119
120        run_blocking(move || {
121            let mut conn = pool
122                .get()
123                .map_err(|e| SessionStoreError::Sink(format!("conn: {e}")))?;
124            diesel::insert_into(jwt_sessions::table)
125                .values(&row)
126                .execute(&mut conn)
127                .map_err(|e| SessionStoreError::Sink(e.to_string()))?;
128            Ok(())
129        })
130        .await
131    }
132
133    async fn is_valid(&self, jti: Uuid, now_us: i64) -> Result<bool, SessionStoreError> {
134        let key = jti.to_string();
135        let pool = self.pool.clone();
136
137        run_blocking(move || {
138            let mut conn = pool
139                .get()
140                .map_err(|e| SessionStoreError::Sink(format!("conn: {e}")))?;
141            let row: Option<SessionRow> = jwt_sessions::table
142                .filter(jwt_sessions::jti.eq(&key))
143                .first(&mut conn)
144                .optional()
145                .map_err(|e| SessionStoreError::Sink(e.to_string()))?;
146            Ok(match row {
147                Some(r) => {
148                    let rec = r.into_record()?;
149                    rec.is_valid_at(now_us)
150                }
151                None => false,
152            })
153        })
154        .await
155    }
156
157    async fn revoke(&self, jti: Uuid, now_us: i64) -> Result<(), SessionStoreError> {
158        let key = jti.to_string();
159        let pool = self.pool.clone();
160
161        run_blocking(move || {
162            let mut conn = pool
163                .get()
164                .map_err(|e| SessionStoreError::Sink(format!("conn: {e}")))?;
165            let n = diesel::update(jwt_sessions::table.filter(jwt_sessions::jti.eq(&key)))
166                .set(jwt_sessions::revoked_at_us.eq(Some(now_us)))
167                .execute(&mut conn)
168                .map_err(|e| SessionStoreError::Sink(e.to_string()))?;
169            if n == 0 {
170                Err(SessionStoreError::NotFound(jti))
171            } else {
172                Ok(())
173            }
174        })
175        .await
176    }
177
178    async fn revoke_all_for_actor(
179        &self,
180        actor_id: &str,
181        now_us: i64,
182    ) -> Result<usize, SessionStoreError> {
183        let key = actor_id.to_string();
184        let pool = self.pool.clone();
185
186        run_blocking(move || {
187            let mut conn = pool
188                .get()
189                .map_err(|e| SessionStoreError::Sink(format!("conn: {e}")))?;
190            let n = diesel::update(
191                jwt_sessions::table
192                    .filter(jwt_sessions::actor_id.eq(&key))
193                    .filter(jwt_sessions::revoked_at_us.is_null()),
194            )
195            .set(jwt_sessions::revoked_at_us.eq(Some(now_us)))
196            .execute(&mut conn)
197            .map_err(|e| SessionStoreError::Sink(e.to_string()))?;
198            Ok(n)
199        })
200        .await
201    }
202
203    async fn prune_expired(&self, now_us: i64) -> Result<usize, SessionStoreError> {
204        let pool = self.pool.clone();
205
206        run_blocking(move || {
207            let mut conn = pool
208                .get()
209                .map_err(|e| SessionStoreError::Sink(format!("conn: {e}")))?;
210            let n =
211                diesel::delete(jwt_sessions::table.filter(jwt_sessions::expires_at_us.le(now_us)))
212                    .execute(&mut conn)
213                    .map_err(|e| SessionStoreError::Sink(e.to_string()))?;
214            Ok(n)
215        })
216        .await
217    }
218}
219
220// Suppress dead_code on SessionRow fields that exist for Queryable derive only.
221#[allow(dead_code)]
222fn _force_use(r: &SessionRow) -> &str {
223    &r.actor_id
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    async fn setup_store() -> SqliteSessionStore {
231        let manager = ConnectionManager::<SqliteConnection>::new(":memory:");
232        let pool = Pool::builder().max_size(1).build(manager).expect("pool");
233        let mut conn = pool.get().unwrap();
234        crate::test_support::migrate(&mut conn);
235        drop(conn);
236        SqliteSessionStore::with_pool(pool)
237    }
238
239    fn record(jti: Uuid, actor: &str) -> SessionRecord {
240        let now = 1_700_000_000_000_000;
241        SessionRecord {
242            jti,
243            actor_id: actor.into(),
244            created_at_us: now,
245            expires_at_us: now + 24 * 3600 * 1_000_000,
246            revoked_at_us: None,
247        }
248    }
249
250    #[tokio::test]
251    async fn test_record_then_is_valid() {
252        let s = setup_store().await;
253        let id = Uuid::new_v4();
254        s.record_session(record(id, "alice")).await.unwrap();
255        assert!(s.is_valid(id, 1_700_000_000_000_001).await.unwrap());
256    }
257
258    #[tokio::test]
259    async fn test_revoke_then_invalid() {
260        let s = setup_store().await;
261        let id = Uuid::new_v4();
262        s.record_session(record(id, "alice")).await.unwrap();
263        s.revoke(id, 1_700_000_000_000_500).await.unwrap();
264        assert!(!s.is_valid(id, 1_700_000_000_000_600).await.unwrap());
265    }
266
267    #[tokio::test]
268    async fn test_revoke_unknown_returns_not_found() {
269        let s = setup_store().await;
270        let err = s.revoke(Uuid::new_v4(), 0).await.unwrap_err();
271        assert!(matches!(err, SessionStoreError::NotFound(_)));
272    }
273
274    #[tokio::test]
275    async fn test_revoke_all_for_actor() {
276        let s = setup_store().await;
277        let a1 = Uuid::new_v4();
278        let a2 = Uuid::new_v4();
279        let b = Uuid::new_v4();
280        s.record_session(record(a1, "alice")).await.unwrap();
281        s.record_session(record(a2, "alice")).await.unwrap();
282        s.record_session(record(b, "bob")).await.unwrap();
283
284        let now = 1_700_000_000_000_500;
285        let n = s.revoke_all_for_actor("alice", now).await.unwrap();
286        assert_eq!(n, 2);
287        assert!(!s.is_valid(a1, now + 1).await.unwrap());
288        assert!(!s.is_valid(a2, now + 1).await.unwrap());
289        assert!(s.is_valid(b, now + 1).await.unwrap());
290    }
291
292    #[tokio::test]
293    async fn test_prune_expired() {
294        let s = setup_store().await;
295        let live = Uuid::new_v4();
296        let dead = Uuid::new_v4();
297        let now = 1_700_000_000_000_000;
298
299        s.record_session(SessionRecord {
300            jti: live,
301            actor_id: "x".into(),
302            created_at_us: now,
303            expires_at_us: now + 1_000_000,
304            revoked_at_us: None,
305        })
306        .await
307        .unwrap();
308        s.record_session(SessionRecord {
309            jti: dead,
310            actor_id: "x".into(),
311            created_at_us: now - 2000,
312            expires_at_us: now - 1000,
313            revoked_at_us: None,
314        })
315        .await
316        .unwrap();
317
318        let n = s.prune_expired(now).await.unwrap();
319        assert_eq!(n, 1);
320        assert!(s.is_valid(live, now + 1).await.unwrap());
321        assert!(!s.is_valid(dead, now + 1).await.unwrap());
322    }
323
324    #[tokio::test]
325    async fn test_indices_used_for_actor_lookup() {
326        let s = setup_store().await;
327        let pool = s.pool.clone();
328        let plan =
329            tokio::task::spawn_blocking(move || -> Result<Vec<String>, SessionStoreError> {
330                let mut conn = pool
331                    .get()
332                    .map_err(|e| SessionStoreError::Sink(e.to_string()))?;
333                let plan: Vec<ExplainRow> = diesel::sql_query(
334                    "EXPLAIN QUERY PLAN SELECT 1 FROM jwt_sessions WHERE actor_id = 'a'",
335                )
336                .load::<ExplainRow>(&mut *conn)
337                .map_err(|e| SessionStoreError::Sink(e.to_string()))?;
338                Ok(plan.into_iter().map(|r| r.detail).collect())
339            })
340            .await
341            .unwrap()
342            .unwrap();
343        assert!(
344            plan.iter().any(|d| d.contains("idx_jwt_sessions_actor_id")),
345            "actor_id query did not use index; plan: {plan:?}"
346        );
347    }
348
349    #[derive(QueryableByName, Debug)]
350    struct ExplainRow {
351        #[diesel(sql_type = diesel::sql_types::Integer)]
352        #[allow(dead_code)]
353        id: i32,
354        #[diesel(sql_type = diesel::sql_types::Integer)]
355        #[allow(dead_code)]
356        parent: i32,
357        #[diesel(sql_type = diesel::sql_types::Integer)]
358        #[allow(dead_code)]
359        notused: i32,
360        #[diesel(sql_type = diesel::sql_types::Text)]
361        detail: String,
362    }
363}