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 OLE_EXT: &[&str] = &["doc", "ppt", "xls"];
11const TEXT_EXT: &[&str] = &["txt", "md", "markdown", "csv", "html", "htm"];
12
13fn head_of(path: &str, n: usize) -> Result<Vec<u8>> {
14    std::fs::read(path)
15        .map(|v| v[..v.len().min(n)].to_vec())
16        .map_err(|e| DocError::corrupt(format!("cannot read {path}: {e}")))
17}
18
19/// 语义（继承 KB 现网经验）：
20///   * OOXML 若文件头是 OLE 容器 → 加密件，归 encrypted；
21///   * legacy（doc/ppt/xls）期望 OLE，不是 OLE → unsupported（伪装）；
22///   * 其余按头/magic 兜底。
23pub fn sniff(path: &str) -> Result<&'static str> {
24    let ext = path
25        .rsplit('.')
26        .next()
27        .map(|e| e.to_ascii_lowercase())
28        .unwrap_or_default();
29    let h = head_of(path, 8)?;
30
31    if h.starts_with(OLE_MAGIC) {
32        if OOXML.contains(&ext.as_str()) {
33            return Err(DocError::encrypted(format!("{ext} 为加密容器（OLE）")));
34        }
35        if OLE_EXT.contains(&ext.as_str()) {
36            return Ok("ole");
37        }
38        return Err(DocError::unsupported(format!(
39            "OLE 容器但扩展名不支持: {ext}"
40        )));
41    }
42    if h.starts_with(ZIP_MAGIC) {
43        // EPUB 也是 ZIP 容器：结构（container.xml/opf/spine）由 epub 后端进一步校验
44        if ext == "epub" {
45            return Ok("epub");
46        }
47        // 细判 ooxml 需要读包内清单，这里仅按扩展名近似（由解析后端进一步校验）
48        if OOXML.contains(&ext.as_str()) {
49            return Ok(match ext.as_str() {
50                "docx" => "docx",
51                "xlsx" => "xlsx",
52                _ => "pptx",
53            });
54        }
55        return Err(DocError::corrupt("ZIP 容器但缺少 OOXML 关键部件"));
56    }
57    if h.starts_with(PDF_MAGIC) {
58        return Ok("pdf");
59    }
60    if OLE_EXT.contains(&ext.as_str()) {
61        return Err(DocError::unsupported(format!(
62            "{ext} 文件头不是 OLE，疑似伪装"
63        )));
64    }
65    if TEXT_EXT.contains(&ext.as_str()) {
66        let kind: &'static str = match ext.as_str() {
67            "txt" => "txt",
68            "md" | "markdown" => "md",
69            "csv" => "csv",
70            "html" | "htm" => "html",
71            _ => "txt",
72        };
73        return Ok(kind);
74    }
75    Ok("txt") // 兜底当纯文本
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81    use std::io::Write;
82
83    #[test]
84    fn sniff_pdf_and_ole() {
85        let d = std::env::temp_dir();
86        let pdf = d.join("docparser_router_test.pdf");
87        let mut f = std::fs::File::create(&pdf).unwrap();
88        f.write_all(b"%PDF-1.7\n%%EOF").unwrap();
89        drop(f);
90        assert_eq!(sniff(pdf.to_str().unwrap()).unwrap(), "pdf");
91        std::fs::remove_file(&pdf).ok();
92
93        // magic 优先：.doc 但头是 %PDF → 按 pdf 路由（与 Python Router 一致）
94        let renamed = d.join("renamed_pdf.doc");
95        let mut f = std::fs::File::create(&renamed).unwrap();
96        f.write_all(b"%PDF-1.7\n%%EOF").unwrap();
97        drop(f);
98        assert_eq!(sniff(renamed.to_str().unwrap()).unwrap(), "pdf");
99        std::fs::remove_file(&renamed).ok();
100
101        // 伪装 doc：非 OLE / ZIP / PDF 头 → unsupported
102        let fake = d.join("fake.doc");
103        let mut f = std::fs::File::create(&fake).unwrap();
104        f.write_all(b"NOTANOFFICEFILE").unwrap();
105        drop(f);
106        let err = sniff(fake.to_str().unwrap()).unwrap_err();
107        assert_eq!(err.code.as_str(), "unsupported");
108        std::fs::remove_file(&fake).ok();
109    }
110
111    #[test]
112    fn ooxml_ole_container_is_encrypted() {
113        let d = std::env::temp_dir();
114        let fx = d.join("enc.docx");
115        let mut f = std::fs::File::create(&fx).unwrap();
116        f.write_all(OLE_MAGIC).unwrap();
117        drop(f);
118        let err = sniff(fx.to_str().unwrap()).unwrap_err();
119        assert_eq!(err.code.as_str(), "encrypted");
120        std::fs::remove_file(&fx).ok();
121    }
122}
docparser_core/router.rs

docparser_core/
router.rs
