sqlite-hardened 0.1.0

//! Correct-by-default sqlx SQLite pool bootstrap.

use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous};
use sqlx::SqlitePool;
use sqlx::{Sqlite, Transaction};
use std::str::FromStr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error(transparent)]
    Sqlx(#[from] sqlx::Error),
}
pub type Result<T> = std::result::Result<T, Error>;

/// Builder for a hardened `SqlitePool`.
pub struct HardenedSqlite {
    url: String,
    journal_mode: SqliteJournalMode,
    foreign_keys: bool,
    busy_timeout: Duration,
    synchronous: SqliteSynchronous,
    create_if_missing: bool,
    max_connections: u32,
    min_connections: Option<u32>,
    memory_name: Option<String>,
}

impl HardenedSqlite {
    /// New builder with the hardened defaults:
    /// WAL, foreign_keys ON, busy_timeout 5s, synchronous NORMAL,
    /// create_if_missing, max_connections 8.
    pub fn new(url: impl Into<String>) -> Self {
        Self {
            url: url.into(),
            journal_mode: SqliteJournalMode::Wal,
            foreign_keys: true,
            busy_timeout: Duration::from_secs(5),
            synchronous: SqliteSynchronous::Normal,
            create_if_missing: true,
            max_connections: 8,
            min_connections: None,
            memory_name: None,
        }
    }

    pub fn journal_mode(mut self, m: SqliteJournalMode) -> Self {
        self.journal_mode = m;
        self
    }
    pub fn foreign_keys(mut self, on: bool) -> Self {
        self.foreign_keys = on;
        self
    }
    pub fn busy_timeout(mut self, d: Duration) -> Self {
        self.busy_timeout = d;
        self
    }
    pub fn synchronous(mut self, s: SqliteSynchronous) -> Self {
        self.synchronous = s;
        self
    }
    pub fn create_if_missing(mut self, on: bool) -> Self {
        self.create_if_missing = on;
        self
    }
    pub fn max_connections(mut self, n: u32) -> Self {
        self.max_connections = n;
        self
    }
    pub fn min_connections(mut self, n: u32) -> Self {
        self.min_connections = Some(n);
        self
    }
    /// Force a specific shared-cache in-memory name (else auto-salted).
    pub fn memory_name(mut self, name: impl Into<String>) -> Self {
        self.memory_name = Some(name.into());
        self
    }

    fn is_memory(&self) -> bool {
        is_memory_url(&self.url)
    }

    /// Build the `SqliteConnectOptions` without connecting.
    pub fn connect_options(&self) -> Result<SqliteConnectOptions> {
        let opts = if self.is_memory() {
            // Explicit unique file:...?mode=memory&cache=shared — NOT a bare
            // ":memory:" URL, which sqlx would turn into SQLITE_OPEN_MEMORY and
            // collapse every pool onto one anonymous shared cache. WAL is
            // meaningless for memory DBs, so it's skipped here.
            let name = self.memory_name.clone().unwrap_or_else(salted_memory_name);
            SqliteConnectOptions::new()
                .filename(format!("file:{name}?mode=memory&cache=shared"))
                .create_if_missing(true)
                .foreign_keys(self.foreign_keys)
                .busy_timeout(self.busy_timeout)
                .synchronous(self.synchronous)
        } else {
            SqliteConnectOptions::from_str(&self.url)?
                .create_if_missing(self.create_if_missing)
                .journal_mode(self.journal_mode)
                .foreign_keys(self.foreign_keys)
                .busy_timeout(self.busy_timeout)
                .synchronous(self.synchronous)
        };
        Ok(opts)
    }

    /// The PRAGMA statements this builder would apply, in order — for callers
    /// on rusqlite or another driver who want the same hardening without the
    /// sqlx pool. WAL is omitted for in-memory configurations.
    pub fn pragma_statements(&self) -> Vec<String> {
        let mut v = Vec::new();
        if !self.is_memory() {
            let jm = match self.journal_mode {
                SqliteJournalMode::Wal => "WAL",
                SqliteJournalMode::Delete => "DELETE",
                SqliteJournalMode::Truncate => "TRUNCATE",
                SqliteJournalMode::Persist => "PERSIST",
                SqliteJournalMode::Memory => "MEMORY",
                SqliteJournalMode::Off => "OFF",
            };
            v.push(format!("PRAGMA journal_mode = {jm};"));
        }
        v.push(format!(
            "PRAGMA foreign_keys = {};",
            if self.foreign_keys { "ON" } else { "OFF" }
        ));
        v.push(format!(
            "PRAGMA busy_timeout = {};",
            self.busy_timeout.as_millis()
        ));
        let sync = match self.synchronous {
            SqliteSynchronous::Off => "OFF",
            SqliteSynchronous::Normal => "NORMAL",
            SqliteSynchronous::Full => "FULL",
            SqliteSynchronous::Extra => "EXTRA",
        };
        v.push(format!("PRAGMA synchronous = {sync};"));
        v
    }

