From 93774ffc3b8cb0b0a87f10cc18072827e2b886dc Mon Sep 17 00:00:00 2001 From: Avinash Sajjanshetty Date: Tue, 12 Aug 2025 21:44:25 +0530 Subject: [PATCH] Add `PRAGMA key` to set the encryption key If set, set the key for the connection --- core/lib.rs | 14 ++++++++ core/pragma.rs | 4 +++ core/translate/pragma.rs | 35 ++++++++++++++++--- vendored/sqlite3-parser/src/parser/ast/mod.rs | 5 +++ 4 files changed, 54 insertions(+), 4 deletions(-) diff --git a/core/lib.rs b/core/lib.rs index e1e5c6ae4..0d0c737ab 100644 --- a/core/lib.rs +++ b/core/lib.rs @@ -76,6 +76,8 @@ use std::{ }; #[cfg(feature = "fs")] use storage::database::DatabaseFile; +#[cfg(feature = "encryption")] +pub use storage::encryption::EncryptionKey; use storage::page_cache::DumbLruPageCache; use storage::pager::{AtomicDbState, DbState}; use storage::sqlite3_ondisk::PageSize; @@ -425,6 +427,8 @@ impl Database { view_transaction_states: RefCell::new(HashMap::new()), metrics: RefCell::new(ConnectionMetrics::new()), is_nested_stmt: Cell::new(false), + #[cfg(feature = "encryption")] + encryption_key: RefCell::new(None), }); let builtin_syms = self.builtin_syms.borrow(); // add built-in extensions symbols to the connection to prevent having to load each time @@ -852,6 +856,8 @@ pub struct Connection { /// Whether the connection is executing a statement initiated by another statement. /// Generally this is only true for ParseSchema. is_nested_stmt: Cell, + #[cfg(feature = "encryption")] + encryption_key: RefCell>, } impl Connection { @@ -1925,6 +1931,14 @@ impl Connection { pub fn get_syms_vtab_mods(&self) -> std::collections::HashSet { self.syms.borrow().vtab_modules.keys().cloned().collect() } + + #[cfg(feature = "encryption")] + pub fn set_encryption_key(&self, key: Option) { + tracing::trace!("setting encryption key for connection"); + *self.encryption_key.borrow_mut() = key.clone(); + let pager = self.pager.borrow(); + pager.set_encryption_key(key); + } } pub struct Statement { diff --git a/core/pragma.rs b/core/pragma.rs index f4b4b3f44..62be7f313 100644 --- a/core/pragma.rs +++ b/core/pragma.rs @@ -107,6 +107,10 @@ pub fn pragma_for(pragma: &PragmaName) -> Pragma { &["query_only"], ), FreelistCount => Pragma::new(PragmaFlags::Result0, &["freelist_count"]), + EncryptionKey => Pragma::new( + PragmaFlags::Result0 | PragmaFlags::SchemaReq | PragmaFlags::NoColumns1, + &["key"], + ), } } diff --git a/core/translate/pragma.rs b/core/translate/pragma.rs index bac0309cd..2af63a977 100644 --- a/core/translate/pragma.rs +++ b/core/translate/pragma.rs @@ -7,11 +7,16 @@ use std::sync::Arc; use turso_sqlite3_parser::ast::{self, ColumnDefinition, Expr, Literal, Name}; use turso_sqlite3_parser::ast::{PragmaName, QualifiedName}; +use super::integrity_check::translate_integrity_check; use crate::pragma::pragma_for; use crate::schema::Schema; +#[cfg(feature = "encryption")] +use crate::storage::encryption::EncryptionKey; use crate::storage::pager::AutoVacuumMode; +use crate::storage::pager::Pager; use crate::storage::sqlite3_ondisk::CacheSize; use crate::storage::wal::CheckpointMode; +use crate::translate::emitter::TransactionMode; use crate::translate::schema::translate_create_table; use crate::util::{normalize_ident, parse_signed_number, parse_string, IOExt as _}; use crate::vdbe::builder::{ProgramBuilder, ProgramBuilderOpts}; @@ -20,10 +25,6 @@ use crate::{bail_parse_error, CaptureDataChangesMode, LimboError, SymbolTable, V use std::str::FromStr; use strum::IntoEnumIterator; -use super::integrity_check::translate_integrity_check; -use crate::storage::pager::Pager; -use crate::translate::emitter::TransactionMode; - fn list_pragmas(program: &mut ProgramBuilder) { for x in PragmaName::iter() { let register = program.emit_string8_new_reg(x.to_string()); @@ -309,6 +310,15 @@ fn update_pragma( connection, program, ), + PragmaName::EncryptionKey => { + #[cfg(feature = "encryption")] + { + let value = parse_string(&value)?; + let key = EncryptionKey::from_string(&value); + connection.set_encryption_key(Some(key)); + } + Ok((program, TransactionMode::None)) + } } } @@ -566,6 +576,23 @@ fn query_pragma( program.add_pragma_result_column(pragma.to_string()); Ok((program, TransactionMode::None)) } + PragmaName::EncryptionKey => { + #[cfg(feature = "encryption")] + { + let msg = { + if connection.encryption_key.borrow().is_some() { + "encryption key is set for this session" + } else { + "encryption key is not set for this session" + } + }; + let register = program.alloc_register(); + program.emit_string8(msg.to_string(), register); + program.emit_result_row(register, 1); + program.add_pragma_result_column(pragma.to_string()); + } + Ok((program, TransactionMode::None)) + } } } diff --git a/vendored/sqlite3-parser/src/parser/ast/mod.rs b/vendored/sqlite3-parser/src/parser/ast/mod.rs index e8d755799..de096827f 100644 --- a/vendored/sqlite3-parser/src/parser/ast/mod.rs +++ b/vendored/sqlite3-parser/src/parser/ast/mod.rs @@ -1829,6 +1829,11 @@ pub enum PragmaName { IntegrityCheck, /// `journal_mode` pragma JournalMode, + /// encryption key for encrypted databases. This is just called `key` because most + /// extensions use this name instead of `encryption_key`. + #[strum(serialize = "key")] + #[cfg_attr(feature = "serde", serde(rename = "key"))] + EncryptionKey, /// Noop as per SQLite docs LegacyFileFormat, /// Set or get the maximum number of pages in the database file.