321 lines
11 KiB
Rust
321 lines
11 KiB
Rust
//! Bounded inputs and outputs; parser allocations and CPU are not process isolated.
|
|
use crate::SourceLocation;
|
|
use anyhow::{bail, Context, Result};
|
|
use quick_xml::{events::Event, name::ResolveResult, NsReader, XmlVersion};
|
|
use sha2::{Digest, Sha256};
|
|
use std::{
|
|
fs::File,
|
|
io::{Cursor, Read},
|
|
path::Path,
|
|
};
|
|
|
|
const MAX_TEXT: usize = 2 * 1024 * 1024;
|
|
const MAX_DOCUMENT: usize = 32 * 1024 * 1024;
|
|
const MAX_XML: usize = 8 * 1024 * 1024;
|
|
const MAX_PDF_STREAM: usize = 4 * 1024 * 1024;
|
|
const MAX_PASSAGES: usize = 10_000;
|
|
|
|
pub(crate) struct Passage {
|
|
pub text: String,
|
|
pub source: SourceLocation,
|
|
}
|
|
|
|
pub(crate) struct SourceInput {
|
|
bytes: Vec<u8>,
|
|
pub hash: String,
|
|
}
|
|
|
|
/// Read once: the fingerprint and parser consume precisely the same bytes.
|
|
pub(crate) fn read_source(path: &Path) -> Result<Option<SourceInput>> {
|
|
let extension = path
|
|
.extension()
|
|
.and_then(|value| value.to_str())
|
|
.unwrap_or("")
|
|
.to_ascii_lowercase();
|
|
let document = matches!(extension.as_str(), "pdf" | "docx");
|
|
if !document
|
|
&& !crate::is_code_path(path)
|
|
&& !matches!(
|
|
extension.as_str(),
|
|
"txt"
|
|
| "md"
|
|
| "json"
|
|
| "toml"
|
|
| "yaml"
|
|
| "yml"
|
|
| "csv"
|
|
| "html"
|
|
| "css"
|
|
| "xml"
|
|
| "log"
|
|
| "sh"
|
|
)
|
|
{
|
|
return Ok(None);
|
|
}
|
|
if !document && File::open(path)?.metadata()?.len() > MAX_TEXT as u64 {
|
|
return Ok(None);
|
|
}
|
|
let bytes = read_bounded(path, if document { MAX_DOCUMENT } else { MAX_TEXT })?;
|
|
let hash = format!("{:x}", Sha256::digest(&bytes));
|
|
Ok(Some(SourceInput { bytes, hash }))
|
|
}
|
|
|
|
pub(crate) fn extract(path: &Path, input: SourceInput) -> Result<Option<Vec<Passage>>> {
|
|
let extension = path
|
|
.extension()
|
|
.and_then(|value| value.to_str())
|
|
.unwrap_or("")
|
|
.to_ascii_lowercase();
|
|
if extension == "pdf" {
|
|
return pdf(&input.bytes).map(Some);
|
|
}
|
|
if extension == "docx" {
|
|
return docx(&input.bytes).map(Some);
|
|
}
|
|
let code = crate::is_code_path(path);
|
|
if input.bytes.contains(&0) {
|
|
return Ok(None);
|
|
}
|
|
let Ok(text) = String::from_utf8(input.bytes) else {
|
|
return Ok(None);
|
|
};
|
|
let lines = text.split_inclusive('\n').collect::<Vec<_>>();
|
|
if code {
|
|
let mut passages = Vec::new();
|
|
for start in (0..lines.len()).step_by(96) {
|
|
if passages.len() >= MAX_PASSAGES {
|
|
return Ok(None);
|
|
}
|
|
let end = (start + 128).min(lines.len());
|
|
passages.push(Passage {
|
|
text: lines[start..end].concat(),
|
|
source: SourceLocation {
|
|
kind: "line".into(),
|
|
number: (start + 1) as u32,
|
|
},
|
|
});
|
|
if end == lines.len() {
|
|
break;
|
|
}
|
|
}
|
|
return Ok(Some(passages));
|
|
}
|
|
Ok(Some(
|
|
lines
|
|
.chunks(64)
|
|
.enumerate()
|
|
.map(|(index, lines)| Passage {
|
|
text: lines.concat(),
|
|
source: SourceLocation {
|
|
kind: "line".into(),
|
|
number: (index * 64 + 1) as u32,
|
|
},
|
|
})
|
|
.collect(),
|
|
))
|
|
}
|
|
|
|
fn read_bounded(path: &Path, maximum: usize) -> Result<Vec<u8>> {
|
|
let file = File::open(path)?;
|
|
if file.metadata()?.len() > maximum as u64 {
|
|
bail!("document input exceeds {maximum} bytes");
|
|
}
|
|
let mut bytes = Vec::new();
|
|
file.take(maximum as u64 + 1).read_to_end(&mut bytes)?;
|
|
if bytes.len() > maximum {
|
|
bail!("document input exceeds {maximum} bytes");
|
|
}
|
|
Ok(bytes)
|
|
}
|
|
|
|
fn pdf(bytes: &[u8]) -> Result<Vec<Passage>> {
|
|
let document = lopdf::Document::load_mem_with_options(
|
|
bytes,
|
|
lopdf::LoadOptions {
|
|
strict: true,
|
|
max_decompressed_size: Some(MAX_PDF_STREAM),
|
|
..Default::default()
|
|
},
|
|
)
|
|
.context("invalid, encrypted, or over-limit PDF")?;
|
|
if document.was_encrypted() || document.is_encrypted() {
|
|
bail!("encrypted PDF is unsupported");
|
|
}
|
|
let pages = document.get_pages();
|
|
if pages.is_empty() {
|
|
bail!("PDF has no pages");
|
|
}
|
|
if pages.len() > 1000 {
|
|
bail!("PDF exceeds 1000 pages");
|
|
}
|
|
let mut passages = Vec::new();
|
|
let mut total = 0;
|
|
for number in pages.keys() {
|
|
let text = document
|
|
.extract_text_with_limit(&[*number], MAX_PDF_STREAM)
|
|
.with_context(|| format!("PDF text extraction failed on page {number}"))?;
|
|
total += text.len();
|
|
if total > MAX_TEXT {
|
|
bail!("PDF text exceeds {MAX_TEXT} bytes");
|
|
}
|
|
if !text.trim().is_empty() {
|
|
passages.push(Passage {
|
|
text,
|
|
source: SourceLocation {
|
|
kind: "page".into(),
|
|
number: *number,
|
|
},
|
|
});
|
|
}
|
|
}
|
|
Ok(passages)
|
|
}
|
|
|
|
fn docx(bytes: &[u8]) -> Result<Vec<Passage>> {
|
|
let mut archive =
|
|
zip::ZipArchive::new(Cursor::new(bytes)).context("invalid DOCX ZIP container")?;
|
|
if archive.len() > 10_000 {
|
|
bail!("DOCX exceeds 10000 archive entries");
|
|
}
|
|
let mut xml = String::new();
|
|
let part = archive
|
|
.by_name("word/document.xml")
|
|
.context("DOCX lacks word/document.xml")?;
|
|
if part.size() > MAX_XML as u64 {
|
|
bail!("DOCX document XML exceeds {MAX_XML} bytes");
|
|
}
|
|
part.take(MAX_XML as u64 + 1)
|
|
.read_to_string(&mut xml)
|
|
.context("DOCX XML must be valid UTF-8")?;
|
|
if xml.len() > MAX_XML {
|
|
bail!("DOCX document XML exceeds {MAX_XML} bytes");
|
|
}
|
|
let mut reader = NsReader::from_str(&xml);
|
|
let mut passages = Vec::new();
|
|
let mut paragraph = String::new();
|
|
let mut number = 0;
|
|
let mut in_paragraph = false;
|
|
let mut in_text = false;
|
|
let mut total = 0;
|
|
let mut depth = 0usize;
|
|
let mut saw_document = false;
|
|
loop {
|
|
let (namespace, event) = reader.read_resolved_event().context("invalid DOCX XML")?;
|
|
let word_namespace = matches!(namespace, ResolveResult::Bound(value) if matches!(value.as_ref(),
|
|
"http://schemas.openxmlformats.org/wordprocessingml/2006/main" | "http://purl.oclc.org/ooxml/wordprocessingml/main"));
|
|
match &event {
|
|
Event::Start(start) => {
|
|
if depth == 0 {
|
|
if saw_document || !word_namespace || start.local_name().as_ref() != "document"
|
|
{
|
|
bail!("invalid DOCX document root");
|
|
}
|
|
saw_document = true;
|
|
}
|
|
depth += 1;
|
|
if depth > 128 {
|
|
bail!("DOCX XML exceeds 128 nesting levels");
|
|
}
|
|
}
|
|
Event::End(_) => depth = depth.checked_sub(1).context("invalid DOCX XML depth")?,
|
|
Event::Empty(_) if depth == 0 => bail!("invalid DOCX document root"),
|
|
_ => {}
|
|
}
|
|
match event {
|
|
Event::Start(event) if word_namespace => match event.local_name().as_ref() {
|
|
"p" => {
|
|
if in_paragraph {
|
|
bail!("nested DOCX paragraphs are unsupported");
|
|
}
|
|
number += 1;
|
|
if number > MAX_PASSAGES {
|
|
bail!("DOCX exceeds 10000 paragraphs");
|
|
}
|
|
in_paragraph = true;
|
|
paragraph.clear();
|
|
}
|
|
"t" => in_text = in_paragraph,
|
|
_ => {}
|
|
},
|
|
Event::End(event) if word_namespace => match event.local_name().as_ref() {
|
|
"t" => in_text = false,
|
|
"p" => {
|
|
in_paragraph = false;
|
|
in_text = false;
|
|
total += paragraph.len();
|
|
if total > MAX_TEXT {
|
|
bail!("DOCX text exceeds {MAX_TEXT} bytes");
|
|
}
|
|
if !paragraph.trim().is_empty() {
|
|
passages.push(Passage {
|
|
text: std::mem::take(&mut paragraph),
|
|
source: SourceLocation {
|
|
kind: "paragraph".into(),
|
|
number: number as u32,
|
|
},
|
|
});
|
|
}
|
|
}
|
|
_ => {}
|
|
},
|
|
Event::Empty(event) if word_namespace => match event.local_name().as_ref() {
|
|
"tab" if in_paragraph => paragraph.push('\t'),
|
|
"br" | "cr" if in_paragraph => paragraph.push('\n'),
|
|
"p" => {
|
|
number += 1;
|
|
if number > MAX_PASSAGES {
|
|
bail!("DOCX exceeds 10000 paragraphs");
|
|
}
|
|
}
|
|
_ => {}
|
|
},
|
|
Event::Text(event) if in_text => {
|
|
paragraph.push_str(&event.xml_content(XmlVersion::Implicit1_0))
|
|
}
|
|
Event::CData(event) if in_text => {
|
|
paragraph.push_str(&event.xml_content(XmlVersion::Implicit1_0))
|
|
}
|
|
Event::GeneralRef(event) if in_text => {
|
|
let reference = format!("&{};", event.as_ref());
|
|
paragraph.push_str(
|
|
&quick_xml::escape::unescape(&reference).context("unsupported DOCX entity")?,
|
|
);
|
|
}
|
|
Event::DocType(_) => bail!("DOCX XML document types are unsupported"),
|
|
Event::Eof => {
|
|
if in_paragraph || depth != 0 || !saw_document {
|
|
bail!("incomplete DOCX document");
|
|
}
|
|
break;
|
|
}
|
|
_ => {}
|
|
}
|
|
if total + paragraph.len() > MAX_TEXT {
|
|
bail!("DOCX text exceeds {MAX_TEXT} bytes");
|
|
}
|
|
}
|
|
Ok(passages)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod provenance_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn parsing_and_fingerprint_share_one_snapshot_even_when_path_changes() -> Result<()> {
|
|
let directory = tempfile::tempdir()?;
|
|
let path = directory.path().join("notes.txt");
|
|
std::fs::write(&path, "original parser input")?;
|
|
let input = read_source(&path)?.unwrap();
|
|
assert_eq!(
|
|
input.hash,
|
|
format!("{:x}", Sha256::digest(b"original parser input"))
|
|
);
|
|
std::fs::write(&path, "replacement on disk")?;
|
|
let passages = extract(&path, input)?.unwrap();
|
|
assert_eq!(passages[0].text, "original parser input");
|
|
Ok(())
|
|
}
|
|
}
|