1#![deny(missing_docs)]
4#![forbid(unsafe_code)]
5
6use std::{
7 error::Error as StdError,
8 fmt,
9 io::{Cursor, Read},
10};
11
12use calamine::{Reader as _, open_workbook_auto_from_rs};
13use quick_xml::{Reader as XmlReader, events::Event as XmlEvent};
14
15#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
17pub enum DocumentFormat {
18 Pdf,
20 Doc,
22 Docx,
24 Spreadsheet,
26 PlainText,
28}
29
30impl DocumentFormat {
31 pub const fn as_str(self) -> &'static str {
33 match self {
34 Self::Pdf => "pdf",
35 Self::Doc => "doc",
36 Self::Docx => "docx",
37 Self::Spreadsheet => "spreadsheet",
38 Self::PlainText => "text",
39 }
40 }
41}
42
43#[derive(Clone, Copy, Debug, Eq, PartialEq)]
45pub enum ErrorKind {
46 InvalidInput,
48 UnsupportedFormat,
50 ExtractionFailed,
52 EmptyText,
54}
55
56#[derive(Debug)]
58pub struct Error {
59 kind: ErrorKind,
60 message: String,
61}
62
63impl Error {
64 fn new(kind: ErrorKind, message: impl Into<String>) -> Self {
65 Self {
66 kind,
67 message: message.into(),
68 }
69 }
70
71 pub const fn kind(&self) -> ErrorKind {
73 self.kind
74 }
75
76 pub fn message(&self) -> &str {
78 &self.message
79 }
80}
81
82impl fmt::Display for Error {
83 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
84 formatter.write_str(&self.message)
85 }
86}
87
88impl StdError for Error {}
89
90pub type Result<T> = std::result::Result<T, Error>;
92
93#[derive(Clone, Debug, Eq, PartialEq)]
95pub struct ExtractionLimits {
96 pub max_bytes: usize,
98 pub max_characters: usize,
100}
101
102impl Default for ExtractionLimits {
103 fn default() -> Self {
104 Self {
105 max_bytes: 20 * 1024 * 1024,
106 max_characters: 1_000_000,
107 }
108 }
109}
110
111#[derive(Clone, Eq, PartialEq)]
113pub struct DocumentInput {
114 pub file_name: String,
116 pub content_type: String,
118 pub data: Vec<u8>,
120}
121
122impl fmt::Debug for DocumentInput {
123 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
124 formatter
125 .debug_struct("DocumentInput")
126 .field("file_name", &self.file_name)
127 .field("content_type", &self.content_type)
128 .field("bytes", &self.data.len())
129 .finish()
130 }
131}
132
133#[derive(Clone, Debug, Eq, PartialEq)]
135pub struct ExtractedDocument {
136 pub file_name: String,
138 pub content_type: String,
140 pub format: DocumentFormat,
142 pub text: String,
144 pub characters: usize,
146 pub truncated: bool,
148}
149
150#[derive(Clone, Debug)]
152pub struct DocumentExtractor {
153 limits: ExtractionLimits,
154}
155
156impl DocumentExtractor {
157 pub fn new(limits: ExtractionLimits) -> Result<Self> {
159 if limits.max_bytes == 0 || limits.max_characters == 0 {
160 return Err(Error::new(
161 ErrorKind::InvalidInput,
162 "document byte and character limits must be greater than zero",
163 ));
164 }
165 Ok(Self { limits })
166 }
167
168 pub fn extract(&self, input: DocumentInput) -> Result<ExtractedDocument> {
170 if input.data.is_empty() || input.data.len() > self.limits.max_bytes {
171 return Err(Error::new(
172 ErrorKind::InvalidInput,
173 format!(
174 "document must contain between 1 and {} bytes",
175 self.limits.max_bytes
176 ),
177 ));
178 }
179 let content_type = input
180 .content_type
181 .split(';')
182 .next()
183 .unwrap_or(&input.content_type)
184 .trim()
185 .to_ascii_lowercase();
186 let format = detect_format(&input.file_name, &content_type).ok_or_else(|| {
187 Error::new(
188 ErrorKind::UnsupportedFormat,
189 "supported documents are PDF, DOC, DOCX, XLSX, XLS, XLSB, ODS, CSV, TSV, and plain text",
190 )
191 })?;
192 let file_name = safe_filename(&input.file_name, format);
193 let extracted = extract_text(format, &input.data).map_err(|message| {
194 Error::new(
195 ErrorKind::ExtractionFailed,
196 format!("document could not be converted to text: {message}"),
197 )
198 })?;
199 if extracted.is_empty() {
200 let message = if format == DocumentFormat::Pdf {
201 "PDF contains no extractable text and may require OCR"
202 } else {
203 "document contains no extractable text"
204 };
205 return Err(Error::new(ErrorKind::EmptyText, message));
206 }
207 let (text, truncated) = truncate_characters(&extracted, self.limits.max_characters);
208 let characters = text.chars().count();
209 Ok(ExtractedDocument {
210 file_name,
211 content_type,
212 format,
213 text,
214 characters,
215 truncated,
216 })
217 }
218}
219
220impl Default for DocumentExtractor {
221 fn default() -> Self {
222 Self::new(ExtractionLimits::default()).expect("default document limits are valid")
223 }
224}
225
226fn detect_format(file_name: &str, content_type: &str) -> Option<DocumentFormat> {
227 let extension = file_name
228 .rsplit_once('.')
229 .map(|(_, extension)| extension.to_ascii_lowercase());
230 match extension.as_deref() {
231 Some("pdf") => Some(DocumentFormat::Pdf),
232 Some("doc") => Some(DocumentFormat::Doc),
233 Some("docx") => Some(DocumentFormat::Docx),
234 Some("xlsx" | "xls" | "xlsb" | "ods") => Some(DocumentFormat::Spreadsheet),
235 Some("csv" | "tsv" | "txt" | "md" | "json" | "yaml" | "yml" | "xml") => {
236 Some(DocumentFormat::PlainText)
237 }
238 _ if content_type == "application/pdf" => Some(DocumentFormat::Pdf),
239 _ if content_type == "application/msword" => Some(DocumentFormat::Doc),
240 _ if content_type
241 == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" =>
242 {
243 Some(DocumentFormat::Docx)
244 }
245 _ if matches!(
246 content_type,
247 "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
248 | "application/vnd.ms-excel"
249 | "application/vnd.ms-excel.sheet.binary.macroenabled.12"
250 | "application/vnd.oasis.opendocument.spreadsheet"
251 ) =>
252 {
253 Some(DocumentFormat::Spreadsheet)
254 }
255 _ if content_type.starts_with("text/")
256 || matches!(
257 content_type,
258 "application/json" | "application/xml" | "application/yaml" | "application/x-yaml"
259 ) =>
260 {
261 Some(DocumentFormat::PlainText)
262 }
263 _ => None,
264 }
265}
266
267fn safe_filename(value: &str, format: DocumentFormat) -> String {
268 let cleaned = value
269 .chars()
270 .filter(|character| {
271 !character.is_control() && !matches!(character, '/' | '\\' | ':' | '\0')
272 })
273 .take(200)
274 .collect::<String>();
275 if cleaned.trim().is_empty() {
276 format!("document.{}", format.as_str())
277 } else {
278 cleaned
279 }
280}
281
282fn local_xml_name(name: &[u8]) -> &[u8] {
283 name.rsplit(|byte| *byte == b':').next().unwrap_or(name)
284}
285
286fn extract_doc_text(bytes: &[u8]) -> std::result::Result<String, String> {
287 let document = match std::panic::catch_unwind(|| unword::parse_doc_with_options(bytes, true)) {
288 Ok(Ok(document)) => document,
289 Ok(Err(_)) => return Err("legacy Word parser rejected the document".to_owned()),
290 Err(_) => {
291 return Err("legacy Word parser panicked while reading the document".to_owned());
292 }
293 };
294
295 let mut output = document.body_text;
296 for text_box in document.textboxes {
297 let text_box = text_box.trim();
298 if text_box.is_empty() {
299 continue;
300 }
301 if !output.trim().is_empty() {
302 output.push_str("\n\n");
303 }
304 output.push_str("Text box:\n");
305 output.push_str(text_box);
306 }
307 Ok(output)
308}
309
310fn extract_docx_text(bytes: &[u8]) -> std::result::Result<String, String> {
311 let mut archive =
312 zip::ZipArchive::new(Cursor::new(bytes)).map_err(|error| error.to_string())?;
313 let mut document = archive
314 .by_name("word/document.xml")
315 .map_err(|error| format!("word/document.xml: {error}"))?;
316 let mut xml = String::new();
317 document
318 .read_to_string(&mut xml)
319 .map_err(|error| error.to_string())?;
320 let mut reader = XmlReader::from_str(&xml);
321 let mut output = String::new();
322 let mut in_text = false;
323 let mut in_cell = false;
324 loop {
325 match reader.read_event() {
326 Ok(XmlEvent::Start(event)) => match local_xml_name(event.name().as_ref()) {
327 b"t" => in_text = true,
328 b"tc" => in_cell = true,
329 _ => {}
330 },
331 Ok(XmlEvent::Empty(event)) => match local_xml_name(event.name().as_ref()) {
332 b"tab" => output.push('\t'),
333 b"br" | b"cr" => output.push('\n'),
334 _ => {}
335 },
336 Ok(XmlEvent::Text(text)) if in_text => {
337 let decoded = text.decode().map_err(|error| error.to_string())?;
338 output.push_str(&decoded);
339 }
340 Ok(XmlEvent::CData(text)) if in_text => {
341 let decoded = text.decode().map_err(|error| error.to_string())?;
342 output.push_str(&decoded);
343 }
344 Ok(XmlEvent::GeneralRef(reference)) if in_text => {
345 let decoded = reference.decode().map_err(|error| error.to_string())?;
346 let escaped = format!("&{decoded};");
347 let resolved =
348 quick_xml::escape::unescape(&escaped).map_err(|error| error.to_string())?;
349 output.push_str(&resolved);
350 }
351 Ok(XmlEvent::End(event)) => match local_xml_name(event.name().as_ref()) {
352 b"t" => in_text = false,
353 b"p" if in_cell => output.push_str(" / "),
354 b"p" => output.push('\n'),
355 b"tc" => {
356 if output.ends_with(" / ") {
357 output.truncate(output.len() - 3);
358 }
359 output.push('\t');
360 in_cell = false;
361 }
362 b"tr" => {
363 if output.ends_with('\t') {
364 output.pop();
365 }
366 output.push('\n');
367 }
368 _ => {}
369 },
370 Ok(XmlEvent::Eof) => break,
371 Err(error) => return Err(error.to_string()),
372 _ => {}
373 }
374 }
375 Ok(output)
376}
377
378fn extract_spreadsheet_text(bytes: &[u8]) -> std::result::Result<String, String> {
379 let mut workbook = open_workbook_auto_from_rs(Cursor::new(bytes.to_vec()))
380 .map_err(|error| error.to_string())?;
381 let sheet_names = workbook.sheet_names().to_vec();
382 let mut output = String::new();
383 for (sheet_index, sheet_name) in sheet_names.iter().enumerate() {
384 if sheet_index > 0 {
385 output.push('\n');
386 }
387 output.push_str("Sheet: ");
388 output.push_str(sheet_name);
389 output.push('\n');
390 let range = workbook
391 .worksheet_range(sheet_name)
392 .map_err(|error| format!("{sheet_name}: {error}"))?;
393 for row in range.rows() {
394 let line = row
395 .iter()
396 .map(ToString::to_string)
397 .map(|cell| cell.replace(['\t', '\r', '\n'], " "))
398 .collect::<Vec<_>>()
399 .join("\t");
400 output.push_str(line.trim_end());
401 output.push('\n');
402 }
403 }
404 Ok(output)
405}
406
407fn normalize_text(value: String) -> String {
408 let normalized = value
409 .replace("\r\n", "\n")
410 .replace('\r', "\n")
411 .replace('\0', "");
412 let mut output = String::with_capacity(normalized.len());
413 let mut blank_lines = 0;
414 for line in normalized.lines() {
415 let line = line.trim_end();
416 if line.trim().is_empty() {
417 blank_lines += 1;
418 if blank_lines > 1 {
419 continue;
420 }
421 } else {
422 blank_lines = 0;
423 }
424 output.push_str(line);
425 output.push('\n');
426 }
427 output.trim().to_owned()
428}
429
430fn extract_text(format: DocumentFormat, bytes: &[u8]) -> std::result::Result<String, String> {
431 let extracted = match format {
432 DocumentFormat::Pdf => {
433 pdf_extract::extract_text_from_mem(bytes).map_err(|error| error.to_string())?
434 }
435 DocumentFormat::Doc => extract_doc_text(bytes)?,
436 DocumentFormat::Docx => extract_docx_text(bytes)?,
437 DocumentFormat::Spreadsheet => extract_spreadsheet_text(bytes)?,
438 DocumentFormat::PlainText => String::from_utf8_lossy(bytes)
439 .trim_start_matches('\u{feff}')
440 .to_owned(),
441 };
442 Ok(normalize_text(extracted))
443}
444
445fn truncate_characters(value: &str, limit: usize) -> (String, bool) {
446 let mut iter = value.char_indices();
447 let Some((boundary, _)) = iter.nth(limit) else {
448 return (value.to_owned(), false);
449 };
450 (value[..boundary].trim_end().to_owned(), true)
451}
452
453#[cfg(test)]
454mod tests {
455 use super::*;
456 use std::io::Write as _;
457
458 const FIXTURE_BODY: &str = "# Legacy heading\rCafé 世界 legacy body\rVisible \u{13} HYPERLINK secret-url \u{14}link label\u{15}\r";
459 const FIXTURE_TEXTBOXES: &str = "First box\rSecond box\r";
460
461 fn put_u16(bytes: &mut [u8], offset: usize, value: u16) {
462 bytes[offset..offset + 2].copy_from_slice(&value.to_le_bytes());
463 }
464
465 fn put_u32(bytes: &mut [u8], offset: usize, value: u32) {
466 bytes[offset..offset + 4].copy_from_slice(&value.to_le_bytes());
467 }
468
469 fn build_doc_fixture(body: &str, textboxes: &str) -> Vec<u8> {
470 assert!(
471 body.chars()
472 .chain(textboxes.chars())
473 .all(|character| character.len_utf16() == 1),
474 "fixture builder supports only single-unit UTF-16 characters"
475 );
476
477 let body_characters = body.chars().count() as u32;
478 let textbox_characters = textboxes.chars().count() as u32;
479 let all_characters = body_characters + textbox_characters;
480 let text_offset = 1_024usize;
481
482 let mut word_document = vec![0u8; 2_048];
483 put_u16(&mut word_document, 0, 0xA5EC);
484 put_u16(&mut word_document, 10, 0);
485 put_u16(&mut word_document, 32, 14);
486 put_u16(&mut word_document, 62, 22);
487
488 let fib_rg_lw_start = 64usize;
489 put_u32(&mut word_document, fib_rg_lw_start + 12, body_characters);
490 put_u32(&mut word_document, fib_rg_lw_start + 16, 0);
491 put_u32(&mut word_document, fib_rg_lw_start + 20, 0);
492 put_u32(&mut word_document, fib_rg_lw_start + 28, 0);
493 put_u32(&mut word_document, fib_rg_lw_start + 32, 0);
494 put_u32(&mut word_document, fib_rg_lw_start + 36, textbox_characters);
495
496 let fib_rg_fc_start = 154usize;
497 put_u32(&mut word_document, fib_rg_fc_start + 8, 32);
498 put_u32(&mut word_document, fib_rg_fc_start + 12, 2);
499 put_u32(&mut word_document, fib_rg_fc_start + 104, 40);
500 put_u32(&mut word_document, fib_rg_fc_start + 108, 4);
501 put_u32(&mut word_document, fib_rg_fc_start + 264, 64);
502 put_u32(&mut word_document, fib_rg_fc_start + 268, 21);
503
504 for (index, character) in body.chars().chain(textboxes.chars()).enumerate() {
505 let code_unit = character as u16;
506 put_u16(&mut word_document, text_offset + index * 2, code_unit);
507 }
508
509 let mut table = vec![0u8; 128];
510 put_u16(&mut table, 32, 0);
511 table[64] = 0x02;
512 put_u32(&mut table, 65, 16);
513 put_u32(&mut table, 69, 0);
514 put_u32(&mut table, 73, all_characters);
515 put_u32(&mut table, 79, text_offset as u32);
516
517 let mut compound =
518 cfb::CompoundFile::create(Cursor::new(Vec::<u8>::new())).expect("create CFB fixture");
519 {
520 let mut stream = compound
521 .create_stream("/WordDocument")
522 .expect("create WordDocument stream");
523 stream
524 .write_all(&word_document)
525 .expect("write WordDocument stream");
526 }
527 {
528 let mut stream = compound
529 .create_stream("/0Table")
530 .expect("create 0Table stream");
531 stream.write_all(&table).expect("write 0Table stream");
532 }
533 compound.flush().expect("flush CFB fixture");
534 compound.into_inner().into_inner()
535 }
536
537 fn doc_fixture() -> Vec<u8> {
538 build_doc_fixture(FIXTURE_BODY, FIXTURE_TEXTBOXES)
539 }
540
541 #[test]
542 fn formats_are_selected_from_extensions_and_media_types() {
543 assert_eq!(
544 detect_format("report.PDF", "application/octet-stream"),
545 Some(DocumentFormat::Pdf)
546 );
547 assert_eq!(
548 detect_format("legacy.DOC", "application/octet-stream"),
549 Some(DocumentFormat::Doc)
550 );
551 assert_eq!(
552 detect_format("legacy", "application/msword"),
553 Some(DocumentFormat::Doc)
554 );
555 assert_eq!(
556 detect_format(
557 "legacy.doc",
558 "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
559 ),
560 Some(DocumentFormat::Doc)
561 );
562 assert_eq!(
563 detect_format("book.xlsx", "application/octet-stream"),
564 Some(DocumentFormat::Spreadsheet)
565 );
566 assert_eq!(
567 detect_format("notes", "text/plain; charset=utf-8"),
568 Some(DocumentFormat::PlainText)
569 );
570 assert_eq!(detect_format("archive.zip", "application/zip"), None);
571 assert_eq!(DocumentFormat::Doc.as_str(), "doc");
572 }
573
574 #[test]
575 fn plain_text_is_normalized_and_bounded() {
576 let extractor = DocumentExtractor::new(ExtractionLimits {
577 max_bytes: 100,
578 max_characters: 5,
579 })
580 .unwrap();
581 let extracted = extractor
582 .extract(DocumentInput {
583 file_name: "notes.txt".into(),
584 content_type: "text/plain".into(),
585 data: b"hello\r\n\r\n\r\nworld\0".to_vec(),
586 })
587 .unwrap();
588 assert_eq!(extracted.text, "hello");
589 assert!(extracted.truncated);
590 }
591
592 #[test]
593 fn doc_extracts_body_unicode_fields_and_textboxes() {
594 let extracted = DocumentExtractor::default()
595 .extract(DocumentInput {
596 file_name: "legacy.DOC".into(),
597 content_type: "application/msword; charset=binary".into(),
598 data: doc_fixture(),
599 })
600 .unwrap();
601
602 assert_eq!(extracted.format, DocumentFormat::Doc);
603 assert_eq!(extracted.content_type, "application/msword");
604 assert!(extracted.text.contains("# Legacy heading"));
605 assert!(extracted.text.contains("Café 世界 legacy body"));
606 assert!(extracted.text.contains("Visible link label"));
607 assert!(!extracted.text.contains("HYPERLINK"));
608 assert!(!extracted.text.contains("secret-url"));
609 assert_eq!(extracted.text.matches("Text box:\n").count(), 2);
610 assert!(
611 extracted
612 .text
613 .contains("Text box:\nFirst box\n\nText box:\nSecond box")
614 );
615 assert_eq!(extracted.characters, extracted.text.chars().count());
616 assert!(!extracted.truncated);
617 }
618
619 #[test]
620 fn doc_uses_common_character_limit() {
621 let extracted = DocumentExtractor::new(ExtractionLimits {
622 max_bytes: 20 * 1024 * 1024,
623 max_characters: 24,
624 })
625 .unwrap()
626 .extract(DocumentInput {
627 file_name: "legacy.doc".into(),
628 content_type: "application/octet-stream".into(),
629 data: doc_fixture(),
630 })
631 .unwrap();
632
633 assert_eq!(extracted.characters, extracted.text.chars().count());
634 assert!(extracted.characters <= 24);
635 assert!(extracted.truncated);
636 }
637
638 #[test]
639 fn byte_limit_is_checked_before_doc_parsing() {
640 let fixture = doc_fixture();
641 let error = DocumentExtractor::new(ExtractionLimits {
642 max_bytes: fixture.len() - 1,
643 max_characters: 1_000_000,
644 })
645 .unwrap()
646 .extract(DocumentInput {
647 file_name: "legacy.doc".into(),
648 content_type: "application/msword".into(),
649 data: fixture,
650 })
651 .unwrap_err();
652
653 assert_eq!(error.kind(), ErrorKind::InvalidInput);
654 }
655
656 #[test]
657 fn empty_doc_text_uses_common_empty_error() {
658 let error = DocumentExtractor::default()
659 .extract(DocumentInput {
660 file_name: "empty.doc".into(),
661 content_type: "application/msword".into(),
662 data: build_doc_fixture("\r", ""),
663 })
664 .unwrap_err();
665
666 assert_eq!(error.kind(), ErrorKind::EmptyText);
667 assert_eq!(error.to_string(), "document contains no extractable text");
668 }
669
670 #[test]
671 fn malformed_doc_inputs_are_contained_and_sanitized() {
672 let fixture = doc_fixture();
673 let cases = vec![
674 b"not a Word document".to_vec(),
675 b"PK\x03\x04DOCX-like bytes".to_vec(),
676 fixture[..fixture.len() / 2].to_vec(),
677 vec![0xA5; 4096],
678 ];
679
680 for data in cases {
681 let input = DocumentInput {
682 file_name: "malformed.doc".into(),
683 content_type: "application/msword".into(),
684 data,
685 };
686 let result = std::panic::catch_unwind(|| DocumentExtractor::default().extract(input));
687 let error = result
688 .expect("legacy parser panic must not escape")
689 .expect_err("malformed DOC must fail");
690 assert_eq!(error.kind(), ErrorKind::ExtractionFailed);
691 assert!(
692 error
693 .to_string()
694 .starts_with("document could not be converted to text: legacy Word parser")
695 );
696 }
697
698 let secret = "raw-document-secret";
699 let input = DocumentInput {
700 file_name: "malformed.doc".into(),
701 content_type: "application/msword".into(),
702 data: secret.as_bytes().to_vec(),
703 };
704 assert!(!format!("{input:?}").contains(secret));
705 let error = DocumentExtractor::default().extract(input).unwrap_err();
706 assert!(!error.to_string().contains(secret));
707 assert!(!format!("{error:?}").contains(secret));
708 }
709}