1use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous};
68use sqlx::SqlitePool;
69use sqlx::{Sqlite, Transaction};
70use std::str::FromStr;
71use std::sync::atomic::{AtomicU64, Ordering};
72use std::time::Duration;
73
74#[derive(Debug, thiserror::Error)]
75pub enum Error {
76 #[error(transparent)]
77 Sqlx(#[from] sqlx::Error),
78}
79pub type Result<T> = std::result::Result<T, Error>;
80
81pub struct HardenedSqlite {
83 url: String,
84 journal_mode: SqliteJournalMode,
85 foreign_keys: bool,
86 busy_timeout: Duration,
87 synchronous: SqliteSynchronous,
88 create_if_missing: bool,
89 max_connections: u32,
90 min_connections: Option<u32>,
91 memory_name: Option<String>,
92}
93
94impl HardenedSqlite {
95 pub fn new(url: impl Into<String>) -> Self {
99 Self {
100 url: url.into(),
101 journal_mode: SqliteJournalMode::Wal,
102 foreign_keys: true,
103 busy_timeout: Duration::from_secs(5),
104 synchronous: SqliteSynchronous::Normal,
105 create_if_missing: true,
106 max_connections: 8,
107 min_connections: None,
108 memory_name: None,
109 }
110 }
111
112 pub fn journal_mode(mut self, m: SqliteJournalMode) -> Self {
113 self.journal_mode = m;
114 self
115 }
116 pub fn foreign_keys(mut self, on: bool) -> Self {
117 self.foreign_keys = on;
118 self
119 }
120 pub fn busy_timeout(mut self, d: Duration) -> Self {
121 self.busy_timeout = d;
122 self
123 }
124 pub fn synchronous(mut self, s: SqliteSynchronous) -> Self {
125 self.synchronous = s;
126 self
127 }
128 pub fn create_if_missing(mut self, on: bool) -> Self {
129 self.create_if_missing = on;
130 self
131 }
132 pub fn max_connections(mut self, n: u32) -> Self {
133 self.max_connections = n;
134 self
135 }
136 pub fn min_connections(mut self, n: u32) -> Self {
137 self.min_connections = Some(n);
138 self
139 }
140 pub fn memory_name(mut self, name: impl Into<String>) -> Self {
142 self.memory_name = Some(name.into());
143 self
144 }
145
146 fn is_memory(&self) -> bool {
147 is_memory_url(&self.url)
148 }
149
150 pub fn connect_options(&self) -> Result<SqliteConnectOptions> {
152 let opts = if self.is_memory() {
153 let name = self.memory_name.clone().unwrap_or_else(salted_memory_name);
158 SqliteConnectOptions::new()
159 .filename(format!("file:{name}?mode=memory&cache=shared"))
160 .create_if_missing(true)
161 .foreign_keys(self.foreign_keys)
162 .busy_timeout(self.busy_timeout)
163 .synchronous(self.synchronous)
164 } else {
165 SqliteConnectOptions::from_str(&self.url)?
166 .create_if_missing(self.create_if_missing)
167 .journal_mode(self.journal_mode)
168 .foreign_keys(self.foreign_keys)
169 .busy_timeout(self.busy_timeout)
170 .synchronous(self.synchronous)
171 };
172 Ok(opts)
173 }
174
175 pub fn pragma_statements(&self) -> Vec<String> {
190 let mut v = Vec::new();
191 if !self.is_memory() {
192 let jm = match self.journal_mode {
193 SqliteJournalMode::Wal => "WAL",
194 SqliteJournalMode::Delete => "DELETE",
195 SqliteJournalMode::Truncate => "TRUNCATE",
196 SqliteJournalMode::Persist => "PERSIST",
197 SqliteJournalMode::Memory => "MEMORY",
198 SqliteJournalMode::Off => "OFF",
199 };
200 v.push(format!("PRAGMA journal_mode = {jm};"));
201 }
202 v.push(format!(
203 "PRAGMA foreign_keys = {};",
204 if self.foreign_keys { "ON" } else { "OFF" }
205 ));
206 v.push(format!(
207 "PRAGMA busy_timeout = {};",
208 self.busy_timeout.as_millis()
209 ));
210 let sync = match self.synchronous {
211 SqliteSynchronous::Off => "OFF",
212 SqliteSynchronous::Normal => "NORMAL",
213 SqliteSynchronous::Full => "FULL",
214 SqliteSynchronous::Extra => "EXTRA",
215 };
216 v.push(format!("PRAGMA synchronous = {sync};"));
217 v
218 }
219
220 pub async fn connect(self) -> Result<SqlitePool> {
222 let opts = self.connect_options()?;
223 let mut pool_opts = SqlitePoolOptions::new().max_connections(self.max_connections);
224 let min = self
227 .min_connections
228 .map(|m| if self.is_memory() { m.max(1) } else { m })
229 .or(if self.is_memory() { Some(1) } else { None });
230 if let Some(m) = min {
231 pool_opts = pool_opts.min_connections(m);
232 }
233 Ok(pool_opts.connect_with(opts).await?)
234 }
235}
236
237pub fn is_memory_url(url: &str) -> bool {
239 matches!(url, ":memory:" | "sqlite::memory:" | "sqlite://:memory:")
240}
241
242static MEM_SERIAL: AtomicU64 = AtomicU64::new(0);
243
244fn salted_memory_name() -> String {
245 let n = MEM_SERIAL.fetch_add(1, Ordering::Relaxed);
246 format!("sqlite-hardened-{}-{n}", std::process::id())
247}
248
249pub async fn begin_immediate(pool: &SqlitePool) -> Result<Transaction<'static, Sqlite>> {
281 Ok(pool.begin_with("BEGIN IMMEDIATE").await?)
282}
283
284#[allow(async_fn_in_trait)]
286pub trait SqlitePoolExt {
287 async fn begin_immediate(&self) -> Result<Transaction<'static, Sqlite>>;
288}
289
290impl SqlitePoolExt for SqlitePool {
291 async fn begin_immediate(&self) -> Result<Transaction<'static, Sqlite>> {
292 begin_immediate(self).await
293 }
294}
295
296#[cfg(test)]
297mod tests {
298 use super::*;
299
300 async fn scalar(pool: &SqlitePool, sql: &str) -> i64 {
301 sqlx::query_scalar::<_, i64>(sql)
302 .fetch_one(pool)
303 .await
304 .unwrap()
305 }
306
307 #[tokio::test]
308 async fn file_pool_applies_hardened_pragmas() {
309 let tmp = tempfile::tempdir().unwrap();
310 let db = tmp.path().join("t.db");
311 let pool = HardenedSqlite::new(format!("sqlite://{}", db.display()))
312 .connect()
313 .await
314 .unwrap();
315 let jm: String = sqlx::query_scalar("PRAGMA journal_mode")
316 .fetch_one(&pool)
317 .await
318 .unwrap();
319 assert_eq!(jm.to_lowercase(), "wal");
320 assert_eq!(scalar(&pool, "PRAGMA foreign_keys").await, 1);
321 assert_eq!(scalar(&pool, "PRAGMA busy_timeout").await, 5000);
322 assert_eq!(scalar(&pool, "PRAGMA synchronous").await, 1); }
324
325 #[tokio::test]
326 async fn overrides_are_honored() {
327 let tmp = tempfile::tempdir().unwrap();
328 let db = tmp.path().join("t.db");
329 let pool = HardenedSqlite::new(format!("sqlite://{}", db.display()))
330 .foreign_keys(false)
331 .busy_timeout(Duration::from_secs(2))
332 .connect()
333 .await
334 .unwrap();
335 assert_eq!(scalar(&pool, "PRAGMA foreign_keys").await, 0);
336 assert_eq!(scalar(&pool, "PRAGMA busy_timeout").await, 2000);
337 }
338
339 #[tokio::test]
340 async fn in_memory_survives_drop_to_zero() {
341 let pool = HardenedSqlite::new(":memory:").connect().await.unwrap();
342 sqlx::query("CREATE TABLE t (id INTEGER PRIMARY KEY)")
343 .execute(&pool)
344 .await
345 .unwrap();
346 sqlx::query("INSERT INTO t (id) VALUES (1)")
347 .execute(&pool)
348 .await
349 .unwrap();
350 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
353 let n = scalar(&pool, "SELECT COUNT(*) FROM t").await;
354 assert_eq!(n, 1);
355 }
356
357 #[tokio::test]
358 async fn two_memory_pools_are_isolated() {
359 let a = HardenedSqlite::new(":memory:").connect().await.unwrap();
360 let b = HardenedSqlite::new(":memory:").connect().await.unwrap();
361 sqlx::query("CREATE TABLE only_in_a (x INTEGER)")
362 .execute(&a)
363 .await
364 .unwrap();
365 let seen: i64 = sqlx::query_scalar(
366 "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='only_in_a'",
367 )
368 .fetch_one(&b)
369 .await
370 .unwrap();
371 assert_eq!(seen, 0);
372 }
373
374 #[tokio::test]
375 async fn begin_immediate_serializes_writers() {
376 let tmp = tempfile::tempdir().unwrap();
377 let db = tmp.path().join("t.db");
378 let pool = HardenedSqlite::new(format!("sqlite://{}", db.display()))
379 .connect()
380 .await
381 .unwrap();
382 sqlx::query("CREATE TABLE c (n INTEGER)")
383 .execute(&pool)
384 .await
385 .unwrap();
386 sqlx::query("INSERT INTO c (n) VALUES (0)")
387 .execute(&pool)
388 .await
389 .unwrap();
390
391 let inc = |p: SqlitePool| async move {
395 let mut tx = p.begin_immediate().await.unwrap();
396 let cur: i64 = sqlx::query_scalar("SELECT n FROM c")
397 .fetch_one(&mut *tx)
398 .await
399 .unwrap();
400 sqlx::query("UPDATE c SET n = ?")
401 .bind(cur + 1)
402 .execute(&mut *tx)
403 .await
404 .unwrap();
405 tx.commit().await.unwrap();
406 };
407 let (a, b) = (pool.clone(), pool.clone());
408 let (ra, rb) = tokio::join!(tokio::spawn(inc(a)), tokio::spawn(inc(b)));
409 ra.unwrap();
410 rb.unwrap();
411
412 let n: i64 = sqlx::query_scalar("SELECT n FROM c")
413 .fetch_one(&pool)
414 .await
415 .unwrap();
416 assert_eq!(n, 2, "both writers committed without a lost update");
417 }
418
419 #[test]
420 fn pragma_statements_reflect_settings() {
421 let stmts = HardenedSqlite::new("sqlite://x.db").pragma_statements();
422 assert!(stmts.iter().any(|s| s == "PRAGMA journal_mode = WAL;"));
423 assert!(stmts.iter().any(|s| s == "PRAGMA foreign_keys = ON;"));
424 assert!(stmts.iter().any(|s| s == "PRAGMA busy_timeout = 5000;"));
425 assert!(stmts.iter().any(|s| s == "PRAGMA synchronous = NORMAL;"));
426 }
427
428 #[test]
429 fn pragma_statements_omit_wal_for_memory() {
430 let stmts = HardenedSqlite::new(":memory:").pragma_statements();
431 assert!(!stmts.iter().any(|s| s.contains("journal_mode")));
432 }
433
434 #[tokio::test]
435 async fn hardened_settings_enforced_with_url_params_present() {
436 let tmp = tempfile::tempdir().unwrap();
437 let db = tmp.path().join("t.db");
438 let url = format!("sqlite://{}?mode=rwc", db.display());
441 let pool = HardenedSqlite::new(url).connect().await.unwrap();
442 assert_eq!(scalar(&pool, "PRAGMA foreign_keys").await, 1);
443 assert_eq!(scalar(&pool, "PRAGMA busy_timeout").await, 5000);
444 }
445}