sqlite_hardened/
lib.rs

1//! Correct-by-default sqlx SQLite pool bootstrap.
2
3use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous};
4use sqlx::SqlitePool;
5use sqlx::{Sqlite, Transaction};
6use std::str::FromStr;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::time::Duration;
9
10#[derive(Debug, thiserror::Error)]
11pub enum Error {
12    #[error(transparent)]
13    Sqlx(#[from] sqlx::Error),
14}
15pub type Result<T> = std::result::Result<T, Error>;
16
17/// Builder for a hardened `SqlitePool`.
18pub struct HardenedSqlite {
19    url: String,
20    journal_mode: SqliteJournalMode,
21    foreign_keys: bool,
22    busy_timeout: Duration,
23    synchronous: SqliteSynchronous,
24    create_if_missing: bool,
25    max_connections: u32,
26    min_connections: Option<u32>,
27    memory_name: Option<String>,
28}
29
30impl HardenedSqlite {
31    /// New builder with the hardened defaults:
32    /// WAL, foreign_keys ON, busy_timeout 5s, synchronous NORMAL,
33    /// create_if_missing, max_connections 8.
34    pub fn new(url: impl Into<String>) -> Self {
35        Self {
36            url: url.into(),
37            journal_mode: SqliteJournalMode::Wal,
38            foreign_keys: true,
39            busy_timeout: Duration::from_secs(5),
40            synchronous: SqliteSynchronous::Normal,
41            create_if_missing: true,
42            max_connections: 8,
43            min_connections: None,
44            memory_name: None,
45        }
46    }
47
48    pub fn journal_mode(mut self, m: SqliteJournalMode) -> Self {
49        self.journal_mode = m;
50        self
51    }
52    pub fn foreign_keys(mut self, on: bool) -> Self {
53        self.foreign_keys = on;
54        self
55    }
56    pub fn busy_timeout(mut self, d: Duration) -> Self {
57        self.busy_timeout = d;
58        self
59    }
60    pub fn synchronous(mut self, s: SqliteSynchronous) -> Self {
61        self.synchronous = s;
62        self
63    }
64    pub fn create_if_missing(mut self, on: bool) -> Self {
65        self.create_if_missing = on;
66        self
67    }
68    pub fn max_connections(mut self, n: u32) -> Self {
69        self.max_connections = n;
70        self
71    }
72    pub fn min_connections(mut self, n: u32) -> Self {
73        self.min_connections = Some(n);
74        self
75    }
76    /// Force a specific shared-cache in-memory name (else auto-salted).
77    pub fn memory_name(mut self, name: impl Into<String>) -> Self {
78        self.memory_name = Some(name.into());
79        self
80    }
81
82    fn is_memory(&self) -> bool {
83        is_memory_url(&self.url)
84    }
85
86    /// Build the `SqliteConnectOptions` without connecting.
87    pub fn connect_options(&self) -> Result<SqliteConnectOptions> {
88        let opts = if self.is_memory() {
89            // Explicit unique file:...?mode=memory&cache=shared — NOT a bare
90            // ":memory:" URL, which sqlx would turn into SQLITE_OPEN_MEMORY and
91            // collapse every pool onto one anonymous shared cache. WAL is
92            // meaningless for memory DBs, so it's skipped here.
93            let name = self.memory_name.clone().unwrap_or_else(salted_memory_name);
94            SqliteConnectOptions::new()
95                .filename(format!("file:{name}?mode=memory&cache=shared"))
96                .create_if_missing(true)
97                .foreign_keys(self.foreign_keys)
98                .busy_timeout(self.busy_timeout)
99                .synchronous(self.synchronous)
100        } else {
101            SqliteConnectOptions::from_str(&self.url)?
102                .create_if_missing(self.create_if_missing)
103                .journal_mode(self.journal_mode)
104                .foreign_keys(self.foreign_keys)
105                .busy_timeout(self.busy_timeout)
106                .synchronous(self.synchronous)
107        };
108        Ok(opts)
109    }
110
111    /// The PRAGMA statements this builder would apply, in order — for callers
112    /// on rusqlite or another driver who want the same hardening without the
113    /// sqlx pool. WAL is omitted for in-memory configurations.
114    pub fn pragma_statements(&self) -> Vec<String> {
115        let mut v = Vec::new();
116        if !self.is_memory() {
117            let jm = match self.journal_mode {
118                SqliteJournalMode::Wal => "WAL",
119                SqliteJournalMode::Delete => "DELETE",
120                SqliteJournalMode::Truncate => "TRUNCATE",
121                SqliteJournalMode::Persist => "PERSIST",
122                SqliteJournalMode::Memory => "MEMORY",
123                SqliteJournalMode::Off => "OFF",
124            };
125            v.push(format!("PRAGMA journal_mode = {jm};"));
126        }
127        v.push(format!(
128            "PRAGMA foreign_keys = {};",
129            if self.foreign_keys { "ON" } else { "OFF" }
130        ));
131        v.push(format!(
132            "PRAGMA busy_timeout = {};",
133            self.busy_timeout.as_millis()
134        ));
135        let sync = match self.synchronous {
136            SqliteSynchronous::Off => "OFF",
137            SqliteSynchronous::Normal => "NORMAL",
138            SqliteSynchronous::Full => "FULL",
139            SqliteSynchronous::Extra => "EXTRA",
140        };
141        v.push(format!("PRAGMA synchronous = {sync};"));
142        v
143    }
144
145    /// Build and connect the pool.
146    pub async fn connect(self) -> Result<SqlitePool> {
147        let opts = self.connect_options()?;
148        let mut pool_opts = SqlitePoolOptions::new().max_connections(self.max_connections);
149        // Pin one connection for in-memory DBs so the shared cache can't be
150        // destroyed when the pool drops to zero live connections.
151        let min = self
152            .min_connections
153            .map(|m| if self.is_memory() { m.max(1) } else { m })
154            .or(if self.is_memory() { Some(1) } else { None });
155        if let Some(m) = min {
156            pool_opts = pool_opts.min_connections(m);
157        }
158        Ok(pool_opts.connect_with(opts).await?)
159    }
160}
161
162/// True for the recognized in-memory URL spellings.
163pub fn is_memory_url(url: &str) -> bool {
164    matches!(url, ":memory:" | "sqlite::memory:" | "sqlite://:memory:")
165}
166
167static MEM_SERIAL: AtomicU64 = AtomicU64::new(0);
168
169fn salted_memory_name() -> String {
170    let n = MEM_SERIAL.fetch_add(1, Ordering::Relaxed);
171    format!("sqlite-hardened-{}-{n}", std::process::id())
172}
173
174/// Begin a WRITE transaction with `BEGIN IMMEDIATE`.
175///
176/// sqlx's `pool.begin()` issues a deferred BEGIN: a transaction that reads
177/// before it writes pins its read snapshot on the first statement, and if
178/// another connection commits in between, the later write cannot upgrade —
179/// SQLite returns a busy-snapshot error even *with* a busy_timeout, because the
180/// conflict is logical, not a lock wait. `BEGIN IMMEDIATE` takes the write lock
181/// up front, so a read-then-write transaction is safe under WAL's single-writer
182/// rule.
183pub async fn begin_immediate(pool: &SqlitePool) -> Result<Transaction<'static, Sqlite>> {
184    Ok(pool.begin_with("BEGIN IMMEDIATE").await?)
185}
186
187/// `pool.begin_immediate()` sugar.
188#[allow(async_fn_in_trait)]
189pub trait SqlitePoolExt {
190    async fn begin_immediate(&self) -> Result<Transaction<'static, Sqlite>>;
191}
192
193impl SqlitePoolExt for SqlitePool {
194    async fn begin_immediate(&self) -> Result<Transaction<'static, Sqlite>> {
195        begin_immediate(self).await
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202
203    async fn scalar(pool: &SqlitePool, sql: &str) -> i64 {
204        sqlx::query_scalar::<_, i64>(sql)
205            .fetch_one(pool)
206            .await
207            .unwrap()
208    }
209
210    #[tokio::test]
211    async fn file_pool_applies_hardened_pragmas() {
212        let tmp = tempfile::tempdir().unwrap();
213        let db = tmp.path().join("t.db");
214        let pool = HardenedSqlite::new(format!("sqlite://{}", db.display()))
215            .connect()
216            .await
217            .unwrap();
218        let jm: String = sqlx::query_scalar("PRAGMA journal_mode")
219            .fetch_one(&pool)
220            .await
221            .unwrap();
222        assert_eq!(jm.to_lowercase(), "wal");
223        assert_eq!(scalar(&pool, "PRAGMA foreign_keys").await, 1);
224        assert_eq!(scalar(&pool, "PRAGMA busy_timeout").await, 5000);
225        assert_eq!(scalar(&pool, "PRAGMA synchronous").await, 1); // NORMAL
226    }
227
228    #[tokio::test]
229    async fn overrides_are_honored() {
230        let tmp = tempfile::tempdir().unwrap();
231        let db = tmp.path().join("t.db");
232        let pool = HardenedSqlite::new(format!("sqlite://{}", db.display()))
233            .foreign_keys(false)
234            .busy_timeout(Duration::from_secs(2))
235            .connect()
236            .await
237            .unwrap();
238        assert_eq!(scalar(&pool, "PRAGMA foreign_keys").await, 0);
239        assert_eq!(scalar(&pool, "PRAGMA busy_timeout").await, 2000);
240    }
241
242    #[tokio::test]
243    async fn in_memory_survives_drop_to_zero() {
244        let pool = HardenedSqlite::new(":memory:").connect().await.unwrap();
245        sqlx::query("CREATE TABLE t (id INTEGER PRIMARY KEY)")
246            .execute(&pool)
247            .await
248            .unwrap();
249        sqlx::query("INSERT INTO t (id) VALUES (1)")
250            .execute(&pool)
251            .await
252            .unwrap();
253        // Give any idle connection time to cycle; the pinned min-connection
254        // keeps the shared-cache memdb alive so the table persists.
255        tokio::time::sleep(std::time::Duration::from_millis(20)).await;
256        let n = scalar(&pool, "SELECT COUNT(*) FROM t").await;
257        assert_eq!(n, 1);
258    }
259
260    #[tokio::test]
261    async fn two_memory_pools_are_isolated() {
262        let a = HardenedSqlite::new(":memory:").connect().await.unwrap();
263        let b = HardenedSqlite::new(":memory:").connect().await.unwrap();
264        sqlx::query("CREATE TABLE only_in_a (x INTEGER)")
265            .execute(&a)
266            .await
267            .unwrap();
268        let seen: i64 = sqlx::query_scalar(
269            "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='only_in_a'",
270        )
271        .fetch_one(&b)
272        .await
273        .unwrap();
274        assert_eq!(seen, 0);
275    }
276
277    #[tokio::test]
278    async fn begin_immediate_serializes_writers() {
279        let tmp = tempfile::tempdir().unwrap();
280        let db = tmp.path().join("t.db");
281        let pool = HardenedSqlite::new(format!("sqlite://{}", db.display()))
282            .connect()
283            .await
284            .unwrap();
285        sqlx::query("CREATE TABLE c (n INTEGER)")
286            .execute(&pool)
287            .await
288            .unwrap();
289        sqlx::query("INSERT INTO c (n) VALUES (0)")
290            .execute(&pool)
291            .await
292            .unwrap();
293
294        // Two concurrent read-then-write transactions. With begin_immediate and
295        // the 5s busy_timeout, the second waits for the first instead of hitting
296        // a busy-snapshot error. Both increments must land → n == 2.
297        let inc = |p: SqlitePool| async move {
298            let mut tx = p.begin_immediate().await.unwrap();
299            let cur: i64 = sqlx::query_scalar("SELECT n FROM c")
300                .fetch_one(&mut *tx)
301                .await
302                .unwrap();
303            sqlx::query("UPDATE c SET n = ?")
304                .bind(cur + 1)
305                .execute(&mut *tx)
306                .await
307                .unwrap();
308            tx.commit().await.unwrap();
309        };
310        let (a, b) = (pool.clone(), pool.clone());
311        let (ra, rb) = tokio::join!(tokio::spawn(inc(a)), tokio::spawn(inc(b)));
312        ra.unwrap();
313        rb.unwrap();
314
315        let n: i64 = sqlx::query_scalar("SELECT n FROM c")
316            .fetch_one(&pool)
317            .await
318            .unwrap();
319        assert_eq!(n, 2, "both writers committed without a lost update");
320    }
321
322    #[test]
323    fn pragma_statements_reflect_settings() {
324        let stmts = HardenedSqlite::new("sqlite://x.db").pragma_statements();
325        assert!(stmts.iter().any(|s| s == "PRAGMA journal_mode = WAL;"));
326        assert!(stmts.iter().any(|s| s == "PRAGMA foreign_keys = ON;"));
327        assert!(stmts.iter().any(|s| s == "PRAGMA busy_timeout = 5000;"));
328        assert!(stmts.iter().any(|s| s == "PRAGMA synchronous = NORMAL;"));
329    }
330
331    #[test]
332    fn pragma_statements_omit_wal_for_memory() {
333        let stmts = HardenedSqlite::new(":memory:").pragma_statements();
334        assert!(!stmts.iter().any(|s| s.contains("journal_mode")));
335    }
336
337    #[tokio::test]
338    async fn hardened_settings_enforced_with_url_params_present() {
339        let tmp = tempfile::tempdir().unwrap();
340        let db = tmp.path().join("t.db");
341        // `mode=rwc` is a pass-through param honored by from_str; the hardened
342        // foreign_keys setting is still enforced (ON) alongside it.
343        let url = format!("sqlite://{}?mode=rwc", db.display());
344        let pool = HardenedSqlite::new(url).connect().await.unwrap();
345        assert_eq!(scalar(&pool, "PRAGMA foreign_keys").await, 1);
346        assert_eq!(scalar(&pool, "PRAGMA busy_timeout").await, 5000);
347    }
348}
sqlite_hardened/lib.rs

sqlite_hardened/
lib.rs
