diff --git a/crates/fetchkit/src/content.rs b/crates/fetchkit/src/content.rs index 0de59cd..7e704f2 100644 --- a/crates/fetchkit/src/content.rs +++ b/crates/fetchkit/src/content.rs @@ -3,6 +3,8 @@ //! Fetchers own retrieval and network policy. Content processors operate only on //! bounded response bytes, so format-specific extraction cannot bypass egress controls. +use std::sync::{Arc, LazyLock}; + use async_trait::async_trait; use bytes::Bytes; use pdf_inspector::process_pdf_mem; @@ -18,7 +20,8 @@ use crate::{PageMetadata, PageQuality}; // pdf-inspector uses Rayon internally; bound concurrent documents so batch fetches // cannot multiply parser thread pools without limit. -static PDF_PROCESSING_LIMIT: Semaphore = Semaphore::const_new(2); +static PDF_PROCESSING_LIMIT: LazyLock> = + LazyLock::new(|| Arc::new(Semaphore::new(2))); /// Bounded response bytes and metadata passed to a [`ContentProcessor`]. pub struct ContentProcessorInput { @@ -268,6 +271,25 @@ impl ContentProcessor for HtmlProcessor { /// Extracts native text PDFs as Markdown using `pdf-inspector`. pub struct PdfProcessor; +async fn run_pdf_task(limit: Arc, task: F) -> Result +where + T: Send + 'static, + F: FnOnce() -> T + Send + 'static, +{ + let permit = limit + .acquire_owned() + .await + .map_err(|error| ContentProcessorError(format!("PDF processing unavailable: {error}")))?; + tokio::task::spawn_blocking(move || { + // Keep the permit in the blocking task: cancelling its async caller must not + // admit another parse while this non-cancellable work is still running. + let _permit = permit; + task() + }) + .await + .map_err(|error| ContentProcessorError(format!("PDF processor task failed: {error}"))) +} + #[async_trait] impl ContentProcessor for PdfProcessor { fn name(&self) -> &'static str { @@ -293,13 +315,11 @@ impl ContentProcessor for PdfProcessor { &self, input: ContentProcessorInput, ) -> Result { - let _permit = PDF_PROCESSING_LIMIT.acquire().await.map_err(|error| { - ContentProcessorError(format!("PDF processing unavailable: {error}")) - })?; - let result = tokio::task::spawn_blocking(move || process_pdf_mem(&input.body)) - .await - .map_err(|error| ContentProcessorError(format!("PDF processor task failed: {error}")))? - .map_err(|error| ContentProcessorError(error.to_string()))?; + let result = run_pdf_task(Arc::clone(&PDF_PROCESSING_LIMIT), move || { + process_pdf_mem(&input.body) + }) + .await? + .map_err(|error| ContentProcessorError(error.to_string()))?; let needs_ocr = !result.pages_needing_ocr.is_empty(); let mut warnings = Vec::new(); @@ -360,6 +380,8 @@ fn comma_separated_pages(pages: &[u32]) -> String { #[cfg(test)] mod tests { + use std::sync::mpsc; + use super::*; struct StubMarkdownConverter; @@ -418,6 +440,30 @@ mod tests { assert!(!processor.matches(&extensionless, Some("application/octet-stream"))); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn cancelled_pdf_caller_does_not_release_running_task_permit() { + let limit = Arc::new(Semaphore::new(1)); + let (started_tx, started_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + let handle = tokio::spawn(run_pdf_task(Arc::clone(&limit), move || { + started_tx.send(()).unwrap(); + release_rx.recv().unwrap(); + })); + + started_rx.recv().unwrap(); + handle.abort(); + assert_eq!(limit.available_permits(), 0); + + release_tx.send(()).unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while limit.available_permits() == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } + #[test] fn registry_uses_first_matching_processor() { let registry = ContentProcessorRegistry::with_defaults(); diff --git a/knowledge/foundations/fetchers.md b/knowledge/foundations/fetchers.md index 59482ab..ed118bf 100644 --- a/knowledge/foundations/fetchers.md +++ b/knowledge/foundations/fetchers.md @@ -82,7 +82,8 @@ allowing callers to replace the default content processor registry. suffix is a fallback only for missing Content-Type or `application/octet-stream`. - Uses `pdf-inspector` to classify and extract Markdown from bounded in-memory bytes. - Runs CPU work through `spawn_blocking`, with at most two PDF documents processed - concurrently per process. + concurrently per process. The blocking task owns its concurrency permit, so caller + cancellation cannot admit replacement work while parsing continues. - Does not perform OCR. Responses identify OCR-required or encoding-problem pages through quality warnings and `suggested_next_action: "use_ocr"`. - Extracted output is capped by the same configured `max_body_size`; partial PDF input diff --git a/knowledge/security/threat-model.md b/knowledge/security/threat-model.md index a4aac0b..8283e36 100644 --- a/knowledge/security/threat-model.md +++ b/knowledge/security/threat-model.md @@ -297,7 +297,7 @@ to protect against changes between validation and write. | TM-DOS-004 | Rapid request flooding via tool | Low | No rate limiting in Fetchkit; caller responsibility | **CALLER RISK** | | TM-DOS-005 | DNS resolution delay | Low | DNS resolution uses system resolver; no explicit timeout on DNS lookup | **ACCEPTED** | | TM-DOS-006 | Memory exhaustion from large HTML conversion | Medium | Conversion input bounded by `max_body_size` (10 MB default) | MITIGATED | -| TM-DOS-007 | CPU/memory exhaustion from adversarial PDF parsing | Medium | Bounded input/output, blocking-task isolation, and two-document concurrency cap; parser-internal expansion remains possible | **ACCEPTED** | +| TM-DOS-007 | CPU/memory exhaustion from adversarial PDF parsing | Medium | Bounded input/output, blocking-task isolation, and a cancellation-safe two-document concurrency cap; parser-internal expansion remains possible | **ACCEPTED** | ### Mitigation Details @@ -321,8 +321,10 @@ against unbounded responses (TM-DOS-001). **TM-DOS-007 — Adversarial PDF parsing (ACCEPTED):** PDF content processors receive only bodies bounded by `max_body_size`; partial bodies are rejected before parsing, extracted output is capped again, and synchronous parsing -runs off the async runtime with at most two documents active per process. A crafted PDF -may still cause expensive parser-internal object-stream expansion or prolonged CPU work. +runs off the async runtime with at most two documents active per process. The blocking +task owns its concurrency permit, preventing cancelled callers from starting replacement +parses while non-cancellable work continues. A crafted PDF may still cause expensive +parser-internal object-stream expansion or prolonged CPU work. Operators handling hostile, high-volume PDFs should isolate the process and apply outer request CPU/memory limits. Accepted because the in-process Rust parser has no reliable cancellation boundary once a blocking parse begins.