    /// Build and connect the pool.
    pub async fn connect(self) -> Result<SqlitePool> {
        let opts = self.connect_options()?;
        let mut pool_opts = SqlitePoolOptions::new().max_connections(self.max_connections);
        // Pin one connection for in-memory DBs so the shared cache can't be
        // destroyed when the pool drops to zero live connections.
        let min = self
            .min_connections
            .map(|m| if self.is_memory() { m.max(1) } else { m })
            .or(if self.is_memory() { Some(1) } else { None });
        if let Some(m) = min {
            pool_opts = pool_opts.min_connections(m);
        }
        Ok(pool_opts.connect_with(opts).await?)
    }
}

/// True for the recognized in-memory URL spellings.
pub fn is_memory_url(url: &str) -> bool {
    matches!(url, ":memory:" | "sqlite::memory:" | "sqlite://:memory:")
}

static MEM_SERIAL: AtomicU64 = AtomicU64::new(0);

fn salted_memory_name() -> String {
    let n = MEM_SERIAL.fetch_add(1, Ordering::Relaxed);
    format!("sqlite-hardened-{}-{n}", std::process::id())
}

/// Begin a WRITE transaction with `BEGIN IMMEDIATE`.
///
/// sqlx's `pool.begin()` issues a deferred BEGIN: a transaction that reads
/// before it writes pins its read snapshot on the first statement, and if
/// another connection commits in between, the later write cannot upgrade —
/// SQLite returns a busy-snapshot error even *with* a busy_timeout, because the
/// conflict is logical, not a lock wait. `BEGIN IMMEDIATE` takes the write lock
/// up front, so a read-then-write transaction is safe under WAL's single-writer
/// rule.
pub async fn begin_immediate(pool: &SqlitePool) -> Result<Transaction<'static, Sqlite>> {
    Ok(pool.begin_with("BEGIN IMMEDIATE").await?)
}

/// `pool.begin_immediate()` sugar.
#[allow(async_fn_in_trait)]
pub trait SqlitePoolExt {
    async fn begin_immediate(&self) -> Result<Transaction<'static, Sqlite>>;
}

impl SqlitePoolExt for SqlitePool {
    async fn begin_immediate(&self) -> Result<Transaction<'static, Sqlite>> {
        begin_immediate(self).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    async fn scalar(pool: &SqlitePool, sql: &str) -> i64 {
        sqlx::query_scalar::<_, i64>(sql)
            .fetch_one(pool)
            .await
            .unwrap()
    }

    #[tokio::test]
    async fn file_pool_applies_hardened_pragmas() {
        let tmp = tempfile::tempdir().unwrap();
        let db = tmp.path().join("t.db");
        let pool = HardenedSqlite::new(format!("sqlite://{}", db.display()))
            .connect()
            .await
            .unwrap();
        let jm: String = sqlx::query_scalar("PRAGMA journal_mode")
            .fetch_one(&pool)
            .await
            .unwrap();
        assert_eq!(jm.to_lowercase(), "wal");
        assert_eq!(scalar(&pool, "PRAGMA foreign_keys").await, 1);
        assert_eq!(scalar(&pool, "PRAGMA busy_timeout").await, 5000);
        assert_eq!(scalar(&pool, "PRAGMA synchronous").await, 1); // NORMAL
    }

    #[tokio::test]
    async fn overrides_are_honored() {
        let tmp = tempfile::tempdir().unwrap();
        let db = tmp.path().join("t.db");
        let pool = HardenedSqlite::new(format!("sqlite://{}", db.display()))
            .foreign_keys(false)
            .busy_timeout(Duration::from_secs(2))
            .connect()
            .await
            .unwrap();
        assert_eq!(scalar(&pool, "PRAGMA foreign_keys").await, 0);
        assert_eq!(scalar(&pool, "PRAGMA busy_timeout").await, 2000);
    }

