Skip to main content

rightkit_db/
lib.rs

1#![forbid(unsafe_code)]
2
3use rusqlite::{Connection, OpenFlags, Transaction, TransactionBehavior};
4use std::ffi::OsString;
5use std::ops::Deref;
6use std::path::{Path, PathBuf};
7use std::sync::{Arc, Condvar, Mutex, MutexGuard, Weak};
8use zeroize::Zeroizing;
9
10mod backup;
11mod migrate;
12mod recover;
13
14pub use backup::{online_backup_to, BackupProgress};
15pub use migrate::{
16    add_column_if_missing, column_exists, migrate, pending_marker, rollback_to, user_version,
17    MigrateOptions, Migration, MigrationReport, PendingMarker, Step,
18};
19#[cfg(feature = "sqlcipher")]
20pub use recover::rekey;
21pub use recover::{quarantine_database, OpenOutcome, RecoveryPolicy};
22
23const BUSY_TIMEOUT_MS: u64 = 5_000;
24const JOURNAL_SIZE_LIMIT: i64 = 67_108_864;
25
26pub type Result<T> = std::result::Result<T, DbError>;
27
28#[derive(Debug, thiserror::Error)]
29pub enum DbError {
30    #[error("SQLite error: {0}")]
31    Sqlite(#[from] rusqlite::Error),
32    #[error("a SQLCipher key requires the rightkit-db `sqlcipher` feature")]
33    SqlcipherFeatureRequired,
34    #[error("SQLCipher key contains characters that are unsafe for PRAGMA interpolation")]
35    InvalidSqlcipherKey,
36    #[error("database integrity check failed: {0}")]
37    Integrity(String),
38    #[error("read-pool capacity must be greater than zero")]
39    InvalidPoolCapacity,
40    #[error("database path has no file name")]
41    InvalidDatabasePath,
42    #[error("database is at schema v{current}, newer than this build understands (v{max}) and not forward compatible")]
43    NewerDatabase { current: i64, max: i64 },
44    #[error("invalid migration list: {0}")]
45    InvalidMigrations(String),
46    #[error("migration v{version} ({name}) failed and was rolled back: {source}")]
47    MigrationFailed {
48        version: i64,
49        name: &'static str,
50        #[source]
51        source: rusqlite::Error,
52    },
53    #[error("cannot roll back: {0}")]
54    Rollback(String),
55    #[error("database backup failed: {0}")]
56    Backup(String),
57    #[error("no candidate key opened the database")]
58    NoKeyOpened,
59    #[error(transparent)]
60    Fs(#[from] rightkit_fs::FsError),
61}
62
63#[derive(Clone, Copy, Debug, Eq, PartialEq)]
64pub enum CacheProfile {
65    Default,
66    LargeWorkingSet,
67}
68
69#[derive(Clone, Copy, Debug, Eq, PartialEq)]
70pub enum DurabilityProfile {
71    Full,
72    Normal,
73}
74
75/// How the journal mode is chosen. WAL needs working shared-memory locking, which network
76/// filesystems often lack and which SQLite will silently downgrade; the default therefore
77/// *verifies* the mode that took effect and falls back to a rollback journal.
78#[derive(Clone, Copy, Debug, Eq, PartialEq)]
79pub enum JournalSelect {
80    /// WAL when the filesystem honors it and is not a detected network mount; else DELETE.
81    WalWithFallback,
82    /// WAL or fail (no fallback).
83    WalStrict,
84    /// Rollback journal.
85    Delete,
86}
87
88/// Environment kill-switch (`wal` | `delete`) for a misclassified mount in the field.
89pub const JOURNAL_ENV: &str = "RIGHTKIT_SQLITE_JOURNAL_MODE";
90
91/// Post-key, pre-pragma per-connection initializer (register functions, collations, ...).
92pub type ConnInit = Arc<dyn Fn(&Connection) -> rusqlite::Result<()> + Send + Sync>;
93
94#[derive(Clone)]
95pub struct DbOptions {
96    pub path: PathBuf,
97    /// SQLCipher passphrase-style key (base64-ish charset). See also [`DbOptions::raw_key`].
98    pub key: Option<Zeroizing<String>>,
99    /// SQLCipher raw 256-bit key, applied as `x'<hex>'` (no KDF). Takes precedence over `key`.
100    pub raw_key: Option<Zeroizing<[u8; 32]>>,
101    pub profile: CacheProfile,
102    pub durability: DurabilityProfile,
103    pub foreign_keys: bool,
104    pub integrity_check_on_open: bool,
105    pub journal: JournalSelect,
106    /// Narrow the database and its sidecars to owner-only (0600) on unix. Default on.
107    pub owner_only: bool,
108    pub init: Option<ConnInit>,
109}
110
111impl DbOptions {
112    pub fn new(path: impl Into<PathBuf>) -> Self {
113        Self {
114            path: path.into(),
115            key: None,
116            raw_key: None,
117            profile: CacheProfile::Default,
118            durability: DurabilityProfile::Normal,
119            foreign_keys: false,
120            integrity_check_on_open: false,
121            journal: JournalSelect::WalWithFallback,
122            owner_only: true,
123            init: None,
124        }
125    }
126
127    pub fn with_passphrase(mut self, key: impl Into<String>) -> Self {
128        self.key = Some(Zeroizing::new(key.into()));
129        self.raw_key = None;
130        self
131    }
132
133    pub fn with_raw_key(mut self, key: [u8; 32]) -> Self {
134        self.raw_key = Some(Zeroizing::new(key));
135        self.key = None;
136        self
137    }
138
139    pub fn is_keyed(&self) -> bool {
140        self.key.is_some() || self.raw_key.is_some()
141    }
142
143    pub fn open_writer(&self) -> Result<Connection> {
144        self.open(ConnectionKind::Writer)
145    }
146
147    pub fn open_reader(&self) -> Result<Connection> {
148        self.open(ConnectionKind::Reader)
149    }
150
151    pub fn open_read_pool(&self, cap: usize) -> Result<ReadPool> {
152        ReadPool::open(self.clone(), cap)
153    }
154
155    fn open(&self, kind: ConnectionKind) -> Result<Connection> {
156        let key = self.key_pragma_value()?;
157        let flags = OpenFlags::SQLITE_OPEN_READ_WRITE
158            | OpenFlags::SQLITE_OPEN_CREATE
159            | OpenFlags::SQLITE_OPEN_NO_MUTEX;
160        let conn = Connection::open_with_flags(&self.path, flags)?;
161        // Narrow the main file before any write so WAL/SHM sidecars inherit owner-only.
162        if self.owner_only {
163            rightkit_fs::restrict_owner_only(&self.path)?;
164        }
165        apply_key_first(&conn, key.as_ref().map(|k| k.as_str()))?;
166        verify_key(&conn)?;
167        if let Some(init) = &self.init {
168            init(&conn)?;
169        }
170        configure_connection(&conn, self, kind)?;
171        if self.integrity_check_on_open {
172            integrity_check(&conn, self.is_keyed())?;
173        }
174        if self.owner_only {
175            harden_sidecars(&self.path)?;
176        }
177        Ok(conn)
178    }
179
180    /// The `PRAGMA key` right-hand side (already quoted), validated against injection.
181    fn key_pragma_value(&self) -> Result<Option<Zeroizing<String>>> {
182        if let Some(raw) = &self.raw_key {
183            #[cfg(not(feature = "sqlcipher"))]
184            {
185                let _ = raw;
186                return Err(DbError::SqlcipherFeatureRequired);
187            }
188            #[cfg(feature = "sqlcipher")]
189            {
190                let hex: String = raw.iter().map(|b| format!("{b:02x}")).collect();
191                return Ok(Some(Zeroizing::new(format!("\"x'{hex}'\""))));
192            }
193        }
194        let key = self.key.as_ref().map(|key| key.as_str());
195        validate_key(key)?;
196        Ok(key.map(|k| Zeroizing::new(format!("'{k}'"))))
197    }
198}
199
200fn harden_sidecars(db: &Path) -> Result<()> {
201    for suffix in ["-wal", "-shm"] {
202        let mut name = OsString::from(db.file_name().ok_or(DbError::InvalidDatabasePath)?);
203        name.push(suffix);
204        rightkit_fs::restrict_owner_only(&db.with_file_name(name))?;
205    }
206    Ok(())
207}
208
209#[derive(Clone, Copy)]
210enum ConnectionKind {
211    Writer,
212    Reader,
213}
214
215fn configure_connection(
216    conn: &Connection,
217    options: &DbOptions,
218    kind: ConnectionKind,
219) -> Result<()> {
220    conn.busy_timeout(std::time::Duration::from_millis(BUSY_TIMEOUT_MS))?;
221    configure_journal(conn, options, kind)?;
222    let synchronous = match options.durability {
223        DurabilityProfile::Full => "FULL",
224        DurabilityProfile::Normal => "NORMAL",
225    };
226    conn.pragma_update(None, "synchronous", synchronous)?;
227    conn.pragma_update(None, "journal_size_limit", JOURNAL_SIZE_LIMIT)?;
228    conn.pragma_update(None, "temp_store", "MEMORY")?;
229    if options.foreign_keys {
230        conn.pragma_update(None, "foreign_keys", true)?;
231    }
232    if options.profile == CacheProfile::LargeWorkingSet {
233        conn.pragma_update(None, "cache_size", -65_536_i64)?;
234        conn.pragma_update(None, "mmap_size", 268_435_456_i64)?;
235    }
236    optimize_on_open(conn)?;
237    if matches!(kind, ConnectionKind::Reader) {
238        conn.pragma_update(None, "query_only", true)?;
239    }
240    Ok(())
241}
242
243fn configure_journal(conn: &Connection, options: &DbOptions, kind: ConnectionKind) -> Result<()> {
244    let wanted = effective_journal(options);
245    if wanted == JournalSelect::Delete {
246        if matches!(kind, ConnectionKind::Writer) {
247            conn.query_row("PRAGMA journal_mode=DELETE", [], |row| {
248                row.get::<_, String>(0)
249            })?;
250        }
251        return Ok(());
252    }
253    let honored = conn
254        .query_row("PRAGMA journal_mode=WAL", [], |row| row.get::<_, String>(0))
255        .map(|mode| mode.eq_ignore_ascii_case("wal"))
256        .unwrap_or(false)
257        // `journal_mode=WAL` can report success before the -shm/-wal sidecars are usable;
258        // one empty write transaction proves the sidecars work before anything depends on them.
259        && (matches!(kind, ConnectionKind::Reader)
260            || conn.execute_batch("BEGIN IMMEDIATE; ROLLBACK;").is_ok());
261    if honored {
262        return Ok(());
263    }
264    if wanted == JournalSelect::WalStrict {
265        return Err(DbError::Sqlite(rusqlite::Error::SqliteFailure(
266            rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_IOERR),
267            Some("WAL journal mode was not honored by this filesystem".into()),
268        )));
269    }
270    let _ = conn.execute_batch("ROLLBACK;");
271    conn.query_row("PRAGMA journal_mode=DELETE", [], |row| {
272        row.get::<_, String>(0)
273    })?;
274    Ok(())
275}
276
277fn effective_journal(options: &DbOptions) -> JournalSelect {
278    match std::env::var(JOURNAL_ENV)
279        .ok()
280        .as_deref()
281        .map(str::to_ascii_lowercase)
282        .as_deref()
283    {
284        Some("wal") => return JournalSelect::WalStrict,
285        Some("delete") | Some("truncate") => return JournalSelect::Delete,
286        _ => {}
287    }
288    if options.journal == JournalSelect::WalWithFallback
289        && rightkit_fs::is_network_filesystem(&options.path)
290    {
291        return JournalSelect::Delete;
292    }
293    options.journal
294}
295
296/// The journal mode SQLite is actually running (`wal`, `delete`, ...).
297pub fn journal_mode(conn: &Connection) -> Result<String> {
298    Ok(conn.query_row("PRAGMA journal_mode", [], |row| row.get::<_, String>(0))?)
299}
300
301/// Start a transaction that reserves SQLite's single writer slot up front.
302/// This keeps read-then-write work out of the unretryable BUSY_SNAPSHOT path.
303pub fn write_txn(conn: &mut Connection) -> Result<Transaction<'_>> {
304    conn.transaction_with_behavior(TransactionBehavior::Immediate)
305        .map_err(Into::into)
306}
307
308#[cfg(feature = "sqlcipher")]
309fn apply_key_first(conn: &Connection, key: Option<&str>) -> Result<()> {
310    if let Some(key) = key {
311        let statement = Zeroizing::new(format!("PRAGMA key = {key};"));
312        conn.execute_batch(statement.as_str())?;
313    }
314    Ok(())
315}
316
317#[cfg(not(feature = "sqlcipher"))]
318fn apply_key_first(_conn: &Connection, key: Option<&str>) -> Result<()> {
319    if key.is_some() {
320        return Err(DbError::SqlcipherFeatureRequired);
321    }
322    Ok(())
323}
324
325fn validate_key(key: Option<&str>) -> Result<()> {
326    #[cfg(not(feature = "sqlcipher"))]
327    if key.is_some() {
328        return Err(DbError::SqlcipherFeatureRequired);
329    }
330    if let Some(key) = key {
331        let safe = !key.is_empty()
332            && key.len() <= 4_096
333            && key.bytes().all(|byte| {
334                byte.is_ascii_alphanumeric() || matches!(byte, b'+' | b'/' | b'=' | b'-' | b'_')
335            });
336        if !safe {
337            return Err(DbError::InvalidSqlcipherKey);
338        }
339    }
340    Ok(())
341}
342
343fn verify_key(conn: &Connection) -> Result<()> {
344    conn.query_row("SELECT count(*) FROM sqlite_master", [], |row| {
345        row.get::<_, i64>(0)
346    })?;
347    Ok(())
348}
349
350fn integrity_check(conn: &Connection, keyed: bool) -> Result<()> {
351    let pragma = if keyed {
352        "PRAGMA cipher_integrity_check"
353    } else {
354        "PRAGMA integrity_check"
355    };
356    let mut statement = conn.prepare(pragma)?;
357    let rows = statement.query_map([], |row| row.get::<_, String>(0))?;
358    let messages = rows.collect::<std::result::Result<Vec<_>, _>>()?;
359    if messages.is_empty()
360        || messages
361            .iter()
362            .all(|message| message.eq_ignore_ascii_case("ok"))
363    {
364        return Ok(());
365    }
366    Err(DbError::Integrity(messages.join("; ")))
367}
368
369#[derive(Clone)]
370pub struct ReadPool {
371    inner: Arc<ReadPoolInner>,
372}
373
374struct ReadPoolInner {
375    options: DbOptions,
376    cap: usize,
377    state: Mutex<ReadPoolState>,
378    available: Condvar,
379}
380
381#[derive(Default)]
382struct ReadPoolState {
383    idle: Vec<Connection>,
384    total: usize,
385}
386
387pub struct PooledReadConnection {
388    conn: Option<Connection>,
389    pool: Weak<ReadPoolInner>,
390}
391
392impl Deref for PooledReadConnection {
393    type Target = Connection;
394
395    fn deref(&self) -> &Self::Target {
396        self.conn.as_ref().expect("pooled connection is present")
397    }
398}
399
400impl Drop for PooledReadConnection {
401    fn drop(&mut self) {
402        let Some(conn) = self.conn.take() else {
403            return;
404        };
405        let Some(pool) = self.pool.upgrade() else {
406            return;
407        };
408        let mut state = lock_recover(&pool.state);
409        state.idle.push(conn);
410        pool.available.notify_one();
411    }
412}
413
414impl ReadPool {
415    fn open(options: DbOptions, cap: usize) -> Result<Self> {
416        if cap == 0 {
417            return Err(DbError::InvalidPoolCapacity);
418        }
419        let pool = Self {
420            inner: Arc::new(ReadPoolInner {
421                options,
422                cap,
423                state: Mutex::new(ReadPoolState::default()),
424                available: Condvar::new(),
425            }),
426        };
427        let conn = pool.inner.options.open_reader()?;
428        let mut state = lock_recover(&pool.inner.state);
429        state.total = 1;
430        state.idle.push(conn);
431        drop(state);
432        Ok(pool)
433    }
434
435    pub fn get(&self) -> Result<PooledReadConnection> {
436        let mut state = lock_recover(&self.inner.state);
437        loop {
438            if let Some(conn) = state.idle.pop() {
439                return Ok(PooledReadConnection {
440                    conn: Some(conn),
441                    pool: Arc::downgrade(&self.inner),
442                });
443            }
444            if state.total < self.inner.cap {
445                state.total += 1;
446                drop(state);
447                return match self.inner.options.open_reader() {
448                    Ok(conn) => Ok(PooledReadConnection {
449                        conn: Some(conn),
450                        pool: Arc::downgrade(&self.inner),
451                    }),
452                    Err(error) => {
453                        let mut state = lock_recover(&self.inner.state);
454                        state.total -= 1;
455                        self.inner.available.notify_one();
456                        Err(error)
457                    }
458                };
459            }
460            state = match self.inner.available.wait(state) {
461                Ok(state) => state,
462                Err(poisoned) => poisoned.into_inner(),
463            };
464        }
465    }
466
467    pub fn capacity(&self) -> usize {
468        self.inner.cap
469    }
470
471    pub fn idle_len(&self) -> usize {
472        lock_recover(&self.inner.state).idle.len()
473    }
474
475    pub fn total_len(&self) -> usize {
476        lock_recover(&self.inner.state).total
477    }
478}
479
480fn lock_recover<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
481    mutex
482        .lock()
483        .unwrap_or_else(|poisoned| poisoned.into_inner())
484}
485
486#[derive(Clone, Copy, Debug, Eq, PartialEq)]
487pub struct CheckpointStatus {
488    pub busy: i64,
489    pub log_frames: i64,
490    pub checkpointed_frames: i64,
491}
492
493pub fn checkpoint_truncate(conn: &Connection) -> Result<CheckpointStatus> {
494    checkpoint(conn, "PRAGMA wal_checkpoint(TRUNCATE)")
495}
496
497pub fn checkpoint_if_wal_larger_than(
498    conn: &Connection,
499    db: &Path,
500    bytes: u64,
501) -> Result<Option<CheckpointStatus>> {
502    let wal = wal_path(db)?;
503    let size = match std::fs::metadata(wal) {
504        Ok(metadata) => metadata.len(),
505        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
506        Err(_) => return Ok(None),
507    };
508    if size <= bytes {
509        return Ok(None);
510    }
511    checkpoint(conn, "PRAGMA wal_checkpoint(PASSIVE)").map(Some)
512}
513
514fn checkpoint(conn: &Connection, pragma: &str) -> Result<CheckpointStatus> {
515    conn.query_row(pragma, [], |row| {
516        Ok(CheckpointStatus {
517            busy: row.get(0)?,
518            log_frames: row.get(1)?,
519            checkpointed_frames: row.get(2)?,
520        })
521    })
522    .map_err(Into::into)
523}
524
525fn wal_path(db: &Path) -> Result<PathBuf> {
526    let file_name = db.file_name().ok_or(DbError::InvalidDatabasePath)?;
527    let mut wal_name = OsString::from(file_name);
528    wal_name.push("-wal");
529    Ok(db.with_file_name(wal_name))
530}
531
532pub fn optimize_on_open(conn: &Connection) -> Result<()> {
533    conn.execute_batch("PRAGMA optimize=0x10002;")?;
534    Ok(())
535}
536
537pub fn optimize_boundary(conn: &Connection) -> Result<()> {
538    conn.execute_batch("PRAGMA optimize;")?;
539    Ok(())
540}