1use crate::{
4 schema::inspect,
5 snapshot::{walk, Snapshot},
6 DatabaseInventory, DatabaseState, KeyFile, ProfileError, ProfileLocation, ProfileReader,
7 ReadOptions, Result, SchemaReport,
8};
9use rusqlite::Connection;
10use sha2::{Digest, Sha256};
11use std::{
12 collections::BTreeMap,
13 fs,
14 io::Read,
15 sync::{Arc, Mutex},
16};
17
18const MAX_CACHED_DATABASES: usize = 32;
20
21struct CachedDatabase {
22 metadata: fs::Metadata,
23 key: String,
24 connection: Arc<Mutex<Connection>>,
25 schema: SchemaReport,
26 schema_version: i64,
27}
28
29fn same_database_file(a: &fs::Metadata, b: &fs::Metadata) -> bool {
30 #[cfg(unix)]
31 {
32 wechat_platform::same_file(a, b)
33 }
34 #[cfg(not(unix))]
35 {
36 a.created().ok() == b.created().ok()
37 && a.modified().ok() == b.modified().ok()
38 && a.len() == b.len()
39 }
40}
41
42#[derive(Default)]
43pub struct OwnedReaderCache {
44 root: Option<std::path::PathBuf>,
45 databases: BTreeMap<String, CachedDatabase>,
46 nested: BTreeMap<String, Arc<Mutex<Connection>>>,
47}
48
49impl OwnedReaderCache {
50 pub fn read_snapshot(
54 &mut self,
55 location: ProfileLocation,
56 keys: KeyFile,
57 options: ReadOptions,
58 ) -> Result<ProfileReader> {
59 let mut reader = self.read(location, keys, options)?;
60 let mut pinned = crate::snapshot::PinnedTransactions {
61 connections: BTreeMap::new(),
62 };
63 let deadline = std::time::Instant::now() + reader.snapshot.options.query_timeout;
64 for entry in reader
65 .snapshot
66 .inventory
67 .iter()
68 .filter(|entry| entry.state == DatabaseState::Readable)
69 {
70 let Some(primary) = reader.snapshot.owned.get(&entry.relative_path) else {
71 return Err(ProfileError::ParallelSnapshotUnavailable);
72 };
73 if !self.nested.contains_key(&entry.relative_path) {
74 self.nested.insert(
75 entry.relative_path.clone(),
76 Arc::new(Mutex::new(
77 reader
78 .snapshot
79 .open_connection(&entry.relative_path, true)?,
80 )),
81 );
82 }
83 let pair = vec![primary.clone(), self.nested[&entry.relative_path].clone()];
84 pinned
85 .connections
86 .insert(entry.relative_path.clone(), pair.clone());
87 for connection in pair {
88 let connection = connection.lock().map_err(|_| ProfileError::Database)?;
89 let mode: String = connection.query_row("PRAGMA journal_mode", [], |r| r.get(0))?;
90 if !mode.eq_ignore_ascii_case("wal") {
91 return Err(ProfileError::ParallelSnapshotUnavailable);
92 }
93 connection
94 .progress_handler(1000, Some(move || std::time::Instant::now() > deadline));
95 connection.set_limit(
96 rusqlite::limits::Limit::SQLITE_LIMIT_LENGTH,
97 reader.snapshot.options.max_blob_bytes as i32,
98 );
99 connection.execute_batch("BEGIN DEFERRED")?;
100 connection.query_row("SELECT count(*) FROM sqlite_schema", [], |r| {
101 r.get::<_, i64>(0)
102 })?;
103 }
104 }
105 self.nested
106 .retain(|relative, _| reader.snapshot.paths.contains_key(relative));
107 reader.snapshot.pinned = Some(pinned);
108 Ok(reader)
109 }
110
111 pub fn read(
115 &mut self,
116 location: ProfileLocation,
117 keys: KeyFile,
118 options: ReadOptions,
119 ) -> Result<ProfileReader> {
120 if !options.stopped_or_consistent_snapshot {
121 return Err(ProfileError::SnapshotRequired);
122 }
123 if keys.account_binding != location.account {
124 return Err(ProfileError::AccountMismatch);
125 }
126 if options.max_blob_bytes == 0
127 || options.max_blob_bytes > 64 * 1024 * 1024
128 || options.max_databases == 0
129 || options.max_databases > 4096
130 || options.query_timeout.is_zero()
131 {
132 return Err(ProfileError::InvalidQuery);
133 }
134 if self.root.as_ref() != Some(&location.root) {
135 self.databases.clear();
136 self.nested.clear();
137 self.root = Some(location.root.clone());
138 }
139 let mut files = Vec::new();
140 walk(&location.storage, &mut files, 0, options.max_databases)?;
141 if files.is_empty() {
142 return Err(ProfileError::InvalidProfile);
143 }
144 let mut snapshot = Snapshot {
145 location,
146 keys,
147 options,
148 generation: String::new(),
149 inventory: Vec::new(),
150 paths: BTreeMap::new(),
151 _temp: None,
152 owned: BTreeMap::new(),
153 pinned: None,
154 };
155 let mut generation = Sha256::new();
156 generation.update(snapshot.location.root.as_os_str().as_encoded_bytes());
157 let mut total = 0u64;
158 let mut schemas = Vec::new();
159 for path in files {
160 let relative = path
161 .strip_prefix(&snapshot.location.storage)
162 .ok()
163 .and_then(|p| p.to_str())
164 .ok_or(ProfileError::UnsupportedFile)?
165 .replace('\\', "/");
166 let metadata = fs::symlink_metadata(&path)?;
167 let mut stamp = Sha256::new();
168 for suffix in ["", "-wal", "-shm"] {
169 let file = std::path::PathBuf::from(format!("{}{suffix}", path.display()));
170 let stat = match fs::symlink_metadata(&file) {
171 Ok(stat) => stat,
172 Err(error)
173 if error.kind() == std::io::ErrorKind::NotFound && !suffix.is_empty() =>
174 {
175 continue
176 }
177 Err(error) => return Err(error.into()),
178 };
179 if !stat.is_file() || wechat_platform::is_link(&stat) {
180 return Err(ProfileError::UnsupportedFile);
181 }
182 total = total
183 .checked_add(stat.len())
184 .ok_or(ProfileError::ResourceLimit)?;
185 if stat.len() > snapshot.options.max_database_bytes
186 || total > snapshot.options.max_total_bytes
187 {
188 return Err(ProfileError::ResourceLimit);
189 }
190 if suffix != "-shm" {
192 stamp.update(suffix.as_bytes());
193 stamp.update(stat.len().to_le_bytes());
194 stamp.update(
195 format!("{:?}:{:?}", stat.modified()?, stat.created().ok()).as_bytes(),
196 );
197 }
198 }
199 let mut prefix = [0u8; 16];
200 fs::File::open(&path)?.read_exact(&mut prefix)?;
201 let plaintext = prefix == *b"SQLite format 3\0";
202 let key = snapshot
203 .keys
204 .databases
205 .get(&relative)
206 .map(|key| {
207 let mut hash = Sha256::new();
208 hash.update(key.salt);
209 hash.update(key.raw_page_key.expose());
210 hex::encode(hash.finalize())
211 })
212 .unwrap_or_default();
213 let entry = DatabaseInventory {
214 relative_path: relative.clone(),
215 bytes: metadata.len(),
216 plaintext,
217 has_wal: path
218 .with_file_name(format!(
219 "{}-wal",
220 path.file_name()
221 .and_then(|s| s.to_str())
222 .ok_or(ProfileError::UnsupportedFile)?
223 ))
224 .try_exists()?,
225 generation: String::new(),
226 state: DatabaseState::Readable,
227 };
228 snapshot.paths.insert(relative.clone(), path);
229 snapshot.inventory.push(entry);
230 let mut uncached = None;
231 let reuse = self.databases.get(&relative).is_some_and(|cached| {
232 same_database_file(&cached.metadata, &metadata)
233 && cached.key == key
234 && (!plaintext || snapshot.options.allow_plaintext)
235 });
236 if !reuse {
237 self.databases.remove(&relative);
238 self.nested.remove(&relative);
239 match snapshot.open_connection(&relative, true) {
240 Ok(connection) => {
241 let schema = inspect(&connection, &relative)
242 .map_err(|_| ProfileError::UnsupportedSchema(relative.clone()))?;
243 let schema_version =
244 connection.query_row("PRAGMA schema_version", [], |r| r.get(0))?;
245 let database = CachedDatabase {
246 metadata,
247 key,
248 connection: Arc::new(Mutex::new(connection)),
249 schema,
250 schema_version,
251 };
252 if self.databases.len() < MAX_CACHED_DATABASES {
253 self.databases.insert(relative.clone(), database);
254 } else {
255 uncached = Some(database);
256 }
257 }
258 Err(ProfileError::KeyUnavailable(_)) => {
259 snapshot.inventory.last_mut().unwrap().state = DatabaseState::KeyUnavailable
260 }
261 Err(ProfileError::KeyMismatch(_)) => {
262 snapshot.inventory.last_mut().unwrap().state = DatabaseState::KeyMismatch
263 }
264 Err(error) => return Err(error),
265 }
266 }
267 let retained = uncached.is_none();
268 if let Some(cached) = self.databases.get_mut(&relative).or(uncached.as_mut()) {
269 let connection = cached
270 .connection
271 .lock()
272 .map_err(|_| ProfileError::Database)?;
273 let deadline = std::time::Instant::now() + snapshot.options.query_timeout;
275 connection
276 .progress_handler(1000, Some(move || std::time::Instant::now() > deadline));
277 connection.set_limit(
278 rusqlite::limits::Limit::SQLITE_LIMIT_LENGTH,
279 snapshot.options.max_blob_bytes as i32,
280 );
281 let schema_version =
282 connection.query_row("PRAGMA schema_version", [], |r| r.get(0))?;
283 if schema_version != cached.schema_version {
284 cached.schema = inspect(&connection, &relative)
285 .map_err(|_| ProfileError::UnsupportedSchema(relative.clone()))?;
286 cached.schema_version = schema_version;
287 }
288 schemas.push(cached.schema.clone());
292 if retained {
293 snapshot
294 .owned
295 .insert(relative.clone(), cached.connection.clone());
296 }
297 }
298 let stamp = hex::encode(stamp.finalize());
299 generation.update(relative.as_bytes());
300 generation.update(stamp.as_bytes());
301 snapshot.inventory.last_mut().unwrap().generation = stamp;
302 }
303 self.databases
304 .retain(|relative, _| snapshot.paths.contains_key(relative));
305 snapshot.generation = hex::encode(generation.finalize());
306 Ok(ProfileReader { snapshot, schemas })
307 }
308}
309
310#[cfg(test)]
311mod tests {
312 use super::*;
313
314 fn fixture() -> (tempfile::TempDir, ProfileLocation, Connection) {
315 let temp = tempfile::tempdir().unwrap();
316 let directory = temp.path().join("account/db_storage");
317 fs::create_dir_all(&directory).unwrap();
318 let writer = Connection::open(directory.join("sample.db")).unwrap();
319 writer.execute_batch("PRAGMA journal_mode=WAL; CREATE TABLE sample(value INTEGER); INSERT INTO sample VALUES(1)").unwrap();
320 let location = ProfileLocation::new(temp.path().join("account")).unwrap();
321 (temp, location, writer)
322 }
323 fn read(cache: &mut OwnedReaderCache, location: &ProfileLocation) -> ProfileReader {
324 cache
325 .read(
326 location.clone(),
327 KeyFile::empty("account"),
328 ReadOptions {
329 stopped_or_consistent_snapshot: true,
330 allow_plaintext: true,
331 ..Default::default()
332 },
333 )
334 .unwrap()
335 }
336 fn value(reader: &ProfileReader) -> i64 {
337 reader
338 .snapshot
339 .open("sample.db")
340 .unwrap()
341 .query_row("SELECT value FROM sample", [], |r| r.get(0))
342 .unwrap()
343 }
344
345 fn pinned(cache: &mut OwnedReaderCache, location: &ProfileLocation) -> ProfileReader {
346 cache
347 .read_snapshot(
348 location.clone(),
349 KeyFile::empty("account"),
350 ReadOptions {
351 stopped_or_consistent_snapshot: true,
352 allow_plaintext: true,
353 ..Default::default()
354 },
355 )
356 .unwrap()
357 }
358 #[test]
359 fn request_snapshots_pin_all_databases_and_nested_queries_at_one_boundary() {
360 let (_temp, location, writer) = fixture();
361 let other = Connection::open(location.storage.join("other.db")).unwrap();
362 other.execute_batch("PRAGMA journal_mode=WAL; CREATE TABLE sample(value INTEGER); INSERT INTO sample VALUES(10)").unwrap();
363 let mut cache = OwnedReaderCache::default();
364 let first = pinned(&mut cache, &location);
365 writer.execute("UPDATE sample SET value=2", []).unwrap();
366 other.execute("UPDATE sample SET value=20", []).unwrap();
367 let held = first.snapshot.open("sample.db").unwrap();
368 assert_eq!(
369 held.query_row("SELECT value FROM sample", [], |r| r.get::<_, i64>(0))
370 .unwrap(),
371 1
372 );
373 assert_eq!(value(&first), 1); assert_eq!(
375 first
376 .snapshot
377 .open("other.db")
378 .unwrap()
379 .query_row("SELECT value FROM sample", [], |r| r.get::<_, i64>(0))
380 .unwrap(),
381 10
382 );
383 drop(held);
384 drop(first);
385 let second = pinned(&mut cache, &location);
386 assert_eq!(value(&second), 2);
387 assert_eq!(
388 second
389 .snapshot
390 .open("other.db")
391 .unwrap()
392 .query_row("SELECT value FROM sample", [], |r| r.get::<_, i64>(0))
393 .unwrap(),
394 20
395 );
396 }
397 #[test]
398 fn independent_pool_connections_agree_on_source_generation() {
399 let (_temp, location, writer) = fixture();
400 let mut a = OwnedReaderCache::default();
401 let before = pinned(&mut a, &location);
402 drop(before);
403 writer.execute("UPDATE sample SET value=2", []).unwrap();
404 let a = pinned(&mut a, &location);
405 let mut b = OwnedReaderCache::default();
406 let b = pinned(&mut b, &location);
407 assert_eq!(a.generation(), b.generation());
408 }
409 #[test]
410 fn rollback_journal_fallback_does_not_change_mode_or_leave_transactions_open() {
411 let (_temp, location, writer) = fixture();
412 writer.execute_batch("PRAGMA journal_mode=DELETE").unwrap();
413 let mut cache = OwnedReaderCache::default();
414 let error = cache.read_snapshot(
415 location.clone(),
416 KeyFile::empty("account"),
417 ReadOptions {
418 stopped_or_consistent_snapshot: true,
419 allow_plaintext: true,
420 ..Default::default()
421 },
422 );
423 assert!(matches!(
424 error,
425 Err(ProfileError::ParallelSnapshotUnavailable)
426 ));
427 writer.execute("UPDATE sample SET value=3", []).unwrap();
428 assert_eq!(value(&read(&mut cache, &location)), 3);
429 assert_eq!(
430 writer
431 .query_row("PRAGMA journal_mode", [], |r| r.get::<_, String>(0))
432 .unwrap(),
433 "delete"
434 );
435 }
436
437 #[test]
438 fn unchanged_reads_reuse_handles_without_a_temporary_copy() {
439 let (_temp, location, _writer) = fixture();
440 let mut cache = OwnedReaderCache::default();
441 let first = read(&mut cache, &location);
442 let second = read(&mut cache, &location);
443 assert!(first.snapshot._temp.is_none());
444 assert!(Arc::ptr_eq(
445 &first.snapshot.owned["sample.db"],
446 &second.snapshot.owned["sample.db"]
447 ));
448 assert_eq!(first.generation(), second.generation());
449 assert_eq!(value(&first), 1);
450 assert!(first
451 .snapshot
452 .open("sample.db")
453 .unwrap()
454 .execute("UPDATE sample SET value=2", [])
455 .is_err());
456 }
457
458 #[test]
459 fn short_read_transactions_are_stable_and_release_before_the_next_read() {
460 let (_temp, location, writer) = fixture();
461 let mut cache = OwnedReaderCache::default();
462 let first = read(&mut cache, &location);
463 let transaction = first.snapshot.open("sample.db").unwrap();
464 let query = || {
465 transaction
466 .query_row("SELECT value FROM sample", [], |r| r.get::<_, i64>(0))
467 .unwrap()
468 };
469 assert_eq!(query(), 1);
470 writer.execute("UPDATE sample SET value=2", []).unwrap();
471 assert_eq!(query(), 1);
472 drop(transaction);
473 let second = read(&mut cache, &location);
474 assert_eq!(value(&second), 2);
475 assert_ne!(first.generation(), second.generation());
476 }
477
478 #[test]
479 fn changed_schemas_and_new_or_removed_databases_refresh_the_inventory() {
480 let (_temp, location, writer) = fixture();
481 let mut cache = OwnedReaderCache::default();
482 let first = read(&mut cache, &location);
483 writer
484 .execute_batch("ALTER TABLE sample ADD COLUMN text_value TEXT")
485 .unwrap();
486 let extra = location.storage.join("extra.db");
487 Connection::open(&extra)
488 .unwrap()
489 .execute_batch("CREATE TABLE other(value INTEGER)")
490 .unwrap();
491 let second = read(&mut cache, &location);
492 assert_ne!(first.generation(), second.generation());
493 assert_eq!(second.inventory().len(), 2);
494 assert!(second
495 .schemas()
496 .iter()
497 .any(|schema| schema.tables.iter().any(|table| table
498 .columns
499 .iter()
500 .any(|column| column.name == "text_value"))));
501 drop(second);
502 cache.databases.remove("extra.db");
503 fs::remove_file(extra).unwrap();
504 assert_eq!(read(&mut cache, &location).inventory().len(), 1);
505 }
506
507 #[test]
508 fn account_binding_and_size_limits_still_apply() {
509 let (_temp, location, _writer) = fixture();
510 let mut cache = OwnedReaderCache::default();
511 let options = ReadOptions {
512 stopped_or_consistent_snapshot: true,
513 allow_plaintext: true,
514 ..Default::default()
515 };
516 assert!(matches!(
517 cache.read(location.clone(), KeyFile::empty("another"), options.clone()),
518 Err(ProfileError::AccountMismatch)
519 ));
520 assert!(matches!(
521 cache.read(
522 location,
523 KeyFile::empty("account"),
524 ReadOptions {
525 max_database_bytes: 1,
526 ..options
527 }
528 ),
529 Err(ProfileError::ResourceLimit)
530 ));
531 }
532 #[test]
533 fn stricter_read_options_apply_to_reused_connections() {
534 let (_temp, location, _writer) = fixture();
535 let mut cache = OwnedReaderCache::default();
536 let original = read(&mut cache, &location);
537 let smaller = cache
538 .read(
539 location.clone(),
540 KeyFile::empty("account"),
541 ReadOptions {
542 stopped_or_consistent_snapshot: true,
543 allow_plaintext: true,
544 max_blob_bytes: 1024,
545 ..Default::default()
546 },
547 )
548 .unwrap();
549 assert!(Arc::ptr_eq(
550 &original.snapshot.owned["sample.db"],
551 &smaller.snapshot.owned["sample.db"]
552 ));
553 assert_eq!(
554 smaller
555 .snapshot
556 .open("sample.db")
557 .unwrap()
558 .limit(rusqlite::limits::Limit::SQLITE_LIMIT_LENGTH),
559 1024
560 );
561 let strict = cache
562 .read(
563 location,
564 KeyFile::empty("account"),
565 ReadOptions {
566 stopped_or_consistent_snapshot: true,
567 ..Default::default()
568 },
569 )
570 .unwrap();
571 assert!(matches!(
572 strict.snapshot.open("sample.db"),
573 Err(ProfileError::KeyUnavailable(_))
574 ));
575 }
576 #[test]
577 fn nested_reads_do_not_wait_on_their_own_checked_out_connection() {
578 let (_temp, location, _writer) = fixture();
579 let reader = read(&mut OwnedReaderCache::default(), &location);
580 let first = reader.snapshot.open("sample.db").unwrap();
581 let second = reader.snapshot.open("sample.db").unwrap();
582 for connection in [&*first, &*second] {
583 assert_eq!(
584 connection
585 .query_row("SELECT value FROM sample", [], |r| r.get::<_, i64>(0))
586 .unwrap(),
587 1
588 );
589 }
590 }
591 #[test]
592 fn large_profiles_do_not_retain_a_handle_for_every_database() {
593 let (_temp, location, _writer) = fixture();
594 for index in 0..MAX_CACHED_DATABASES + 2 {
595 Connection::open(location.storage.join(format!("extra-{index:02}.db")))
596 .unwrap()
597 .execute_batch("CREATE TABLE extra(value INTEGER); INSERT INTO extra VALUES(7)")
598 .unwrap();
599 }
600 let mut cache = OwnedReaderCache::default();
601 let reader = read(&mut cache, &location);
602 assert_eq!(cache.databases.len(), MAX_CACHED_DATABASES);
603 assert_eq!(reader.snapshot.owned.len(), MAX_CACHED_DATABASES);
604 assert_eq!(reader.inventory().len(), MAX_CACHED_DATABASES + 3);
605 assert!(!reader.snapshot.owned.contains_key("sample.db"));
607 assert_eq!(value(&reader), 1);
608 assert!(reader
609 .snapshot
610 .open("sample.db")
611 .unwrap()
612 .execute("UPDATE sample SET value=2", [])
613 .is_err());
614 }
615 #[test]
616 fn changed_encryption_keys_cannot_reuse_an_already_unlocked_handle() {
617 let temp = tempfile::tempdir().unwrap();
618 let (location, keys, _) =
619 crate::initialize_profile(temp.path().join("account"), &[7; 32]).unwrap();
620 let options = ReadOptions {
621 stopped_or_consistent_snapshot: true,
622 ..Default::default()
623 };
624 let mut cache = OwnedReaderCache::default();
625 let original = cache
626 .read(location.clone(), keys.clone(), options.clone())
627 .unwrap();
628 assert!(original.contacts(&crate::ListQuery::default()).is_ok());
629 let relative = "contact/contact.db";
630 let mut changed = keys.clone();
631 changed.databases.insert(
632 relative.into(),
633 crate::DatabaseKey::new(relative.into(), keys.databases[relative].salt, [9; 32])
634 .unwrap(),
635 );
636 let rejected = cache
637 .read(location.clone(), changed, options.clone())
638 .unwrap();
639 assert!(matches!(
640 rejected.contacts(&crate::ListQuery::default()),
641 Err(ProfileError::KeyMismatch(_))
642 ));
643 let restored = cache.read(location, keys, options).unwrap();
644 assert!(restored.contacts(&crate::ListQuery::default()).is_ok());
645 assert!(!Arc::ptr_eq(
646 &original.snapshot.owned[relative],
647 &restored.snapshot.owned[relative]
648 ));
649 }
650 #[cfg(unix)]
651 #[test]
652 fn replaced_files_reopen_instead_of_serving_an_old_inode() {
653 let (_temp, location, writer) = fixture();
654 let mut cache = OwnedReaderCache::default();
655 let original = read(&mut cache, &location);
656 let old_handle = Arc::downgrade(&original.snapshot.owned["sample.db"]);
657 drop(original);
658 writer
661 .execute_batch("PRAGMA wal_checkpoint(TRUNCATE)")
662 .unwrap();
663 drop(writer);
664 let replacement = location.storage.join("replacement");
665 let db = Connection::open(&replacement).unwrap();
666 db.execute_batch("CREATE TABLE sample(value INTEGER); INSERT INTO sample VALUES(9)")
667 .unwrap();
668 drop(db);
669 fs::rename(replacement, location.storage.join("sample.db")).unwrap();
670 let current = read(&mut cache, &location);
671 assert!(old_handle.upgrade().is_none());
672 assert_eq!(value(¤t), 9);
673 }
674}