teamcraft/
pdftext.rs

1//! Text of PDF reports for the full-text search and the assistant
2//! (`read_file_text`); see `docs/src/research/pdf-text-extraction.md`.
3//!
4//! A PDF is untrusted input, so it is never parsed in the server process: the
5//! hidden `teamcraft text-extract` child process does it, with a time limit,
6//! the memory limit of the 3D previews and a size cap before the file is
7//! opened (in tests, where there is no such binary, on a blocking thread).
8//! The result is bounded as well – [`MAX_PAGES`] pages, [`MAX_CHARS`]
9//! characters. There is no OCR: a scanned PDF has no text.
10//!
11//! The same child reads HTML reports (`--format html`, see [`crate::htmltext`]);
12//! `extract_isolated` runs either.
13
14use std::{
15    path::Path,
16    sync::{
17        atomic::{AtomicU64, Ordering},
18        LazyLock,
19    },
20    time::Duration,
21};
22
23use tokio::sync::Semaphore;
24
25/// A PDF above this size is not opened (CFD reports are small; big PDFs are scans).
26pub const MAX_PDF_BYTES: u64 = 32 * 1024 * 1024;
27
28/// Pages read.
29pub const MAX_PAGES: usize = 100;
30
31/// Characters kept (the cap of every extracted text).
32pub const MAX_CHARS: usize = 200_000;
33
34/// Time limit of one extraction.
35pub const TIME_LIMIT: Duration = Duration::from_secs(60);
36
37/// Extractions at the same time.
38const PARALLEL: usize = 2;
39
40static SLOTS: LazyLock<Semaphore> = LazyLock::new(|| Semaphore::new(PARALLEL));
41
42/// Address space of the child, set at start-up from the preview limits.
43static MEMORY_BYTES: AtomicU64 = AtomicU64::new(1024 * 1024 * 1024);
44
45/// Takes the memory limit of the child process (0 = unlimited).
46pub fn configure(memory_bytes: u64) {
47    MEMORY_BYTES.store(memory_bytes, Ordering::Relaxed);
48}
49
50/// Whether a file name is a PDF.
51pub fn is_pdf(filename: &str) -> bool {
52    filename
53        .rsplit_once('.')
54        .is_some_and(|(stem, ext)| !stem.is_empty() && ext.eq_ignore_ascii_case("pdf"))
55}
56
57/// The text of a PDF held in memory: at most [`MAX_PAGES`] pages and
58/// [`MAX_CHARS`] characters; pages that cannot be read are skipped. A parser
59/// panic is an error, not a crash.
60pub fn extract(bytes: &[u8]) -> anyhow::Result<String> {
61    let bytes = bytes.to_vec();
62    let result = std::panic::catch_unwind(move || -> anyhow::Result<String> {
63        let doc = pdf_oxide::PdfDocument::from_bytes(bytes)
64            .map_err(|e| anyhow::anyhow!("not a readable PDF: {e}"))?;
65        let pages = doc
66            .page_count()
67            .map_err(|e| anyhow::anyhow!("not a readable PDF: {e}"))?;
68        let mut out = String::new();
69        for page in 0..pages.min(MAX_PAGES) {
70            if let Ok(text) = doc.extract_text(page) {
71                out.push_str(&text);
72                out.push('\n');
73            }
74            if out.chars().count() > MAX_CHARS {
75                break;
76            }
77        }
78        Ok(tidy(&out))
79    });
80    match result {
81        Ok(r) => r,
82        Err(_) => Err(anyhow::anyhow!("the PDF could not be read")),
83    }
84}
85
86/// Control characters (NUL included) become spaces, line breaks stay (tables
87/// keep their rows), runs of blank lines collapse, the length is capped.
88pub(crate) fn tidy(text: &str) -> String {
89    let mut out = String::with_capacity(text.len().min(MAX_CHARS + 16));
90    let mut blank_lines = 0;
91    for line in text.lines() {
92        let line: String = line
93            .chars()
94            .map(|c| if c.is_control() { ' ' } else { c })
95            .collect();
96        let line = line.trim_end();
97        if line.trim().is_empty() {
98            blank_lines += 1;
99            if blank_lines > 1 {
100                continue;
101            }
102        } else {
103            blank_lines = 0;
104        }
105        out.push_str(line);
106        out.push('\n');
107        if out.chars().count() > MAX_CHARS {
108            break;
109        }
110    }
111    out.trim().chars().take(MAX_CHARS).collect()
112}
113
114/// One job of the child process: the text of the PDF read from `input` (stdin)
115/// is written to `output` (stdout). At most `max_input_bytes` are read; more is
116/// refused. Returns the number of characters.
117pub fn run_job(
118    input: impl std::io::Read,
119    mut output: impl std::io::Write,
120    max_input_bytes: u64,
121) -> anyhow::Result<usize> {
122    use std::io::Read as _;
123    let mut bytes = Vec::new();
124    input
125        .take(max_input_bytes.saturating_add(1))
126        .read_to_end(&mut bytes)?;
127    if bytes.len() as u64 > max_input_bytes {
128        anyhow::bail!(
129            "This PDF is too large to read (more than {} MB).",
130            max_input_bytes / (1024 * 1024)
131        );
132    }
133    let text = extract(&bytes)?;
134    output.write_all(text.as_bytes())?;
135    output.flush()?;
136    Ok(text.chars().count())
137}
138
139/// What the `text-extract` child reads: a PDF or an HTML report.
140#[derive(Clone, Copy, Debug, PartialEq, Eq)]
141pub enum TextFormat {
142    Pdf,
143    /// See [`crate::htmltext`].
144    Html,
145}
146
147impl TextFormat {
148    /// The `--format` value of the child.
149    pub fn as_str(self) -> &'static str {
150        match self {
151            TextFormat::Pdf => "pdf",
152            TextFormat::Html => "html",
153        }
154    }
155
156    pub fn parse(s: &str) -> Option<Self> {
157        match s {
158            "pdf" => Some(TextFormat::Pdf),
159            "html" => Some(TextFormat::Html),
160            _ => None,
161        }
162    }
163
164    /// Bytes the child reads: a bigger PDF is refused, a bigger HTML file is cut.
165    pub fn max_input_bytes(self) -> u64 {
166        match self {
167            TextFormat::Pdf => MAX_PDF_BYTES,
168            TextFormat::Html => crate::htmltext::MAX_HTML_BYTES,
169        }
170    }
171
172    fn name(self) -> &'static str {
173        match self {
174            TextFormat::Pdf => "PDF",
175            TextFormat::Html => "HTML file",
176        }
177    }
178}
179
180/// One job of the child process in the given format (see [`run_job`] and
181/// [`crate::htmltext::run_job`]).
182pub fn run_job_as(
183    format: TextFormat,
184    input: impl std::io::Read,
185    output: impl std::io::Write,
186    max_input_bytes: u64,
187) -> anyhow::Result<usize> {
188    match format {
189        TextFormat::Pdf => run_job(input, output, max_input_bytes),
190        TextFormat::Html => crate::htmltext::run_job(input, output, max_input_bytes),
191    }
192}
193
194/// The text of a stored PDF; empty if it has none or cannot be read (a
195/// scanned or damaged file never fails an upload). The opened file is handed to the
196/// child on its standard input and the text comes back on its standard
197/// output, so no path travels anywhere.
198pub async fn extract_file(file: tokio::fs::File) -> String {
199    // size before reading: a big PDF is a scan, not a report
200    match file.metadata().await {
201        Ok(meta) if meta.len() <= MAX_PDF_BYTES => {}
202        _ => return String::new(),
203    }
204    extract_isolated(file, TextFormat::Pdf).await
205}
206
207/// Runs one extraction in the `text-extract` child (or, in tests, on a
208/// blocking thread): one of the [`PARALLEL`] slots, the time limit, the
209/// memory limit, a cleared environment. Empty if there is no text.
210pub(crate) async fn extract_isolated(file: tokio::fs::File, format: TextFormat) -> String {
211    let Ok(_slot) = SLOTS.acquire().await else {
212        return String::new();
213    };
214    let file = file.into_std().await;
215    let outcome = match crate::preview::worker_exe() {
216        Some(exe) => run_worker(&exe, file, format).await,
217        None => run_in_process(file, format).await,
218    };
219    match outcome {
220        Ok(text) => text.chars().take(MAX_CHARS).collect(),
221        Err(message) => {
222            tracing::debug!(%message, "no text from the {}", format.name());
223            String::new()
224        }
225    }
226}
227
228async fn run_worker(
229    exe: &Path,
230    input: std::fs::File,
231    format: TextFormat,
232) -> Result<String, String> {
233    let child = tokio::process::Command::new(exe)
234        .arg("text-extract")
235        .arg("--format")
236        .arg(format.as_str())
237        .arg("--max-input-bytes")
238        .arg(format.max_input_bytes().to_string())
239        .arg("--max-memory-bytes")
240        .arg(MEMORY_BYTES.load(Ordering::Relaxed).to_string())
241        .env_clear()
242        .env("RUST_LOG", "error")
243        .stdin(std::process::Stdio::from(input))
244        .stdout(std::process::Stdio::piped())
245        .stderr(std::process::Stdio::piped())
246        .kill_on_drop(true)
247        .spawn()
248        .map_err(|e| format!("the text extractor could not be started: {e}"))?;
249    match tokio::time::timeout(TIME_LIMIT, child.wait_with_output()).await {
250        Err(_) => Err(format!("no result within {} s", TIME_LIMIT.as_secs())),
251        Ok(Err(e)) => Err(e.to_string()),
252        Ok(Ok(out)) if out.status.success() => {
253            Ok(String::from_utf8_lossy(&out.stdout).into_owned())
254        }
255        Ok(Ok(out)) => Err(String::from_utf8_lossy(&out.stderr)
256            .lines()
257            .rev()
258            .find(|l| !l.trim().is_empty())
259            .unwrap_or("the extractor stopped unexpectedly")
260            .to_string()),
261    }
262}
263
264async fn run_in_process(input: std::fs::File, format: TextFormat) -> Result<String, String> {
265    let task = tokio::task::spawn_blocking(move || {
266        let mut text = Vec::new();
267        run_job_as(format, input, &mut text, format.max_input_bytes())
268            .map(|_| String::from_utf8_lossy(&text).into_owned())
269            .map_err(|e| e.to_string())
270    });
271    match tokio::time::timeout(TIME_LIMIT, task).await {
272        Err(_) => Err("timed out".into()),
273        Ok(Err(e)) => Err(e.to_string()),
274        Ok(Ok(r)) => r,
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281
282    /// A one-page PDF with the given lines of text (Helvetica).
283    pub(crate) fn sample_pdf(lines: &[&str]) -> Vec<u8> {
284        let mut content = String::from("BT /F1 12 Tf 14 TL 50 740 Td\n");
285        for line in lines {
286            let escaped = line
287                .replace('\\', "\\\\")
288                .replace('(', "\\(")
289                .replace(')', "\\)");
290            content.push_str(&format!("({escaped}) Tj T*\n"));
291        }
292        content.push_str("ET");
293        let objects = [
294            "<< /Type /Catalog /Pages 2 0 R >>".to_string(),
295            "<< /Type /Pages /Kids [3 0 R] /Count 1 >>".to_string(),
296            "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] /Contents 4 0 R \
297             /Resources << /Font << /F1 5 0 R >> >> >>"
298                .to_string(),
299            format!(
300                "<< /Length {} >>\nstream\n{content}\nendstream",
301                content.len()
302            ),
303            "<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>".to_string(),
304        ];
305        let mut pdf = String::from("%PDF-1.4\n");
306        let mut offsets = Vec::new();
307        for (i, body) in objects.iter().enumerate() {
308            offsets.push(pdf.len());
309            pdf.push_str(&format!("{} 0 obj\n{body}\nendobj\n", i + 1));
310        }
311        let xref = pdf.len();
312        pdf.push_str(&format!(
313            "xref\n0 {}\n0000000000 65535 f \n",
314            objects.len() + 1
315        ));
316        for offset in offsets {
317            pdf.push_str(&format!("{offset:010} 00000 n \n"));
318        }
319        pdf.push_str(&format!(
320            "trailer\n<< /Size {} /Root 1 0 R >>\nstartxref\n{xref}\n%%EOF\n",
321            objects.len() + 1
322        ));
323        pdf.into_bytes()
324    }
325
326    #[test]
327    fn the_text_of_a_pdf_is_found() {
328        let pdf = sample_pdf(&[
329            "Pressure drop: 388.25 Pa",
330            "Max. von Mises stress 182.5 MPa",
331        ]);
332        let text = extract(&pdf).unwrap();
333        assert!(text.contains("388.25"), "{text:?}");
334        assert!(text.contains("von Mises"), "{text:?}");
335    }
336
337    #[test]
338    fn damaged_and_hostile_files_fail_cleanly() {
339        let pdf = sample_pdf(&["Pressure drop: 388.25 Pa"]);
340        for cut in [0, 5, 40, pdf.len() / 2, pdf.len() - 10] {
341            // an error or whatever text survives – never a panic
342            let _ = extract(&pdf[..cut]);
343        }
344        assert!(extract(b"not a pdf at all").is_err());
345        assert!(extract(&[]).is_err());
346        let mut flipped = pdf.clone();
347        for i in (0..flipped.len()).step_by(7) {
348            flipped[i] ^= 0xFF;
349        }
350        let _ = extract(&flipped);
351    }
352
353    #[test]
354    fn tidy_keeps_rows_and_drops_noise() {
355        assert_eq!(tidy("a\u{0}b  \n\n\n\nc\t d\n"), "a b\n\nc  d");
356        let long = "x".repeat(MAX_CHARS * 2);
357        assert_eq!(tidy(&long).chars().count(), MAX_CHARS);
358    }
359
360    #[test]
361    fn names_and_size_cap() {
362        assert!(is_pdf("report.PDF") && is_pdf("a.b.pdf") && !is_pdf("pdf") && !is_pdf(".pdf"));
363        let dir = std::env::temp_dir().join(format!("tc-pdf-{}", uuid::Uuid::new_v4()));
364        std::fs::create_dir_all(&dir).unwrap();
365        let big = dir.join("big.pdf");
366        std::fs::File::create(&big)
367            .unwrap()
368            .set_len(5 * 1024 * 1024)
369            .unwrap();
370        let mut sink = Vec::new();
371        let e = run_job(std::fs::File::open(&big).unwrap(), &mut sink, 1024 * 1024).unwrap_err();
372        assert!(e.to_string().contains("too large to read"), "{e}");
373        assert!(sink.is_empty());
374        std::fs::remove_dir_all(&dir).ok();
375    }
376
377    #[tokio::test]
378    async fn stored_pdfs_are_read_in_process_during_tests() {
379        let dir = std::env::temp_dir().join(format!("tc-pdf-{}", uuid::Uuid::new_v4()));
380        std::fs::create_dir_all(&dir).unwrap();
381        let path = dir.join("report.pdf");
382        std::fs::write(&path, sample_pdf(&["Mass flow 2.25 kg/s"])).unwrap();
383        assert!(extract_file(tokio::fs::File::open(&path).await.unwrap())
384            .await
385            .contains("2.25"));
386        std::fs::write(&path, b"junk").unwrap();
387        assert_eq!(
388            extract_file(tokio::fs::File::open(&path).await.unwrap()).await,
389            ""
390        );
391        std::fs::remove_dir_all(&dir).ok();
392    }
393}
teamcraft/pdftext.rs

teamcraft/
pdftext.rs
