kcode-doc-extraction 0.2.0

//! Local PDF, DOC, DOCX, spreadsheet, and plain-text extraction.

#![deny(missing_docs)]
#![forbid(unsafe_code)]

use std::{
    error::Error as StdError,
    fmt,
    io::{Cursor, Read},
};

use calamine::{Reader as _, open_workbook_auto_from_rs};
use quick_xml::{Reader as XmlReader, events::Event as XmlEvent};

/// Supported local document representation.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum DocumentFormat {
    /// Searchable PDF text extraction without OCR.
    Pdf,
    /// Microsoft Word 97-2003 binary document.
    Doc,
    /// Microsoft Word Open XML document.
    Docx,
    /// XLSX, XLS, XLSB, or ODS workbook.
    Spreadsheet,
    /// UTF-8-family text, including CSV, TSV, Markdown, JSON, YAML, and XML.
    PlainText,
}

impl DocumentFormat {
    /// Returns the stable format label.
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Pdf => "pdf",
            Self::Doc => "doc",
            Self::Docx => "docx",
            Self::Spreadsheet => "spreadsheet",
            Self::PlainText => "text",
        }
    }
}

/// Stable document-extraction failure classification.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ErrorKind {
    /// Input metadata, size, or bytes were invalid.
    InvalidInput,
    /// The filename and media type do not identify a supported format.
    UnsupportedFormat,
    /// A local parser could not extract the document.
    ExtractionFailed,
    /// The parser found no readable text.
    EmptyText,
}

/// A local document validation or extraction failure.
#[derive(Debug)]
pub struct Error {
    kind: ErrorKind,
    message: String,
}

impl Error {
    fn new(kind: ErrorKind, message: impl Into<String>) -> Self {
        Self {
            kind,
            message: message.into(),
        }
    }

    /// Returns the stable failure classification.
    pub const fn kind(&self) -> ErrorKind {
        self.kind
    }

    /// Returns the parser or validation detail.
    pub fn message(&self) -> &str {
        &self.message
    }
}

impl fmt::Display for Error {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.message)
    }
}

impl StdError for Error {}

/// Result type returned by document extraction.
pub type Result<T> = std::result::Result<T, Error>;

/// Input and output resource limits.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ExtractionLimits {
    /// Maximum in-memory document bytes.
    pub max_bytes: usize,
    /// Maximum returned Unicode characters.
    pub max_characters: usize,
}

impl Default for ExtractionLimits {
    fn default() -> Self {
        Self {
            max_bytes: 20 * 1024 * 1024,
            max_characters: 1_000_000,
        }
    }
}

/// One complete in-memory document.
#[derive(Clone, Eq, PartialEq)]
pub struct DocumentInput {
    /// Caller-supplied filename used for format detection and returned safely.
    pub file_name: String,
    /// Caller-supplied media type used as a format-detection fallback.
    pub content_type: String,
    /// Complete document bytes.
    pub data: Vec<u8>,
}

impl fmt::Debug for DocumentInput {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("DocumentInput")
            .field("file_name", &self.file_name)
            .field("content_type", &self.content_type)
            .field("bytes", &self.data.len())
            .finish()
    }
}

/// Successfully extracted normalized text and metadata.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ExtractedDocument {
    /// Safe display filename.
    pub file_name: String,
    /// Original normalized media type.
    pub content_type: String,
    /// Detected document format.
    pub format: DocumentFormat,
    /// Normalized, non-empty extracted text.
    pub text: String,
    /// Returned Unicode character count.
    pub characters: usize,
    /// Whether the character limit truncated the text.
    pub truncated: bool,
}

/// Stateless local document extractor with fixed limits.
#[derive(Clone, Debug)]
pub struct DocumentExtractor {
    limits: ExtractionLimits,
}

impl DocumentExtractor {
    /// Constructs an extractor after validating nonzero limits.
    pub fn new(limits: ExtractionLimits) -> Result<Self> {
        if limits.max_bytes == 0 || limits.max_characters == 0 {
            return Err(Error::new(
                ErrorKind::InvalidInput,
                "document byte and character limits must be greater than zero",
            ));
        }
        Ok(Self { limits })
    }

