1use 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#[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#[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}