ragfs_extract/
pdf.rs

1//! PDF content extractor.
2//!
3//! Uses pdf-extract to extract text content and lopdf for embedded images.
4
5use async_trait::async_trait;
6use flate2::read::ZlibDecoder;
7use lopdf::Document;
8use ragfs_core::{
9    ContentElement, ContentExtractor, ContentMetadataInfo, ExtractError, ExtractedContent,
10    ExtractedImage,
11};
12use std::io::Read;
13use std::path::Path;
14use tracing::{debug, warn};
15
16/// Extractor for PDF files.
17pub struct PdfExtractor;
18
19impl PdfExtractor {
20    /// Create a new PDF extractor.
21    #[must_use]
22    pub fn new() -> Self {
23        Self
24    }
25}
26
27impl Default for PdfExtractor {
28    fn default() -> Self {
29        Self::new()
30    }
31}
32
33#[async_trait]
34impl ContentExtractor for PdfExtractor {
35    fn supported_types(&self) -> &[&str] {
36        &["application/pdf"]
37    }
38
39    fn can_extract_by_extension(&self, path: &Path) -> bool {
40        path.extension()
41            .and_then(|ext| ext.to_str())
42            .is_some_and(|ext| ext.eq_ignore_ascii_case("pdf"))
43    }
44
45    async fn extract(&self, path: &Path) -> Result<ExtractedContent, ExtractError> {
46        debug!("Extracting PDF: {:?}", path);
47
48        // Read PDF file
49        let bytes = tokio::fs::read(path).await?;
50
51        // Extract text using pdf-extract (blocking operation)
52        let text = tokio::task::spawn_blocking({
53            let bytes = bytes.clone();
54            move || extract_pdf_text(&bytes)
55        })
56        .await
57        .map_err(|e| ExtractError::Failed(format!("Task join error: {e}")))?
58        .map_err(|e| ExtractError::Failed(format!("PDF extraction failed: {e}")))?;
59
60        // Extract images using lopdf (blocking operation)
61        let images = tokio::task::spawn_blocking(move || extract_pdf_images(&bytes))
62            .await
63            .map_err(|e| ExtractError::Failed(format!("Image extraction task error: {e}")))?;
64
65        // Split into pages/paragraphs for elements
66        let elements = build_elements(&text);
67
68        // Estimate page count from page breaks or text length
69        let page_count = estimate_page_count(&text);
70
71        Ok(ExtractedContent {
72            text,
73            elements,
74            images,
75            metadata: ContentMetadataInfo {
76                page_count: Some(page_count),
77                ..Default::default()
78            },
79        })
80    }
81}
82
83/// Extract text from PDF bytes using pdf-extract.
84fn extract_pdf_text(bytes: &[u8]) -> Result<String, String> {
85    pdf_extract::extract_text_from_mem(bytes).map_err(|e| e.to_string())
86}
87
88/// Configuration for image extraction limits.
89const MAX_IMAGES: usize = 100;
90const MAX_TOTAL_BYTES: usize = 50 * 1024 * 1024; // 50MB
91const MIN_DIMENSION: u32 = 50; // Skip tiny images (icons, etc.)
92
93/// Extract images from PDF document using lopdf.
94fn extract_pdf_images(bytes: &[u8]) -> Vec<ExtractedImage> {
95    let doc = match Document::load_mem(bytes) {
96        Ok(d) => d,
97        Err(e) => {
98            warn!("Failed to load PDF for image extraction: {}", e);
99            return vec![];
100        }
101    };
102
103    let mut images = Vec::new();
104    let mut total_bytes = 0usize;
105
106    let pages = doc.get_pages();
107    for (page_num, page_id) in pages {
108        if images.len() >= MAX_IMAGES {
109            debug!(
110                "Reached maximum image count ({}), stopping extraction",
111                MAX_IMAGES
112            );
113            break;
114        }
115
116        match doc.get_page_images(page_id) {
117            Ok(page_images) => {
118                for pdf_image in page_images {
119                    if images.len() >= MAX_IMAGES || total_bytes >= MAX_TOTAL_BYTES {
120                        break;
121                    }
122
123                    // Skip tiny images
124                    if pdf_image.width < i64::from(MIN_DIMENSION)
125                        || pdf_image.height < i64::from(MIN_DIMENSION)
126                    {
127                        debug!(
128                            "Skipping small image: {}x{}",
129                            pdf_image.width, pdf_image.height
130                        );
131                        continue;
132                    }
133
134                    // Never let a single image, or the running total, exceed the
135                    // cap: a small FlateDecode stream can inflate without bound.
136                    let budget = MAX_TOTAL_BYTES - total_bytes;
137                    if let Some(extracted) = decode_pdf_image(&pdf_image, page_num, budget) {
138                        total_bytes += extracted.data.len();
139                        images.push(extracted);
140                    }
141                }
142            }
143            Err(e) => {
144                debug!("Failed to get images from page {}: {}", page_num, e);
145            }
146        }
147    }
148
149    debug!(
150        "Extracted {} images from PDF ({} bytes total)",
151        images.len(),
152        total_bytes
153    );
154    images
155}
156
157/// Decode a PDF image into `ExtractedImage` format.
158///
159/// `budget` is the number of bytes still available under `MAX_TOTAL_BYTES`;
160/// anything larger is rejected rather than allocated.
161fn decode_pdf_image(
162    pdf_image: &lopdf::xobject::PdfImage<'_>,
163    page_num: u32,
164    budget: usize,
165) -> Option<ExtractedImage> {
166    let filters = pdf_image.filters.as_ref()?;
167
168    // Determine MIME type and decode based on filter
169    let (data, mime_type) = if filters.iter().any(|f| f == "DCTDecode") {
170        // JPEG - can use raw content directly
171        (pdf_image.content.to_vec(), "image/jpeg".to_string())
172    } else if filters.iter().any(|f| f == "FlateDecode") {
173        // Compressed raw image data - decompress and convert to PNG
174        match decode_flate_image(pdf_image, budget) {
175            Ok((data, mime)) => (data, mime),
176            Err(e) => {
177                debug!("Failed to decode FlateDecode image: {}", e);
178                return None;
179            }
180        }
181    } else if filters.iter().any(|f| f == "JPXDecode") {
182        // JPEG 2000 - use raw content
183        (pdf_image.content.to_vec(), "image/jp2".to_string())
184    } else {
185        // Unsupported filter
186        debug!("Unsupported image filter: {:?}", filters);
187        return None;
188    };
189
190    // The DCT/JPX paths are bounded by the file size, but they still have to
191    // respect the remaining budget.
192    if data.len() > budget {
193        debug!(
194            "Skipping image: {} bytes exceeds the remaining budget of {} bytes",
195            data.len(),
196            budget
197        );
198        return None;
199    }
200
201    Some(ExtractedImage {
202        data,
203        mime_type,
204        caption: None, // Will be filled by vision model in future
205        page: Some(page_num),
206    })
207}
208
209/// Decode `FlateDecode` compressed image to PNG.
210fn decode_flate_image(
211    pdf_image: &lopdf::xobject::PdfImage<'_>,
212    budget: usize,
213) -> Result<(Vec<u8>, String), String> {
214    // Decompress with a hard bound. `read_to_end` on a hostile stream is a
215    // decompression bomb; read at most one byte past the budget so an oversized
216    // image is rejected instead of silently truncated.
217    let mut decoder = ZlibDecoder::new(pdf_image.content);
218    let mut decompressed = Vec::new();
219    decoder
220        .by_ref()
221        .take(budget as u64 + 1)
222        .read_to_end(&mut decompressed)
223        .map_err(|e| format!("Decompression failed: {e}"))?;
224
225    if decompressed.len() > budget {
226        return Err(format!(
227            "Decompressed image exceeds the remaining budget of {budget} bytes"
228        ));
229    }
230
231    // Determine color space and create image
232    let color_space = pdf_image.color_space.as_deref().unwrap_or("DeviceRGB");
233    let width = pdf_image.width as u32;
234    let height = pdf_image.height as u32;
235
236    let img = match color_space {
237        "DeviceRGB" | "RGB" => image::RgbImage::from_raw(width, height, decompressed)
238            .map(image::DynamicImage::ImageRgb8),
239        "DeviceGray" | "Gray" => image::GrayImage::from_raw(width, height, decompressed)
240            .map(image::DynamicImage::ImageLuma8),
241        "DeviceCMYK" | "CMYK" => {
242            // Convert CMYK to RGB
243            let rgb_data = cmyk_to_rgb(&decompressed);
244            image::RgbImage::from_raw(width, height, rgb_data).map(image::DynamicImage::ImageRgb8)
245        }
246        _ => {
247            // Attempt RGB as fallback
248            debug!("Unknown color space '{}', attempting RGB", color_space);
249            image::RgbImage::from_raw(width, height, decompressed)
250                .map(image::DynamicImage::ImageRgb8)
251        }
252    };
253
254    let img = img.ok_or_else(|| "Failed to create image from raw data".to_string())?;
255
256    // Encode to PNG
257    let mut png_data = Vec::new();
258    img.write_to(
259        &mut std::io::Cursor::new(&mut png_data),
260        image::ImageFormat::Png,
261    )
262    .map_err(|e| format!("PNG encoding failed: {e}"))?;
263
264    Ok((png_data, "image/png".to_string()))
265}
266
267/// Convert CMYK bytes to RGB.
268#[allow(clippy::many_single_char_names)]
269fn cmyk_to_rgb(cmyk: &[u8]) -> Vec<u8> {
270    let mut rgb = Vec::with_capacity((cmyk.len() / 4) * 3);
271    for chunk in cmyk.as_chunks::<4>().0 {
272        let c = f32::from(chunk[0]) / 255.0;
273        let m = f32::from(chunk[1]) / 255.0;
274        let y = f32::from(chunk[2]) / 255.0;
275        let k = f32::from(chunk[3]) / 255.0;
276
277        let r = 255.0 * (1.0 - c) * (1.0 - k);
278        let g = 255.0 * (1.0 - m) * (1.0 - k);
279        let b = 255.0 * (1.0 - y) * (1.0 - k);
280
281        #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
282        {
283            rgb.push(r as u8);
284            rgb.push(g as u8);
285            rgb.push(b as u8);
286        }
287    }
288    rgb
289}
290
291/// Build `ContentElements` from extracted text.
292fn build_elements(text: &str) -> Vec<ContentElement> {
293    let mut elements = Vec::new();
294    let mut current_offset = 0u64;
295
296    // Split by double newlines to get paragraphs
297    for paragraph in text.split("\n\n") {
298        let trimmed = paragraph.trim();
299        if trimmed.is_empty() {
300            current_offset += paragraph.len() as u64 + 2; // +2 for \n\n
301            continue;
302        }
303
304        // Check if it looks like a heading (short, possibly capitalized)
305        if looks_like_heading(trimmed) {
306            elements.push(ContentElement::Heading {
307                level: 1,
308                text: trimmed.to_string(),
309                byte_offset: current_offset,
310            });
311        } else {
312            elements.push(ContentElement::Paragraph {
313                text: trimmed.to_string(),
314                byte_offset: current_offset,
315            });
316        }
317
318        current_offset += paragraph.len() as u64 + 2;
319    }
320
321    elements
322}
323
324/// Heuristic to detect if text looks like a heading.
325fn looks_like_heading(text: &str) -> bool {
326    // Short text (likely a title/heading)
327    if text.len() > 100 {
328        return false;
329    }
330
331    // No period at end (headings typically don't end with periods)
332    if text.ends_with('.') {
333        return false;
334    }
335
336    // Single line
337    if text.contains('\n') {
338        return false;
339    }
340
341    // All caps or title case with few words
342    let words: Vec<&str> = text.split_whitespace().collect();
343    if words.len() <= 8 {
344        // Check if mostly capitalized
345        let caps_count = words
346            .iter()
347            .filter(|w| w.chars().next().is_some_and(char::is_uppercase))
348            .count();
349        return caps_count >= words.len() / 2;
350    }
351
352    false
353}
354
355/// Estimate page count from text.
356fn estimate_page_count(text: &str) -> u32 {
357    // Look for form feed characters (page breaks)
358    let form_feeds = text.matches('\x0C').count();
359    if form_feeds > 0 {
360        return (form_feeds + 1) as u32;
361    }
362
363    // Estimate based on character count (~3000 chars per page average)
364    let chars = text.len();
365    std::cmp::max(1, (chars / 3000) as u32)
366}
367
368#[cfg(test)]
369mod tests {
370    use super::*;
371
372    #[test]
373    fn test_looks_like_heading() {
374        assert!(looks_like_heading("Chapter 1"));
375        assert!(looks_like_heading("INTRODUCTION"));
376        assert!(looks_like_heading("The Quick Brown Fox"));
377        assert!(!looks_like_heading("This is a normal sentence."));
378        assert!(!looks_like_heading(
379            "This is a very long paragraph that goes on and on and definitely is not a heading"
380        ));
381    }
382
383    #[test]
384    fn test_estimate_page_count() {
385        assert_eq!(estimate_page_count("short"), 1);
386        assert_eq!(estimate_page_count(&"x".repeat(6000)), 2);
387        assert_eq!(estimate_page_count("page1\x0Cpage2\x0Cpage3"), 3);
388    }
389
390    #[test]
391    fn test_build_elements() {
392        let text = "Title\n\nFirst paragraph here.\n\nSecond paragraph.";
393        let elements = build_elements(text);
394        assert_eq!(elements.len(), 3);
395    }
396
397    /// A hostile PDF can make a tiny `FlateDecode` stream inflate without
398    /// bound, so the decoder must refuse anything above the remaining budget.
399    #[test]
400    fn test_decode_flate_image_respects_budget() {
401        use flate2::Compression;
402        use flate2::write::ZlibEncoder;
403        use std::io::Write;
404
405        // 4x4 DeviceRGB = 48 bytes of pixel data.
406        let pixels = vec![7u8; 48];
407        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
408        encoder.write_all(&pixels).unwrap();
409        let compressed = encoder.finish().unwrap();
410
411        let dict = lopdf::Dictionary::new();
412        let image = lopdf::xobject::PdfImage {
413            id: (1, 0),
414            width: 4,
415            height: 4,
416            color_space: Some("DeviceRGB".to_string()),
417            filters: Some(vec!["FlateDecode".to_string()]),
418            bits_per_component: Some(8),
419            content: &compressed,
420            origin_dict: &dict,
421        };
422
423        // A budget below the decompressed size must be refused, not truncated.
424        assert!(decode_flate_image(&image, 16).is_err());
425
426        // A sufficient budget decodes to a PNG.
427        let (png, mime) = decode_flate_image(&image, 4096).unwrap();
428        assert_eq!(mime, "image/png");
429        assert!(!png.is_empty());
430    }
431}
ragfs_extract/pdf.rs

ragfs_extract/
pdf.rs