    #[tokio::test]
    async fn in_memory_survives_drop_to_zero() {
        let pool = HardenedSqlite::new(":memory:").connect().await.unwrap();
        sqlx::query("CREATE TABLE t (id INTEGER PRIMARY KEY)")
            .execute(&pool)
            .await
            .unwrap();
        sqlx::query("INSERT INTO t (id) VALUES (1)")
            .execute(&pool)
            .await
            .unwrap();
        // Give any idle connection time to cycle; the pinned min-connection
        // keeps the shared-cache memdb alive so the table persists.
        tokio::time::sleep(std::time::Duration::from_millis(20)).await;
        let n = scalar(&pool, "SELECT COUNT(*) FROM t").await;
        assert_eq!(n, 1);
    }

    #[tokio::test]
    async fn two_memory_pools_are_isolated() {
        let a = HardenedSqlite::new(":memory:").connect().await.unwrap();
        let b = HardenedSqlite::new(":memory:").connect().await.unwrap();
        sqlx::query("CREATE TABLE only_in_a (x INTEGER)")
            .execute(&a)
            .await
            .unwrap();
        let seen: i64 = sqlx::query_scalar(
            "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='only_in_a'",
        )
        .fetch_one(&b)
        .await
        .unwrap();
        assert_eq!(seen, 0);
    }

    #[tokio::test]
    async fn begin_immediate_serializes_writers() {
        let tmp = tempfile::tempdir().unwrap();
        let db = tmp.path().join("t.db");
        let pool = HardenedSqlite::new(format!("sqlite://{}", db.display()))
            .connect()
            .await
            .unwrap();
        sqlx::query("CREATE TABLE c (n INTEGER)")
            .execute(&pool)
            .await
            .unwrap();
        sqlx::query("INSERT INTO c (n) VALUES (0)")
            .execute(&pool)
            .await
            .unwrap();

        // Two concurrent read-then-write transactions. With begin_immediate and
        // the 5s busy_timeout, the second waits for the first instead of hitting
        // a busy-snapshot error. Both increments must land → n == 2.
        let inc = |p: SqlitePool| async move {
            let mut tx = p.begin_immediate().await.unwrap();
            let cur: i64 = sqlx::query_scalar("SELECT n FROM c")
                .fetch_one(&mut *tx)
                .await
                .unwrap();
            sqlx::query("UPDATE c SET n = ?")
                .bind(cur + 1)
                .execute(&mut *tx)
                .await
                .unwrap();
            tx.commit().await.unwrap();
        };
        let (a, b) = (pool.clone(), pool.clone());
        let (ra, rb) = tokio::join!(tokio::spawn(inc(a)), tokio::spawn(inc(b)));
        ra.unwrap();
        rb.unwrap();

        let n: i64 = sqlx::query_scalar("SELECT n FROM c")
            .fetch_one(&pool)
            .await
            .unwrap();
        assert_eq!(n, 2, "both writers committed without a lost update");
    }

    #[test]
    fn pragma_statements_reflect_settings() {
        let stmts = HardenedSqlite::new("sqlite://x.db").pragma_statements();
        assert!(stmts.iter().any(|s| s == "PRAGMA journal_mode = WAL;"));
        assert!(stmts.iter().any(|s| s == "PRAGMA foreign_keys = ON;"));
        assert!(stmts.iter().any(|s| s == "PRAGMA busy_timeout = 5000;"));
        assert!(stmts.iter().any(|s| s == "PRAGMA synchronous = NORMAL;"));
    }

    #[test]
    fn pragma_statements_omit_wal_for_memory() {
        let stmts = HardenedSqlite::new(":memory:").pragma_statements();
        assert!(!stmts.iter().any(|s| s.contains("journal_mode")));
    }

    #[tokio::test]
    async fn hardened_settings_enforced_with_url_params_present() {
        let tmp = tempfile::tempdir().unwrap();
        let db = tmp.path().join("t.db");
        // `mode=rwc` is a pass-through param honored by from_str; the hardened
        // foreign_keys setting is still enforced (ON) alongside it.
        let url = format!("sqlite://{}?mode=rwc", db.display());
        let pool = HardenedSqlite::new(url).connect().await.unwrap();
        assert_eq!(scalar(&pool, "PRAGMA foreign_keys").await, 1);
        assert_eq!(scalar(&pool, "PRAGMA busy_timeout").await, 5000);
    }
}