docparser_core/
router.rs

1//! 路由：按 magic + 扩展名判定格式（与 Python `engine.Router` 语义对齐）。
2
3use crate::error::{DocError, Result};
4
5const OLE_MAGIC: &[u8] = b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1";
6const ZIP_MAGIC: &[u8] = b"PK\x03\x04";
7const PDF_MAGIC: &[u8] = b"%PDF";
8
9const OOXML: &[&str] = &["docx", "xlsx", "pptx"];
10const ODF: &[&str] = &["odt", "ods", "odp"];
11const OLE_EXT: &[&str] = &["doc", "ppt", "xls"];
12const TEXT_EXT: &[&str] = &["txt", "md", "markdown", "csv", "html", "htm"];
13/// 字幕：没有可靠 magic（.srt 首行是序号），按扩展名直连后端，内容合法性由后端判
14const SUB_EXT: &[&str] = &["srt", "vtt", "ass", "ssa"];
15/// RTF：有 magic（`{\rtf`），但按扩展名直连，头部合法性由后端判（避免伪装件被当纯文本）
16const RTF_EXT: &[&str] = &["rtf"];
17
18fn head_of(path: &str, n: usize) -> Result<Vec<u8>> {
19    std::fs::read(path)
20        .map(|v| v[..v.len().min(n)].to_vec())
21        .map_err(|e| DocError::corrupt(format!("cannot read {path}: {e}")))
22}
23
24/// 语义（继承 KB 现网经验）：
25///   * OOXML 若文件头是 OLE 容器 → 加密件，归 encrypted；
26///   * legacy（doc/ppt/xls）期望 OLE，不是 OLE → unsupported（伪装）；
27///   * 其余按头/magic 兜底。
28pub fn sniff(path: &str) -> Result<&'static str> {
29    let ext = path
30        .rsplit('.')
31        .next()
32        .map(|e| e.to_ascii_lowercase())
33        .unwrap_or_default();
34    let h = head_of(path, 8)?;
35
36    if h.starts_with(OLE_MAGIC) {
37        if OOXML.contains(&ext.as_str()) {
38            return Err(DocError::encrypted(format!("{ext} 为加密容器（OLE）")));
39        }
40        if OLE_EXT.contains(&ext.as_str()) {
41            return Ok("ole");
42        }
43        return Err(DocError::unsupported(format!(
44            "OLE 容器但扩展名不支持: {ext}"
45        )));
46    }
47    if h.starts_with(ZIP_MAGIC) {
48        // EPUB 也是 ZIP 容器：结构（container.xml/opf/spine）由 epub 后端进一步校验
49        if ext == "epub" {
50            return Ok("epub");
51        }
52        // ODF（ODT/ODS/ODP）同样是 ZIP 容器（内层 content.xml），按扩展名分流；
53        // 真实类型由 docparser-odf 后端读 mimetype/content.xml 进一步校验
54        if ODF.contains(&ext.as_str()) {
55            return Ok(match ext.as_str() {
56                "odt" => "odt",
57                "ods" => "ods",
58                _ => "odp",
59            });
60        }
61        // 细判 ooxml 需要读包内清单，这里仅按扩展名近似（由解析后端进一步校验）
62        if OOXML.contains(&ext.as_str()) {
63            return Ok(match ext.as_str() {
64                "docx" => "docx",
65                "xlsx" => "xlsx",
66                _ => "pptx",
67            });
68        }
69        return Err(DocError::corrupt("ZIP 容器但缺少 OOXML 关键部件"));
70    }
71    if h.starts_with(PDF_MAGIC) {
72        return Ok("pdf");
73    }
74    if RTF_EXT.contains(&ext.as_str()) {
75        // 内容合法性（`{\rtf` 头）由 rtf 后端判；这里不抢判，避免伪装件被当纯文本
76        return Ok("rtf");
77    }
78    if OLE_EXT.contains(&ext.as_str()) {
79        return Err(DocError::unsupported(format!(
80            "{ext} 文件头不是 OLE，疑似伪装"
81        )));
82    }
83    if TEXT_EXT.contains(&ext.as_str()) {
84        let kind: &'static str = match ext.as_str() {
85            "txt" => "txt",
86            "md" | "markdown" => "md",
87            "csv" => "csv",
88            "html" | "htm" => "html",
89            _ => "txt",
90        };
91        return Ok(kind);
92    }
93    if SUB_EXT.contains(&ext.as_str()) {
94        return Ok(subtitle_format(&ext));
95    }
96    Ok("txt") // 兜底当纯文本
97}
98
99fn subtitle_format(ext: &str) -> &'static str {
100    match ext {
101        "vtt" => "vtt",
102        "ass" => "ass",
103        "ssa" => "ssa",
104        _ => "srt",
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111    use std::io::Write;
112
113    #[test]
114    fn sniff_pdf_and_ole() {
115        let d = std::env::temp_dir();
116        let pdf = d.join("docparser_router_test.pdf");
117        let mut f = std::fs::File::create(&pdf).unwrap();
118        f.write_all(b"%PDF-1.7\n%%EOF").unwrap();
119        drop(f);
120        assert_eq!(sniff(pdf.to_str().unwrap()).unwrap(), "pdf");
121        std::fs::remove_file(&pdf).ok();
122
123        // magic 优先：.doc 但头是 %PDF → 按 pdf 路由（与 Python Router 一致）
124        let renamed = d.join("renamed_pdf.doc");
125        let mut f = std::fs::File::create(&renamed).unwrap();
126        f.write_all(b"%PDF-1.7\n%%EOF").unwrap();
127        drop(f);
128        assert_eq!(sniff(renamed.to_str().unwrap()).unwrap(), "pdf");
129        std::fs::remove_file(&renamed).ok();
130
131        // 伪装 doc：非 OLE / ZIP / PDF 头 → unsupported
132        let fake = d.join("fake.doc");
133        let mut f = std::fs::File::create(&fake).unwrap();
134        f.write_all(b"NOTANOFFICEFILE").unwrap();
135        drop(f);
136        let err = sniff(fake.to_str().unwrap()).unwrap_err();
137        assert_eq!(err.code.as_str(), "unsupported");
138        std::fs::remove_file(&fake).ok();
139    }
140
141    #[test]
142    fn ooxml_ole_container_is_encrypted() {
143        let d = std::env::temp_dir();
144        let fx = d.join("enc.docx");
145        let mut f = std::fs::File::create(&fx).unwrap();
146        f.write_all(OLE_MAGIC).unwrap();
147        drop(f);
148        let err = sniff(fx.to_str().unwrap()).unwrap_err();
149        assert_eq!(err.code.as_str(), "encrypted");
150        std::fs::remove_file(&fx).ok();
151    }
152
153    /// 字幕没有可靠 magic：按扩展名直连（内容合法性由 subtitle 后端判）。
154    #[test]
155    fn subtitle_extensions_route_to_subtitle_formats() {
156        let d = std::env::temp_dir();
157        for (name, want) in [
158            ("r.srt", "srt"),
159            ("r.vtt", "vtt"),
160            ("r.ass", "ass"),
161            ("r.ssa", "ssa"),
162        ] {
163            let p = d.join(format!("docparser_router_{name}"));
164            let mut f = std::fs::File::create(&p).unwrap();
165            f.write_all(b"1\n00:00:01,000 --> 00:00:02,000\nhi\n")
166                .unwrap();
167            drop(f);
168            assert_eq!(sniff(p.to_str().unwrap()).unwrap(), want, "{name}");
169            std::fs::remove_file(&p).ok();
170        }
171    }
172
173    /// RTF 虽有 `{\rtf` magic，仍按扩展名直连：头部合法性交给 rtf 后端判。
174    #[test]
175    fn rtf_extension_routes_to_rtf() {
176        let d = std::env::temp_dir();
177        let p = d.join("docparser_router_r.rtf");
178        std::fs::write(&p, br"{\rtf1\ansi hi\par}").unwrap();
179        assert_eq!(sniff(p.to_str().unwrap()).unwrap(), "rtf");
180        std::fs::remove_file(&p).ok();
181    }
182
183    /// ODF 也是 ZIP 容器：按扩展名分流到 odt / ods / odp（内容由 odf 后端再校验）。
184    #[test]
185    fn zip_container_routes_odf_family() {
186        let d = std::env::temp_dir();
187        for (name, want) in [
188            ("r.odt", "odt"),
189            ("r.ods", "ods"),
190            ("r.odp", "odp"),
191            ("r.docx", "docx"),
192        ] {
193            let p = d.join(format!("docparser_router_{name}"));
194            let mut f = std::fs::File::create(&p).unwrap();
195            f.write_all(ZIP_MAGIC).unwrap();
196            drop(f);
197            assert_eq!(sniff(p.to_str().unwrap()).unwrap(), want, "{name}");
198            std::fs::remove_file(&p).ok();
199        }
200    }
201}
docparser_core/router.rs

docparser_core/
router.rs
