1use crate::{
2 model::*,
3 schema::{inspect, message_table, quote_identifier},
4 snapshot::Snapshot,
5 KeyFile, ProfileError, ProfileLocation, ReadOptions, Result, SchemaReport,
6};
7use base64::{
8 engine::general_purpose::{STANDARD, URL_SAFE_NO_PAD},
9 Engine,
10};
11use md5::Md5;
12use rusqlite::{params, types::ValueRef, Connection, OptionalExtension, Row};
13use serde::{Deserialize, Serialize};
14use sha2::{Digest, Sha256};
15use std::{
16 collections::{BTreeMap, BTreeSet},
17 io::Read,
18};
19#[path = "thread_previews.rs"]
20mod thread_previews;
21
22pub struct ProfileReader {
23 pub(crate) snapshot: Snapshot,
24 pub(crate) schemas: Vec<SchemaReport>,
25}
26#[derive(Serialize, Deserialize)]
27#[serde(deny_unknown_fields)]
28struct Cursor {
29 version: u32,
30 generation: String,
31 query: String,
32 key: CursorKey,
33}
34#[derive(Clone, Serialize, Deserialize)]
35#[serde(tag = "type", deny_unknown_fields)]
36enum CursorKey {
37 Name {
38 name: String,
39 },
40 Thread {
41 time: i64,
42 name: String,
43 },
44 Message {
45 time: i64,
46 database: String,
47 table: String,
48 local: i64,
49 },
50}
51
52fn limit(n: Option<usize>) -> Result<usize> {
53 let n = n.unwrap_or(50);
54 if !(1..=200).contains(&n) {
55 Err(ProfileError::InvalidQuery)
56 } else {
57 Ok(n)
58 }
59}
60fn query_hash<T: Serialize>(kind: &str, q: &T) -> Result<String> {
61 let b = serde_json::to_vec(&(kind, q)).map_err(|_| ProfileError::InvalidQuery)?;
62 Ok(hex::encode(Sha256::digest(b)))
63}
64fn encode_cursor(generation: &str, query: &str, key: CursorKey) -> Result<String> {
65 Ok(URL_SAFE_NO_PAD.encode(
66 serde_json::to_vec(&Cursor {
67 version: 1,
68 generation: generation.into(),
69 query: query.into(),
70 key,
71 })
72 .map_err(|_| ProfileError::InvalidCursor)?,
73 ))
74}
75fn decode_cursor(s: Option<&str>, generation: &str, query: &str) -> Result<Option<CursorKey>> {
76 let Some(s) = s else {
77 return Ok(None);
78 };
79 if s.len() > 8192 {
80 return Err(ProfileError::InvalidCursor);
81 }
82 let b = URL_SAFE_NO_PAD
83 .decode(s)
84 .map_err(|_| ProfileError::InvalidCursor)?;
85 let c: Cursor = serde_json::from_slice(&b).map_err(|_| ProfileError::InvalidCursor)?;
86 if c.version != 1 || c.query != query {
87 return Err(ProfileError::InvalidCursor);
88 }
89 if c.generation != generation {
90 return Err(ProfileError::CursorExpired);
91 }
92 Ok(Some(c.key))
93}
94pub(crate) fn raw(row: &Row<'_>) -> rusqlite::Result<RawFields> {
95 let mut out = BTreeMap::new();
96 for (i, name) in row.as_ref().column_names().iter().enumerate() {
97 let value = match row.get_ref(i)? {
98 ValueRef::Null => SqlValue::Null,
99 ValueRef::Integer(n) => SqlValue::Integer(n.to_string()),
100 ValueRef::Real(n) => SqlValue::Real(n.to_string()),
101 ValueRef::Blob(b) => SqlValue::Blob(STANDARD.encode(b)),
102 ValueRef::Text(b) => match std::str::from_utf8(b) {
103 Ok(s) => SqlValue::Text(s.into()),
104 Err(_) => SqlValue::TextBytes(STANDARD.encode(b)),
105 },
106 };
107 out.insert((*name).into(), value);
108 }
109 Ok(out)
110}
111fn text(r: &RawFields, k: &str) -> Option<String> {
112 match r.get(k) {
113 Some(SqlValue::Text(s)) => Some(s.clone()),
114 _ => None,
115 }
116}
117pub(crate) fn num(r: &RawFields, k: &str) -> Option<i64> {
118 match r.get(k) {
119 Some(SqlValue::Integer(s)) => s.parse().ok(),
120 _ => None,
121 }
122}
123pub(crate) fn contact_not_deleted(fields: &RawFields) -> bool {
127 matches!(fields.get("delete_flag"), Some(SqlValue::Null))
128 || num(fields, "delete_flag") == Some(0)
129}
130fn blob(r: &RawFields, k: &str) -> Result<Option<Vec<u8>>> {
131 match r.get(k) {
132 Some(SqlValue::Text(s)) => Ok(Some(s.as_bytes().to_vec())),
133 Some(SqlValue::Blob(s) | SqlValue::TextBytes(s)) => Ok(Some(
134 STANDARD
135 .decode(s)
136 .map_err(|_| ProfileError::InvalidMetadata)?,
137 )),
138 _ => Ok(None),
139 }
140}
141fn label(r: &RawFields, username: &str) -> String {
142 text(r, "remark")
143 .filter(|s| !s.is_empty())
144 .or_else(|| text(r, "nick_name").filter(|s| !s.is_empty()))
145 .unwrap_or_else(|| username.into())
146}
147fn contact(r: RawFields) -> Result<Contact> {
148 let username = text(&r, "username").ok_or(ProfileError::InvalidMetadata)?;
149 let id = num(&r, "id").ok_or(ProfileError::InvalidMetadata)?;
150 Ok(Contact {
151 label: label(&r, &username),
152 is_group: username.ends_with("@chatroom"),
153 username,
154 local_id: id.to_string(),
155 nickname: text(&r, "nick_name"),
156 remark: text(&r, "remark"),
157 alias: text(&r, "alias"),
158 local_type: num(&r, "local_type"),
159 raw_fields: r,
160 })
161}
162fn group_member(contact: Contact, display: Option<String>) -> GroupMember {
163 let label = contact
164 .remark
165 .clone()
166 .filter(|s| !s.is_empty())
167 .or_else(|| display.clone().filter(|s| !s.is_empty()))
168 .or_else(|| contact.nickname.clone().filter(|s| !s.is_empty()))
169 .unwrap_or_else(|| contact.username.clone());
170 GroupMember {
171 username: contact.username.clone(),
172 local_id: contact.local_id.clone(),
173 label,
174 room_display_name: display,
175 contact,
176 }
177}
178pub(crate) fn decode(r: &RawFields, column: &str, flag: &str, max: usize) -> Result<DecodedField> {
179 let flag = num(r, flag);
180 let Some(mut bytes) = blob(r, column)? else {
181 return Ok(DecodedField {
182 text: None,
183 state: DecodeState::Null,
184 compression_flag: flag,
185 });
186 };
187 if bytes.len() > max {
188 return Err(ProfileError::ResourceLimit);
189 }
190 let mut state = DecodeState::Utf8;
191 if flag.is_some_and(|f| f != 0 && f != 4) {
192 return Ok(DecodedField {
193 text: None,
194 state: DecodeState::UnsupportedCompression,
195 compression_flag: flag,
196 });
197 }
198 if bytes.starts_with(&[0x28, 0xb5, 0x2f, 0xfd]) {
199 let result = (|| {
200 let mut decoder = zstd::stream::read::Decoder::new(bytes.as_slice())
201 .ok()?
202 .take((max + 1) as u64);
203 let mut output = Vec::new();
204 decoder.read_to_end(&mut output).ok()?;
205 if output.len() > max {
206 return None;
207 }
208 Some(output)
209 })();
210 let Some(decoded) = result else {
211 return Ok(DecodedField {
212 text: None,
213 state: DecodeState::InvalidCompressedData,
214 compression_flag: flag,
215 });
216 };
217 bytes = decoded;
218 state = DecodeState::Zstd;
219 }
220 match String::from_utf8(bytes) {
221 Ok(s) => Ok(DecodedField {
222 text: Some(s),
223 state,
224 compression_flag: flag,
225 }),
226 Err(_) => Ok(DecodedField {
227 text: None,
228 state: DecodeState::InvalidUtf8,
229 compression_flag: flag,
230 }),
231 }
232}
233fn name_map(c: &Connection) -> Result<BTreeMap<i64, String>> {
234 let present: bool = c.query_row(
235 "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type='table' AND name='Name2Id')",
236 [],
237 |r| r.get(0),
238 )?;
239 if !present {
240 let messages: bool = c.query_row(
243 "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type='table' AND name GLOB 'Msg_*')",
244 [],
245 |r| r.get(0),
246 )?;
247 if messages {
248 return Err(ProfileError::UnsupportedSchema(
249 "message sender namespace".into(),
250 ));
251 }
252 return Ok(BTreeMap::new());
253 }
254 let mut s = c.prepare("SELECT rowid,user_name FROM Name2Id")?;
255 let map = s
256 .query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?
257 .collect::<std::result::Result<BTreeMap<_, _>, _>>()?;
258 if map.len() > 1_000_000 {
259 return Err(ProfileError::ResourceLimit);
260 }
261 Ok(map)
262}
263fn shard(name: &str) -> bool {
264 let b = name.rsplit('/').next().unwrap_or("");
265 let stem = b.strip_suffix(".db").unwrap_or("");
266 ["message_", "biz_message_"].iter().any(|p| {
267 stem.strip_prefix(p)
268 .is_some_and(|n| !n.is_empty() && n.bytes().all(|b| b.is_ascii_digit()))
269 })
270}
271fn validate_name(s: &str) -> Result<()> {
272 if s.is_empty() || s.len() > 4096 || s.contains('\0') {
273 Err(ProfileError::InvalidQuery)
274 } else {
275 Ok(())
276 }
277}
278
279impl ProfileReader {
280 pub fn open(location: ProfileLocation, keys: KeyFile, options: ReadOptions) -> Result<Self> {
281 let snapshot = Snapshot::create(location, keys, options)?;
282 let mut schemas = Vec::new();
283 for db in &snapshot.inventory {
284 if db.state == DatabaseState::Readable {
285 let connection = snapshot.open(&db.relative_path)?;
286 schemas.push(
287 inspect(&connection, &db.relative_path)
288 .map_err(|_| ProfileError::UnsupportedSchema(db.relative_path.clone()))?,
289 );
290 }
291 }
292 Ok(Self { snapshot, schemas })
293 }
294 pub fn generation(&self) -> &str {
295 &self.snapshot.generation
296 }
297 pub fn schemas(&self) -> &[SchemaReport] {
298 &self.schemas
299 }
300 pub fn coverage(&self) -> Coverage {
301 Coverage {
302 partial: self
303 .snapshot
304 .inventory
305 .iter()
306 .any(|d| d.state != DatabaseState::Readable),
307 databases: self.snapshot.inventory.clone(),
308 }
309 }
310 pub fn inventory(&self) -> &[DatabaseInventory] {
311 &self.snapshot.inventory
312 }
313 pub fn require_writable(&self) -> Result<()> {
314 Err(ProfileError::WritesUnsupported)
315 }
316 fn role(&self, basename: &str, capability: &str) -> Result<&str> {
317 let candidates: Vec<_> = self
318 .snapshot
319 .inventory
320 .iter()
321 .filter(|v| v.relative_path.rsplit('/').next() == Some(basename))
322 .collect();
323 if candidates.len() != 1 {
324 return Err(ProfileError::UnsupportedSchema(basename.into()));
325 }
326 let db = candidates[0];
327 match db.state {
328 DatabaseState::KeyUnavailable => {
329 return Err(ProfileError::KeyUnavailable(db.relative_path.clone()))
330 }
331 DatabaseState::KeyMismatch => {
332 return Err(ProfileError::KeyMismatch(db.relative_path.clone()))
333 }
334 _ => {}
335 }
336 if !self
337 .schemas
338 .iter()
339 .any(|s| s.database == db.relative_path && s.supports(capability))
340 {
341 return Err(ProfileError::UnsupportedSchema(db.relative_path.clone()));
342 }
343 Ok(&db.relative_path)
344 }
345 fn page<T>(
346 &self,
347 mut items: Vec<T>,
348 n: usize,
349 q: &str,
350 key: impl Fn(&T) -> CursorKey,
351 ) -> Result<Page<T>> {
352 let more = items.len() > n;
353 items.truncate(n);
354 let next_cursor = if more {
355 items
356 .last()
357 .map(|v| encode_cursor(self.generation(), q, key(v)))
358 .transpose()?
359 } else {
360 None
361 };
362 Ok(Page {
363 items,
364 next_cursor,
365 source_generation: self.generation().into(),
366 coverage: self.coverage(),
367 })
368 }
369 pub fn contacts(&self, q: &ListQuery) -> Result<Page<Contact>> {
370 let n = limit(q.limit)?;
371 let hash = query_hash("contacts", &q.search)?;
372 let after = match decode_cursor(q.cursor.as_deref(), self.generation(), &hash)? {
373 None => String::new(),
374 Some(CursorKey::Name { name }) => name,
375 _ => return Err(ProfileError::InvalidCursor),
376 };
377 let c = self.snapshot.open(self.role("contact.db", "contacts")?)?;
378 let mut s=c.prepare("SELECT * FROM contact WHERE username>?1 AND (?2 IS NULL OR instr(username,?2)>0 OR instr(coalesce(nick_name,''),?2)>0 OR instr(coalesce(remark,''),?2)>0 OR instr(coalesce(alias,''),?2)>0) ORDER BY username LIMIT ?3")?;
379 let records = s
380 .query_map(params![after, q.search, n + 1], raw)?
381 .collect::<std::result::Result<Vec<_>, _>>()?;
382 let items = records
383 .into_iter()
384 .map(contact)
385 .collect::<Result<Vec<_>>>()?;
386 self.page(items, n, &hash, |v| CursorKey::Name {
387 name: v.username.clone(),
388 })
389 }
390 pub fn thread_exists(&self, username: &str) -> Result<bool> {
391 validate_name(username)?;
392 let c = self.snapshot.open(self.role("session.db", "threads")?)?;
393 Ok(c.query_row(
394 "SELECT EXISTS(SELECT 1 FROM SessionTable WHERE username=?1)",
395 [username],
396 |row| row.get(0),
397 )?)
398 }
399 pub fn threads(&self, q: &ListQuery) -> Result<Page<Thread>> {
400 let n = limit(q.limit)?;
401 let hash = query_hash("threads", &q.search)?;
402 let after = match decode_cursor(q.cursor.as_deref(), self.generation(), &hash)? {
403 None => None,
404 Some(CursorKey::Thread { time, name }) => Some((time, name)),
405 _ => return Err(ProfileError::InvalidCursor),
406 };
407 let c = self.snapshot.open(self.role("session.db", "threads")?)?;
408 let mut s=c.prepare("SELECT * FROM SessionTable WHERE (?1 IS NULL OR coalesce(sort_timestamp,0)<?1 OR (coalesce(sort_timestamp,0)=?1 AND username>?2)) AND (?3 IS NULL OR instr(username,?3)>0) ORDER BY coalesce(sort_timestamp,0) DESC,username ASC LIMIT ?4")?;
409 let records = s
410 .query_map(
411 params![
412 after.as_ref().map(|a| a.0),
413 after.as_ref().map(|a| &a.1),
414 q.search,
415 n + 1
416 ],
417 raw,
418 )?
419 .collect::<std::result::Result<Vec<_>, _>>()?;
420 let labels = self.labels()?;
421 let items = records
422 .into_iter()
423 .map(|r| {
424 let username = text(&r, "username").ok_or(ProfileError::InvalidMetadata)?;
425 Ok(Thread {
426 label: labels
427 .get(&username)
428 .cloned()
429 .unwrap_or_else(|| username.clone()),
430 username,
431 unread_count: num(&r, "unread_count"),
432 hidden: num(&r, "is_hidden").map(|n| n != 0),
433 sort_timestamp: num(&r, "sort_timestamp"),
434 last_timestamp: num(&r, "last_timestamp"),
435 raw_fields: r,
436 })
437 })
438 .collect::<Result<Vec<_>>>()?;
439 self.page(items, n, &hash, |v| CursorKey::Thread {
440 time: v.sort_timestamp.unwrap_or(0),
441 name: v.username.clone(),
442 })
443 }
444 pub fn labels(&self) -> Result<BTreeMap<String, String>> {
445 let Ok(rel) = self.role("contact.db", "contacts") else {
447 return Ok(BTreeMap::new());
448 };
449 let c = self.snapshot.open(rel)?;
450 let mut s = c.prepare("SELECT username,remark,nick_name FROM contact LIMIT 1000001")?;
451 let mut out = BTreeMap::new();
452 let mut rows = s.query([])?;
453 while let Some(r) = rows.next()? {
454 let name: String = r.get(0)?;
455 let remark: Option<String> = r.get(1)?;
456 let nick: Option<String> = r.get(2)?;
457 out.insert(
458 name.clone(),
459 remark
460 .filter(|s| !s.is_empty())
461 .or_else(|| nick.filter(|s| !s.is_empty()))
462 .unwrap_or(name),
463 );
464 if out.len() > 1_000_000 {
465 return Err(ProfileError::ResourceLimit);
466 }
467 }
468 Ok(out)
469 }
470 pub fn message_references(
475 &self,
476 chat: Option<&str>,
477 max_items: usize,
478 ) -> Result<Vec<MessageReference>> {
479 if max_items == 0 || max_items > 100_000 {
480 return Err(ProfileError::InvalidQuery);
481 }
482 if let Some(chat) = chat {
483 validate_name(chat)?;
484 }
485 let target = chat.map(|s| format!("Msg_{:x}", Md5::digest(s.as_bytes())));
486 let mut names: BTreeSet<String> = self.labels()?.into_keys().collect();
487 if let Some(chat) = chat {
488 names.insert(chat.to_owned());
489 }
490 let mut records = Vec::new();
491 let mut found = false;
492 for db in self
493 .snapshot
494 .inventory
495 .iter()
496 .filter(|v| shard(&v.relative_path))
497 {
498 if db.state != DatabaseState::Readable {
499 continue;
500 }
501 let schema = self
502 .schemas
503 .iter()
504 .find(|s| s.database == db.relative_path)
505 .ok_or(ProfileError::InvalidMetadata)?;
506 let c = self.snapshot.open(&db.relative_path)?;
507 names.extend(name_map(&c)?.into_values());
508 for table in schema.tables.iter().filter(|t| {
509 message_table(&t.name)
510 && target
511 .as_ref()
512 .is_none_or(|wanted| wanted.eq_ignore_ascii_case(&t.name))
513 }) {
514 require_message_columns(table, &db.relative_path)?;
515 found = true;
516 let sql = format!(
517 "SELECT local_id,create_time FROM {} LIMIT ?1",
518 quote_identifier(&table.name)
519 );
520 let mut statement = c.prepare(&sql)?;
521 let mut rows = statement.query([max_items + 1 - records.len()])?;
522 while let Some(row) = rows.next()? {
523 let local: i64 = row.get(0)?;
524 let time: i64 = row.get(1)?;
525 records.push((time, db.relative_path.clone(), table.name.clone(), local));
526 if records.len() > max_items {
527 return Err(ProfileError::ResourceLimit);
528 }
529 }
530 }
531 }
532 if !found
533 && self
534 .snapshot
535 .inventory
536 .iter()
537 .any(|d| shard(&d.relative_path) && d.state != DatabaseState::Readable)
538 {
539 return Err(ProfileError::KeyUnavailable("message shards".into()));
540 }
541 let mut hashes: BTreeMap<String, Vec<String>> = BTreeMap::new();
542 for name in names {
543 hashes
544 .entry(format!("{:x}", Md5::digest(name.as_bytes())))
545 .or_default()
546 .push(name);
547 }
548 records.sort_unstable_by(|a, b| b.cmp(a));
549 Ok(records
550 .into_iter()
551 .map(|(_, database, table, local)| {
552 let peer = hashes
553 .get(&table[4..].to_ascii_lowercase())
554 .filter(|v| v.len() == 1)
555 .map(|v| v[0].clone());
556 MessageReference {
557 id: SourceMessageId {
558 database,
559 table,
560 local_id: local.to_string(),
561 },
562 peer,
563 }
564 })
565 .collect())
566 }
567
568 pub fn latest_messages(
572 &self,
573 chats: &[String],
574 own_username: Option<&str>,
575 ) -> Result<BTreeMap<String, Message>> {
576 if chats.len() > 200 {
577 return Err(ProfileError::InvalidQuery);
578 }
579 for chat in chats {
580 validate_name(chat)?;
581 }
582 if let Some(own) = own_username {
583 validate_name(own)?;
584 }
585 if chats.is_empty() {
586 return Ok(BTreeMap::new());
587 }
588 let shards: Vec<_> = self
589 .inventory()
590 .iter()
591 .filter(|db| shard(&db.relative_path))
592 .collect();
593 if shards.is_empty() || shards.iter().any(|db| db.state != DatabaseState::Readable) {
594 return Err(ProfileError::KeyUnavailable("message shards".into()));
595 }
596 let labels = BTreeMap::new();
599 let mut result: BTreeMap<String, Message> = BTreeMap::new();
600 for db in shards {
601 let schema = self
602 .schemas
603 .iter()
604 .find(|s| s.database == db.relative_path)
605 .ok_or(ProfileError::InvalidMetadata)?;
606 let c = self.snapshot.open(&db.relative_path)?;
607 let names = name_map(&c)?;
608 let mut selected = Vec::new();
609 let mut columns = BTreeSet::new();
610 for chat in chats {
611 let table = format!("Msg_{:x}", Md5::digest(chat.as_bytes()));
612 let Some(table) = schema
613 .tables
614 .iter()
615 .find(|t| t.name.eq_ignore_ascii_case(&table))
616 else {
617 continue;
618 };
619 if ![
620 "local_id",
621 "server_id",
622 "local_type",
623 "create_time",
624 "real_sender_id",
625 "message_content",
626 "source",
627 "WCDB_CT_message_content",
628 "WCDB_CT_source",
629 ]
630 .iter()
631 .all(|name| table.columns.iter().any(|c| c.name == *name))
632 {
633 return Err(ProfileError::UnsupportedSchema(db.relative_path.clone()));
634 }
635 columns.extend(table.columns.iter().map(|column| column.name.clone()));
636 selected.push((chat, table));
637 }
638 if selected.is_empty() {
639 continue;
640 }
641 let mut discriminator = "__agent_wechat_preview".to_string();
644 while columns.contains(&discriminator) {
645 discriminator.push('_');
646 }
647 let sql = selected.iter().enumerate().map(|(index, (_, table))| {
648 let projection = columns.iter().map(|column| {
649 if table.columns.iter().any(|c| c.name == *column) {
650 quote_identifier(column)
651 } else {
652 format!("NULL AS {}", quote_identifier(column))
653 }
654 }).collect::<Vec<_>>().join(",");
655 format!("SELECT {index} AS {},{projection} FROM (SELECT * FROM {} ORDER BY create_time DESC,local_id DESC LIMIT 1)", quote_identifier(&discriminator), quote_identifier(&table.name))
656 }).collect::<Vec<_>>().join(" UNION ALL ");
657 #[cfg(test)]
658 batch_lookup_tests::record_batch_query();
659 let mut statement = c.prepare(&sql)?;
660 let mut rows = statement.query([])?;
661 while let Some(row) = rows.next()? {
662 let index: usize = row.get(0)?;
663 let (chat, table) = selected.get(index).ok_or(ProfileError::InvalidMetadata)?;
664 let mut fields = raw(row)?;
665 fields.retain(|name, _| table.columns.iter().any(|column| column.name == *name));
666 let local = num(&fields, "local_id")
667 .filter(|id| *id > 0)
668 .ok_or(ProfileError::InvalidMetadata)?;
669 let message = self.message_from_fields_with_peer(
670 fields,
671 &SourceMessageId {
672 database: db.relative_path.clone(),
673 table: table.name.clone(),
674 local_id: local.to_string(),
675 },
676 own_username,
677 &labels,
678 &names,
679 Some(chat),
680 )?;
681 if result
682 .get(*chat)
683 .is_none_or(|old| message_order(&message) > message_order(old))
684 {
685 result.insert((*chat).clone(), message);
686 }
687 }
688 }
689 Ok(result)
690 }
691
692 pub fn messages(&self, q: &MessageQuery) -> Result<Page<Message>> {
693 let n = limit(q.limit)?;
694 if let Some(chat) = &q.chat {
695 validate_name(chat)?;
696 }
697 if let Some(own) = &q.own_username {
698 validate_name(own)?;
699 }
700 let hash = query_hash("messages", &(&q.chat, &q.own_username, q.message_type))?;
701 let after = match decode_cursor(q.cursor.as_deref(), self.generation(), &hash)? {
702 None => None,
703 Some(CursorKey::Message {
704 time,
705 database,
706 table,
707 local,
708 }) if message_table(&table) => Some((time, database, table, local)),
709 _ => return Err(ProfileError::InvalidCursor),
710 };
711 let target = q
712 .chat
713 .as_ref()
714 .map(|s| format!("Msg_{:x}", Md5::digest(s.as_bytes())));
715 let mut all_names: BTreeSet<String> = self.labels()?.into_keys().collect();
716 for db in self
717 .snapshot
718 .inventory
719 .iter()
720 .filter(|v| shard(&v.relative_path) && v.state == DatabaseState::Readable)
721 {
722 let c = self.snapshot.open(&db.relative_path)?;
723 all_names.extend(name_map(&c)?.into_values());
724 }
725 if let Some(chat) = &q.chat {
726 all_names.insert(chat.clone());
727 }
728 let mut hashes: BTreeMap<String, Vec<String>> = BTreeMap::new();
729 for name in all_names {
730 hashes
731 .entry(format!("{:x}", Md5::digest(name.as_bytes())))
732 .or_default()
733 .push(name);
734 }
735 let mut items = Vec::new();
736 let mut found = false;
737 for db in self
738 .snapshot
739 .inventory
740 .iter()
741 .filter(|v| shard(&v.relative_path))
742 {
743 if db.state != DatabaseState::Readable {
744 continue;
745 }
746 let schema = self
747 .schemas
748 .iter()
749 .find(|s| s.database == db.relative_path)
750 .ok_or(ProfileError::InvalidMetadata)?;
751 let c = self.snapshot.open(&db.relative_path)?;
752 let names = name_map(&c)?;
753 for table in schema.tables.iter().filter(|t| {
754 message_table(&t.name)
755 && target
756 .as_ref()
757 .is_none_or(|wanted| wanted.eq_ignore_ascii_case(&t.name))
758 }) {
759 let table_name = &table.name;
760 require_message_columns(table, &db.relative_path)?;
761 found = true;
762 let (max_time, equal_mode, local_bound) =
763 if let Some((time, ref database, ref last_table, local)) = after {
764 let cmp = (&db.relative_path, table_name).cmp(&(database, last_table));
765 (
766 Some(time),
767 if cmp.is_lt() {
768 1
769 } else if cmp.is_gt() {
770 -1
771 } else {
772 0
773 },
774 Some(local),
775 )
776 } else {
777 (None, 0, None)
778 };
779 let sql=format!("SELECT * FROM {} WHERE (?1 IS NULL OR create_time<?1 OR (create_time=?1 AND (?2=1 OR (?2=0 AND local_id<?3)))) AND (?4 IS NULL OR (local_type & 2147483647)=?4) ORDER BY create_time DESC,local_id DESC LIMIT ?5",quote_identifier(table_name));
780 let mut s = c.prepare(&sql)?;
781 let records = s
782 .query_map(
783 params![max_time, equal_mode, local_bound, q.message_type, n + 1],
784 raw,
785 )?
786 .collect::<std::result::Result<Vec<_>, _>>()?;
787 let peer = hashes
788 .get(&table_name[4..].to_ascii_lowercase())
789 .filter(|v| v.len() == 1)
790 .map(|v| v[0].clone());
791 for r in records {
792 let local = num(&r, "local_id").ok_or(ProfileError::InvalidMetadata)?;
793 let signed = num(&r, "server_id");
794 let sender = num(&r, "real_sender_id").and_then(|id| names.get(&id).cloned());
795 let outgoing = q
796 .own_username
797 .as_ref()
798 .and_then(|own| sender.as_ref().map(|s| s == own));
799 let local_type = num(&r, "local_type").ok_or(ProfileError::InvalidMetadata)?;
800 items.push(Message {
801 id: SourceMessageId {
802 database: db.relative_path.clone(),
803 table: table_name.clone(),
804 local_id: local.to_string(),
805 },
806 server_message_id: signed
807 .filter(|v| *v != 0)
808 .map(|v| (v as u64).to_string()),
809 server_id_signed: signed.map(|n| n.to_string()),
810 peer: peer.clone(),
811 sender,
812 outgoing,
813 message_type: local_type & 0x7fff_ffff,
814 local_type: local_type.to_string(),
815 create_time: num(&r, "create_time").ok_or(ProfileError::InvalidMetadata)?,
816 sort_seq: num(&r, "sort_seq").map(|n| n.to_string()),
817 content: decode(
818 &r,
819 "message_content",
820 "WCDB_CT_message_content",
821 self.snapshot.options.max_blob_bytes,
822 )?,
823 source: decode(
824 &r,
825 "source",
826 "WCDB_CT_source",
827 self.snapshot.options.max_blob_bytes,
828 )?,
829 raw_fields: r,
830 });
831 }
832 items.sort_by(|a, b| message_order(b).cmp(&message_order(a)));
833 items.truncate(n + 1);
834 }
835 }
836 if !found
837 && self
838 .snapshot
839 .inventory
840 .iter()
841 .any(|d| shard(&d.relative_path) && d.state != DatabaseState::Readable)
842 {
843 return Err(ProfileError::KeyUnavailable("message shards".into()));
844 }
845 self.page(items, n, &hash, |v| CursorKey::Message {
846 time: v.create_time,
847 database: v.id.database.clone(),
848 table: v.id.table.clone(),
849 local: v.id.local_id.parse().unwrap_or(0),
850 })
851 }
852
853 pub fn message(&self, id: &SourceMessageId, own_username: Option<&str>) -> Result<Message> {
856 if !message_table(&id.table) || !shard(&id.database) {
857 return Err(ProfileError::InvalidQuery);
858 }
859 let c = self.snapshot.open(&id.database)?;
860 let names = name_map(&c)?;
861 let labels = self.labels()?;
862 self.message_from_connection(&c, id, own_username, &labels, &names)
863 }
864
865 pub fn messages_by_ids(
869 &self,
870 ids: &[SourceMessageId],
871 own_username: Option<&str>,
872 ) -> Result<Vec<Message>> {
873 if ids.len() > 200 {
874 return Err(ProfileError::InvalidQuery);
875 }
876 if ids.is_empty() {
877 return Ok(Vec::new());
878 }
879 if let Some(own) = own_username {
880 validate_name(own)?;
881 }
882 let mut groups: BTreeMap<&str, Vec<&SourceMessageId>> = BTreeMap::new();
883 for id in ids {
884 if !message_table(&id.table) || !shard(&id.database) {
885 return Err(ProfileError::InvalidQuery);
886 }
887 groups.entry(&id.database).or_default().push(id);
888 }
889 let labels = self.labels()?;
890 let mut result = Vec::with_capacity(ids.len());
891 for (database, group) in groups {
892 let c = self.snapshot.open(database)?;
893 let names = name_map(&c)?;
894 for id in group {
895 match self.message_from_connection(&c, id, own_username, &labels, &names) {
896 Ok(message) => result.push(message),
897 Err(ProfileError::NotCached) => {}
898 Err(error) => return Err(error),
899 }
900 }
901 }
902 Ok(result)
903 }
904
905 pub fn messages_by_server_id(
908 &self,
909 chat: &str,
910 server_id: &str,
911 own_username: Option<&str>,
912 ) -> Result<Vec<Message>> {
913 validate_name(chat)?;
914 let server = server_id
915 .parse::<u64>()
916 .map_err(|_| ProfileError::InvalidQuery)?;
917 if server == 0 || server.to_string() != server_id {
918 return Err(ProfileError::InvalidQuery);
919 }
920 let table_name = format!("Msg_{:x}", Md5::digest(chat.as_bytes()));
921 let labels = self.labels()?;
922 let mut result = Vec::new();
923 let shards: Vec<_> = self.schemas.iter().filter(|s| shard(&s.database)).collect();
924 if shards.len() > 128 {
925 return Err(ProfileError::InvalidQuery);
926 }
927 if self
928 .coverage()
929 .databases
930 .iter()
931 .any(|d| shard(&d.relative_path) && d.state != DatabaseState::Readable)
932 {
933 return Err(ProfileError::KeyUnavailable("message shards".into()));
934 }
935 for schema in shards {
936 let Some(table) = schema.tables.iter().find(|t| t.name == table_name) else {
937 continue;
938 };
939 if !["local_id", "server_id"]
940 .iter()
941 .all(|name| table.columns.iter().any(|c| c.name == *name))
942 {
943 return Err(ProfileError::UnsupportedSchema(schema.database.clone()));
944 }
945 let c = self.snapshot.open(&schema.database)?;
946 let names = name_map(&c)?;
947 let mut statement = c.prepare(&format!(
948 "SELECT local_id FROM {} WHERE server_id=?1 LIMIT 2",
949 quote_identifier(&table_name)
950 ))?;
951 let ids = statement
952 .query_map([server as i64], |row| row.get::<_, i64>(0))?
953 .collect::<std::result::Result<Vec<_>, _>>()?;
954 for local in ids {
955 let mut message = self.message_from_connection(
956 &c,
957 &SourceMessageId {
958 database: schema.database.clone(),
959 table: table_name.clone(),
960 local_id: local.to_string(),
961 },
962 own_username,
963 &labels,
964 &names,
965 )?;
966 message.peer = Some(chat.into());
969 result.push(message);
970 if result.len() == 2 {
971 return Ok(result);
972 }
973 }
974 }
975 Ok(result)
976 }
977
978 pub fn visit_messages_by_server_ids<E>(
983 &self,
984 targets: &[(String, String)],
985 own_username: Option<&str>,
986 mut visit: impl FnMut(Message) -> std::result::Result<(), E>,
987 ) -> std::result::Result<(), E>
988 where
989 E: From<ProfileError> + From<rusqlite::Error>,
990 {
991 if targets.len() > 200 {
992 return Err(ProfileError::InvalidQuery.into());
993 }
994 if targets.is_empty() {
995 return Ok(());
996 }
997 if let Some(own) = own_username {
998 validate_name(own)?;
999 }
1000 let mut groups: BTreeMap<&str, BTreeSet<i64>> = BTreeMap::new();
1001 for (chat, server) in targets {
1002 validate_name(chat)?;
1003 let number = server
1004 .parse::<u64>()
1005 .map_err(|_| ProfileError::InvalidQuery)?;
1006 if number == 0 || number.to_string() != *server {
1007 return Err(ProfileError::InvalidQuery.into());
1008 }
1009 groups.entry(chat).or_default().insert(number as i64);
1010 }
1011 let shards: Vec<_> = self.schemas.iter().filter(|s| shard(&s.database)).collect();
1012 if shards.len() > 128 {
1013 return Err(ProfileError::InvalidQuery.into());
1014 }
1015 if self
1016 .coverage()
1017 .databases
1018 .iter()
1019 .any(|d| shard(&d.relative_path) && d.state != DatabaseState::Readable)
1020 {
1021 return Err(ProfileError::KeyUnavailable("message shards".into()).into());
1022 }
1023 let labels = self.labels()?;
1024 let mut counts: BTreeMap<(&str, i64), usize> = BTreeMap::new();
1025 for schema in shards {
1026 let tables: Vec<_> = groups
1027 .iter()
1028 .filter_map(|(&chat, servers)| {
1029 let table_name = format!("Msg_{:x}", Md5::digest(chat.as_bytes()));
1030 let table = schema.tables.iter().find(|t| t.name == table_name)?;
1031 let pending: Vec<_> = servers
1032 .iter()
1033 .copied()
1034 .filter(|server| counts.get(&(chat, *server)).copied().unwrap_or(0) < 2)
1035 .collect();
1036 (!pending.is_empty()).then_some((chat, table, pending))
1037 })
1038 .collect();
1039 if tables.is_empty() {
1040 continue;
1041 }
1042 let c = self.snapshot.open(&schema.database)?;
1043 let names = name_map(&c)?;
1044 for (chat, table, servers) in tables {
1045 require_message_columns(table, &schema.database)?;
1046 let table_name = quote_identifier(&table.name);
1047 let placeholders = vec!["?"; servers.len()].join(",");
1048 let sql = format!("SELECT * FROM {table_name} WHERE local_id IN (SELECT local_id FROM (SELECT local_id,ROW_NUMBER() OVER (PARTITION BY server_id ORDER BY local_id) AS quote_rank FROM {table_name} WHERE server_id IN ({placeholders})) WHERE quote_rank<=2) ORDER BY server_id,local_id");
1051 let mut statement = c.prepare(&sql)?;
1052 #[cfg(test)]
1053 batch_lookup_tests::record_batch_query();
1054 let mut rows = statement.query(rusqlite::params_from_iter(servers))?;
1055 while let Some(row) = rows.next()? {
1056 let server: i64 = row.get("server_id")?;
1057 let count = counts.entry((chat, server)).or_default();
1058 if *count == 2 {
1059 continue;
1060 }
1061 let local: i64 = row.get("local_id")?;
1062 if local <= 0 {
1063 return Err(ProfileError::InvalidQuery.into());
1064 }
1065 let id = SourceMessageId {
1066 database: schema.database.clone(),
1067 table: table.name.clone(),
1068 local_id: local.to_string(),
1069 };
1070 let mut message =
1071 self.message_from_fields(raw(row)?, &id, own_username, &labels, &names)?;
1072 message.peer = Some(chat.into());
1073 visit(message)?;
1074 *count += 1;
1075 }
1076 }
1077 }
1078 Ok(())
1079 }
1080
1081 fn message_from_connection(
1082 &self,
1083 c: &Connection,
1084 id: &SourceMessageId,
1085 own_username: Option<&str>,
1086 labels: &BTreeMap<String, String>,
1087 names: &BTreeMap<i64, String>,
1088 ) -> Result<Message> {
1089 if !message_table(&id.table) || !shard(&id.database) {
1090 return Err(ProfileError::InvalidQuery);
1091 }
1092 if let Some(own) = own_username {
1093 validate_name(own)?;
1094 }
1095 let local = id
1096 .local_id
1097 .parse::<i64>()
1098 .map_err(|_| ProfileError::InvalidQuery)?;
1099 if local <= 0 || local.to_string() != id.local_id {
1100 return Err(ProfileError::InvalidQuery);
1101 }
1102 let table = self
1103 .schemas
1104 .iter()
1105 .find(|s| s.database == id.database)
1106 .and_then(|s| s.tables.iter().find(|t| t.name == id.table))
1107 .ok_or(ProfileError::InvalidQuery)?;
1108 if ![
1109 "local_id",
1110 "server_id",
1111 "local_type",
1112 "create_time",
1113 "real_sender_id",
1114 "message_content",
1115 "source",
1116 "WCDB_CT_message_content",
1117 "WCDB_CT_source",
1118 ]
1119 .iter()
1120 .all(|n| table.columns.iter().any(|c| c.name == *n))
1121 {
1122 return Err(ProfileError::UnsupportedSchema(id.database.clone()));
1123 }
1124 #[cfg(test)]
1125 batch_lookup_tests::record_point_query();
1126 let r = c
1127 .query_row(
1128 &format!(
1129 "SELECT * FROM {} WHERE local_id=?1",
1130 quote_identifier(&id.table)
1131 ),
1132 [local],
1133 raw,
1134 )
1135 .optional()?
1136 .ok_or(ProfileError::NotCached)?;
1137 self.message_from_fields(r, id, own_username, labels, names)
1138 }
1139
1140 fn message_from_fields(
1141 &self,
1142 r: RawFields,
1143 id: &SourceMessageId,
1144 own_username: Option<&str>,
1145 labels: &BTreeMap<String, String>,
1146 names: &BTreeMap<i64, String>,
1147 ) -> Result<Message> {
1148 self.message_from_fields_with_peer(r, id, own_username, labels, names, None)
1149 }
1150
1151 fn message_from_fields_with_peer(
1152 &self,
1153 r: RawFields,
1154 id: &SourceMessageId,
1155 own_username: Option<&str>,
1156 labels: &BTreeMap<String, String>,
1157 names: &BTreeMap<i64, String>,
1158 known_peer: Option<&str>,
1159 ) -> Result<Message> {
1160 let peer = if let Some(peer) = known_peer {
1161 Some(peer.to_owned())
1162 } else {
1163 let mut candidates = labels
1164 .keys()
1165 .cloned()
1166 .chain(names.values().cloned())
1167 .filter(|name| {
1168 format!("{:x}", Md5::digest(name.as_bytes()))
1169 .eq_ignore_ascii_case(&id.table[4..])
1170 })
1171 .collect::<BTreeSet<_>>();
1172 if candidates.len() == 1 {
1173 candidates.pop_first()
1174 } else {
1175 None
1176 }
1177 };
1178 let signed = num(&r, "server_id");
1179 let sender = num(&r, "real_sender_id").and_then(|n| names.get(&n).cloned());
1180 let outgoing = own_username.and_then(|own| sender.as_ref().map(|s| s == own));
1181 let local_type = num(&r, "local_type").ok_or(ProfileError::InvalidMetadata)?;
1182 Ok(Message {
1183 id: id.clone(),
1184 server_message_id: signed.filter(|v| *v != 0).map(|v| (v as u64).to_string()),
1185 server_id_signed: signed.map(|n| n.to_string()),
1186 peer,
1187 sender,
1188 outgoing,
1189 message_type: local_type & 0x7fff_ffff,
1190 local_type: local_type.to_string(),
1191 create_time: num(&r, "create_time").ok_or(ProfileError::InvalidMetadata)?,
1192 sort_seq: num(&r, "sort_seq").map(|n| n.to_string()),
1193 content: decode(
1194 &r,
1195 "message_content",
1196 "WCDB_CT_message_content",
1197 self.snapshot.options.max_blob_bytes,
1198 )?,
1199 source: decode(
1200 &r,
1201 "source",
1202 "WCDB_CT_source",
1203 self.snapshot.options.max_blob_bytes,
1204 )?,
1205 raw_fields: r,
1206 })
1207 }
1208 pub fn groups(&self, q: &ListQuery) -> Result<Page<Group>> {
1209 let n = limit(q.limit)?;
1210 let hash = query_hash("groups", &q.search)?;
1211 let after = match decode_cursor(q.cursor.as_deref(), self.generation(), &hash)? {
1212 None => String::new(),
1213 Some(CursorKey::Name { name }) => name,
1214 _ => return Err(ProfileError::InvalidCursor),
1215 };
1216 let c = self.snapshot.open(self.role("contact.db", "groups")?)?;
1217 let mut s=c.prepare("SELECT * FROM chat_room WHERE username>?1 AND (?2 IS NULL OR instr(username,?2)>0) ORDER BY username LIMIT ?3")?;
1218 let records = s
1219 .query_map(params![after, q.search, n + 1], raw)?
1220 .collect::<std::result::Result<Vec<_>, _>>()?;
1221 let labels = self.labels()?;
1222 let mut items = Vec::new();
1223 for r in records {
1224 let username = text(&r, "username").ok_or(ProfileError::InvalidMetadata)?;
1225 let id = num(&r, "id").ok_or(ProfileError::InvalidMetadata)?;
1226 let announcement = c
1227 .query_row(
1228 "SELECT username_,announcement_ FROM chat_room_info_detail WHERE room_id_=?1",
1229 [id],
1230 |r| Ok((r.get::<_, String>(0)?, r.get::<_, Option<String>>(1)?)),
1231 )
1232 .optional()?;
1233 if announcement.as_ref().is_some_and(|a| a.0 != username) {
1234 return Err(ProfileError::InvalidMetadata);
1235 }
1236 items.push(Group {
1237 label: labels
1238 .get(&username)
1239 .cloned()
1240 .unwrap_or_else(|| username.clone()),
1241 username,
1242 local_id: id.to_string(),
1243 owner: text(&r, "owner"),
1244 announcement: announcement.and_then(|v| v.1),
1245 raw_fields: r,
1246 });
1247 }
1248 self.page(items, n, &hash, |v| CursorKey::Name {
1249 name: v.username.clone(),
1250 })
1251 }
1252 pub fn group_mention_permissions(
1253 &self,
1254 room: &str,
1255 own: Option<&str>,
1256 ) -> Result<GroupMentionPermissions> {
1257 validate_name(room)?;
1258 let c = self.snapshot.open(self.role("contact.db", "groups")?)?;
1259 let (id, owner, packed, membership) = c
1260 .query_row(
1261 "SELECT r.id,r.owner,r.ext_buffer,c.is_in_chat_room FROM chat_room r LEFT JOIN contact c ON c.username=r.username WHERE r.username=?1",
1262 [room],
1263 |r| Ok((r.get::<_, i64>(0)?, r.get::<_, Option<String>>(1)?, r.get::<_, Option<Vec<u8>>>(2)?, r.get::<_, Option<i64>>(3)?)),
1264 )
1265 .optional()?
1266 .ok_or(ProfileError::NotCached)?;
1267 let administrator_ids = packed.as_deref().map(room_administrators).transpose()?;
1268 let can_mention_all = if membership == Some(2) {
1269 Some(false)
1271 } else if let Some(own) = own {
1272 validate_name(own)?;
1273 let is_member: bool = c.query_row(
1274 "SELECT EXISTS(SELECT 1 FROM chatroom_member m JOIN contact c ON c.id=m.member_id WHERE m.room_id=?1 AND c.username=?2)",
1275 params![id, own],
1276 |r| r.get(0),
1277 )?;
1278 if !is_member {
1279 None
1280 } else if owner.as_deref() == Some(own) {
1281 Some(true)
1282 } else {
1283 administrator_ids.as_ref().and_then(|ids| {
1284 if ids.iter().any(|id| id == own) {
1285 Some(true)
1286 } else {
1287 owner.as_ref().filter(|id| !id.is_empty()).map(|_| false)
1288 }
1289 })
1290 }
1291 } else {
1292 None
1293 };
1294 Ok(GroupMentionPermissions {
1295 administrator_ids,
1296 can_mention_all,
1297 })
1298 }
1299 pub fn group_members(&self, room: &str, q: &ListQuery) -> Result<Page<GroupMember>> {
1300 validate_name(room)?;
1301 let n = limit(q.limit)?;
1302 let hash = query_hash("members", &(room, &q.search))?;
1303 let after = match decode_cursor(q.cursor.as_deref(), self.generation(), &hash)? {
1304 None => String::new(),
1305 Some(CursorKey::Name { name }) => name,
1306 _ => return Err(ProfileError::InvalidCursor),
1307 };
1308 let c = self.snapshot.open(self.role("contact.db", "groups")?)?;
1309 let room_row = c
1310 .query_row(
1311 "SELECT id,ext_buffer FROM chat_room WHERE username=?1",
1312 [room],
1313 |r| Ok((r.get::<_, i64>(0)?, r.get::<_, Option<Vec<u8>>>(1)?)),
1314 )
1315 .optional()?
1316 .ok_or(ProfileError::NotCached)?;
1317 let displays = room_displays(room_row.1.as_deref().unwrap_or(&[]))?;
1318 let mut sql_members = BTreeSet::new();
1319 let mut stmt=c.prepare("SELECT contact.username FROM chatroom_member JOIN contact ON contact.id=chatroom_member.member_id WHERE chatroom_member.room_id=?1")?;
1320 for row in stmt.query_map([room_row.0], |r| r.get::<_, String>(0))? {
1321 sql_members.insert(row?);
1322 if sql_members.len() > 100000 {
1323 return Err(ProfileError::ResourceLimit);
1324 }
1325 }
1326 if !displays.is_empty() && displays.keys().cloned().collect::<BTreeSet<_>>() != sql_members
1327 {
1328 return Err(ProfileError::InvalidMetadata);
1329 }
1330 let mut stmt=c.prepare("SELECT c.* FROM contact c JOIN chatroom_member m ON m.member_id=c.id WHERE m.room_id=?1 AND c.username>?2 AND (?3 IS NULL OR instr(c.username,?3)>0 OR instr(coalesce(c.nick_name,''),?3)>0 OR instr(coalesce(c.remark,''),?3)>0) ORDER BY c.username LIMIT ?4")?;
1331 let records = stmt
1332 .query_map(params![room_row.0, after, q.search, n + 1], raw)?
1333 .collect::<std::result::Result<Vec<_>, _>>()?;
1334 let mut items = Vec::new();
1335 for r in records {
1336 let contact = contact(r)?;
1337 let display = displays.get(&contact.username).cloned().flatten();
1338 items.push(group_member(contact, display));
1339 }
1340 self.page(items, n, &hash, |v| CursorKey::Name {
1341 name: v.username.clone(),
1342 })
1343 }
1344 pub fn avatar(&self, username: &str) -> Result<CachedBytes> {
1345 validate_name(username)?;
1346 let rel = self.role("head_image.db", "avatars")?;
1347 let c = self.snapshot.open(rel)?;
1348 let bytes = c
1349 .query_row(
1350 "SELECT image_buffer FROM head_image WHERE username=?1",
1351 [username],
1352 |r| r.get::<_, Vec<u8>>(0),
1353 )
1354 .optional()?
1355 .ok_or(ProfileError::NotCached)?;
1356 let (media_type, _, _) = super::images::validate_image(&bytes)?;
1357 Ok(CachedBytes {
1358 media_type,
1359 sha256: hex::encode(Sha256::digest(&bytes)),
1360 bytes,
1361 source_database: rel.into(),
1362 source_table: "head_image".into(),
1363 })
1364 }
1365 pub fn cached_voice(
1368 &self,
1369 database: &str,
1370 chat: &str,
1371 local_id: &str,
1372 server_id: Option<&str>,
1373 ) -> Result<CachedBytes> {
1374 validate_name(chat)?;
1375 let local: i64 = local_id.parse().map_err(|_| ProfileError::InvalidQuery)?;
1376 let sid = server_id
1377 .map(|s| s.parse::<u64>().map(|n| n as i64))
1378 .transpose()
1379 .map_err(|_| ProfileError::InvalidQuery)?;
1380 let c = self.snapshot.open(database)?;
1381 if !self
1382 .schemas
1383 .iter()
1384 .any(|s| s.database == database && s.supports("voice"))
1385 {
1386 return Err(ProfileError::UnsupportedSchema(database.into()));
1387 }
1388 let mut s=c.prepare("SELECT voice_data FROM VoiceInfo JOIN Name2Id n ON n.rowid=VoiceInfo.chat_name_id WHERE n.user_name=?1 AND VoiceInfo.local_id=?2 AND (?3 IS NULL OR VoiceInfo.svr_id=?3) AND length(voice_data)>0 LIMIT 2")?;
1389 let mut found = s
1390 .query_map(params![chat, local, sid], |r| r.get::<_, Vec<u8>>(0))?
1391 .collect::<std::result::Result<Vec<_>, _>>()?;
1392 if found.len() != 1 {
1393 return Err(if found.is_empty() {
1394 ProfileError::NotCached
1395 } else {
1396 ProfileError::InvalidQuery
1397 });
1398 }
1399 let bytes = found.remove(0);
1400 Ok(CachedBytes {
1401 media_type: "application/octet-stream".into(),
1402 sha256: hex::encode(Sha256::digest(&bytes)),
1403 bytes,
1404 source_database: database.into(),
1405 source_table: "VoiceInfo".into(),
1406 })
1407 }
1408 pub fn decode_dat(&self, bytes: &[u8]) -> Result<(Vec<u8>, crate::DatLayout)> {
1409 crate::decode_dat_v2(
1410 bytes,
1411 self.snapshot
1412 .keys
1413 .image_dat
1414 .as_ref()
1415 .ok_or(ProfileError::ImageKeyUnavailable)?,
1416 )
1417 }
1418}
1419fn message_order(m: &Message) -> (i64, &str, &str, i64) {
1420 (
1421 m.create_time,
1422 &m.id.database,
1423 &m.id.table,
1424 m.id.local_id.parse().unwrap_or(0),
1425 )
1426}
1427fn require_message_columns(table: &crate::TableSchema, database: &str) -> Result<()> {
1428 let required = [
1429 "local_id",
1430 "server_id",
1431 "local_type",
1432 "create_time",
1433 "real_sender_id",
1434 "message_content",
1435 "source",
1436 "WCDB_CT_message_content",
1437 "WCDB_CT_source",
1438 ];
1439 if required
1440 .iter()
1441 .all(|name| table.columns.iter().any(|c| c.name == *name))
1442 {
1443 Ok(())
1444 } else {
1445 Err(ProfileError::UnsupportedSchema(database.to_owned()))
1446 }
1447}
1448
1449pub(crate) fn protobuf_fields(b: &[u8]) -> Result<Vec<(u32, u8, &[u8])>> {
1450 fn var(b: &[u8], i: &mut usize) -> Result<u64> {
1451 let mut n = 0u64;
1452 for k in 0..10 {
1453 let x = *b.get(*i).ok_or(ProfileError::InvalidMetadata)?;
1454 *i += 1;
1455 if k == 9 && x > 1 {
1456 return Err(ProfileError::InvalidMetadata);
1457 }
1458 n |= u64::from(x & 127) << (7 * k);
1459 if x < 128 {
1460 return Ok(n);
1461 }
1462 }
1463 Err(ProfileError::InvalidMetadata)
1464 }
1465 if b.len() > 16 * 1024 * 1024 {
1466 return Err(ProfileError::ResourceLimit);
1467 }
1468 let mut out = Vec::new();
1469 let mut i = 0;
1470 while i < b.len() {
1471 let tag = var(b, &mut i)?;
1472 let n = (tag >> 3) as u32;
1473 let w = (tag & 7) as u8;
1474 if tag >> 3 == 0 || tag >> 3 >= 1 << 29 {
1475 return Err(ProfileError::InvalidMetadata);
1476 }
1477 let (start, end) = match w {
1478 0 => {
1479 let start = i;
1480 var(b, &mut i)?;
1481 (start, i)
1482 }
1483 1 => (i, i.checked_add(8).ok_or(ProfileError::InvalidMetadata)?),
1484 5 => (i, i.checked_add(4).ok_or(ProfileError::InvalidMetadata)?),
1485 2 => {
1486 let z =
1487 usize::try_from(var(b, &mut i)?).map_err(|_| ProfileError::InvalidMetadata)?;
1488 (i, i.checked_add(z).ok_or(ProfileError::InvalidMetadata)?)
1489 }
1490 _ => return Err(ProfileError::InvalidMetadata),
1491 };
1492 let value = b.get(start..end).ok_or(ProfileError::InvalidMetadata)?;
1493 i = end;
1494 out.push((n, w, value));
1495 if out.len() > 100000 {
1496 return Err(ProfileError::ResourceLimit);
1497 }
1498 }
1499 Ok(out)
1500}
1501
1502fn room_administrators(b: &[u8]) -> Result<Vec<String>> {
1503 let mut ids = Vec::new();
1504 let mut seen = BTreeSet::new();
1505 for (n, w, v) in protobuf_fields(b)? {
1508 if n != 5 {
1509 continue;
1510 }
1511 if w != 2 {
1512 return Err(ProfileError::InvalidMetadata);
1513 }
1514 let id = std::str::from_utf8(v).map_err(|_| ProfileError::InvalidMetadata)?;
1515 validate_name(id)?;
1516 if !seen.insert(id.to_owned()) || ids.len() >= 10000 {
1517 return Err(ProfileError::InvalidMetadata);
1518 }
1519 ids.push(id.to_owned());
1520 }
1521 Ok(ids)
1522}
1523
1524fn room_displays(b: &[u8]) -> Result<BTreeMap<String, Option<String>>> {
1525 let mut out = BTreeMap::new();
1526 for (n, w, v) in protobuf_fields(b)? {
1527 if n != 1 {
1528 continue;
1529 }
1530 if w != 2 {
1531 return Err(ProfileError::InvalidMetadata);
1532 }
1533 let fields = protobuf_fields(v)?;
1534 let scalar = |number| -> Result<Option<String>> {
1535 let mut values = fields.iter().filter(|(n, _, _)| *n == number);
1536 let one = values.next();
1537 if values.next().is_some() {
1538 return Err(ProfileError::InvalidMetadata);
1539 }
1540 one.map(|(_, w, v)| {
1541 if *w == 2 {
1542 String::from_utf8(v.to_vec()).map_err(|_| ProfileError::InvalidMetadata)
1543 } else {
1544 Err(ProfileError::InvalidMetadata)
1545 }
1546 })
1547 .transpose()
1548 };
1549 let username = scalar(1)?.ok_or(ProfileError::InvalidMetadata)?;
1550 let display = scalar(2)?;
1551 if out.insert(username, display).is_some() {
1552 return Err(ProfileError::InvalidMetadata);
1553 }
1554 }
1555 Ok(out)
1556}
1557
1558#[cfg(test)]
1559mod batch_lookup_tests {
1560 use super::*;
1561 use std::cell::Cell;
1562
1563 thread_local! {
1564 static BATCH_QUERIES: Cell<usize> = const { Cell::new(0) };
1565 static POINT_QUERIES: Cell<usize> = const { Cell::new(0) };
1566 }
1567 pub(super) fn record_batch_query() {
1568 BATCH_QUERIES.with(|n| n.set(n.get() + 1));
1569 }
1570 pub(super) fn record_point_query() {
1571 POINT_QUERIES.with(|n| n.set(n.get() + 1));
1572 }
1573
1574 fn fixture() -> (tempfile::TempDir, ProfileReader) {
1575 let temp = tempfile::tempdir().unwrap();
1576 let directory = temp.path().join("account/db_storage/message");
1577 std::fs::create_dir_all(&directory).unwrap();
1578 let table = format!("Msg_{:x}", Md5::digest(b"peer"));
1579 for shard in 0..2 {
1580 let db = Connection::open(directory.join(format!("message_{shard}.db"))).unwrap();
1581 db.execute_batch("CREATE TABLE Name2Id(user_name TEXT PRIMARY KEY,is_session INTEGER)")
1582 .unwrap();
1583 for user in if shard == 0 {
1584 ["self", "peer"]
1585 } else {
1586 ["peer", "self"]
1587 } {
1588 db.execute("INSERT INTO Name2Id VALUES(?1,1)", [user])
1589 .unwrap();
1590 }
1591 db.execute_batch(&format!("CREATE TABLE {table}(local_id INTEGER PRIMARY KEY,server_id INTEGER,local_type INTEGER,create_time INTEGER,real_sender_id INTEGER,message_content TEXT,source TEXT,WCDB_CT_message_content INTEGER,WCDB_CT_source INTEGER)")).unwrap();
1592 for local in 1..=25 {
1593 db.execute(
1594 &format!("INSERT INTO {table} VALUES(?1,?2,1,100,?3,?4,'',0,0)"),
1595 params![
1596 local,
1597 shard * 25 + local,
1598 if shard == 0 { 2 } else { 1 },
1599 format!("Shard {shard} row {local}")
1600 ],
1601 )
1602 .unwrap();
1603 }
1604 db.execute(
1605 &format!("INSERT INTO {table} VALUES(1000,-1,1,100,1,'Unsigned ID','',0,0)"),
1606 [],
1607 )
1608 .unwrap();
1609 for local in 2000..2003 {
1610 db.execute(
1611 &format!("INSERT INTO {table} VALUES(?1,99,1,100,1,'Duplicate','',0,0)"),
1612 [local],
1613 )
1614 .unwrap();
1615 }
1616 }
1617 let reader = ProfileReader::open(
1618 ProfileLocation::new(temp.path().join("account")).unwrap(),
1619 KeyFile::empty("account"),
1620 ReadOptions {
1621 allow_plaintext: true,
1622 stopped_or_consistent_snapshot: true,
1623 ..Default::default()
1624 },
1625 )
1626 .unwrap();
1627 (temp, reader)
1628 }
1629
1630 #[test]
1631 fn notification_preferences_use_one_query_for_a_full_page() {
1632 let (temp, original) = fixture();
1633 drop(original);
1634 let directory = temp.path().join("account/db_storage/contact");
1635 std::fs::create_dir_all(&directory).unwrap();
1636 let db = Connection::open(directory.join("contact.db")).unwrap();
1637 db.execute_batch(include_str!("layout/contact.db.sql"))
1638 .unwrap();
1639 let chats: Vec<_> = (0..200).map(|n| format!("peer_{n}")).collect();
1640 for (id, chat) in chats.iter().enumerate() {
1641 db.execute(
1642 "INSERT INTO contact(id,username,flag,local_type,delete_flag) VALUES(?1,?2,?3,1,?4)",
1643 params![id, chat, if id % 2 == 0 { 3 } else { 515 }, if id % 3 == 0 { None } else { Some(0) }],
1644 )
1645 .unwrap();
1646 }
1647 let reader = ProfileReader::open(
1648 ProfileLocation::new(temp.path().join("account")).unwrap(),
1649 KeyFile::empty("account"),
1650 ReadOptions {
1651 allow_plaintext: true,
1652 stopped_or_consistent_snapshot: true,
1653 ..Default::default()
1654 },
1655 )
1656 .unwrap();
1657 BATCH_QUERIES.with(|n| n.set(0));
1658 let states = reader.thread_muted(&chats).unwrap();
1659 assert_eq!(states.len(), 200);
1660 for (id, chat) in chats.iter().enumerate() {
1661 assert_eq!(states[chat], Some(id % 2 != 0));
1662 }
1663 BATCH_QUERIES.with(|n| assert_eq!(n.get(), 1));
1664 assert!(reader.thread_muted(&[]).unwrap().is_empty());
1665 assert!(matches!(
1666 reader.thread_muted(&vec!["peer".into(); 201]),
1667 Err(ProfileError::InvalidQuery)
1668 ));
1669 }
1670
1671 #[test]
1672 fn pin_preferences_use_one_query_and_keep_incomplete_or_special_peers_unknown() {
1673 let (temp, original) = fixture();
1674 drop(original);
1675 let directory = temp.path().join("account/db_storage/contact");
1676 std::fs::create_dir_all(&directory).unwrap();
1677 let db = Connection::open(directory.join("contact.db")).unwrap();
1678 db.execute_batch(include_str!("layout/contact.db.sql"))
1679 .unwrap();
1680 let chats: Vec<_> = (0..200)
1681 .map(|n| {
1682 if n % 2 == 0 {
1683 format!("peer_{n}")
1684 } else {
1685 format!("room_{n}@chatroom")
1686 }
1687 })
1688 .collect();
1689 for (id, chat) in chats.iter().enumerate() {
1690 let kind = if id % 2 == 0 { 1 } else { 2 };
1691 let flag = kind | 0x200 | if id % 4 < 2 { 0x800 } else { 0 };
1693 db.execute(
1694 "INSERT INTO contact(id,username,flag,local_type,delete_flag,is_in_chat_room) VALUES(?1,?2,?3,?4,?5,1)",
1695 params![id, chat, flag, kind, if id % 3 == 0 { None } else { Some(0) }],
1696 )
1697 .unwrap();
1698 }
1699 db.execute_batch(
1700 "INSERT INTO contact(id,username,flag,local_type,delete_flag,is_in_chat_room) VALUES
1701 (201,'missing_flag',NULL,1,0,1),
1702 (202,'negative_flag',-1,1,0,1),
1703 (203,'wide_flag',4294967296,1,0,1),
1704 (204,'malformed_flag','not a flag',1,0,1),
1705 (205,'not_friend',2048,1,0,1),
1706 (206,'deleted',2049,1,1,1),
1707 (207,'wrong_type',2049,7,0,1),
1708 (208,'left@chatroom',2050,2,0,2),
1709 (209,'not_group@chatroom',2049,1,0,1),
1710 (210,'filehelper',2049,1,0,1),
1711 (211,'weixin',2049,1,0,1),
1712 (212,'gh_service',2049,1,0,1),
1713 (213,'notification_messages',2049,1,0,1),
1714 (214,'malformed_delete',2049,1,'bad',1);",
1715 )
1716 .unwrap();
1717 let reader = ProfileReader::open(
1718 ProfileLocation::new(temp.path().join("account")).unwrap(),
1719 KeyFile::empty("account"),
1720 ReadOptions {
1721 allow_plaintext: true,
1722 stopped_or_consistent_snapshot: true,
1723 ..Default::default()
1724 },
1725 )
1726 .unwrap();
1727 BATCH_QUERIES.with(|n| n.set(0));
1728 let states = reader.thread_pinned(&chats).unwrap();
1729 assert_eq!(states.len(), 200);
1730 for (id, chat) in chats.iter().enumerate() {
1731 assert_eq!(states[chat], Some(id % 4 < 2));
1732 }
1733 BATCH_QUERIES.with(|n| assert_eq!(n.get(), 1));
1734 let unknown = [
1735 "missing",
1736 "missing_flag",
1737 "negative_flag",
1738 "wide_flag",
1739 "malformed_flag",
1740 "not_friend",
1741 "deleted",
1742 "wrong_type",
1743 "left@chatroom",
1744 "not_group@chatroom",
1745 "filehelper",
1746 "weixin",
1747 "gh_service",
1748 "notification_messages",
1749 "malformed_delete",
1750 ]
1751 .map(str::to_owned);
1752 let states = reader.thread_pinned(&unknown).unwrap();
1753 for name in unknown {
1754 assert_eq!(states.get(&name).copied().flatten(), None, "{name}");
1755 }
1756 assert!(reader.thread_pinned(&[]).unwrap().is_empty());
1757 assert!(matches!(
1758 reader.thread_pinned(&vec!["peer".into(); 201]),
1759 Err(ProfileError::InvalidQuery)
1760 ));
1761 }
1762
1763 fn collect(
1764 reader: &ProfileReader,
1765 targets: &[(String, String)],
1766 own: Option<&str>,
1767 ) -> Result<Vec<Message>> {
1768 let mut messages = Vec::new();
1769 reader.visit_messages_by_server_ids::<ProfileError>(targets, own, |message| {
1770 messages.push(message);
1771 Ok(())
1772 })?;
1773 Ok(messages)
1774 }
1775
1776 #[test]
1777 fn sidebar_previews_use_one_compound_query_per_shard_for_a_full_page() {
1778 let (temp, original) = fixture();
1779 drop(original);
1780 let source = format!("Msg_{:x}", Md5::digest(b"peer"));
1781 let chats: Vec<_> = (0..50).map(|n| format!("peer_{n}")).collect();
1782 for shard in 0..2 {
1783 let db = Connection::open(
1784 temp.path()
1785 .join(format!("account/db_storage/message/message_{shard}.db")),
1786 )
1787 .unwrap();
1788 for chat in &chats {
1789 let table = format!("Msg_{:x}", Md5::digest(chat.as_bytes()));
1790 db.execute_batch(&format!("CREATE TABLE {table} AS SELECT * FROM {source}; CREATE INDEX {table}_time ON {table}(create_time DESC,local_id DESC)")).unwrap();
1791 }
1792 let table = format!("Msg_{:x}", Md5::digest(chats[0].as_bytes()));
1793 db.execute_batch(&format!("ALTER TABLE {table} ADD COLUMN __agent_wechat_preview BLOB; UPDATE {table} SET __agent_wechat_preview=X'0102'")).unwrap();
1794 }
1795 let reader = ProfileReader::open(
1796 ProfileLocation::new(temp.path().join("account")).unwrap(),
1797 KeyFile::empty("account"),
1798 ReadOptions {
1799 allow_plaintext: true,
1800 stopped_or_consistent_snapshot: true,
1801 ..Default::default()
1802 },
1803 )
1804 .unwrap();
1805 BATCH_QUERIES.with(|n| n.set(0));
1806 POINT_QUERIES.with(|n| n.set(0));
1807 let latest = reader.latest_messages(&chats, Some("self")).unwrap();
1808 assert_eq!(latest.len(), chats.len());
1809 assert_eq!(BATCH_QUERIES.with(Cell::get), 2);
1810 assert_eq!(POINT_QUERIES.with(Cell::get), 0);
1811 for chat in &chats {
1812 let expected = reader.message(&latest[chat].id, Some("self")).unwrap();
1813 assert_eq!(latest[chat].raw_fields, expected.raw_fields);
1814 assert_eq!(latest[chat].sender, expected.sender);
1815 }
1816 assert!(!latest[&chats[1]]
1817 .raw_fields
1818 .contains_key("__agent_wechat_preview"));
1819 }
1820
1821 #[test]
1822 fn a_page_uses_one_bulk_lookup_per_shard_and_no_original_point_reads() {
1823 let (_temp, reader) = fixture();
1824 let targets: Vec<_> = (1..=50).map(|n| ("peer".into(), n.to_string())).collect();
1825 BATCH_QUERIES.with(|n| n.set(0));
1826 POINT_QUERIES.with(|n| n.set(0));
1827 let batch = collect(&reader, &targets, Some("self")).unwrap();
1828 assert_eq!(BATCH_QUERIES.with(Cell::get), 2);
1829 assert_eq!(POINT_QUERIES.with(Cell::get), 0);
1830 assert_eq!(batch.len(), 50);
1831 assert!(batch
1832 .iter()
1833 .all(|m| m.sender.as_deref() == Some("peer") && m.outgoing == Some(false)));
1834 let individual: Vec<_> = targets
1835 .iter()
1836 .flat_map(|(thread, server)| {
1837 reader
1838 .messages_by_server_id(thread, server, Some("self"))
1839 .unwrap()
1840 })
1841 .collect();
1842 assert_eq!(POINT_QUERIES.with(Cell::get), 50);
1843 let by_id = |messages: Vec<Message>| {
1844 messages
1845 .into_iter()
1846 .map(|m| (format!("{:?}", m.id), serde_json::to_value(m).unwrap()))
1847 .collect::<BTreeMap<_, _>>()
1848 };
1849 assert_eq!(by_id(batch), by_id(individual));
1850 }
1851
1852 #[test]
1853 fn batch_lookup_keeps_unsigned_ids_duplicates_and_conversations_separate() {
1854 let (_temp, reader) = fixture();
1855 let targets = vec![
1856 ("peer".into(), u64::MAX.to_string()),
1857 ("peer".into(), "99".into()),
1858 ("peer".into(), "99".into()),
1859 ("other".into(), "1".into()),
1860 ("peer".into(), "99999".into()),
1861 ];
1862 let found = collect(&reader, &targets, None).unwrap();
1863 assert_eq!(found.len(), 4);
1864 assert_eq!(
1865 found
1866 .iter()
1867 .filter(|m| m.server_message_id.as_deref() == Some("99"))
1868 .count(),
1869 2
1870 );
1871 assert_eq!(
1872 found
1873 .iter()
1874 .filter(|m| m.server_message_id.as_deref() == Some(u64::MAX.to_string().as_str()))
1875 .count(),
1876 2
1877 );
1878 assert!(found.iter().all(|m| m.peer.as_deref() == Some("peer")));
1879 for invalid in ["0", "01", "-1", "18446744073709551616"] {
1880 assert!(matches!(
1881 collect(&reader, &[("peer".into(), invalid.into())], None),
1882 Err(ProfileError::InvalidQuery)
1883 ));
1884 }
1885 assert!(matches!(
1886 collect(&reader, &vec![("peer".into(), "1".into()); 201], None),
1887 Err(ProfileError::InvalidQuery)
1888 ));
1889 assert!(collect(&reader, &[], None).unwrap().is_empty());
1890 }
1891
1892 #[test]
1893 fn batch_visitor_stops_on_a_consumer_error_without_buffering_more_originals() {
1894 let (_temp, reader) = fixture();
1895 let targets: Vec<_> = (1..=50).map(|n| ("peer".into(), n.to_string())).collect();
1896 let mut visited = 0;
1897 let error = reader.visit_messages_by_server_ids::<ProfileError>(&targets, None, |_| {
1898 visited += 1;
1899 Err(ProfileError::ResourceLimit)
1900 });
1901 assert!(matches!(error, Err(ProfileError::ResourceLimit)));
1902 assert_eq!(visited, 1);
1903 }
1904}