1use crate::{
2 CachedBytes, CachedImage, ImageQuality, ProfileError, ProfileReader, Result, SourceMessageId,
3 VerifiedImageQuality,
4};
5use image::{ImageFormat, ImageReader, Limits};
6use md5::{Digest, Md5};
7use quick_xml::{events::Event, Reader};
8use rusqlite::OptionalExtension;
9use sha2::Sha256;
10use std::{
11 collections::{BTreeMap, BTreeSet},
12 fs,
13 io::{Cursor, Read},
14 path::{Component, Path},
15};
16
17pub(crate) fn validate_image(bytes: &[u8]) -> Result<(String, u32, u32)> {
18 if bytes.starts_with(b"wxgf")
19 || bytes.starts_with(b"GIF8")
20 || (bytes.starts_with(b"RIFF") && bytes.get(8..12) == Some(b"WEBP"))
21 {
22 let inspection =
23 wechat_media::inspect_image(bytes).map_err(|_| ProfileError::InvalidMetadata)?;
24 return Ok((inspection.content_type, inspection.width, inspection.height));
25 }
26 let mut reader = ImageReader::new(Cursor::new(bytes))
27 .with_guessed_format()
28 .map_err(|_| ProfileError::InvalidMetadata)?;
29 let format = reader.format().ok_or(ProfileError::InvalidMetadata)?;
30 let mime = match format {
31 ImageFormat::Jpeg => "image/jpeg",
32 ImageFormat::Png => "image/png",
33 _ => return Err(ProfileError::InvalidMetadata),
34 };
35 let mut limits = Limits::default();
36 limits.max_image_width = Some(20000);
37 limits.max_image_height = Some(20000);
38 limits.max_alloc = Some(128 * 1024 * 1024);
39 reader.limits(limits);
40 let decoded = reader.decode().map_err(|_| ProfileError::InvalidMetadata)?;
41 Ok((mime.into(), decoded.width(), decoded.height()))
42}
43fn component(s: &str) -> bool {
44 !s.is_empty()
45 && s.len() <= 255
46 && Path::new(s).components().count() == 1
47 && Path::new(s)
48 .components()
49 .all(|c| matches!(c, Component::Normal(_)))
50 && !s.contains(['\\', '\0'])
51}
52fn stem(s: &str) -> Option<String> {
53 let s = s.strip_suffix(".dat").unwrap_or(s);
54 let s = s
55 .strip_suffix("_h")
56 .or_else(|| s.strip_suffix("_t"))
57 .unwrap_or(s);
58 if s.len() == 32 && s.bytes().all(|c| c.is_ascii_hexdigit()) {
59 Some(s.to_ascii_lowercase())
60 } else {
61 None
62 }
63}
64pub(crate) fn image_md5(xml: &str) -> Result<Option<String>> {
65 if xml.len() > 16 * 1024 * 1024 {
66 return Err(ProfileError::ResourceLimit);
67 }
68 let mut reader = Reader::from_str(xml);
69 let mut value = None;
70 let mut count = 0;
71 loop {
72 match reader.read_event() {
73 Ok(Event::DocType(_)) => return Err(ProfileError::InvalidMetadata),
74 Ok(Event::Start(e) | Event::Empty(e)) => {
75 count += 1;
76 if count > 100000 {
77 return Err(ProfileError::ResourceLimit);
78 }
79 if e.name().as_ref() == b"img" {
80 for a in e.attributes() {
81 let a = a.map_err(|_| ProfileError::InvalidMetadata)?;
82 if a.key.as_ref() == b"md5" {
83 let s = std::str::from_utf8(a.value.as_ref())
84 .map_err(|_| ProfileError::InvalidMetadata)?
85 .to_ascii_lowercase();
86 if s.len() != 32
87 || !s.bytes().all(|c| c.is_ascii_hexdigit())
88 || value.replace(s).is_some()
89 {
90 return Err(ProfileError::InvalidMetadata);
91 }
92 }
93 }
94 }
95 }
96 Ok(Event::Eof) => break,
97 Ok(_) => {}
98 Err(_) => return Err(ProfileError::InvalidMetadata),
99 }
100 }
101 Ok(value)
102}
103struct ImageResources {
104 base: String,
105 peer_directory: String,
106 details: BTreeMap<String, u64>,
107}
108impl ProfileReader {
109 fn image_resources(
110 &self,
111 id: &SourceMessageId,
112 local: i64,
113 row: &crate::RawFields,
114 ) -> Result<Option<ImageResources>> {
115 let databases = self
116 .snapshot
117 .inventory
118 .iter()
119 .filter(|d| d.relative_path.rsplit('/').next() == Some("message_resource.db"))
120 .collect::<Vec<_>>();
121 if databases.is_empty() {
122 return Ok(None);
123 }
124 if databases.len() != 1 {
125 return Err(ProfileError::InvalidMetadata);
126 }
127 let message = self.message(id, None)?;
128 let peer = message.peer.ok_or(ProfileError::InvalidMetadata)?;
129 let time = crate::reader::num(row, "create_time").ok_or(ProfileError::InvalidMetadata)?;
130 let server = crate::reader::num(row, "server_id").ok_or(ProfileError::InvalidMetadata)?;
131 let c = self.snapshot.open(&databases[0].relative_path)?;
132 let rows=c.prepare("SELECT r.message_id,r.packed_info FROM MessageResourceInfo r JOIN ChatName2Id n ON n.rowid=r.chat_id WHERE n.user_name=?1 AND r.message_local_id=?2 AND r.message_create_time=?3 AND r.message_svr_id=?4 LIMIT 2")?.query_map(rusqlite::params![peer,local,time,server],|r|Ok((r.get::<_,i64>(0)?,r.get::<_,Vec<u8>>(1)?)))?.collect::<std::result::Result<Vec<_>,_>>()?;
133 if rows.is_empty() {
134 return Ok(None);
135 }
136 if rows.len() != 1 {
137 return Err(ProfileError::InvalidMetadata);
138 }
139 let (resource, packed) = &rows[0];
140 let mut bases = BTreeSet::new();
141 for (field, wire, inner) in crate::reader::protobuf_fields(packed)? {
142 if field == 2 && wire == 2 {
143 for (f, w, b) in crate::reader::protobuf_fields(inner)? {
144 if f == 1 && w == 2 {
145 let base =
146 std::str::from_utf8(b).map_err(|_| ProfileError::InvalidMetadata)?;
147 if !crate::media::digest_hex(base, 32) {
148 return Err(ProfileError::InvalidMetadata);
149 }
150 bases.insert(base.to_ascii_lowercase());
151 }
152 }
153 }
154 }
155 if bases.len() != 1 {
156 return Err(ProfileError::InvalidMetadata);
157 }
158 let mut details = BTreeMap::new();
159 for row in c
160 .prepare("SELECT type,size,status FROM MessageResourceDetail WHERE message_id=?1")?
161 .query_map([resource], |r| {
162 Ok((
163 r.get::<_, i64>(0)?,
164 r.get::<_, i64>(1)?,
165 r.get::<_, i64>(2)?,
166 ))
167 })?
168 {
169 let (kind, size, status) = row?;
170 let suffix = match kind {
171 65537 => "_h",
172 131073 => "",
173 262145 => "_t",
174 _ => continue,
175 };
176 if status == 1 && size > 0 && details.insert(suffix.to_owned(), size as u64).is_some() {
177 return Err(ProfileError::InvalidMetadata);
178 }
179 }
180 Ok(Some(ImageResources {
181 base: bases.pop_first().ok_or(ProfileError::InvalidMetadata)?,
182 peer_directory: format!("{:x}", Md5::digest(peer.as_bytes())),
183 details,
184 }))
185 }
186 pub fn cached_image(&self, id: &SourceMessageId, quality: ImageQuality) -> Result<CachedImage> {
191 if !crate::schema::message_table(&id.table) {
192 return Err(ProfileError::InvalidQuery);
193 }
194 let local = id
195 .local_id
196 .parse::<i64>()
197 .map_err(|_| ProfileError::InvalidQuery)?;
198 let c = self.snapshot.open(&id.database)?;
199 let q = crate::schema::quote_identifier(&id.table);
200 let row = c
201 .query_row(
202 &format!("SELECT * FROM {q} WHERE local_id=?1"),
203 [local],
204 super::reader::raw,
205 )
206 .optional()?
207 .ok_or(ProfileError::NotCached)?;
208 if super::reader::num(&row, "local_type").map(|n| n & 0x7fffffff) != Some(3) {
209 return Err(ProfileError::InvalidQuery);
210 }
211 let decoded = super::reader::decode(
212 &row,
213 "message_content",
214 "WCDB_CT_message_content",
215 self.snapshot.options.max_blob_bytes,
216 )?;
217 let hash = image_md5(
218 decoded
219 .text
220 .as_deref()
221 .ok_or(ProfileError::InvalidMetadata)?,
222 )?
223 .ok_or(ProfileError::NotCached)?;
224 let resource = self.image_resources(id, local, &row)?;
225 if matches!(quality, ImageQuality::Standard | ImageQuality::Thumbnail) && resource.is_none()
226 {
227 return Err(ProfileError::ImageQualityUnavailable);
228 }
229 let hard = self
230 .snapshot
231 .inventory
232 .iter()
233 .filter(|d| d.relative_path.rsplit('/').next() == Some("hardlink.db"))
234 .collect::<Vec<_>>();
235 if hard.len() != 1 {
236 return Err(ProfileError::NotCached);
237 }
238 let c = self.snapshot.open(&hard[0].relative_path)?;
239 let query = if resource.is_some() {
240 "SELECT a.username,b.username,h.file_name FROM image_hardlink_info_v4 h JOIN dir2id a ON a.rowid=h.dir1 JOIN dir2id b ON b.rowid=h.dir2 WHERE a.username=?1 AND h.file_name IN(?2,?3,?4) LIMIT 201"
241 } else {
242 "SELECT a.username,b.username,h.file_name FROM image_hardlink_info_v4 h JOIN dir2id a ON a.rowid=h.dir1 JOIN dir2id b ON b.rowid=h.dir2 WHERE lower(h.md5)=?1 LIMIT 201"
243 };
244 let mut s = c.prepare(query)?;
245 let args = if let Some(r) = &resource {
246 vec![
247 r.peer_directory.clone(),
248 format!("{}.dat", r.base),
249 format!("{}_h.dat", r.base),
250 format!("{}_t.dat", r.base),
251 ]
252 } else {
253 vec![hash.clone()]
254 };
255 let rows = s
256 .query_map(rusqlite::params_from_iter(args), |r| {
257 Ok((
258 r.get::<_, String>(0)?,
259 r.get::<_, String>(1)?,
260 r.get::<_, String>(2)?,
261 ))
262 })?
263 .collect::<std::result::Result<Vec<_>, _>>()?;
264 if rows.len() > 200 {
265 return Err(ProfileError::ResourceLimit);
266 }
267 let mut groups = BTreeSet::new();
268 for (d1, d2, file) in rows {
269 if !component(&d1) || !component(&d2) || !component(&file) {
270 return Err(ProfileError::InvalidMetadata);
271 }
272 if let Some(base) = stem(&file) {
273 groups.insert((d1, d2, base));
274 }
275 }
276 if groups.len() != 1 {
277 return Err(if groups.is_empty() {
278 ProfileError::NotCached
279 } else {
280 ProfileError::InvalidMetadata
281 });
282 }
283 let (d1, d2, base) = groups.into_iter().next().unwrap();
284 let directory = self
285 .snapshot
286 .location
287 .root
288 .join("msg/attach")
289 .join(d1)
290 .join(d2)
291 .join("Img");
292 let mut current = self.snapshot.location.root.clone();
293 let parts = directory
294 .strip_prefix(¤t)
295 .map_err(|_| ProfileError::InvalidMetadata)?
296 .components()
297 .map(|p| p.as_os_str().to_owned())
298 .collect::<Vec<_>>();
299 for part in parts {
300 current.push(part);
301 let m = fs::symlink_metadata(¤t).map_err(|_| ProfileError::NotCached)?;
302 if !m.is_dir() || wechat_platform::is_link(&m) {
303 return Err(ProfileError::UnsupportedFile);
304 }
305 }
306 let mut variants = Vec::new();
307 for suffix in ["", "_h", "_t"] {
308 let p = directory.join(format!("{base}{suffix}.dat"));
309 let m = match fs::symlink_metadata(&p) {
310 Ok(m) => m,
311 Err(e) if e.kind() == std::io::ErrorKind::NotFound => continue,
312 Err(e) => return Err(e.into()),
313 };
314 if !m.is_file() || wechat_platform::is_link(&m) {
315 return Err(ProfileError::UnsupportedFile);
316 }
317 if m.len() > crate::MAX_DAT_PLAINTEXT_BYTES as u64 + 31 {
318 return Err(ProfileError::ResourceLimit);
319 }
320 let before = m.modified()?;
321 let mut data = Vec::with_capacity(m.len() as usize);
322 fs::File::open(&p)?
323 .take(crate::MAX_DAT_PLAINTEXT_BYTES as u64 + 32)
324 .read_to_end(&mut data)?;
325 if data.len() > crate::MAX_DAT_PLAINTEXT_BYTES + 31 {
326 return Err(ProfileError::ResourceLimit);
327 }
328 let after = fs::metadata(&p)?;
329 if before != after.modified()? || m.len() != after.len() {
330 return Err(ProfileError::SourceChanged);
331 }
332 let (plain, _) = self.decode_dat(&data)?;
333 let (mime, width, height) = validate_image(&plain)?;
334 let original = format!("{:x}", Md5::digest(&plain)) == hash;
335 let role = if resource.as_ref().is_some_and(|r| {
336 r.base == base && r.details.get(suffix).copied() == Some(data.len() as u64)
337 }) {
338 match suffix {
339 "" => VerifiedImageQuality::Standard,
340 "_t" => VerifiedImageQuality::Thumbnail,
341 _ => VerifiedImageQuality::Unknown,
342 }
343 } else {
344 VerifiedImageQuality::Unknown
345 };
346 let verified = match quality {
349 ImageQuality::Standard if role == VerifiedImageQuality::Standard => role,
350 ImageQuality::Thumbnail if role == VerifiedImageQuality::Thumbnail => role,
351 _ if original => VerifiedImageQuality::Original,
352 _ => role,
353 };
354 if match quality {
355 ImageQuality::Original => !original,
356 ImageQuality::Standard => verified != VerifiedImageQuality::Standard,
357 ImageQuality::Thumbnail => verified != VerifiedImageQuality::Thumbnail,
358 ImageQuality::Best => false,
359 } {
360 continue;
361 }
362 variants.push(CachedImage {
363 media: CachedBytes {
364 media_type: mime,
365 sha256: hex::encode(Sha256::digest(&plain)),
366 bytes: plain,
367 source_database: id.database.clone(),
368 source_table: id.table.clone(),
369 },
370 quality: verified,
371 width,
372 height,
373 cache_variant: if suffix.is_empty() {
374 "unsuffixed".into()
375 } else {
376 suffix.into()
377 },
378 });
379 }
380 variants.sort_by_key(|v| {
381 (
382 v.quality == VerifiedImageQuality::Original,
383 u64::from(v.width) * u64::from(v.height),
384 v.media.bytes.len(),
385 )
386 });
387 variants.pop().ok_or(ProfileError::NotCached)
388 }
389}
390
391#[cfg(test)]
392mod codec_tests {
393 use super::*;
394
395 #[test]
396 fn legacy_jpeg_png_profile_dimension_limits_are_unchanged() {
397 for format in [ImageFormat::Jpeg, ImageFormat::Png] {
398 let mut encoded = Cursor::new(Vec::new());
399 image::DynamicImage::new_rgb8(5000, 1)
400 .write_to(&mut encoded, format)
401 .unwrap();
402 let (_, width, height) = validate_image(encoded.get_ref()).unwrap();
403 assert_eq!((width, height), (5000, 1));
404 }
405 }
406
407 #[test]
408 fn additional_image_formats_use_shared_bounded_inspection() {
409 let webp = include_bytes!("../tests/fixtures/two-frame.webp");
410 assert_eq!(validate_image(webp).unwrap(), ("image/webp".into(), 12, 8));
411 let wxgf = include_bytes!("../tests/fixtures/static.wxgf");
412 if std::process::Command::new("ffmpeg")
413 .arg("-version")
414 .stdout(std::process::Stdio::null())
415 .stderr(std::process::Stdio::null())
416 .status()
417 .is_ok()
418 {
419 assert_eq!(
420 validate_image(wxgf).unwrap(),
421 ("image/x-wxgf".into(), 32, 24)
422 );
423 } else {
424 assert!(validate_image(wxgf).is_err());
425 }
426 }
427}