diff --git a/Cargo.lock b/Cargo.lock index de294c4..e4a48ac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3203,6 +3203,7 @@ dependencies = [ "serde_json", "tempfile", "thiserror 2.0.18", + "wx-decrypt", "zstd", ] diff --git a/Cargo.toml b/Cargo.toml index 8729538..8e60b1e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,3 +8,18 @@ edition = "2021" license = "MIT" repository = "https://github.com/pandorafuture/wx-cli" description = "WeChat macOS database decryption and query tool" + +# PBKDF2 intentionally runs 256k rounds. Keep the crypto crate optimized in +# dev/test builds so parallel integration tests and local debug binaries do not +# spend seconds per database deriving SQLCipher keys. +[profile.dev.package.wx-decrypt] +opt-level = 3 + +[profile.dev.package.pbkdf2] +opt-level = 3 + +[profile.dev.package.sha2] +opt-level = 3 + +[profile.dev.package.hmac] +opt-level = 3 diff --git a/crates/wx-cli/src/cmd/query.rs b/crates/wx-cli/src/cmd/query.rs index 1999bd5..543d5dd 100644 --- a/crates/wx-cli/src/cmd/query.rs +++ b/crates/wx-cli/src/cmd/query.rs @@ -77,24 +77,23 @@ pub fn cmd_query( if options.is_enabled() && !preserve_local_warning { let client = ThinClient::new(options.clone()); - match client.probe_health().and_then(|_| { - fetch_remote_query( - &client, - contact, - since, - until, - msg_type.clone(), - effective_limit, - offset, - order.clone(), - around_sort_seq, - around_server_id, - context, - after_sort_seq, - show_hidden, - ) - }) { - Ok(envelope) => { + match client.probe_health() { + Ok(()) => { + let envelope = fetch_remote_query( + &client, + contact, + since, + until, + msg_type.clone(), + effective_limit, + offset, + order.clone(), + around_sort_seq, + around_server_id, + context, + after_sort_seq, + show_hidden, + )?; let is_group = envelope .items .first() diff --git a/crates/wx-cli/src/cmd/search.rs b/crates/wx-cli/src/cmd/search.rs index b5db28e..dcccf1d 100644 --- a/crates/wx-cli/src/cmd/search.rs +++ b/crates/wx-cli/src/cmd/search.rs @@ -1,8 +1,6 @@ use std::path::PathBuf; -use wx_context::{ - open_fts_connection_with_key, AccountContext, ContactResolver, ResolveParams, -}; +use wx_context::{register_mm_fts_tokenizer, AccountContext, ContactResolver, ResolveParams}; use super::thin_client::{ThinClient, ThinClientCliArgs, ThinClientOptions}; use crate::output::{JsonEnvelope, PagingMeta, StatsMeta}; @@ -64,7 +62,10 @@ fn load_local_search( // --- Native FTS search → fallback to scan --- let use_fallback = match db.message_fts_path.as_deref() { - Some(fts_path) => match open_fts_connection_with_key(fts_path, acct.raw_key.as_ref()) { + Some(fts_path) => match db.open_related_readonly(fts_path).and_then(|conn| { + register_mm_fts_tokenizer(&conn).map_err(wx_db::DbError::FtsInit)?; + Ok(conn) + }) { Ok(conn) => { match wx_db::native_fts::search_message_fts( &conn, diff --git a/crates/wx-cli/src/cmd/serve/mod.rs b/crates/wx-cli/src/cmd/serve/mod.rs index e812cdc..d4cb1c1 100644 --- a/crates/wx-cli/src/cmd/serve/mod.rs +++ b/crates/wx-cli/src/cmd/serve/mod.rs @@ -19,8 +19,8 @@ use tokio::signal::unix::SignalKind; use tokio::sync::{broadcast, mpsc, watch}; use tokio_util::sync::CancellationToken; use wx_context::{ - open_fts_connection_with_key, register_mm_fts_tokenizer, write_shard_metadata_sidecar, - AccountContext, ContactResolver, DecryptRequest, PersistentCache, ResolveParams, + register_mm_fts_tokenizer, write_shard_metadata_sidecar, AccountContext, ContactResolver, + DecryptRequest, PersistentCache, ResolveParams, }; use crate::util::{print_cache_stats, print_detection_note}; @@ -122,7 +122,10 @@ pub async fn cmd_serve( // 3b. Open independent FTS connection (outside WechatDb Mutex) let fts_conn = db.message_fts_path.as_deref().and_then(|fts_path| { - match open_fts_connection_with_key(fts_path, acct.raw_key.as_ref()) { + match db.open_related_readonly(fts_path).and_then(|conn| { + register_mm_fts_tokenizer(&conn).map_err(wx_db::DbError::FtsInit)?; + Ok(conn) + }) { Ok(conn) => { if let Ok(mode) = conn.query_row("PRAGMA journal_mode", [], |r| r.get::<_, String>(0)) @@ -196,7 +199,7 @@ pub async fn cmd_serve( // 3c. Open hardlink.db connection (pooled, outside WechatDb Mutex) let hardlink_db_conn = if hardlink_db_path.exists() { - match wx_db::open_readonly_connection(&hardlink_db_path, acct.raw_key.as_ref()) { + match db.open_related_readonly(&hardlink_db_path) { Ok(conn) => { eprintln!("server/hardlink: opened pooled connection"); Some(conn) @@ -222,9 +225,14 @@ pub async fn cmd_serve( } let watch_mode = resolve_watch_mode(poll, fsnotify); + let monitor_derived_keys = wx_context::persisted_derived_keys(&acct)?; let config = wx_monitor::MonitorConfig { encrypted_session_dir, - key_material: acct.key_material.clone(), + key_material: if monitor_derived_keys.is_empty() { + acct.key_material.clone() + } else { + wx_decrypt::KeyMaterial::EncKeys(monitor_derived_keys) + }, params, watch_mode: watch_mode.clone(), poll_interval: Duration::from_millis(poll_ms), @@ -239,7 +247,6 @@ pub async fn cmd_serve( // Capture values before moving db into Mutex let fts_path_for_refresh = db.message_fts_path.clone(); - let raw_key_for_refresh = acct.raw_key; // 5. Create refresh task channels let (refresh_tx, refresh_rx) = mpsc::channel::(64); @@ -398,7 +405,6 @@ pub async fn cmd_serve( shutdown_bg.clone(), ) .with_fts(bg_state.fts_conn.clone(), fts_path_for_refresh) - .with_raw_key(raw_key_for_refresh) .with_caches( Some(Arc::clone(&bg_state.name2id_cache)), Some(Arc::clone(&bg_state.media_db_paths)), diff --git a/crates/wx-cli/src/cmd/serve/refresh.rs b/crates/wx-cli/src/cmd/serve/refresh.rs index 5acac90..a4aadbd 100644 --- a/crates/wx-cli/src/cmd/serve/refresh.rs +++ b/crates/wx-cli/src/cmd/serve/refresh.rs @@ -6,7 +6,7 @@ use rusqlite::Connection; use tokio::sync::{mpsc, watch}; use tokio_util::sync::CancellationToken; use wx_context::{ - open_fts_connection, open_fts_connection_with_key, DecryptProgress, DecryptRequest, + open_fts_connection, register_mm_fts_tokenizer, DecryptProgress, DecryptRequest, PersistentCache, }; use wx_db::WechatDb; @@ -35,8 +35,6 @@ pub struct RefreshTask { fts_conn: Option>>, /// Path to FTS DB for reopening. fts_path: Option, - /// Raw key for encrypted FTS reopen. - raw_key: Option<[u8; 32]>, /// Cache of name2id mapping — cleared when FTS is reopened. name2id_cache: Option>>>>, /// Cache of media DB paths — cleared on every refresh. @@ -61,18 +59,12 @@ impl RefreshTask { shutdown, fts_conn: None, fts_path: None, - raw_key: None, name2id_cache: None, media_db_paths: None, hardlink_db_conn: None, } } - pub fn with_raw_key(mut self, raw_key: Option<[u8; 32]>) -> Self { - self.raw_key = raw_key; - self - } - /// Set the independent FTS connection and path for refresh reopening. pub fn with_fts( mut self, @@ -128,7 +120,6 @@ impl RefreshTask { let cache = self.cache.clone(); let fts_conn = self.fts_conn.clone(); let fts_path = self.fts_path.clone(); - let raw_key = self.raw_key; let success = tokio::task::spawn_blocking(move || { if let Some(cache) = cache { // Decrypt-cache mode: decrypt then selective reopen @@ -264,7 +255,10 @@ impl RefreshTask { // Reopen independent FTS connection if let (Some(fts_mutex), Some(path)) = (&fts_conn, &fts_path) { - match open_fts_connection_with_key(path, raw_key.as_ref()) { + match guard.open_related_readonly(path).and_then(|conn| { + register_mm_fts_tokenizer(&conn).map_err(wx_db::DbError::FtsInit)?; + Ok(conn) + }) { Ok(new_conn) => { if let Ok(mut fts_guard) = fts_mutex.lock() as Result, _> diff --git a/crates/wx-cli/src/cmd/watch.rs b/crates/wx-cli/src/cmd/watch.rs index ea9b2dd..dbddea2 100644 --- a/crates/wx-cli/src/cmd/watch.rs +++ b/crates/wx-cli/src/cmd/watch.rs @@ -318,9 +318,14 @@ pub async fn cmd_watch( } let watch_mode = resolve_watch_mode(poll, fsnotify); + let monitor_derived_keys = wx_context::persisted_derived_keys(&acct)?; let config = wx_monitor::MonitorConfig { encrypted_session_dir, - key_material: acct.key_material.clone(), + key_material: if monitor_derived_keys.is_empty() { + acct.key_material.clone() + } else { + wx_decrypt::KeyMaterial::EncKeys(monitor_derived_keys) + }, params, watch_mode: watch_mode.clone(), poll_interval: Duration::from_millis(poll_ms), diff --git a/crates/wx-cli/src/util.rs b/crates/wx-cli/src/util.rs index a05c9f7..8c94627 100644 --- a/crates/wx-cli/src/util.rs +++ b/crates/wx-cli/src/util.rs @@ -10,7 +10,7 @@ pub fn open_db_core( ) -> Result<(wx_db::WechatDb, Option), Box> { if acct.raw_key.is_some() { eprintln!("Direct encrypted open (SQLCipher)"); - let db = wx_context::open_encrypted_db(acct)?; + let db = wx_context::open_encrypted_db_core(acct)?; Ok((db, None)) } else { let params = &wx_decrypt::MACOS_4_1_7_31; @@ -18,7 +18,7 @@ pub fn open_db_core( let stats = DecryptRequest::new() .core() .execute_with_progress(&cache, progress)?; - let db = wx_db::WechatDb::open(cache.decrypted_root())?; + let db = wx_db::WechatDb::open_core(cache.decrypted_root())?; Ok((db, Some(stats))) } } @@ -139,9 +139,10 @@ pub fn effective_limit_all(all: bool, limit: usize) -> usize { } } -/// Attempt a remote API call via ThinClient; on connection/auth failure fall back to the -/// local path. This encapsulates the `probe_health → remote_fn → should_fallback → local_fn` -/// pattern shared by `search`, `contacts`, and `sessions`. +/// Attempt a remote API call via ThinClient. In auto mode, fall back locally only when +/// the initial health probe cannot reach/authenticate with a usable server. Once health +/// succeeds, a failed business request is returned to the caller instead of launching an +/// expensive local SQLCipher query after waiting for the remote timeout. pub fn try_remote_or_local( options: &ThinClientOptions, remote_fn: impl FnOnce(&ThinClient) -> Result, @@ -150,8 +151,8 @@ pub fn try_remote_or_local( ) -> Result> { if options.is_enabled() { let client = ThinClient::new(options.clone()); - match client.probe_health().and_then(|_| remote_fn(&client)) { - Ok(result) => return Ok(result), + match client.probe_health() { + Ok(()) => return remote_fn(&client).map_err(Into::into), Err(err) if err.should_fallback(options.mode) => { eprintln!( "note: remote server unavailable, falling back to local {label} ({})", diff --git a/crates/wx-cli/tests/thin_client_cli.rs b/crates/wx-cli/tests/thin_client_cli.rs index 4ae928d..a37ef54 100644 --- a/crates/wx-cli/tests/thin_client_cli.rs +++ b/crates/wx-cli/tests/thin_client_cli.rs @@ -4,6 +4,8 @@ use std::process::Command; use std::thread; use std::time::{Duration, Instant}; +use tempfile::TempDir; + fn bin() -> &'static str { env!("CARGO_BIN_EXE_wx-cli") } @@ -278,7 +280,8 @@ fn server_only_fails_when_remote_unavailable() { #[test] fn unavailable_remote_falls_back_to_local() { - let output = Command::new(bin()) + let (mut command, _home) = command_without_local_account(); + let output = command .args(["sessions", "--server-url", "http://127.0.0.1:9"]) .output() .expect("run sessions fallback"); @@ -314,7 +317,8 @@ fn no_server_bypasses_remote_probe() { } }); - let output = Command::new(bin()) + let (mut command, _home) = command_without_local_account(); + let output = command .args([ "sessions", "--no-server", @@ -333,6 +337,39 @@ fn no_server_bypasses_remote_probe() { ); } +#[test] +fn healthy_server_business_transport_failure_does_not_fall_back() { + let (base_url, handle) = spawn_sequence_server(2, |_request, index| match index { + 0 => http_response("200 OK", "{\"ready\":true}"), + // Close the second connection without a response. This is classified as an + // unavailable transport error, but health already proved the server was selected. + 1 => String::new(), + _ => unreachable!(), + }); + + let output = Command::new(bin()) + .args(["sessions", "--server-url", &base_url]) + .output() + .expect("run sessions with failed business request"); + + assert!(!output.status.success()); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!(stderr.contains("error:")); + assert!(!stderr.contains("falling back to local")); + handle.join().unwrap(); +} + +fn command_without_local_account() -> (Command, TempDir) { + let home = TempDir::new().expect("create isolated home"); + let mut command = Command::new(bin()); + command + .env("HOME", home.path()) + .env_remove("WECHAT_CLI_DATA_DIR") + .env_remove("WECHAT_CLI_ACCOUNT") + .env_remove("WECHAT_CLI_KEY"); + (command, home) +} + fn spawn_sequence_server( expected_requests: usize, responder: impl Fn(String, usize) -> String + Send + 'static, diff --git a/crates/wx-context/src/lib.rs b/crates/wx-context/src/lib.rs index 6e50381..cf1e4a7 100644 --- a/crates/wx-context/src/lib.rs +++ b/crates/wx-context/src/lib.rs @@ -36,6 +36,25 @@ pub use progress::{DecryptProgress, DecryptStats}; pub use shard_routing::{route_shards_for_query, write_shard_metadata_sidecar}; pub use visibility::VisibilityIndex; +/// Read persisted per-database derived keys for this account. +/// Ephemeral `--key` contexts deliberately ignore the store because the supplied +/// raw key may not match its cached entries. +pub fn persisted_derived_keys( + account: &AccountContext, +) -> Result, ContextError> { + if !account.writeback_enabled { + return Ok(Vec::new()); + } + let store = wx_keychain::KeyStore::load_default()?; + Ok(match store.resolve_key_material(&account.account_id) { + Some(wx_decrypt::KeyMaterial::EncKeys(pairs)) => pairs, + Some(wx_decrypt::KeyMaterial::EncKey { key, salt }) => { + vec![wx_decrypt::EncKeyPair { key, salt }] + } + _ => Vec::new(), + }) +} + /// Open encrypted WeChat DB directory directly (no pool, no FTS). /// For one-shot commands: contacts, sessions, query, search, export. pub fn open_encrypted_db(account: &AccountContext) -> Result { @@ -43,7 +62,25 @@ pub fn open_encrypted_db(account: &AccountContext) -> Result Result { + let raw_key = account + .raw_key + .ok_or_else(|| ContextError::Cache("raw_key required for encrypted direct open".into()))?; + let encrypted_root = account.data_dir.join("db_storage"); + let derived_keys = persisted_derived_keys(account)?; + let db = wx_db::WechatDb::open_encrypted_core_with_key_cache( + &encrypted_root, + raw_key, + &derived_keys, + )?; Ok(db) } @@ -56,9 +93,11 @@ pub fn open_encrypted_db_with_pool( .raw_key .ok_or_else(|| ContextError::Cache("raw_key required for encrypted direct open".into()))?; let encrypted_root = account.data_dir.join("db_storage"); - let db = wx_db::WechatDb::open_encrypted_with_pool( + let derived_keys = persisted_derived_keys(account)?; + let db = wx_db::WechatDb::open_encrypted_with_pool_and_key_cache( &encrypted_root, raw_key, + &derived_keys, register_mm_fts_tokenizer, )?; Ok(db) diff --git a/crates/wx-db/Cargo.toml b/crates/wx-db/Cargo.toml index d1cf4cd..43ef091 100644 --- a/crates/wx-db/Cargo.toml +++ b/crates/wx-db/Cargo.toml @@ -12,6 +12,7 @@ thiserror = "2" serde = { version = "1", features = ["derive"] } serde_json = "1" hex = "0.4" +wx-decrypt = { path = "../wx-decrypt" } [dev-dependencies] insta = { version = "1", features = ["yaml"] } diff --git a/crates/wx-db/src/fts.rs b/crates/wx-db/src/fts.rs index 124b982..8e73c20 100644 --- a/crates/wx-db/src/fts.rs +++ b/crates/wx-db/src/fts.rs @@ -414,7 +414,8 @@ impl WechatDb { )?; for shard in &self.shards { - let shard_conn = WechatDb::open_shard_with_key(shard, self.raw_key.as_ref())?; + let shard_conn = + WechatDb::open_shard_with_key(shard, self.sqlcipher_key.as_ref())?; // List Msg_* tables in this shard let mut table_stmt = shard_conn.prepare( diff --git a/crates/wx-db/src/messages.rs b/crates/wx-db/src/messages.rs index f652481..8844f47 100644 --- a/crates/wx-db/src/messages.rs +++ b/crates/wx-db/src/messages.rs @@ -12,7 +12,7 @@ use crate::model::{ effective_limit, split_local_type, AnchorMode, Message, MessageQuery, MessageQueryResult, QueryStats, SortOrder, }; -use crate::open::{MessageShard, WechatDb}; +use crate::open::{MessageShard, SqlcipherKey, WechatDb}; /// Dispatch mode for regular (non-anchor) queries. enum RegularQueryMode { @@ -56,13 +56,13 @@ fn prepare_shard_query<'a>( table_name: &str, warnings: &mut Vec, pooled_conn: Option<&'a Connection>, - raw_key: Option<&[u8; 32]>, + sqlcipher_key: Option<&SqlcipherKey>, ) -> Option> { let shard_path = shard.path.display().to_string(); let conn = match pooled_conn { Some(conn) => ShardConnection::Borrowed(conn), - None => match WechatDb::open_shard_with_key(shard, raw_key) { + None => match WechatDb::open_shard_with_key(shard, sqlcipher_key) { Ok(c) => ShardConnection::Owned(c), Err(e) => { warnings.push(ShardWarning { @@ -176,7 +176,7 @@ impl WechatDb { &table_name, &mut shard_warnings, self.pool().and_then(|pool| pool.get(&shard.path)), - self.raw_key.as_ref(), + self.sqlcipher_key.as_ref(), ) { Some(p) => p, None => continue, @@ -296,12 +296,12 @@ impl WechatDb { for shard in &shards { let count = if let Some(pool) = self.pool() { if let Some(conn) = pool.get(&shard.path) { - Self::count_shard(&conn, &sql, start_time, end_time, msg_type_filter) + Self::count_shard(conn, &sql, start_time, end_time, msg_type_filter) } else { continue; } } else { - match crate::open::open_connection(&shard.path, self.raw_key.as_ref()) { + match crate::open::open_connection(&shard.path, self.sqlcipher_key.as_ref()) { Ok(conn) => { Self::count_shard(&conn, &sql, start_time, end_time, msg_type_filter) } @@ -385,7 +385,7 @@ impl WechatDb { table_name, &mut shard_warnings, self.pool().and_then(|pool| pool.get(&shard.path)), - self.raw_key.as_ref(), + self.sqlcipher_key.as_ref(), ) { Some(p) => p, None => continue, @@ -476,7 +476,7 @@ impl WechatDb { table_name, &mut shard_warnings, self.pool().and_then(|pool| pool.get(&shard.path)), - self.raw_key.as_ref(), + self.sqlcipher_key.as_ref(), ) { Some(p) => p, None => continue, @@ -607,7 +607,7 @@ impl WechatDb { table_name, &mut shard_warnings, self.pool().and_then(|pool| pool.get(&shard.path)), - self.raw_key.as_ref(), + self.sqlcipher_key.as_ref(), ) { Some(p) => p, None => continue, @@ -683,7 +683,7 @@ impl WechatDb { table_name, &mut shard_warnings, self.pool().and_then(|pool| pool.get(&shard.path)), - self.raw_key.as_ref(), + self.sqlcipher_key.as_ref(), ) { Some(p) => p, None => continue, @@ -812,16 +812,21 @@ impl WechatDb { known_usernames.iter().map(|u| (u.clone(), 0)).collect(); for shard in self.all_shards() { - let conn = match WechatDb::open_shard_with_key(shard, self.raw_key.as_ref()) { - Ok(c) => c, - Err(e) => { - eprintln!( - "warn: bulk_max_sort_seq: open shard {} failed: {e}", - shard.path.display() - ); - continue; + let conn = if let Some(conn) = self.pool().and_then(|pool| pool.get(&shard.path)) { + ShardConnection::Borrowed(conn) + } else { + match WechatDb::open_shard_with_key(shard, self.sqlcipher_key.as_ref()) { + Ok(conn) => ShardConnection::Owned(conn), + Err(e) => { + eprintln!( + "warn: bulk_max_sort_seq: open shard {} failed: {e}", + shard.path.display() + ); + continue; + } } }; + let conn = conn.as_conn(); // Discover Msg_* tables in this shard let mut stmt = match conn diff --git a/crates/wx-db/src/open.rs b/crates/wx-db/src/open.rs index 94dbc42..d93167b 100644 --- a/crates/wx-db/src/open.rs +++ b/crates/wx-db/src/open.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use std::fmt; use std::os::raw::c_void; use std::path::{Path, PathBuf}; -use std::sync::{Arc, RwLock}; +use std::sync::{Arc, Mutex, RwLock}; use rusqlite::Connection; @@ -18,6 +18,111 @@ pub(crate) struct MessageShard { pub end_unix: i64, } +/// A raw WeChat key plus an in-process cache of SQLCipher's derived keys. +/// +/// SQLCipher normally runs its 256k-round PBKDF2 every time a connection is +/// opened. WeChat uses a different salt per database, but the same database is +/// often opened several times during one command (metadata scan, query, count, +/// refresh). Passing SQLCipher's raw keyspec lets us derive once per salt and +/// reuse the result for every subsequent connection. +#[derive(Clone)] +pub(crate) struct SqlcipherKey { + raw_key: [u8; 32], + derived_keys: Arc>>, +} + +#[derive(Clone, Copy)] +struct CachedKey { + key: [u8; 32], + /// Preloaded keys come from the persisted key store and get one raw-key + /// fallback if validation fails. Keys derived in this process are trusted. + preloaded: bool, +} + +impl SqlcipherKey { + fn new(raw_key: [u8; 32]) -> Self { + Self::with_preloaded(raw_key, &[]) + } + + fn with_preloaded(raw_key: [u8; 32], pairs: &[wx_decrypt::EncKeyPair]) -> Self { + let derived_keys = pairs + .iter() + .map(|pair| { + ( + pair.salt, + CachedKey { + key: pair.key, + preloaded: true, + }, + ) + }) + .collect(); + Self { + raw_key, + derived_keys: Arc::new(Mutex::new(derived_keys)), + } + } + + fn keyspec_for_path(&self, path: &Path) -> Result<(Vec, [u8; 16], bool), DbError> { + let salt = wx_decrypt::read_db_salt(path) + .map_err(|e| DbError::EncryptionKey(format!("failed to read database salt: {e}")))?; + + let cached = { + let mut cache = self.derived_keys.lock().map_err(|_| { + DbError::EncryptionKey("derived-key cache lock was poisoned".into()) + })?; + *cache.entry(salt).or_insert_with(|| { + let key = wx_decrypt::kdf::derive_enc_key( + &self.raw_key, + &salt, + &wx_decrypt::MACOS_4_1_7_31, + ); + CachedKey { + key, + preloaded: false, + } + }) + }; + + // SQLCipher raw-key syntax includes the original 16-byte database salt. + // Supplying this ASCII keyspec to sqlite3_key() skips SQLCipher's PBKDF2. + let keyspec = format!("x'{}{}'", hex::encode(cached.key), hex::encode(salt)).into_bytes(); + Ok((keyspec, salt, cached.preloaded)) + } + + fn mark_verified(&self, salt: [u8; 16]) -> Result<(), DbError> { + let mut cache = self + .derived_keys + .lock() + .map_err(|_| DbError::EncryptionKey("derived-key cache lock was poisoned".into()))?; + if let Some(entry) = cache.get_mut(&salt) { + entry.preloaded = false; + } + Ok(()) + } + + fn rederive_keyspec(&self, salt: [u8; 16]) -> Result, DbError> { + let key = + wx_decrypt::kdf::derive_enc_key(&self.raw_key, &salt, &wx_decrypt::MACOS_4_1_7_31); + self.derived_keys + .lock() + .map_err(|_| DbError::EncryptionKey("derived-key cache lock was poisoned".into()))? + .insert( + salt, + CachedKey { + key, + preloaded: false, + }, + ); + Ok(format!("x'{}{}'", hex::encode(key), hex::encode(salt)).into_bytes()) + } + + #[cfg(test)] + fn cached_salt_count(&self) -> usize { + self.derived_keys.lock().unwrap().len() + } +} + /// Handle to an opened (decrypted) WeChat database directory. /// /// Holds connections to contact/session databases and metadata about @@ -34,8 +139,8 @@ pub struct WechatDb { pub contact_fts_path: Option, /// Optional pre-opened connection pool for serve mode. pub(crate) pool: Option, - /// Raw key for encrypted direct open. Stored for reopen operations. - pub(crate) raw_key: Option<[u8; 32]>, + /// Shared raw/derived key state for encrypted direct open and reopen operations. + pub(crate) sqlcipher_key: Option, /// Lazily initialized cache of label_id -> label_name from contact_label table. /// Cleared on `reopen_contacts()` so label changes are visible. pub(crate) label_cache: RwLock>>, @@ -54,30 +159,59 @@ pub fn open_readonly_connection( path: &Path, raw_key: Option<&[u8; 32]>, ) -> Result { - let conn = Connection::open_with_flags(path, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)?; - if let Some(key) = raw_key { - unsafe { - let rc = rusqlite::ffi::sqlite3_key(conn.handle(), key.as_ptr() as *const c_void, 32); - if rc != 0 { - return Err(DbError::EncryptionKey(format!( - "sqlite3_key failed: rc={rc}" - ))); - } - } - conn.query_row("SELECT count(*) FROM sqlite_master", [], |r| { - r.get::<_, i64>(0) - }) - .map_err(|_| DbError::EncryptionKey("incorrect key or not an encrypted database".into()))?; - conn.execute_batch("PRAGMA query_only = ON")?; - } - Ok(conn) + let key = raw_key.copied().map(SqlcipherKey::new); + open_connection(path, key.as_ref()) } pub(crate) fn open_connection( path: &Path, - raw_key: Option<&[u8; 32]>, + sqlcipher_key: Option<&SqlcipherKey>, ) -> Result { - open_readonly_connection(path, raw_key) + if let Some(key) = sqlcipher_key { + let (keyspec, salt, preloaded) = key.keyspec_for_path(path)?; + match open_connection_with_keyspec(path, &keyspec) { + Ok(conn) => { + if preloaded { + key.mark_verified(salt)?; + } + Ok(conn) + } + Err(DbError::EncryptionKey(_)) if preloaded => { + // Persisted entries are an optimization, never a single point of + // failure. Re-derive once from the raw key if an entry is stale. + let keyspec = key.rederive_keyspec(salt)?; + open_connection_with_keyspec(path, &keyspec) + } + Err(err) => Err(err), + } + } else { + Ok(Connection::open_with_flags( + path, + rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY, + )?) + } +} + +fn open_connection_with_keyspec(path: &Path, keyspec: &[u8]) -> Result { + let conn = Connection::open_with_flags(path, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)?; + unsafe { + let rc = rusqlite::ffi::sqlite3_key( + conn.handle(), + keyspec.as_ptr() as *const c_void, + keyspec.len() as i32, + ); + if rc != 0 { + return Err(DbError::EncryptionKey(format!( + "sqlite3_key failed: rc={rc}" + ))); + } + } + conn.query_row("SELECT count(*) FROM sqlite_master", [], |r| { + r.get::<_, i64>(0) + }) + .map_err(|_| DbError::EncryptionKey("incorrect key or not an encrypted database".into()))?; + conn.execute_batch("PRAGMA query_only = ON")?; + Ok(conn) } impl WechatDb { @@ -87,7 +221,13 @@ impl WechatDb { /// does not exist. Message shards are optional here; message queries will /// return `DbError::NoShards` if no numbered shard is available. pub fn open(path: impl AsRef) -> Result { - Self::open_internal(path.as_ref(), None) + Self::open_internal(path.as_ref(), None, true) + } + + /// Open only contact.db and session.db, without scanning message shards. + /// Useful for contacts, sessions, and monitoring commands that never read messages. + pub fn open_core(path: impl AsRef) -> Result { + Self::open_internal(path.as_ref(), None, false) } /// Open a decrypted WeChat database directory with a pre-opened @@ -104,7 +244,39 @@ impl WechatDb { /// Open an encrypted WeChat database directory directly using `sqlite3_key()`. pub fn open_encrypted(path: impl AsRef, raw_key: [u8; 32]) -> Result { - Self::open_internal(path.as_ref(), Some(raw_key)) + Self::open_internal(path.as_ref(), Some(SqlcipherKey::new(raw_key)), true) + } + + /// Open an encrypted directory and seed the per-salt cache with persisted + /// derived keys, falling back to the raw key for missing or stale entries. + pub fn open_encrypted_with_key_cache( + path: impl AsRef, + raw_key: [u8; 32], + pairs: &[wx_decrypt::EncKeyPair], + ) -> Result { + Self::open_internal( + path.as_ref(), + Some(SqlcipherKey::with_preloaded(raw_key, pairs)), + true, + ) + } + + /// Open only encrypted contact.db and session.db, without scanning message shards. + pub fn open_encrypted_core(path: impl AsRef, raw_key: [u8; 32]) -> Result { + Self::open_internal(path.as_ref(), Some(SqlcipherKey::new(raw_key)), false) + } + + /// Core-only variant of [`WechatDb::open_encrypted_with_key_cache`]. + pub fn open_encrypted_core_with_key_cache( + path: impl AsRef, + raw_key: [u8; 32], + pairs: &[wx_decrypt::EncKeyPair], + ) -> Result { + Self::open_internal( + path.as_ref(), + Some(SqlcipherKey::with_preloaded(raw_key, pairs)), + false, + ) } /// Open an encrypted WeChat database directory with a pre-opened @@ -114,35 +286,53 @@ impl WechatDb { raw_key: [u8; 32], fts_init: impl Fn(&Connection) -> Result<(), String> + Send + Sync + 'static, ) -> Result { - Self::open_with_pool_internal(path, Some(raw_key), fts_init) + Self::open_with_pool_internal(path, Some(SqlcipherKey::new(raw_key)), fts_init) } - fn open_internal(path: &Path, raw_key: Option<[u8; 32]>) -> Result { + /// Pool variant seeded with persisted per-salt derived keys. + pub fn open_encrypted_with_pool_and_key_cache( + path: impl AsRef, + raw_key: [u8; 32], + pairs: &[wx_decrypt::EncKeyPair], + fts_init: impl Fn(&Connection) -> Result<(), String> + Send + Sync + 'static, + ) -> Result { + Self::open_with_pool_internal( + path, + Some(SqlcipherKey::with_preloaded(raw_key, pairs)), + fts_init, + ) + } + + fn open_internal( + path: &Path, + sqlcipher_key: Option, + scan_message_shards: bool, + ) -> Result { if !path.exists() { return Err(DbError::NotFound(path.display().to_string())); } - let key_ref = raw_key.as_ref(); + let key_ref = sqlcipher_key.as_ref(); // Open contact.db let contact_path = path.join("contact").join("contact.db"); if !contact_path.exists() { return Err(DbError::NotFound(contact_path.display().to_string())); } - let contact_conn = open_readonly_connection(&contact_path, key_ref)?; + let contact_conn = open_connection(&contact_path, key_ref)?; // Open session.db let session_path = path.join("session").join("session.db"); if !session_path.exists() { return Err(DbError::NotFound(session_path.display().to_string())); } - let session_conn = open_readonly_connection(&session_path, key_ref)?; + let session_conn = open_connection(&session_path, key_ref)?; // Scan message shards let msg_dir = path.join("message"); let mut shards = Vec::new(); - if msg_dir.is_dir() { + if scan_message_shards && msg_dir.is_dir() { let mut entries: Vec = std::fs::read_dir(&msg_dir)? .filter_map(|e| e.ok()) .map(|e| e.path()) @@ -196,23 +386,23 @@ impl WechatDb { } }, pool: None, - raw_key, + sqlcipher_key, label_cache: RwLock::new(None), }) } fn open_with_pool_internal( path: impl AsRef, - raw_key: Option<[u8; 32]>, + sqlcipher_key: Option, fts_init: impl Fn(&Connection) -> Result<(), String> + Send + Sync + 'static, ) -> Result { - let mut db = Self::open_internal(path.as_ref(), raw_key)?; + let mut db = Self::open_internal(path.as_ref(), sqlcipher_key.clone(), true)?; let fts_init_arc: Arc = Arc::new(fts_init); let pool = ShardPool::open( &db.shards, db.message_fts_path.as_deref(), Some(fts_init_arc), - raw_key, + sqlcipher_key, )?; db.pool = Some(pool); Ok(db) @@ -220,14 +410,14 @@ impl WechatDb { /// Re-open the session.db connection to pick up external changes. pub fn reopen_sessions(&mut self) -> Result<(), DbError> { - self.session_conn = open_connection(&self.session_path, self.raw_key.as_ref())?; + self.session_conn = open_connection(&self.session_path, self.sqlcipher_key.as_ref())?; Ok(()) } /// Re-open the contact.db connection to pick up external changes. /// Also invalidates the label cache so it is reloaded on next query. pub fn reopen_contacts(&mut self) -> Result<(), DbError> { - self.contact_conn = open_connection(&self.contact_path, self.raw_key.as_ref())?; + self.contact_conn = open_connection(&self.contact_path, self.sqlcipher_key.as_ref())?; *self.label_cache.write().unwrap() = None; Ok(()) } @@ -270,6 +460,13 @@ impl WechatDb { self.pool.as_ref() } + /// Open another database from the same encrypted account while reusing this + /// handle's derived-key cache. This is used by serve-mode auxiliary FTS and + /// media connections so refreshes do not re-run PBKDF2. + pub fn open_related_readonly(&self, path: &Path) -> Result { + open_connection(path, self.sqlcipher_key.as_ref()) + } + /// Return shards whose time range overlaps `[start, end]`. pub(crate) fn shards_for_range(&self, start: i64, end: i64) -> Vec<&MessageShard> { self.shards @@ -307,9 +504,9 @@ impl WechatDb { /// Open a SQLite connection to a specific shard, optionally encrypted. pub(crate) fn open_shard_with_key( shard: &MessageShard, - raw_key: Option<&[u8; 32]>, + sqlcipher_key: Option<&SqlcipherKey>, ) -> Result { - open_readonly_connection(&shard.path, raw_key) + open_connection(&shard.path, sqlcipher_key) } } @@ -341,8 +538,8 @@ fn is_numbered_message_shard(path: &Path) -> bool { /// Try to read the timestamp from a message shard's Timestamp table. /// Returns 0 if the table does not exist or is empty. -fn read_shard_timestamp(path: &Path, raw_key: Option<&[u8; 32]>) -> i64 { - let conn = match open_connection(path, raw_key) { +fn read_shard_timestamp(path: &Path, sqlcipher_key: Option<&SqlcipherKey>) -> i64 { + let conn = match open_connection(path, sqlcipher_key) { Ok(c) => c, Err(_) => return 0, }; @@ -425,9 +622,35 @@ mod tests { build_encrypted_db_storage(&root, &raw_key); let mut db = WechatDb::open_encrypted(&root, raw_key).unwrap(); + let key = db.sqlcipher_key.clone().unwrap(); + let cached_before = key.cached_salt_count(); + assert_eq!( + cached_before, 3, + "contact, session, and message salts cached" + ); // Reopen should succeed (re-applies sqlite3_key) db.reopen_sessions().unwrap(); db.reopen_contacts().unwrap(); + assert_eq!( + key.cached_salt_count(), + cached_before, + "reopen must reuse derived keys instead of deriving again" + ); + } + + #[test] + fn open_core_does_not_scan_message_shards() { + let tmp = TempDir::new().unwrap(); + let root = tmp.path().join("db_storage"); + std::fs::create_dir_all(root.join("contact")).unwrap(); + std::fs::create_dir_all(root.join("session")).unwrap(); + std::fs::create_dir_all(root.join("message")).unwrap(); + Connection::open(root.join("contact/contact.db")).unwrap(); + Connection::open(root.join("session/session.db")).unwrap(); + std::fs::write(root.join("message/message_0.db"), b"not a sqlite database").unwrap(); + + let db = WechatDb::open_core(&root).unwrap(); + assert!(db.shards.is_empty()); } #[test] @@ -457,7 +680,57 @@ mod tests { "CREATE TABLE t (id INTEGER); INSERT INTO t VALUES (42);", ); - let conn = open_connection(&path, Some(&raw_key)).unwrap(); + let key = SqlcipherKey::new(raw_key); + let conn = open_connection(&path, Some(&key)).unwrap(); + let val: i64 = conn + .query_row("SELECT id FROM t", [], |r| r.get(0)) + .unwrap(); + assert_eq!(val, 42); + } + + #[test] + fn preloaded_derived_key_opens_encrypted_database() { + let tmp = TempDir::new().unwrap(); + let path = tmp.path().join("enc.db"); + let raw_key = [0xAB_u8; 32]; + create_encrypted_db( + &path, + &raw_key, + "CREATE TABLE t (id INTEGER); INSERT INTO t VALUES (42);", + ); + let salt = wx_decrypt::read_db_salt(&path).unwrap(); + let enc_key = wx_decrypt::kdf::derive_enc_key(&raw_key, &salt, &wx_decrypt::MACOS_4_1_7_31); + let key = + SqlcipherKey::with_preloaded(raw_key, &[wx_decrypt::EncKeyPair { key: enc_key, salt }]); + + let conn = open_connection(&path, Some(&key)).unwrap(); + let val: i64 = conn + .query_row("SELECT id FROM t", [], |r| r.get(0)) + .unwrap(); + assert_eq!(val, 42); + assert_eq!(key.cached_salt_count(), 1); + } + + #[test] + fn stale_preloaded_key_falls_back_to_raw_key() { + let tmp = TempDir::new().unwrap(); + let path = tmp.path().join("enc.db"); + let raw_key = [0xAB_u8; 32]; + create_encrypted_db( + &path, + &raw_key, + "CREATE TABLE t (id INTEGER); INSERT INTO t VALUES (42);", + ); + let salt = wx_decrypt::read_db_salt(&path).unwrap(); + let key = SqlcipherKey::with_preloaded( + raw_key, + &[wx_decrypt::EncKeyPair { + key: [0xCD; 32], + salt, + }], + ); + + let conn = open_connection(&path, Some(&key)).unwrap(); let val: i64 = conn .query_row("SELECT id FROM t", [], |r| r.get(0)) .unwrap(); diff --git a/crates/wx-db/src/pool.rs b/crates/wx-db/src/pool.rs index 99294a0..8f22b46 100644 --- a/crates/wx-db/src/pool.rs +++ b/crates/wx-db/src/pool.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use rusqlite::Connection; use crate::error::DbError; -use crate::open::MessageShard; +use crate::open::{MessageShard, SqlcipherKey}; pub(crate) type FtsInitFn = dyn Fn(&Connection) -> Result<(), String> + Send + Sync; @@ -19,7 +19,7 @@ pub struct ShardPool { fts_conn: Option, fts_path: Option, fts_init: Option>, - raw_key: Option<[u8; 32]>, + sqlcipher_key: Option, } impl std::fmt::Debug for ShardPool { @@ -41,22 +41,22 @@ impl ShardPool { shards: &[MessageShard], fts_path: Option<&Path>, fts_init: Option>, - raw_key: Option<[u8; 32]>, + sqlcipher_key: Option, ) -> Result { let mut conns = HashMap::with_capacity(shards.len()); for shard in shards { - let conn = crate::open::open_connection(&shard.path, raw_key.as_ref())?; + let conn = crate::open::open_connection(&shard.path, sqlcipher_key.as_ref())?; conns.insert(shard.path.clone(), conn); } let fts_conn = match (fts_path, &fts_init) { (Some(path), Some(init)) => { - let conn = crate::open::open_connection(path, raw_key.as_ref())?; + let conn = crate::open::open_connection(path, sqlcipher_key.as_ref())?; init(&conn).map_err(DbError::FtsInit)?; Some(conn) } (Some(path), None) => { - let conn = crate::open::open_connection(path, raw_key.as_ref())?; + let conn = crate::open::open_connection(path, sqlcipher_key.as_ref())?; Some(conn) } _ => None, @@ -67,7 +67,7 @@ impl ShardPool { fts_conn, fts_path: fts_path.map(|p| p.to_path_buf()), fts_init, - raw_key, + sqlcipher_key, }) } @@ -79,7 +79,7 @@ impl ShardPool { /// Close and reopen one shard connection. pub fn reopen_shard(&mut self, path: &Path) -> Result<(), DbError> { if self.conns.contains_key(path) { - let conn = crate::open::open_connection(path, self.raw_key.as_ref())?; + let conn = crate::open::open_connection(path, self.sqlcipher_key.as_ref())?; self.conns.insert(path.to_path_buf(), conn); } Ok(()) @@ -93,7 +93,7 @@ impl ShardPool { /// Close and reopen the FTS connection, re-registering the tokenizer. pub fn reopen_fts(&mut self) -> Result<(), DbError> { if let Some(path) = &self.fts_path { - let conn = crate::open::open_connection(path, self.raw_key.as_ref())?; + let conn = crate::open::open_connection(path, self.sqlcipher_key.as_ref())?; if let Some(init) = &self.fts_init { init(&conn).map_err(DbError::FtsInit)?; } @@ -106,7 +106,7 @@ impl ShardPool { pub fn reopen_all(&mut self) -> Result<(), DbError> { let paths: Vec = self.conns.keys().cloned().collect(); for path in paths { - let conn = crate::open::open_connection(&path, self.raw_key.as_ref())?; + let conn = crate::open::open_connection(&path, self.sqlcipher_key.as_ref())?; self.conns.insert(path, conn); } self.reopen_fts()?; diff --git a/crates/wx-monitor/src/monitor.rs b/crates/wx-monitor/src/monitor.rs index a705eb3..22aec25 100644 --- a/crates/wx-monitor/src/monitor.rs +++ b/crates/wx-monitor/src/monitor.rs @@ -5,7 +5,7 @@ use std::time::Duration; use futures_core::Stream; use wx_db::{SessionQuery, WechatDb}; -use wx_decrypt::{CryptoParams, KeyMaterial}; +use wx_decrypt::{CryptoParams, EncKeyPair, KeyMaterial}; use crate::cache::{DecryptCache, UpdateKind}; use crate::error::MonitorError; @@ -100,7 +100,19 @@ impl WechatMonitor { let (db, cache) = if let (Some(raw_key), Some(ref encrypted_root)) = (config.raw_key, &config.encrypted_root) { - let db = WechatDb::open_encrypted(encrypted_root, raw_key)?; + let derived_keys: Vec = match &config.key_material { + KeyMaterial::EncKeys(pairs) => pairs.clone(), + KeyMaterial::EncKey { key, salt } => vec![EncKeyPair { + key: *key, + salt: *salt, + }], + KeyMaterial::RawKey(_) => Vec::new(), + }; + let db = WechatDb::open_encrypted_core_with_key_cache( + encrypted_root, + raw_key, + &derived_keys, + )?; (db, None) } else { let mut cache = DecryptCache::new( @@ -109,7 +121,7 @@ impl WechatMonitor { config.params, )?; cache.initial_decrypt()?; - let db = WechatDb::open(cache.decrypted_root())?; + let db = WechatDb::open_core(cache.decrypted_root())?; (db, Some(cache)) };