Skip to main content

wechat_profile/
reader.rs

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}
123// Older sparse contact projections left the nullable deletion column unset.
124// Only SQL NULL and integer zero establish a non-deleted row; absent or
125// malformed values must not turn an incomplete identity into an eligible peer.
126pub(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        // The stock client can create an empty next-ordinal database before
241        // initializing its message tables. It has no sender namespace yet.
242        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        // Contact coverage can be unavailable without preventing thread reads.
446        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    /// Capture complete history membership in one pass per shard. This reads
471    /// only IDs/timestamps; callers resolve delivered content with
472    /// `messages_by_ids` against this same snapshot. No state is cached across
473    /// snapshots, and incompatible message tables remain errors.
474    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    /// Read the latest saved message for each requested conversation in one
569    /// snapshot. One compound query per shard selects one indexed row per
570    /// conversation; no point lookups or older-body decoding are needed.
571    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        // The selected conversation already supplies the peer identity, so do
597        // not rescan the contact namespace for every preview's table hash.
598        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            // Preserve unknown columns even when conversation tables differ.
642            // A private discriminator must not collide with a profile column.
643            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    /// Read one source-qualified observation without paging through history.
854    /// Resolve sender numbers only through this database's own Name2Id map.
855    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    /// Resolve at most one page without decoding unrelated history. Missing rows
866    /// are omitted; unreadable databases and incompatible records remain errors.
867    /// Source-qualified IDs preserve each shard's own sender namespace.
868    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    /// Resolve an unsigned server ID in one conversation only. Return up to two
906    /// observations so callers can reject ambiguous IDs across shards.
907    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                // The exact table hash proves the peer even when Name2Id no
967                // longer contains this conversation's label.
968                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    /// Read quoted originals for one page from this captured snapshot. Group by
979    /// conversation/shard, decode only matching rows, and retain at most two
980    /// observations per identity just like the single-reference lookup. The
981    /// visitor consumes each row before the next, avoiding a page of large bodies.
982    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                // The window bounds duplicate observations without fetching or
1049                // decoding unrelated bodies, even if a profile has duplicate IDs.
1050                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            // Official self-leave can retain the former owner/admin bytes.
1270            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    /// Read one exact cached voice record. The supplied media-shard namespace is
1366    /// explicit; its Name2Id is not interchangeable with the message shard.
1367    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    // The supported Linux room model stores the administrator vector as
1506    // repeated string field5 of chat_room.ext_buffer, separate from member flags.
1507    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            // Other flag bits must not affect the independently saved pin bit.
1692            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}