    /// Extracts and normalizes one complete document in the current thread.
    pub fn extract(&self, input: DocumentInput) -> Result<ExtractedDocument> {
        if input.data.is_empty() || input.data.len() > self.limits.max_bytes {
            return Err(Error::new(
                ErrorKind::InvalidInput,
                format!(
                    "document must contain between 1 and {} bytes",
                    self.limits.max_bytes
                ),
            ));
        }
        let content_type = input
            .content_type
            .split(';')
            .next()
            .unwrap_or(&input.content_type)
            .trim()
            .to_ascii_lowercase();
        let format = detect_format(&input.file_name, &content_type).ok_or_else(|| {
            Error::new(
                ErrorKind::UnsupportedFormat,
                "supported documents are PDF, DOC, DOCX, XLSX, XLS, XLSB, ODS, CSV, TSV, and plain text",
            )
        })?;
        let file_name = safe_filename(&input.file_name, format);
        let extracted = extract_text(format, &input.data).map_err(|message| {
            Error::new(
                ErrorKind::ExtractionFailed,
                format!("document could not be converted to text: {message}"),
            )
        })?;
        if extracted.is_empty() {
            let message = if format == DocumentFormat::Pdf {
                "PDF contains no extractable text and may require OCR"
            } else {
                "document contains no extractable text"
            };
            return Err(Error::new(ErrorKind::EmptyText, message));
        }
        let (text, truncated) = truncate_characters(&extracted, self.limits.max_characters);
        let characters = text.chars().count();
        Ok(ExtractedDocument {
            file_name,
            content_type,
            format,
            text,
            characters,
            truncated,
        })
    }
}

impl Default for DocumentExtractor {
    fn default() -> Self {
        Self::new(ExtractionLimits::default()).expect("default document limits are valid")
    }
}

fn detect_format(file_name: &str, content_type: &str) -> Option<DocumentFormat> {
    let extension = file_name
        .rsplit_once('.')
        .map(|(_, extension)| extension.to_ascii_lowercase());
    match extension.as_deref() {
        Some("pdf") => Some(DocumentFormat::Pdf),
        Some("doc") => Some(DocumentFormat::Doc),
        Some("docx") => Some(DocumentFormat::Docx),
        Some("xlsx" | "xls" | "xlsb" | "ods") => Some(DocumentFormat::Spreadsheet),
        Some("csv" | "tsv" | "txt" | "md" | "json" | "yaml" | "yml" | "xml") => {
            Some(DocumentFormat::PlainText)
        }
        _ if content_type == "application/pdf" => Some(DocumentFormat::Pdf),
        _ if content_type == "application/msword" => Some(DocumentFormat::Doc),
        _ if content_type
            == "application/vnd.openxmlformats-officedocument.wordprocessingml.document" =>
        {
            Some(DocumentFormat::Docx)
        }
        _ if matches!(
            content_type,
            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                | "application/vnd.ms-excel"
                | "application/vnd.ms-excel.sheet.binary.macroenabled.12"
                | "application/vnd.oasis.opendocument.spreadsheet"
        ) =>
        {
            Some(DocumentFormat::Spreadsheet)
        }
        _ if content_type.starts_with("text/")
            || matches!(
                content_type,
                "application/json" | "application/xml" | "application/yaml" | "application/x-yaml"
            ) =>
        {
            Some(DocumentFormat::PlainText)
        }
        _ => None,
    }
}

fn safe_filename(value: &str, format: DocumentFormat) -> String {
    let cleaned = value
        .chars()
        .filter(|character| {
            !character.is_control() && !matches!(character, '/' | '\\' | ':' | '\0')
        })
        .take(200)
        .collect::<String>();
    if cleaned.trim().is_empty() {
        format!("document.{}", format.as_str())
    } else {
        cleaned
    }
}

fn local_xml_name(name: &[u8]) -> &[u8] {
    name.rsplit(|byte| *byte == b':').next().unwrap_or(name)
}

fn extract_doc_text(bytes: &[u8]) -> std::result::Result<String, String> {
    let document = match std::panic::catch_unwind(|| unword::parse_doc_with_options(bytes, true)) {
        Ok(Ok(document)) => document,
        Ok(Err(_)) => return Err("legacy Word parser rejected the document".to_owned()),
        Err(_) => {
            return Err("legacy Word parser panicked while reading the document".to_owned());
        }
    };

    let mut output = document.body_text;
    for text_box in document.textboxes {
        let text_box = text_box.trim();
        if text_box.is_empty() {
            continue;
        }
        if !output.trim().is_empty() {
            output.push_str("\n\n");
        }
        output.push_str("Text box:\n");
        output.push_str(text_box);
    }
    Ok(output)
}

