repair malformed PDF containers

src/bin/detect_pdf.rs

@@ -1,8 +1,12 @@
 //! CLI tool for detecting PDF type (text-based vs scanned)
 
-use pdf_inspector::{detect_pdf_type, process_pdf_with_options, PdfOptions, PdfType, ProcessMode};
+use pdf_inspector::{
+    detect_pdf_type, detector::estimate_page_count_from_bytes, process_pdf_with_options,
+    PdfOptions, PdfType, ProcessMode,
+};
 use std::env;
 use std::fmt::Write;
+use std::fs;
 use std::process;
 use std::time::Instant;
 
@@ -64,6 +68,32 @@ fn pdf_type_str(pdf_type: &PdfType) -> &'static str {
     }
 }
 
+fn page_count_hint(pdf_path: &str) -> Option<u32> {
+    fs::read(pdf_path)
+        .ok()
+        .map(|bytes| estimate_page_count_from_bytes(&bytes))
+        .filter(|&count| count > 0)
+}
+
+fn print_error(e: &pdf_inspector::PdfError, pdf_path: &str, json_output: bool) {
+    if json_output {
+        if let Some(count) = page_count_hint(pdf_path) {
+            println!(
+                r#"{{"error":"{}","page_count_hint":{}}}"#,
+                json_escape(&e.to_string()),
+                count
+            );
+        } else {
+            println!(r#"{{"error":"{}"}}"#, json_escape(&e.to_string()));
+        }
+    } else {
+        eprintln!("Error: {}", e);
+        if let Some(count) = page_count_hint(pdf_path) {
+            eprintln!("Page count hint: {}", count);
+        }
+    }
+}
+
 fn run_analyze(pdf_path: &str, json_output: bool, start: Instant) {
     match process_pdf_with_options(pdf_path, PdfOptions::new().mode(ProcessMode::Analyze)) {
         Ok(result) => {
@@ -135,11 +165,7 @@ fn run_analyze(pdf_path: &str, json_output: bool, start: Instant) {
             }
         }
         Err(e) => {
-            if json_output {
-                println!(r#"{{"error":"{}"}}"#, e);
-            } else {
-                eprintln!("Error: {}", e);
-            }
+            print_error(&e, pdf_path, json_output);
             process::exit(1);
         }
     }
@@ -236,11 +262,7 @@ fn run_detect_only(pdf_path: &str, json_output: bool, start: Instant) {
             }
         }
         Err(e) => {
-            if json_output {
-                println!(r#"{{"error":"{}"}}"#, e);
-            } else {
-                eprintln!("Error: {}", e);
-            }
+            print_error(&e, pdf_path, json_output);
             process::exit(1);
         }
     }

src/detector.rs

@@ -97,26 +97,9 @@ pub fn detect_pdf_type_with_config<P: AsRef<Path>>(
 ) -> Result<PdfTypeResult, PdfError> {
     crate::validate_pdf_file(&path)?;
 
-    // First, load metadata only (fast operation)
-    let metadata = match Document::load_metadata(&path) {
-        Ok(m) => m,
-        Err(ref e) if crate::is_encrypted_lopdf_error(e) => {
-            Document::load_metadata_with_password(&path, "")?
-        }
-        Err(e) => return Err(e.into()),
-    };
+    let (doc, page_count) = crate::load_document_from_path(&path)?;
 
-    // Then load the full document for content inspection
-    // We use filtered loading to skip heavy objects we don't need
-    let doc = match Document::load(&path) {
-        Ok(d) => d,
-        Err(ref e) if crate::is_encrypted_lopdf_error(e) => {
-            Document::load_with_password(&path, "")?
-        }
-        Err(e) => return Err(e.into()),
-    };
-
-    detect_from_document(&doc, metadata.page_count, &config)
+    detect_from_document(&doc, page_count, &config)
 }
 
 /// Detect PDF type from memory buffer
@@ -131,25 +114,64 @@ pub fn detect_pdf_type_mem_with_config(
 ) -> Result<PdfTypeResult, PdfError> {
     crate::validate_pdf_bytes(buffer)?;
 
-    // Load metadata first (fast)
-    let metadata = match Document::load_metadata_mem(buffer) {
-        Ok(m) => m,
-        Err(ref e) if crate::is_encrypted_lopdf_error(e) => {
-            Document::load_metadata_mem_with_password(buffer, "")?
-        }
-        Err(e) => return Err(e.into()),
-    };
+    let (doc, page_count) = crate::load_document_from_mem(buffer)?;
 
-    // Load document for inspection
-    let doc = match Document::load_mem(buffer) {
-        Ok(d) => d,
-        Err(ref e) if crate::is_encrypted_lopdf_error(e) => {
-            Document::load_mem_with_options(buffer, lopdf::LoadOptions::with_password(""))?
-        }
-        Err(e) => return Err(e.into()),
-    };
+    detect_from_document(&doc, page_count, &config)
+}
 
-    detect_from_document(&doc, metadata.page_count, &config)
+/// Heuristic page-count fallback for malformed PDFs that cannot be parsed.
+///
+/// This scans raw bytes for page dictionaries (`/Type /Page`) while excluding
+/// the page tree node (`/Type /Pages`). It is intended as a low-confidence hint
+/// for diagnostics; parsed page-tree counts remain authoritative.
+pub fn estimate_page_count_from_bytes(buffer: &[u8]) -> u32 {
+    let mut count = 0u32;
+    let mut pos = 0usize;
+
+    while let Some(rel_idx) = find_bytes(&buffer[pos..], b"/Type") {
+        let mut value_pos = pos + rel_idx + b"/Type".len();
+        value_pos = skip_pdf_whitespace(buffer, value_pos);
+
+        if buffer.get(value_pos) == Some(&b'/') {
+            let name_start = value_pos + 1;
+            let name_end = name_start + b"Page".len();
+            if name_end <= buffer.len()
+                && &buffer[name_start..name_end] == b"Page"
+                && buffer
+                    .get(name_end)
+                    .is_none_or(|b| is_pdf_name_delimiter(*b))
+            {
+                count += 1;
+            }
+        }
+
+        pos += rel_idx + b"/Type".len();
+    }
+
+    count
+}
+
+fn find_bytes(haystack: &[u8], needle: &[u8]) -> Option<usize> {
+    haystack.windows(needle.len()).position(|w| w == needle)
+}
+
+fn skip_pdf_whitespace(buffer: &[u8], mut pos: usize) -> usize {
+    while pos < buffer.len() && is_pdf_whitespace(buffer[pos]) {
+        pos += 1;
+    }
+    pos
+}
+
+fn is_pdf_whitespace(byte: u8) -> bool {
+    matches!(byte, b'\0' | b'\t' | b'\n' | 0x0C | b'\r' | b' ')
+}
+
+fn is_pdf_name_delimiter(byte: u8) -> bool {
+    is_pdf_whitespace(byte)
+        || matches!(
+            byte,
+            b'(' | b')' | b'<' | b'>' | b'[' | b']' | b'{' | b'}' | b'/' | b'%'
+        )
 }
 
 /// Detection logic on a pre-loaded document.

src/extractor/mod.rs

@@ -36,26 +36,14 @@ pub(crate) use layout::ColumnRegion;
 /// Extract text from PDF file as plain string
 pub fn extract_text<P: AsRef<Path>>(path: P) -> Result<String, PdfError> {
     crate::validate_pdf_file(&path)?;
-    let doc = match Document::load(&path) {
-        Ok(d) => d,
-        Err(ref e) if crate::is_encrypted_lopdf_error(e) => {
-            Document::load_with_password(&path, "")?
-        }
-        Err(e) => return Err(e.into()),
-    };
+    let (doc, _) = crate::load_document_from_path(&path)?;
     extract_text_from_doc(&doc)
 }
 
 /// Extract text from PDF memory buffer
 pub fn extract_text_mem(buffer: &[u8]) -> Result<String, PdfError> {
     crate::validate_pdf_bytes(buffer)?;
-    let doc = match Document::load_mem(buffer) {
-        Ok(d) => d,
-        Err(ref e) if crate::is_encrypted_lopdf_error(e) => {
-            Document::load_mem_with_options(buffer, lopdf::LoadOptions::with_password(""))?
-        }
-        Err(e) => return Err(e.into()),
-    };
+    let (doc, _) = crate::load_document_from_mem(buffer)?;
     extract_text_from_doc(&doc)
 }
 
@@ -91,13 +79,7 @@ pub(crate) fn extract_text_with_positions_and_rects<P: AsRef<Path>>(
     page_filter: Option<&HashSet<u32>>,
 ) -> Result<PageExtraction, PdfError> {
     crate::validate_pdf_file(&path)?;
-    let doc = match Document::load(&path) {
-        Ok(d) => d,
-        Err(ref e) if crate::is_encrypted_lopdf_error(e) => {
-            Document::load_with_password(&path, "")?
-        }
-        Err(e) => return Err(e.into()),
-    };
+    let (doc, _) = crate::load_document_from_path(&path)?;
     let font_cmaps = FontCMaps::from_doc(&doc);
     let (extraction, _thresholds, _gid_pages) =
         extract_positioned_text_from_doc(&doc, &font_cmaps, page_filter)?;
@@ -124,13 +106,7 @@ pub(crate) fn extract_text_with_positions_mem_and_rects(
     page_filter: Option<&HashSet<u32>>,
 ) -> Result<PageExtraction, PdfError> {
     crate::validate_pdf_bytes(buffer)?;
-    let doc = match Document::load_mem(buffer) {
-        Ok(d) => d,
-        Err(ref e) if crate::is_encrypted_lopdf_error(e) => {
-            Document::load_mem_with_options(buffer, lopdf::LoadOptions::with_password(""))?
-        }
-        Err(e) => return Err(e.into()),
-    };
+    let (doc, _) = crate::load_document_from_mem(buffer)?;
     let font_cmaps = FontCMaps::from_doc(&doc);
     let (extraction, _thresholds, _gid_pages) =
         extract_positioned_text_from_doc(&doc, &font_cmaps, page_filter)?;

src/lib.rs

@@ -1381,30 +1381,133 @@ fn tsr_region_contains_item(item: &TextItem, bounds: RegionBounds) -> bool {
 /// `Document::load_metadata` for page count + `Document::load` for content
 /// are combined here, but lopdf loads the full doc in `load()` so we extract
 /// page count from it directly to avoid the metadata-only round-trip.
-fn load_document_from_path<P: AsRef<Path>>(path: P) -> Result<(Document, u32), PdfError> {
+pub(crate) fn load_document_from_path<P: AsRef<Path>>(
+    path: P,
+) -> Result<(Document, u32), PdfError> {
     let buffer = std::fs::read(&path)?;
     load_document_from_mem(&buffer)
 }
 
 /// Load a PDF from a memory buffer.
-fn load_document_from_mem(buffer: &[u8]) -> Result<(Document, u32), PdfError> {
+pub(crate) fn load_document_from_mem(buffer: &[u8]) -> Result<(Document, u32), PdfError> {
     // Fix malformed struct element names before parsing. Some PDF generators
     // write bare names (/S Code) instead of proper PDF names (/S /Code), which
     // causes lopdf to silently drop the entire object.
     let fixed = structure_tree::fix_bare_struct_names(buffer);
     let buf = fixed.as_ref();
 
-    let doc = match Document::load_mem(buf) {
-        Ok(d) => d,
-        Err(ref e) if is_encrypted_lopdf_error(e) => {
-            Document::load_mem_with_options(buf, lopdf::LoadOptions::with_password(""))?
+    let doc = match load_document_bytes(buf) {
+        Ok(doc) => doc,
+        Err(first_err) => {
+            for repaired in repair_pdf_container_candidates(buf) {
+                match load_document_bytes(&repaired) {
+                    Ok(doc) => {
+                        log::debug!("loaded PDF after repairing malformed container bytes");
+                        let page_count = doc.get_pages().len() as u32;
+                        return Ok((doc, page_count));
+                    }
+                    Err(e) => {
+                        if is_encrypted_lopdf_error(&e) {
+                            return Err(e.into());
+                        }
+                    }
+                }
+            }
+            return Err(first_err.into());
         }
-        Err(e) => return Err(e.into()),
     };
     let page_count = doc.get_pages().len() as u32;
     Ok((doc, page_count))
 }
 
+fn load_document_bytes(buf: &[u8]) -> Result<Document, lopdf::Error> {
+    match Document::load_mem(buf) {
+        Ok(doc) => Ok(doc),
+        Err(ref e) if is_encrypted_lopdf_error(e) => {
+            Document::load_mem_with_options(buf, lopdf::LoadOptions::with_password(""))
+        }
+        Err(e) => Err(e),
+    }
+}
+
+fn repair_pdf_container_candidates(buf: &[u8]) -> Vec<Vec<u8>> {
+    let mut candidates = Vec::new();
+
+    add_repair_candidate(&mut candidates, append_missing_eof_marker(buf), buf);
+
+    let stripped = strip_leading_pdf_container_bytes(buf);
+    if let Some(stripped_buf) = stripped.as_deref() {
+        add_repair_candidate(&mut candidates, Some(stripped_buf.to_vec()), buf);
+        add_repair_candidate(
+            &mut candidates,
+            append_missing_eof_marker(stripped_buf),
+            buf,
+        );
+    }
+
+    candidates
+}
+
+fn add_repair_candidate(
+    candidates: &mut Vec<Vec<u8>>,
+    candidate: Option<Vec<u8>>,
+    original: &[u8],
+) {
+    let Some(candidate) = candidate else {
+        return;
+    };
+    if candidate.as_slice() == original {
+        return;
+    }
+    if candidates.iter().any(|existing| existing == &candidate) {
+        return;
+    }
+    candidates.push(candidate);
+}
+
+fn append_missing_eof_marker(buf: &[u8]) -> Option<Vec<u8>> {
+    if contains_recent_eof_marker(buf) {
+        return None;
+    }
+
+    let mut end = buf.len();
+    while end > 0 && buf[end - 1].is_ascii_whitespace() {
+        end -= 1;
+    }
+
+    if !buf[..end].ends_with(b"%%EO") {
+        return None;
+    }
+
+    let mut repaired = Vec::with_capacity(end + 2);
+    repaired.extend_from_slice(&buf[..end]);
+    repaired.extend_from_slice(b"F\n");
+    Some(repaired)
+}
+
+fn contains_recent_eof_marker(buf: &[u8]) -> bool {
+    let start = buf.len().saturating_sub(1024);
+    buf[start..].windows(b"%%EOF".len()).any(|w| w == b"%%EOF")
+}
+
+fn strip_leading_pdf_container_bytes(buf: &[u8]) -> Option<Vec<u8>> {
+    let mut start = if buf.starts_with(&[0xEF, 0xBB, 0xBF]) {
+        3
+    } else {
+        0
+    };
+
+    while start < buf.len() && buf[start].is_ascii_whitespace() {
+        start += 1;
+    }
+
+    if start > 0 && buf[start..].starts_with(b"%PDF-") {
+        Some(buf[start..].to_vec())
+    } else {
+        None
+    }
+}
+
 /// Core processing pipeline operating on a pre-loaded document.
 fn process_document(
     doc: Document,

tests/integration_tests.rs

@@ -1,6 +1,6 @@
 //! Integration tests for pdf-to-markdown library
 
-use pdf_inspector::detector::{DetectionConfig, ScanStrategy};
+use pdf_inspector::detector::{estimate_page_count_from_bytes, DetectionConfig, ScanStrategy};
 use pdf_inspector::extractor::group_into_lines;
 use pdf_inspector::types::TextLine;
 use pdf_inspector::{
@@ -11,6 +11,83 @@ use pdf_inspector::{
 };
 use std::collections::HashSet;
 
+fn make_minimal_text_pdf() -> Vec<u8> {
+    let mut pdf = b"%PDF-1.4\n".to_vec();
+    let mut offsets = vec![0usize];
+
+    fn add_object(pdf: &mut Vec<u8>, offsets: &mut Vec<usize>, id: usize, body: &str) {
+        offsets.push(pdf.len());
+        pdf.extend_from_slice(format!("{id} 0 obj\n").as_bytes());
+        pdf.extend_from_slice(body.as_bytes());
+        pdf.extend_from_slice(b"\nendobj\n");
+    }
+
+    add_object(
+        &mut pdf,
+        &mut offsets,
+        1,
+        "<< /Type /Catalog /Pages 2 0 R >>",
+    );
+    add_object(
+        &mut pdf,
+        &mut offsets,
+        2,
+        "<< /Type /Pages /Kids [3 0 R] /Count 1 >>",
+    );
+    add_object(
+        &mut pdf,
+        &mut offsets,
+        3,
+        "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Resources << /Font << /F1 5 0 R >> >> /Contents 4 0 R >>",
+    );
+
+    let content = "BT /F1 12 Tf 100 700 Td (Hello World) Tj 0 -14 Td (Second Line) Tj 0 -14 Td (Third Line) Tj ET";
+    add_object(
+        &mut pdf,
+        &mut offsets,
+        4,
+        &format!(
+            "<< /Length {} >>\nstream\n{}\nendstream",
+            content.len(),
+            content
+        ),
+    );
+    add_object(
+        &mut pdf,
+        &mut offsets,
+        5,
+        "<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>",
+    );
+
+    let xref_start = pdf.len();
+    pdf.extend_from_slice(format!("xref\n0 {}\n", offsets.len()).as_bytes());
+    pdf.extend_from_slice(b"0000000000 65535 f \n");
+    for offset in offsets.iter().skip(1) {
+        pdf.extend_from_slice(format!("{offset:010} 00000 n \n").as_bytes());
+    }
+    pdf.extend_from_slice(
+        format!(
+            "trailer\n<< /Size {} /Root 1 0 R >>\nstartxref\n{}\n%%EOF",
+            offsets.len(),
+            xref_start
+        )
+        .as_bytes(),
+    );
+
+    pdf
+}
+
+fn truncate_eof_marker(mut pdf: Vec<u8>) -> Vec<u8> {
+    assert!(pdf.ends_with(b"%%EOF"));
+    pdf.pop();
+    pdf
+}
+
+fn add_leading_tab(mut pdf: Vec<u8>) -> Vec<u8> {
+    pdf.insert(0, b'\t');
+    pdf
+}
+
 // Helper to create test TextItems
 fn make_text_item(text: &str, x: f32, y: f32, font_size: f32, page: u32) -> TextItem {
     use pdf_inspector::types::ItemType;
@@ -826,6 +903,73 @@ fn test_bom_prefixed_pdf_header_not_rejected() {
     }
 }
 
+#[test]
+fn test_process_pdf_mem_repairs_truncated_eof_marker() {
+    let pdf = truncate_eof_marker(make_minimal_text_pdf());
+
+    let result = process_pdf_mem(&pdf).expect("truncated %%EO marker should be repaired");
+
+    assert_eq!(result.pdf_type, PdfType::TextBased);
+    assert_eq!(result.page_count, 1);
+    assert!(
+        result
+            .markdown
+            .as_deref()
+            .unwrap_or_default()
+            .contains("Hello World"),
+        "repaired PDF should still extract text"
+    );
+}
+
+#[test]
+fn test_process_pdf_mem_repairs_leading_tab_and_truncated_eof() {
+    let pdf = add_leading_tab(truncate_eof_marker(make_minimal_text_pdf()));
+
+    let result = process_pdf_mem(&pdf).expect("leading whitespace + %%EO should be repaired");
+
+    assert_eq!(result.pdf_type, PdfType::TextBased);
+    assert_eq!(result.page_count, 1);
+    assert!(
+        result
+            .markdown
+            .as_deref()
+            .unwrap_or_default()
+            .contains("Hello World"),
+        "repaired PDF should still extract text"
+    );
+}
+
+#[test]
+fn test_detect_pdf_type_repairs_container_from_path() {
+    let pdf = add_leading_tab(truncate_eof_marker(make_minimal_text_pdf()));
+    let dir = tempfile::tempdir().unwrap();
+    let path = dir.path().join("broken-container.pdf");
+    std::fs::write(&path, pdf).unwrap();
+
+    let result = detect_pdf_type(&path).expect("detector should use shared repair loader");
+
+    assert_eq!(result.pdf_type, PdfType::TextBased);
+    assert_eq!(result.page_count, 1);
+    assert_eq!(result.pages_with_text, 1);
+}
+
+#[test]
+fn test_extract_text_mem_uses_container_repair() {
+    let pdf = truncate_eof_marker(make_minimal_text_pdf());
+
+    let text = pdf_inspector::extractor::extract_text_mem(&pdf)
+        .expect("plain text extraction should use shared repair loader");
+
+    assert!(text.contains("Hello World"));
+}
+
+#[test]
+fn test_estimate_page_count_from_bytes_excludes_pages_tree() {
+    let pdf = add_leading_tab(truncate_eof_marker(make_minimal_text_pdf()));
+
+    assert_eq!(estimate_page_count_from_bytes(&pdf), 1);
+}
+
 #[test]
 fn test_not_a_pdf_detect_pdf_type_mem() {
     // Verify detect_pdf_type_mem is also guarded