tdm_server_rust/utils/
archive_extract.rs

1//! 图源解析工具 (Archive / Image Extraction)
2//!
3//! 把话数图源（单图或压缩包）解析为按页序排序的图片字节序列。
4//!
5//! 支持：单图（jpg/png/webp/gif/avif/bmp/tiff）、zip(cbz)、rar(cbr)、7z、tar、
6//! tar.gz(tgz)、tar.bz2、单层 gz、单层 bz2。
7//!
8//! 解压优先使用 Rust 库：`zip` / `tar` / `flate2` / `bzip2-rs` / `sevenz-rust` / `unrar`。
9//! 格式判定：**魔数优先于扩展名**（图源常被误标为 `.zip` 实为 `.rar` 等）。
10//! 提取后按文件名自然顺序（数字感知）排序，保证 1,2,...,10 而非 1,10,2。
11
12use crate::error::{ApiResult, AppError};
13use std::io::Read;
14use std::path::Path;
15
16/// 提取出的单张图片
17#[derive(Debug, Clone)]
18pub struct ExtractedImage {
19    /// 来源文件名（压缩包内相对路径或原图名）
20    pub name: String,
21    /// 图片字节
22    pub data: Vec<u8>,
23}
24
25/// 支持的图片扩展名
26const IMAGE_EXTS: &[&str] = &[
27    "jpg", "jpeg", "png", "webp", "gif", "avif", "bmp", "tif", "tiff",
28];
29
30/// 判断文件名是否为图片
31pub fn is_image_name(name: &str) -> bool {
32    let lower = name.to_ascii_lowercase();
33    IMAGE_EXTS.iter().any(|e| lower.ends_with(&format!(".{e}")))
34}
35
36/// 自然顺序比较（数字感知）：用于页序排序
37fn natural_cmp(a: &str, b: &str) -> std::cmp::Ordering {
38    natord::compare(&a.to_ascii_lowercase(), &b.to_ascii_lowercase())
39}
40
41/// 压缩包/图源格式种类
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43enum ArchiveKind {
44    /// tar.gz / tgz
45    TarGz,
46    /// tar.bz2 / tbz2 / tbz
47    TarBz2,
48    /// 纯 tar
49    Tar,
50    /// zip / cbz
51    Zip,
52    /// 7z
53    SevenZ,
54    /// rar / cbr
55    Rar,
56    /// 单层 gz（非 tar）
57    Gz,
58    /// 单层 bz2（非 tar）
59    Bz2,
60    /// 单张图片
61    SingleImage,
62}
63
64/// 按文件头魔数识别格式；无法识别时返回 `None`
65fn sniff_archive_kind(bytes: &[u8]) -> Option<ArchiveKind> {
66    if bytes.starts_with(b"Rar!") {
67        return Some(ArchiveKind::Rar);
68    }
69    if bytes.starts_with(b"PK") {
70        return Some(ArchiveKind::Zip);
71    }
72    if bytes.len() >= 6 && bytes[..6] == [0x37, 0x7a, 0xbc, 0xaf, 0x27, 0x1c] {
73        return Some(ArchiveKind::SevenZ);
74    }
75    if bytes.starts_with(b"BZh") {
76        return Some(ArchiveKind::Bz2);
77    }
78    if bytes.starts_with(&[0x1f, 0x8b]) {
79        return Some(ArchiveKind::Gz);
80    }
81    if bytes.len() > 262 && bytes[257..262] == *b"ustar" {
82        return Some(ArchiveKind::Tar);
83    }
84    None
85}
86
87/// 由扩展名推断格式
88fn archive_kind_from_filename(lower: &str) -> Option<ArchiveKind> {
89    if lower.ends_with(".tar.gz") || lower.ends_with(".tgz") {
90        Some(ArchiveKind::TarGz)
91    } else if lower.ends_with(".tar.bz2") || lower.ends_with(".tbz2") || lower.ends_with(".tbz") {
92        Some(ArchiveKind::TarBz2)
93    } else if lower.ends_with(".tar") {
94        Some(ArchiveKind::Tar)
95    } else if lower.ends_with(".zip") || lower.ends_with(".cbz") {
96        Some(ArchiveKind::Zip)
97    } else if lower.ends_with(".7z") {
98        Some(ArchiveKind::SevenZ)
99    } else if lower.ends_with(".rar") || lower.ends_with(".cbr") {
100        Some(ArchiveKind::Rar)
101    } else if lower.ends_with(".gz") {
102        Some(ArchiveKind::Gz)
103    } else if lower.ends_with(".bz2") {
104        Some(ArchiveKind::Bz2)
105    } else {
106        None
107    }
108}
109
110/// 综合魔数与扩展名解析格式；魔数与扩展名冲突时以魔数为准
111fn resolve_archive_kind(filename: &str, bytes: &[u8]) -> Result<ArchiveKind, AppError> {
112    let lower = filename.to_ascii_lowercase();
113    if is_image_name(&lower) || infer_is_image(bytes) {
114        return Ok(ArchiveKind::SingleImage);
115    }
116
117    let sniffed = sniff_archive_kind(bytes);
118    let from_name = archive_kind_from_filename(&lower);
119
120    let kind = match (sniffed, from_name) {
121        (Some(ArchiveKind::Gz), Some(ArchiveKind::TarGz) | None) if is_tar_gz_name(&lower) => {
122            ArchiveKind::TarGz
123        }
124        (Some(ArchiveKind::Bz2), Some(ArchiveKind::TarBz2) | None) if is_tar_bz2_name(&lower) => {
125            ArchiveKind::TarBz2
126        }
127        (Some(k), _) => k,
128        (None, Some(k)) => k,
129        (None, None) => {
130            return Err(AppError::business(format!(
131                "无法识别的图源格式：{filename}（支持 图片/zip/rar/7z/tar/gz/bz2）"
132            )));
133        }
134    };
135    Ok(kind)
136}
137
138/// 是否为 tar.gz / tgz 扩展名
139fn is_tar_gz_name(lower: &str) -> bool {
140    lower.ends_with(".tar.gz") || lower.ends_with(".tgz")
141}
142
143/// 是否为 tar.bz2 类扩展名
144fn is_tar_bz2_name(lower: &str) -> bool {
145    lower.ends_with(".tar.bz2") || lower.ends_with(".tbz2") || lower.ends_with(".tbz")
146}
147
148/// 按已判定格式解压图源
149fn extract_by_kind(
150    kind: ArchiveKind,
151    filename: &str,
152    bytes: &[u8],
153) -> ApiResult<Vec<ExtractedImage>> {
154    let lower = filename.to_ascii_lowercase();
155    match kind {
156        ArchiveKind::TarGz => collect_tar(flate2::read::GzDecoder::new(bytes)),
157        ArchiveKind::TarBz2 => collect_tar(bzip2_rs::DecoderReader::new(bytes)),
158        ArchiveKind::Tar => collect_tar(bytes),
159        ArchiveKind::Zip => collect_zip(bytes),
160        ArchiveKind::SevenZ => collect_7z(bytes),
161        ArchiveKind::Rar => collect_rar(bytes),
162        ArchiveKind::Gz => collect_single_gz(&lower, bytes),
163        ArchiveKind::Bz2 => collect_single_bz2(&lower, bytes),
164        ArchiveKind::SingleImage => Ok(vec![ExtractedImage {
165            name: filename.to_string(),
166            data: bytes.to_vec(),
167        }]),
168    }
169}
170
171/// 解析图源字节为按页序排序的图片序列。
172///
173/// ## 参数
174/// - `filename`: 图源原始文件名（用于判定压缩格式）
175/// - `bytes`: 图源字节
176///
177/// ## 返回
178/// - `Ok(Vec<ExtractedImage>)`: 至少一张图片，已按自然顺序排序
179/// - `Err`: 无法识别格式或解压失败或不含图片
180pub fn extract_images(filename: &str, bytes: &[u8]) -> ApiResult<Vec<ExtractedImage>> {
181    let kind = resolve_archive_kind(filename, bytes)?;
182    let mut images = extract_by_kind(kind, filename, bytes)?;
183
184    images.retain(|img| is_image_name(&img.name) || infer_is_image(&img.data));
185    if images.is_empty() {
186        return Err(AppError::business("图源中未找到任何图片喵"));
187    }
188    images.sort_by(|a, b| natural_cmp(&a.name, &b.name));
189    Ok(images)
190}
191
192/// infer 按魔数判断是否图片
193fn infer_is_image(bytes: &[u8]) -> bool {
194    infer::get(bytes)
195        .map(|t| t.matcher_type() == infer::MatcherType::Image)
196        .unwrap_or(false)
197}
198
199/// 解析 zip / cbz
200fn collect_zip(bytes: &[u8]) -> ApiResult<Vec<ExtractedImage>> {
201    let reader = std::io::Cursor::new(bytes);
202    let mut archive = zip::ZipArchive::new(reader)
203        .map_err(|e| AppError::business(format!("打开 zip 失败：{e}")))?;
204    let mut out = Vec::new();
205    for i in 0..archive.len() {
206        let mut file = archive
207            .by_index(i)
208            .map_err(|e| AppError::business(format!("读取 zip 条目失败：{e}")))?;
209        if !file.is_file() {
210            continue;
211        }
212        let name = file.name().to_string();
213        if !is_image_name(&name) {
214            continue;
215        }
216        let mut data = Vec::with_capacity(file.size() as usize);
217        file.read_to_end(&mut data)
218            .map_err(|e| AppError::business(format!("解压 zip 条目失败：{e}")))?;
219        out.push(ExtractedImage { name, data });
220    }
221    Ok(out)
222}
223
224/// 解析 tar 流（含 tar.gz / tar.bz2 已包装好的 Reader）
225fn collect_tar<R: Read>(reader: R) -> ApiResult<Vec<ExtractedImage>> {
226    let mut archive = tar::Archive::new(reader);
227    let mut out = Vec::new();
228    let entries = archive
229        .entries()
230        .map_err(|e| AppError::business(format!("读取 tar 失败：{e}")))?;
231    for entry in entries {
232        let mut entry = entry.map_err(|e| AppError::business(format!("读取 tar 条目失败：{e}")))?;
233        let path = entry
234            .path()
235            .map(|p| p.to_string_lossy().to_string())
236            .unwrap_or_default();
237        if !is_image_name(&path) {
238            continue;
239        }
240        let mut data = Vec::new();
241        entry
242            .read_to_end(&mut data)
243            .map_err(|e| AppError::business(format!("解压 tar 条目失败：{e}")))?;
244        out.push(ExtractedImage { name: path, data });
245    }
246    Ok(out)
247}
248
249/// 解析单层 gz（非 tar）：解出单文件并视为一张图
250fn collect_single_gz(lower_name: &str, bytes: &[u8]) -> ApiResult<Vec<ExtractedImage>> {
251    let mut decoder = flate2::read::GzDecoder::new(bytes);
252    let mut data = Vec::new();
253    decoder
254        .read_to_end(&mut data)
255        .map_err(|e| AppError::business(format!("解压 gz 失败：{e}")))?;
256    let name = lower_name.trim_end_matches(".gz").to_string();
257    Ok(vec![ExtractedImage { name, data }])
258}
259
260/// 解析单层 bz2（非 tar）：解出单文件并视为一张图
261fn collect_single_bz2(lower_name: &str, bytes: &[u8]) -> ApiResult<Vec<ExtractedImage>> {
262    let mut decoder = bzip2_rs::DecoderReader::new(bytes);
263    let mut data = Vec::new();
264    decoder
265        .read_to_end(&mut data)
266        .map_err(|e| AppError::business(format!("解压 bz2 失败：{e}")))?;
267    let name = lower_name.trim_end_matches(".bz2").to_string();
268    Ok(vec![ExtractedImage { name, data }])
269}
270
271/// 解析 7z：解压到临时目录后回收图片
272fn collect_7z(bytes: &[u8]) -> ApiResult<Vec<ExtractedImage>> {
273    let dir =
274        tempfile::tempdir().map_err(|e| AppError::business(format!("创建临时目录失败：{e}")))?;
275    let src = dir.path().join("source.7z");
276    std::fs::write(&src, bytes)
277        .map_err(|e| AppError::business(format!("写入临时 7z 失败：{e}")))?;
278    let out = dir.path().join("out");
279    std::fs::create_dir_all(&out)
280        .map_err(|e| AppError::business(format!("创建解压目录失败：{e}")))?;
281    sevenz_rust::decompress_file(&src, &out)
282        .map_err(|e| AppError::business(format!("解压 7z 失败：{e}")))?;
283    walk_dir_images(&out)
284}
285
286/// 解析 rar：解压到临时目录后回收图片
287fn collect_rar(bytes: &[u8]) -> ApiResult<Vec<ExtractedImage>> {
288    let dir =
289        tempfile::tempdir().map_err(|e| AppError::business(format!("创建临时目录失败：{e}")))?;
290    let src = dir.path().join("source.rar");
291    std::fs::write(&src, bytes)
292        .map_err(|e| AppError::business(format!("写入临时 rar 失败：{e}")))?;
293    let out = dir.path().join("out");
294    std::fs::create_dir_all(&out)
295        .map_err(|e| AppError::business(format!("创建解压目录失败：{e}")))?;
296    let out_str = out.to_string_lossy().to_string();
297
298    let mut archive = unrar::Archive::new(&src)
299        .open_for_processing()
300        .map_err(|e| AppError::business(format!("打开 rar 失败：{e}")))?;
301    while let Some(header) = archive
302        .read_header()
303        .map_err(|e| AppError::business(format!("读取 rar 头失败：{e}")))?
304    {
305        archive = if header.entry().is_file() {
306            header
307                .extract_with_base(&out_str)
308                .map_err(|e| AppError::business(format!("解压 rar 条目失败：{e}")))?
309        } else {
310            header
311                .skip()
312                .map_err(|e| AppError::business(format!("跳过 rar 目录失败：{e}")))?
313        };
314    }
315    walk_dir_images(&out)
316}
317
318/// 递归遍历目录收集图片（用于 7z/rar 解压结果）
319fn walk_dir_images(dir: &Path) -> ApiResult<Vec<ExtractedImage>> {
320    let mut out = Vec::new();
321    for entry in walkdir::WalkDir::new(dir).into_iter().flatten() {
322        if !entry.file_type().is_file() {
323            continue;
324        }
325        let path = entry.path();
326        let name = path
327            .strip_prefix(dir)
328            .unwrap_or(path)
329            .to_string_lossy()
330            .replace('\\', "/");
331        if !is_image_name(&name) {
332            continue;
333        }
334        let data = std::fs::read(path)
335            .map_err(|e| AppError::business(format!("读取解压图片失败：{e}")))?;
336        out.push(ExtractedImage { name, data });
337    }
338    Ok(out)
339}
340
341#[cfg(test)]
342mod tests {
343    use super::*;
344
345    /// 自然顺序排序：数字感知
346    #[test]
347    fn natural_order_sorts_numerically() {
348        let mut names = ["10.jpg", "2.jpg", "1.jpg"];
349        names.sort_by(|a, b| natural_cmp(a, b));
350        assert_eq!(names, ["1.jpg", "2.jpg", "10.jpg"]);
351    }
352
353    /// 图片扩展名识别
354    #[test]
355    fn detects_image_names() {
356        assert!(is_image_name("a/b/01.PNG"));
357        assert!(!is_image_name("note.txt"));
358    }
359
360    /// zip 导入：乱序图片按自然序还原，非图片条目被忽略
361    #[test]
362    fn extract_zip_orders_and_filters_non_image() {
363        use std::io::Write;
364        use zip::write::SimpleFileOptions;
365
366        let mut buf = Vec::new();
367        {
368            let cursor = std::io::Cursor::new(&mut buf);
369            let mut writer = zip::ZipWriter::new(cursor);
370            // Stored 不依赖压缩 feature，保证测试稳定
371            let opts =
372                SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored);
373            for name in ["10.png", "2.png", "readme.txt", "1.png"] {
374                writer.start_file(name, opts).unwrap();
375                writer.write_all(b"fake-bytes").unwrap();
376            }
377            writer.finish().unwrap();
378        }
379
380        let images = extract_images("pages.zip", &buf).unwrap();
381        let names: Vec<&str> = images.iter().map(|i| i.name.as_str()).collect();
382        assert_eq!(names, ["1.png", "2.png", "10.png"]);
383    }
384
385    /// 不支持的图源格式应报错
386    #[test]
387    fn unknown_format_errors() {
388        let err = extract_images("note.txt", b"plain text").unwrap_err();
389        assert!(format!("{err:?}").contains("无法识别"));
390    }
391
392    /// 魔数为 rar、扩展名为 zip 时按 rar 解析（常见误标）
393    #[test]
394    fn rar_magic_overrides_zip_extension() {
395        assert_eq!(
396            sniff_archive_kind(b"Rar!\x1a\x07\x01\x00"),
397            Some(ArchiveKind::Rar)
398        );
399        let kind = resolve_archive_kind("fake.zip", b"Rar!\x1a\x07\x01\x00").unwrap();
400        assert_eq!(kind, ArchiveKind::Rar);
401    }
402}
tdm_server_rust/utils/archive_extract.rs

tdm_server_rust/utils/
archive_extract.rs