fn extract_docx_text(bytes: &[u8]) -> std::result::Result<String, String> {
    let mut archive =
        zip::ZipArchive::new(Cursor::new(bytes)).map_err(|error| error.to_string())?;
    let mut document = archive
        .by_name("word/document.xml")
        .map_err(|error| format!("word/document.xml: {error}"))?;
    let mut xml = String::new();
    document
        .read_to_string(&mut xml)
        .map_err(|error| error.to_string())?;
    let mut reader = XmlReader::from_str(&xml);
    let mut output = String::new();
    let mut in_text = false;
    let mut in_cell = false;
    loop {
        match reader.read_event() {
            Ok(XmlEvent::Start(event)) => match local_xml_name(event.name().as_ref()) {
                b"t" => in_text = true,
                b"tc" => in_cell = true,
                _ => {}
            },
            Ok(XmlEvent::Empty(event)) => match local_xml_name(event.name().as_ref()) {
                b"tab" => output.push('\t'),
                b"br" | b"cr" => output.push('\n'),
                _ => {}
            },
            Ok(XmlEvent::Text(text)) if in_text => {
                let decoded = text.decode().map_err(|error| error.to_string())?;
                output.push_str(&decoded);
            }
            Ok(XmlEvent::CData(text)) if in_text => {
                let decoded = text.decode().map_err(|error| error.to_string())?;
                output.push_str(&decoded);
            }
            Ok(XmlEvent::GeneralRef(reference)) if in_text => {
                let decoded = reference.decode().map_err(|error| error.to_string())?;
                let escaped = format!("&{decoded};");
                let resolved =
                    quick_xml::escape::unescape(&escaped).map_err(|error| error.to_string())?;
                output.push_str(&resolved);
            }
            Ok(XmlEvent::End(event)) => match local_xml_name(event.name().as_ref()) {
                b"t" => in_text = false,
                b"p" if in_cell => output.push_str(" / "),
                b"p" => output.push('\n'),
                b"tc" => {
                    if output.ends_with(" / ") {
                        output.truncate(output.len() - 3);
                    }
                    output.push('\t');
                    in_cell = false;
                }
                b"tr" => {
                    if output.ends_with('\t') {
                        output.pop();
                    }
                    output.push('\n');
                }
                _ => {}
            },
            Ok(XmlEvent::Eof) => break,
            Err(error) => return Err(error.to_string()),
            _ => {}
        }
    }
    Ok(output)
}

fn extract_spreadsheet_text(bytes: &[u8]) -> std::result::Result<String, String> {
    let mut workbook = open_workbook_auto_from_rs(Cursor::new(bytes.to_vec()))
        .map_err(|error| error.to_string())?;
    let sheet_names = workbook.sheet_names().to_vec();
    let mut output = String::new();
    for (sheet_index, sheet_name) in sheet_names.iter().enumerate() {
        if sheet_index > 0 {
            output.push('\n');
        }
        output.push_str("Sheet: ");
        output.push_str(sheet_name);
        output.push('\n');
        let range = workbook
            .worksheet_range(sheet_name)
            .map_err(|error| format!("{sheet_name}: {error}"))?;
        for row in range.rows() {
            let line = row
                .iter()
                .map(ToString::to_string)
                .map(|cell| cell.replace(['\t', '\r', '\n'], " "))
                .collect::<Vec<_>>()
                .join("\t");
            output.push_str(line.trim_end());
            output.push('\n');
        }
    }
    Ok(output)
}

fn normalize_text(value: String) -> String {
    let normalized = value
        .replace("\r\n", "\n")
        .replace('\r', "\n")
        .replace('\0', "");
    let mut output = String::with_capacity(normalized.len());
    let mut blank_lines = 0;
    for line in normalized.lines() {
        let line = line.trim_end();
        if line.trim().is_empty() {
            blank_lines += 1;
            if blank_lines > 1 {
                continue;
            }
        } else {
            blank_lines = 0;
        }
        output.push_str(line);
        output.push('\n');
    }
    output.trim().to_owned()
}

fn extract_text(format: DocumentFormat, bytes: &[u8]) -> std::result::Result<String, String> {
    let extracted = match format {
        DocumentFormat::Pdf => {
            pdf_extract::extract_text_from_mem(bytes).map_err(|error| error.to_string())?
        }
        DocumentFormat::Doc => extract_doc_text(bytes)?,
        DocumentFormat::Docx => extract_docx_text(bytes)?,
        DocumentFormat::Spreadsheet => extract_spreadsheet_text(bytes)?,
        DocumentFormat::PlainText => String::from_utf8_lossy(bytes)
            .trim_start_matches('\u{feff}')
            .to_owned(),
    };
    Ok(normalize_text(extracted))
}

