Skip to main content

wechat_profile/
images.rs

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    /// Resolve one exact local message through its advertised original hash and
187    /// hardlink namespace. Ambiguous groups are rejected. Standard/Thumbnail
188    /// require the exact resource basename/type/status/encoded-size join. A suffix
189    /// alone is not a quality proof; Original also requires the XML content MD5.
190    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(&current)
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(&current).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            // Equal bytes can legitimately occupy multiple advertised roles. Honor
347            // an explicit resource-bound role before comparing original content.
348            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}