use anyhow::Result;
use lopdf::{
content::{Content, Operation},
dictionary, Document, Object, Stream,
};
use searchi_core::SearchIndex;
use std::{fs, io::Write, path::Path};
fn pdf(path: &Path) -> Result<()> {
pdf_with_text(
path,
&[
"First introductory page",
"Supplier cancellation clause on second page",
],
)
}
fn pdf_with_text(path: &Path, texts: &[&str]) -> Result<()> {
let mut document = Document::with_version("1.5");
let pages_id = document.new_object_id();
let font = document.add_object(dictionary! {"Type"=>"Font", "Subtype"=>"Type1", "BaseFont"=>"Helvetica", "Encoding"=>"WinAnsiEncoding"});
let resources = document.add_object(dictionary! {"Font"=>dictionary! {"F1"=>font}});
let mut pages = Vec::new();
for text in texts {
let content = Content {
operations: vec![
Operation::new("BT", vec![]),
Operation::new("Tf", vec![Object::Name(b"F1".to_vec()), 12.into()]),
Operation::new("Td", vec![50.into(), 700.into()]),
Operation::new("Tj", vec![Object::string_literal(*text)]),
Operation::new("ET", vec![]),
],
};
let stream = document.add_object(Stream::new(dictionary! {}, content.encode()?));
let page = document
.add_object(dictionary! {"Type"=>"Page", "Parent"=>pages_id, "Contents"=>stream});
pages.push(Object::Reference(page));
}
document.objects.insert(pages_id, Object::Dictionary(dictionary! {"Type"=>"Pages", "Kids"=>pages,
"Count"=>texts.len() as i64, "Resources"=>resources, "MediaBox"=>vec![0.into(),0.into(),595.into(),842.into()]}));
let catalog = document.add_object(dictionary! {"Type"=>"Catalog", "Pages"=>pages_id});
document.trailer.set("Root", catalog);
document.compress();
document.save(path)?;
Ok(())
}
#[test]
fn compressed_document_limits_are_enforced() -> Result<()> {
let root = tempfile::tempdir()?;
let text = "a".repeat(4 * 1024 * 1024 + 100);
pdf_with_text(&root.path().join("bomb.pdf"), &[&text])?;
let xml = format!(
r#"{}"#,
"b".repeat(2 * 1024 * 1024 + 1)
);
docx(&root.path().join("over-limit.docx"), &xml)?;
let mut index = SearchIndex::open(":memory:")?;
assert_eq!(index.index_folder(root.path())?.failures, 2);
assert_eq!(index.stats()?.content_files, 0);
assert!(index.search("bomb.pdf", 10)?[0].extraction_error.is_some());
assert!(index.search("over-limit.docx", 10)?[0]
.extraction_error
.as_ref()
.unwrap()
.contains("exceeds"));
Ok(())
}
#[test]
fn encrypted_pdf_is_explicitly_unsupported() -> Result<()> {
let root = tempfile::tempdir()?;
let path = root.path().join("protected.pdf");
pdf(&path)?;
let mut document = Document::load(&path)?;
document.trailer.set(
"ID",
vec![
Object::string_literal("fixture-id"),
Object::string_literal("fixture-id"),
],
);
let state = lopdf::EncryptionState::try_from(lopdf::EncryptionVersion::V1 {
document: &document,
owner_password: "owner",
user_password: "reader",
permissions: lopdf::Permissions::all(),
})?;
document.encrypt(&state)?;
document.save(&path)?;
let mut index = SearchIndex::open(":memory:")?;
assert_eq!(index.index_folder(root.path())?.failures, 1);
assert!(index.search("protected.pdf", 10)?[0]
.extraction_error
.as_ref()
.unwrap()
.contains("encrypted"));
assert!(index.search("cancellation", 10)?.is_empty());
Ok(())
}
fn docx(path: &Path, xml: &str) -> Result<()> {
let mut writer = zip::ZipWriter::new(fs::File::create(path)?);
let options = zip::write::SimpleFileOptions::default()
.compression_method(zip::CompressionMethod::Deflated);
for (name, content) in [
(
"[Content_Types].xml",
r#""#,
),
(
"_rels/.rels",
r#""#,
),
("word/document.xml", xml),
] {
writer.start_file(name, options)?;
writer.write_all(content.as_bytes())?;
}
writer.finish()?;
Ok(())
}
#[test]
fn real_pdf_and_docx_return_matching_locations_and_persist() -> Result<()> {
let root = tempfile::tempdir()?;
pdf(&root.path().join("agreement.pdf"))?;
docx(
&root.path().join("notes.docx"),
r#"IntroductionResearch & development budget"#,
)?;
let db = root.path().join("index.db");
let mut index = SearchIndex::open(&db)?;
assert_eq!(index.index_folder(root.path())?.indexed, 2);
let result = index.search("cancellation", 10)?;
assert_eq!(result[0].source.as_ref().unwrap().kind, "page");
assert_eq!(result[0].source.as_ref().unwrap().number, 2);
assert!(result[0].snippet.contains("cancellation"));
let result = index.search("development", 10)?;
assert_eq!(result[0].source.as_ref().unwrap().kind, "paragraph");
assert_eq!(result[0].source.as_ref().unwrap().number, 3);
assert!(result[0].snippet.contains("Research & development"));
let batch = index.passage_batch(0, 2)?;
assert_eq!(batch.len(), 2);
let remainder = index.passage_batch(batch[1].id, 256)?;
assert_eq!(remainder.len(), 2);
// Directory traversal order differs by OS/filesystem. Pagination must
// preserve every passage exactly once, regardless of which file was first.
assert!(batch
.iter()
.all(|first| remainder.iter().all(|next| first.id != next.id)));
let all = batch.iter().chain(&remainder).collect::>();
assert!(all
.iter()
.any(|passage| passage.source.kind == "paragraph" && passage.source.number == 3));
assert!(all
.iter()
.any(|passage| passage.source.kind == "page" && passage.source.number == 2));
assert!(index
.passage_batch(remainder.last().unwrap().id, 256)?
.is_empty());
assert!(index.contains_path(&result[0].path)?);
assert!(!index.contains_path("unknown")?);
drop(index);
let index = SearchIndex::open(&db)?;
assert_eq!(
index.search("budget", 10)?[0]
.source
.as_ref()
.unwrap()
.number,
3
);
Ok(())
}
#[test]
fn corrupt_document_reports_failure_and_preserves_filename_search() -> Result<()> {
let root = tempfile::tempdir()?;
fs::write(root.path().join("notes.txt"), "persistent evidence")?;
let mut index = SearchIndex::open(":memory:")?;
index.index_folder(root.path())?;
fs::write(root.path().join("broken.pdf"), b"not a PDF")?;
assert_eq!(index.index_folder(root.path())?.failures, 1);
let result = index.search("broken.pdf", 10)?;
assert!(result[0].extraction_error.as_ref().unwrap().contains("PDF"));
assert_eq!(index.search("evidence", 10)?.len(), 1);
fs::remove_file(root.path().join("broken.pdf"))?;
docx(
&root.path().join("broken.docx"),
"broken",
)?;
assert_eq!(index.index_folder(root.path())?.failures, 1);
assert!(index.search("broken.docx", 10)?[0]
.extraction_error
.is_some());
pdf(&root.path().join("agreement.pdf"))?;
index.index_folder(root.path())?;
assert_eq!(index.search("cancellation", 10)?.len(), 1);
fs::write(root.path().join("agreement.pdf"), "corrupted replacement")?;
assert_eq!(index.index_folder(root.path())?.failures, 1);
assert!(index.search("cancellation", 10)?.is_empty());
assert!(index.search("agreement.pdf", 10)?[0]
.extraction_error
.is_some());
Ok(())
}