fn truncate_characters(value: &str, limit: usize) -> (String, bool) {
    let mut iter = value.char_indices();
    let Some((boundary, _)) = iter.nth(limit) else {
        return (value.to_owned(), false);
    };
    (value[..boundary].trim_end().to_owned(), true)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write as _;

    const FIXTURE_BODY: &str = "# Legacy heading\rCafé 世界 legacy body\rVisible \u{13} HYPERLINK secret-url \u{14}link label\u{15}\r";
    const FIXTURE_TEXTBOXES: &str = "First box\rSecond box\r";

    fn put_u16(bytes: &mut [u8], offset: usize, value: u16) {
        bytes[offset..offset + 2].copy_from_slice(&value.to_le_bytes());
    }

    fn put_u32(bytes: &mut [u8], offset: usize, value: u32) {
        bytes[offset..offset + 4].copy_from_slice(&value.to_le_bytes());
    }

    fn build_doc_fixture(body: &str, textboxes: &str) -> Vec<u8> {
        assert!(
            body.chars()
                .chain(textboxes.chars())
                .all(|character| character.len_utf16() == 1),
            "fixture builder supports only single-unit UTF-16 characters"
        );

        let body_characters = body.chars().count() as u32;
        let textbox_characters = textboxes.chars().count() as u32;
        let all_characters = body_characters + textbox_characters;
        let text_offset = 1_024usize;

        let mut word_document = vec![0u8; 2_048];
        put_u16(&mut word_document, 0, 0xA5EC);
        put_u16(&mut word_document, 10, 0);
        put_u16(&mut word_document, 32, 14);
        put_u16(&mut word_document, 62, 22);

        let fib_rg_lw_start = 64usize;
        put_u32(&mut word_document, fib_rg_lw_start + 12, body_characters);
        put_u32(&mut word_document, fib_rg_lw_start + 16, 0);
        put_u32(&mut word_document, fib_rg_lw_start + 20, 0);
        put_u32(&mut word_document, fib_rg_lw_start + 28, 0);
        put_u32(&mut word_document, fib_rg_lw_start + 32, 0);
        put_u32(&mut word_document, fib_rg_lw_start + 36, textbox_characters);

        let fib_rg_fc_start = 154usize;
        put_u32(&mut word_document, fib_rg_fc_start + 8, 32);
        put_u32(&mut word_document, fib_rg_fc_start + 12, 2);
        put_u32(&mut word_document, fib_rg_fc_start + 104, 40);
        put_u32(&mut word_document, fib_rg_fc_start + 108, 4);
        put_u32(&mut word_document, fib_rg_fc_start + 264, 64);
        put_u32(&mut word_document, fib_rg_fc_start + 268, 21);

        for (index, character) in body.chars().chain(textboxes.chars()).enumerate() {
            let code_unit = character as u16;
            put_u16(&mut word_document, text_offset + index * 2, code_unit);
        }

        let mut table = vec![0u8; 128];
        put_u16(&mut table, 32, 0);
        table[64] = 0x02;
        put_u32(&mut table, 65, 16);
        put_u32(&mut table, 69, 0);
        put_u32(&mut table, 73, all_characters);
        put_u32(&mut table, 79, text_offset as u32);

        let mut compound =
            cfb::CompoundFile::create(Cursor::new(Vec::<u8>::new())).expect("create CFB fixture");
        {
            let mut stream = compound
                .create_stream("/WordDocument")
                .expect("create WordDocument stream");
            stream
                .write_all(&word_document)
                .expect("write WordDocument stream");
        }
        {
            let mut stream = compound
                .create_stream("/0Table")
                .expect("create 0Table stream");
            stream.write_all(&table).expect("write 0Table stream");
        }
        compound.flush().expect("flush CFB fixture");
        compound.into_inner().into_inner()
    }

    fn doc_fixture() -> Vec<u8> {
        build_doc_fixture(FIXTURE_BODY, FIXTURE_TEXTBOXES)
    }

    #[test]
    fn formats_are_selected_from_extensions_and_media_types() {
        assert_eq!(
            detect_format("report.PDF", "application/octet-stream"),
            Some(DocumentFormat::Pdf)
        );
        assert_eq!(
            detect_format("legacy.DOC", "application/octet-stream"),
            Some(DocumentFormat::Doc)
        );
        assert_eq!(
            detect_format("legacy", "application/msword"),
            Some(DocumentFormat::Doc)
        );
        assert_eq!(
            detect_format(
                "legacy.doc",
                "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
            ),
            Some(DocumentFormat::Doc)
        );
        assert_eq!(
            detect_format("book.xlsx", "application/octet-stream"),
            Some(DocumentFormat::Spreadsheet)
        );
        assert_eq!(
            detect_format("notes", "text/plain; charset=utf-8"),
            Some(DocumentFormat::PlainText)
        );
        assert_eq!(detect_format("archive.zip", "application/zip"), None);
        assert_eq!(DocumentFormat::Doc.as_str(), "doc");
    }

    #[test]
    fn plain_text_is_normalized_and_bounded() {
        let extractor = DocumentExtractor::new(ExtractionLimits {
            max_bytes: 100,
            max_characters: 5,
        })
        .unwrap();
        let extracted = extractor
            .extract(DocumentInput {
                file_name: "notes.txt".into(),
                content_type: "text/plain".into(),
                data: b"hello\r\n\r\n\r\nworld\0".to_vec(),
            })
            .unwrap();
        assert_eq!(extracted.text, "hello");
        assert!(extracted.truncated);
    }

    #[test]
    fn doc_extracts_body_unicode_fields_and_textboxes() {
        let extracted = DocumentExtractor::default()
            .extract(DocumentInput {
                file_name: "legacy.DOC".into(),
                content_type: "application/msword; charset=binary".into(),
                data: doc_fixture(),
            })
            .unwrap();

        assert_eq!(extracted.format, DocumentFormat::Doc);
        assert_eq!(extracted.content_type, "application/msword");
        assert!(extracted.text.contains("# Legacy heading"));
        assert!(extracted.text.contains("Café 世界 legacy body"));
        assert!(extracted.text.contains("Visible link label"));
        assert!(!extracted.text.contains("HYPERLINK"));
        assert!(!extracted.text.contains("secret-url"));
        assert_eq!(extracted.text.matches("Text box:\n").count(), 2);
        assert!(
            extracted
                .text
                .contains("Text box:\nFirst box\n\nText box:\nSecond box")
        );
        assert_eq!(extracted.characters, extracted.text.chars().count());
        assert!(!extracted.truncated);
    }

    #[test]
    fn doc_uses_common_character_limit() {
        let extracted = DocumentExtractor::new(ExtractionLimits {
            max_bytes: 20 * 1024 * 1024,
            max_characters: 24,
        })
        .unwrap()
        .extract(DocumentInput {
            file_name: "legacy.doc".into(),
            content_type: "application/octet-stream".into(),
            data: doc_fixture(),
        })
        .unwrap();

        assert_eq!(extracted.characters, extracted.text.chars().count());
        assert!(extracted.characters <= 24);
        assert!(extracted.truncated);
    }

    #[test]
    fn byte_limit_is_checked_before_doc_parsing() {
        let fixture = doc_fixture();
        let error = DocumentExtractor::new(ExtractionLimits {
            max_bytes: fixture.len() - 1,
            max_characters: 1_000_000,
        })
        .unwrap()
        .extract(DocumentInput {
            file_name: "legacy.doc".into(),
            content_type: "application/msword".into(),
            data: fixture,
        })
        .unwrap_err();

        assert_eq!(error.kind(), ErrorKind::InvalidInput);
    }

    #[test]
    fn empty_doc_text_uses_common_empty_error() {
        let error = DocumentExtractor::default()
            .extract(DocumentInput {
                file_name: "empty.doc".into(),
                content_type: "application/msword".into(),
                data: build_doc_fixture("\r", ""),
            })
            .unwrap_err();

        assert_eq!(error.kind(), ErrorKind::EmptyText);
        assert_eq!(error.to_string(), "document contains no extractable text");
    }

    #[test]
    fn malformed_doc_inputs_are_contained_and_sanitized() {
        let fixture = doc_fixture();
        let cases = vec![
            b"not a Word document".to_vec(),
            b"PK\x03\x04DOCX-like bytes".to_vec(),
            fixture[..fixture.len() / 2].to_vec(),
            vec![0xA5; 4096],
        ];

        for data in cases {
            let input = DocumentInput {
                file_name: "malformed.doc".into(),
                content_type: "application/msword".into(),
                data,
            };
            let result = std::panic::catch_unwind(|| DocumentExtractor::default().extract(input));
            let error = result
                .expect("legacy parser panic must not escape")
                .expect_err("malformed DOC must fail");
            assert_eq!(error.kind(), ErrorKind::ExtractionFailed);
            assert!(
                error
                    .to_string()
                    .starts_with("document could not be converted to text: legacy Word parser")
            );
        }

        let secret = "raw-document-secret";
        let input = DocumentInput {
            file_name: "malformed.doc".into(),
            content_type: "application/msword".into(),
            data: secret.as_bytes().to_vec(),
        };
        assert!(!format!("{input:?}").contains(secret));
        let error = DocumentExtractor::default().extract(input).unwrap_err();
        assert!(!error.to_string().contains(secret));
        assert!(!format!("{error:?}").contains(secret));
    }
}