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 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 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 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 pub fn require_writable(&self) -> Result<()> {
276 Err(ProfileError::WritesUnsupported)
277 }
278}