Skip to main content

wechat_profile/
schema.rs

1use crate::{ProfileError, Result};
2use rusqlite::{Connection, OptionalExtension};
3use serde::Serialize;
4use sha2::{Digest, Sha256};
5use std::collections::{BTreeMap, BTreeSet};
6
7#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
8#[serde(rename_all = "camelCase")]
9pub struct ColumnSchema {
10    pub name: String,
11    pub declared_type: String,
12    pub not_null: bool,
13    pub default_sql: Option<String>,
14    pub primary_key_position: i64,
15    pub hidden: i64,
16}
17#[derive(Clone, Debug, Serialize)]
18#[serde(rename_all = "camelCase")]
19pub struct IndexSchema {
20    pub name: String,
21    pub unique: bool,
22    pub origin: String,
23    pub partial: bool,
24    pub columns: Vec<Option<String>>,
25    pub definition: Option<String>,
26}
27#[derive(Clone, Debug, Serialize)]
28#[serde(rename_all = "camelCase")]
29pub struct TableSchema {
30    pub name: String,
31    pub definition: Option<String>,
32    pub columns: Vec<ColumnSchema>,
33    pub indexes: Vec<IndexSchema>,
34    pub triggers: Vec<(String, Option<String>)>,
35}
36#[derive(Clone, Debug, Serialize)]
37#[serde(rename_all = "camelCase")]
38pub struct SchemaReport {
39    pub database: String,
40    pub fingerprint: String,
41    pub user_version: i64,
42    pub application_id: i64,
43    /// SQLite's schema counter is diagnostic; it is not a WeChat release.
44    pub sqlite_schema_counter: i64,
45    pub tables: Vec<TableSchema>,
46    pub readable_capabilities: Vec<String>,
47    pub writes_supported: bool,
48}
49pub(crate) fn quote_identifier(s: &str) -> String {
50    format!("\"{}\"", s.replace('"', "\"\""))
51}
52pub(crate) fn message_table(s: &str) -> bool {
53    s.len() == 36 && s.starts_with("Msg_") && s.as_bytes()[4..].iter().all(u8::is_ascii_hexdigit)
54}
55
56const MAX_SCHEMA_TABLES: usize = 16384;
57
58pub(crate) fn inspect_named_table(c: &Connection, name: &str) -> Result<Option<TableSchema>> {
59    let count = c.query_row(
60        "SELECT count(*) FROM sqlite_master WHERE type='table'",
61        [],
62        |row| row.get::<_, usize>(0),
63    )?;
64    if count > MAX_SCHEMA_TABLES {
65        return Err(ProfileError::ResourceLimit);
66    }
67    let definition = c
68        .query_row(
69            "SELECT sql FROM sqlite_master WHERE type='table' AND name=?1",
70            [name],
71            |row| row.get::<_, Option<String>>(0),
72        )
73        .optional()?;
74    definition
75        .map(|sql| inspect_table(c, name.into(), sql))
76        .transpose()
77}
78
79fn inspect_table(c: &Connection, name: String, definition: Option<String>) -> Result<TableSchema> {
80    // Custom WCDB tokenizers/modules are not part of this reader. Retain
81    // their SQL in the fingerprint without asking SQLite to instantiate it.
82    if definition.as_ref().is_some_and(|s| {
83        s.trim_start()
84            .to_ascii_uppercase()
85            .starts_with("CREATE VIRTUAL TABLE")
86    }) {
87        return Ok(TableSchema {
88            name,
89            definition,
90            columns: Vec::new(),
91            indexes: Vec::new(),
92            triggers: Vec::new(),
93        });
94    }
95    let q = quote_identifier(&name);
96    let mut s = c.prepare(&format!("PRAGMA table_xinfo({q})"))?;
97    let columns = s
98        .query_map([], |r| {
99            Ok(ColumnSchema {
100                name: r.get(1)?,
101                declared_type: r.get(2)?,
102                not_null: r.get::<_, i64>(3)? != 0,
103                default_sql: r.get(4)?,
104                primary_key_position: r.get(5)?,
105                hidden: r.get(6)?,
106            })
107        })?
108        .collect::<std::result::Result<Vec<_>, _>>()?;
109    let mut s = c.prepare(&format!("PRAGMA index_list({q})"))?;
110    let raw = s
111        .query_map([], |r| {
112            Ok((
113                r.get::<_, String>(1)?,
114                r.get::<_, i64>(2)? != 0,
115                r.get::<_, String>(3)?,
116                r.get::<_, i64>(4)? != 0,
117            ))
118        })?
119        .collect::<std::result::Result<Vec<_>, _>>()?;
120    let mut indexes = Vec::new();
121    for (name, unique, origin, partial) in raw {
122        let mut s = c.prepare(&format!("PRAGMA index_info({})", quote_identifier(&name)))?;
123        let columns = s
124            .query_map([], |r| r.get::<_, Option<String>>(2))?
125            .collect::<std::result::Result<Vec<_>, _>>()?;
126        let definition = c
127            .query_row(
128                "SELECT sql FROM sqlite_master WHERE type='index' AND name=?1",
129                [&name],
130                |r| r.get::<_, Option<String>>(0),
131            )
132            .optional()?
133            .flatten();
134        indexes.push(IndexSchema {
135            name,
136            unique,
137            origin,
138            partial,
139            columns,
140            definition,
141        });
142    }
143    indexes.sort_by(|a, b| a.name.cmp(&b.name));
144    let mut s = c.prepare(
145        "SELECT name,sql FROM sqlite_master WHERE type='trigger' AND tbl_name=?1 ORDER BY name",
146    )?;
147    let triggers = s
148        .query_map([&name], |r| Ok((r.get(0)?, r.get(1)?)))?
149        .collect::<std::result::Result<Vec<_>, _>>()?;
150    Ok(TableSchema {
151        name,
152        definition,
153        columns,
154        indexes,
155        triggers,
156    })
157}
158
159pub(crate) fn inspect(c: &Connection, database: &str) -> Result<SchemaReport> {
160    let user_version = c.query_row("PRAGMA user_version", [], |r| r.get(0))?;
161    let application_id = c.query_row("PRAGMA application_id", [], |r| r.get(0))?;
162    let counter = c.query_row("PRAGMA schema_version", [], |r| r.get(0))?;
163    let mut stm =
164        c.prepare("SELECT name,sql FROM sqlite_master WHERE type='table' ORDER BY name")?;
165    let names = stm
166        .query_map([], |r| {
167            Ok((r.get::<_, String>(0)?, r.get::<_, Option<String>>(1)?))
168        })?
169        .collect::<std::result::Result<Vec<_>, _>>()?;
170    if names.len() > MAX_SCHEMA_TABLES {
171        return Err(ProfileError::ResourceLimit);
172    }
173    let tables = names
174        .into_iter()
175        .map(|(name, definition)| inspect_table(c, name, definition))
176        .collect::<Result<Vec<_>>>()?;
177    // Counter excluded: rebuilding an identical schema need not change the
178    // compatibility fingerprint. The actual definitions/columns/indexes bind it.
179    let canonical = serde_json::to_vec(&(user_version, application_id, &tables))
180        .map_err(|_| ProfileError::InvalidMetadata)?;
181    let fingerprint = hex::encode(Sha256::digest(&canonical));
182    let map: BTreeMap<_, BTreeSet<_>> = tables
183        .iter()
184        .map(|t| {
185            (
186                t.name.as_str(),
187                t.columns.iter().map(|v| v.name.as_str()).collect(),
188            )
189        })
190        .collect();
191    let has = |table: &str, columns: &[&str]| {
192        map.get(table)
193            .is_some_and(|s| columns.iter().all(|c| s.contains(c)))
194    };
195    let mut caps = Vec::new();
196    if has(
197        "contact",
198        &[
199            "id",
200            "username",
201            "nick_name",
202            "remark",
203            "alias",
204            "local_type",
205        ],
206    ) {
207        caps.push("contacts".into());
208    }
209    if has("chat_room", &["id", "username", "owner", "ext_buffer"])
210        && has("chatroom_member", &["room_id", "member_id"])
211    {
212        caps.push("groups".into());
213    }
214    if has(
215        "SessionTable",
216        &[
217            "username",
218            "unread_count",
219            "sort_timestamp",
220            "last_timestamp",
221            "is_hidden",
222        ],
223    ) {
224        caps.push("threads".into());
225    }
226    if has(
227        "head_image",
228        &["username", "md5", "image_buffer", "update_time"],
229    ) {
230        caps.push("avatars".into());
231    }
232    if has(
233        "VoiceInfo",
234        &["chat_name_id", "local_id", "svr_id", "voice_data"],
235    ) && has("Name2Id", &["user_name"])
236    {
237        caps.push("voice".into());
238    }
239    if has("Name2Id", &["user_name"])
240        && tables.iter().filter(|t| message_table(&t.name)).any(|t| {
241            [
242                "local_id",
243                "server_id",
244                "local_type",
245                "create_time",
246                "real_sender_id",
247                "message_content",
248                "source",
249                "WCDB_CT_message_content",
250                "WCDB_CT_source",
251            ]
252            .iter()
253            .all(|n| t.columns.iter().any(|c| c.name == *n))
254        })
255    {
256        caps.push("messages".into());
257    }
258    Ok(SchemaReport {
259        database: database.into(),
260        fingerprint,
261        user_version,
262        application_id,
263        sqlite_schema_counter: counter,
264        tables,
265        readable_capabilities: caps,
266        writes_supported: false,
267    })
268}
269
270impl SchemaReport {
271    pub fn supports(&self, capability: &str) -> bool {
272        self.readable_capabilities.iter().any(|v| v == capability)
273    }
274    /// No generic writer is authorized merely by a readable or additive schema.
275    pub fn require_writable(&self) -> Result<()> {
276        Err(ProfileError::WritesUnsupported)
277    }
278}