Skip to main content

oxidize_pdf/recovery/
scanner.rs

1//! PDF object scanner for recovery operations
2
3use crate::error::Result;
4use std::collections::HashMap;
5use std::fs::File;
6use std::io::{BufReader, Read};
7use std::path::Path;
8
9/// Scanner for finding valid PDF objects
10pub struct ObjectScanner {
11    /// Found objects indexed by ID
12    objects: HashMap<u32, ScannedObject>,
13    /// Current scan statistics
14    stats: ScanStats,
15}
16
17/// A scanned PDF object
18#[derive(Debug, Clone)]
19pub struct ScannedObject {
20    /// Object ID
21    pub id: u32,
22    /// Generation number
23    pub generation: u16,
24    /// File offset
25    pub offset: u64,
26    /// Object type if detected
27    pub object_type: Option<ObjectType>,
28    /// Whether object appears valid
29    pub is_valid: bool,
30}
31
32/// Types of PDF objects
33#[derive(Debug, Clone, PartialEq)]
34pub enum ObjectType {
35    Page,
36    Pages,
37    Catalog,
38    Font,
39    Image,
40    Stream,
41    Dictionary,
42    Array,
43    Other(String),
44}
45
46/// Scan statistics
47#[derive(Debug, Default, Clone)]
48pub struct ScanStats {
49    /// Total bytes scanned
50    pub bytes_scanned: u64,
51    /// Number of objects found
52    pub objects_found: usize,
53    /// Number of valid objects
54    pub valid_objects: usize,
55    /// Number of pages found
56    pub pages_found: usize,
57    /// Scan duration in milliseconds
58    pub duration_ms: u64,
59}
60
61/// Result of scanning operation
62#[derive(Debug)]
63pub struct ScanResult {
64    /// All found objects
65    pub objects: Vec<ScannedObject>,
66    /// Total objects found
67    pub total_objects: usize,
68    /// Valid objects
69    pub valid_objects: usize,
70    /// Estimated page count
71    pub estimated_pages: u32,
72    /// Scan statistics
73    pub stats: ScanStats,
74}
75
76impl Default for ObjectScanner {
77    fn default() -> Self {
78        Self::new()
79    }
80}
81
82impl ObjectScanner {
83    /// Create a new object scanner
84    pub fn new() -> Self {
85        Self {
86            objects: HashMap::new(),
87            stats: ScanStats::default(),
88        }
89    }
90
91    /// Scan a file for PDF objects
92    pub fn scan_file<P: AsRef<Path>>(&mut self, path: P) -> Result<ScanResult> {
93        let start_time = std::time::Instant::now();
94
95        let mut file = File::open(path)?;
96        let mut reader = BufReader::new(&mut file);
97
98        // Read file in chunks
99        let mut buffer = vec![0u8; 1024 * 1024]; // 1MB chunks
100        let mut file_offset = 0u64;
101
102        loop {
103            match reader.read(&mut buffer) {
104                Ok(0) => break, // EOF
105                Ok(n) => {
106                    self.scan_buffer(&buffer[..n], file_offset)?;
107                    file_offset += n as u64;
108                    self.stats.bytes_scanned += n as u64;
109                }
110                Err(e) => return Err(e.into()),
111            }
112        }
113
114        self.stats.duration_ms = start_time.elapsed().as_millis() as u64;
115
116        // Build result
117        let mut objects: Vec<_> = self.objects.values().cloned().collect();
118        objects.sort_by_key(|obj| obj.id);
119
120        let result = ScanResult {
121            total_objects: objects.len(),
122            valid_objects: objects.iter().filter(|o| o.is_valid).count(),
123            estimated_pages: self.stats.pages_found as u32,
124            objects,
125            stats: self.stats.clone(),
126        };
127
128        Ok(result)
129    }
130
131    /// Scan a buffer for objects
132    fn scan_buffer(&mut self, buffer: &[u8], base_offset: u64) -> Result<()> {
133        let mut pos = 0;
134
135        while pos < buffer.len() {
136            // Look for object pattern: "N G obj"
137            if let Some(obj_start) = find_object_start(&buffer[pos..]) {
138                let absolute_pos = pos + obj_start;
139
140                // Try to parse object header
141                if let Some((id, gen)) = parse_object_header(&buffer[pos..absolute_pos]) {
142                    let object_offset = base_offset + pos as u64;
143
144                    // Scan object content
145                    let object_type = detect_object_type(&buffer[absolute_pos..]);
146                    let is_valid = validate_object(&buffer[absolute_pos..]);
147
148                    if object_type == Some(ObjectType::Page) {
149                        self.stats.pages_found += 1;
150                    }
151
152                    let scanned_obj = ScannedObject {
153                        id,
154                        generation: gen,
155                        offset: object_offset,
156                        object_type,
157                        is_valid,
158                    };
159
160                    self.objects.insert(id, scanned_obj);
161                    self.stats.objects_found += 1;
162
163                    if is_valid {
164                        self.stats.valid_objects += 1;
165                    }
166                }
167
168                pos = absolute_pos + 4; // Skip past " obj"
169            } else {
170                break;
171            }
172        }
173
174        Ok(())
175    }
176
177    /// Get scan statistics
178    pub fn stats(&self) -> &ScanStats {
179        &self.stats
180    }
181
182    /// Reset scanner state
183    pub fn reset(&mut self) {
184        self.objects.clear();
185        self.stats = ScanStats::default();
186    }
187}
188
189fn find_object_start(buffer: &[u8]) -> Option<usize> {
190    buffer
191        .windows(4)
192        .position(|window| window == b" obj")
193        .map(|pos| pos + 1) // Return position after the space in " obj"
194}
195
196fn parse_object_header(buffer: &[u8]) -> Option<(u32, u16)> {
197    // Look backwards for "ID GEN obj" pattern
198    let text = std::str::from_utf8(buffer).ok()?;
199    let parts: Vec<&str> = text.split_whitespace().collect();
200
201    if parts.len() >= 2 {
202        let id = parts[parts.len() - 2].parse().ok()?;
203        let gen = parts[parts.len() - 1].parse().ok()?;
204        Some((id, gen))
205    } else {
206        None
207    }
208}
209
210fn detect_object_type(content: &[u8]) -> Option<ObjectType> {
211    // Simple type detection based on content
212    if content.is_empty() {
213        return None;
214    }
215
216    let text = String::from_utf8_lossy(&content[..content.len().min(200)]);
217
218    if text.contains("/Type /Page") && !text.contains("/Type /Pages") {
219        Some(ObjectType::Page)
220    } else if text.contains("/Type /Pages") {
221        Some(ObjectType::Pages)
222    } else if text.contains("/Type /Catalog") {
223        Some(ObjectType::Catalog)
224    } else if text.contains("/Type /Font") {
225        Some(ObjectType::Font)
226    } else if text.contains("/Subtype /Image") {
227        Some(ObjectType::Image)
228    } else if text.contains("stream") {
229        Some(ObjectType::Stream)
230    } else if text.starts_with("<<") {
231        Some(ObjectType::Dictionary)
232    } else if text.starts_with('[') {
233        Some(ObjectType::Array)
234    } else {
235        None
236    }
237}
238
239fn validate_object(content: &[u8]) -> bool {
240    // Check if object has proper ending
241    if content.len() < 10 {
242        return false;
243    }
244
245    // Look for endobj
246    content.windows(6).any(|window| window == b"endobj")
247}
248
249/// Quick scan for basic file info
250pub fn quick_scan<P: AsRef<Path>>(path: P) -> Result<ScanResult> {
251    let mut scanner = ObjectScanner::new();
252    scanner.scan_file(path)
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258
259    #[test]
260    fn test_object_scanner_creation() {
261        let scanner = ObjectScanner::new();
262        assert_eq!(scanner.stats.objects_found, 0);
263        assert_eq!(scanner.stats.valid_objects, 0);
264    }
265
266    #[test]
267    fn test_find_object_start() {
268        let buffer = b"some text 1 0 obj content";
269        assert_eq!(find_object_start(buffer), Some(14));
270
271        let buffer = b"this has no pdf data";
272        assert_eq!(find_object_start(buffer), None);
273    }
274
275    #[test]
276    fn test_parse_object_header() {
277        let buffer = b"1 0";
278        assert_eq!(parse_object_header(buffer), Some((1, 0)));
279
280        let buffer = b"123 5";
281        assert_eq!(parse_object_header(buffer), Some((123, 5)));
282    }
283
284    #[test]
285    fn test_detect_object_type() {
286        let page = b"<< /Type /Page /Parent 2 0 R >>";
287        assert_eq!(detect_object_type(page), Some(ObjectType::Page));
288
289        let catalog = b"<< /Type /Catalog /Pages 2 0 R >>";
290        assert_eq!(detect_object_type(catalog), Some(ObjectType::Catalog));
291
292        let stream = b"<< /Length 100 >> stream";
293        assert_eq!(detect_object_type(stream), Some(ObjectType::Stream));
294    }
295
296    #[test]
297    fn test_validate_object() {
298        let valid = b"<< /Type /Page >> endobj";
299        assert!(validate_object(valid));
300
301        let invalid = b"<< /Type /Page >>";
302        assert!(!validate_object(invalid));
303    }
304
305    #[test]
306    fn test_object_scanner_default() {
307        let scanner = ObjectScanner::default();
308        assert_eq!(scanner.stats.objects_found, 0);
309        assert_eq!(scanner.stats.valid_objects, 0);
310        assert_eq!(scanner.stats.bytes_scanned, 0);
311        assert_eq!(scanner.stats.pages_found, 0);
312        assert_eq!(scanner.stats.duration_ms, 0);
313    }
314
315    #[test]
316    fn test_scanned_object_creation() {
317        let obj = ScannedObject {
318            id: 42,
319            generation: 0,
320            offset: 1024,
321            object_type: Some(ObjectType::Page),
322            is_valid: true,
323        };
324
325        assert_eq!(obj.id, 42);
326        assert_eq!(obj.generation, 0);
327        assert_eq!(obj.offset, 1024);
328        assert_eq!(obj.object_type, Some(ObjectType::Page));
329        assert!(obj.is_valid);
330    }
331
332    #[test]
333    fn test_scanned_object_debug_clone() {
334        let obj = ScannedObject {
335            id: 123,
336            generation: 5,
337            offset: 2048,
338            object_type: Some(ObjectType::Font),
339            is_valid: false,
340        };
341
342        let debug_str = format!("{obj:?}");
343        assert!(debug_str.contains("ScannedObject"));
344        assert!(debug_str.contains("123"));
345
346        let cloned = obj.clone();
347        assert_eq!(cloned.id, obj.id);
348        assert_eq!(cloned.generation, obj.generation);
349        assert_eq!(cloned.offset, obj.offset);
350        assert_eq!(cloned.object_type, obj.object_type);
351        assert_eq!(cloned.is_valid, obj.is_valid);
352    }
353
354    #[test]
355    fn test_object_type_variants() {
356        let types = vec![
357            ObjectType::Page,
358            ObjectType::Pages,
359            ObjectType::Catalog,
360            ObjectType::Font,
361            ObjectType::Image,
362            ObjectType::Stream,
363            ObjectType::Dictionary,
364            ObjectType::Array,
365            ObjectType::Other("Custom".to_string()),
366        ];
367
368        for obj_type in types {
369            let debug_str = format!("{obj_type:?}");
370            assert!(!debug_str.is_empty());
371
372            let cloned = obj_type.clone();
373            assert_eq!(obj_type, cloned);
374        }
375    }
376
377    #[test]
378    fn test_object_type_equality() {
379        assert_eq!(ObjectType::Page, ObjectType::Page);
380        assert_ne!(ObjectType::Page, ObjectType::Pages);
381        assert_ne!(ObjectType::Font, ObjectType::Image);
382
383        assert_eq!(
384            ObjectType::Other("Test".to_string()),
385            ObjectType::Other("Test".to_string())
386        );
387        assert_ne!(
388            ObjectType::Other("Test1".to_string()),
389            ObjectType::Other("Test2".to_string())
390        );
391    }
392
393    #[test]
394    fn test_scan_stats_default() {
395        let stats = ScanStats::default();
396        assert_eq!(stats.bytes_scanned, 0);
397        assert_eq!(stats.objects_found, 0);
398        assert_eq!(stats.valid_objects, 0);
399        assert_eq!(stats.pages_found, 0);
400        assert_eq!(stats.duration_ms, 0);
401    }
402
403    #[test]
404    fn test_scan_stats_debug_clone() {
405        let stats = ScanStats {
406            bytes_scanned: 1024 * 1024,
407            objects_found: 100,
408            valid_objects: 95,
409            pages_found: 10,
410            duration_ms: 250,
411        };
412
413        let debug_str = format!("{stats:?}");
414        assert!(debug_str.contains("ScanStats"));
415
416        let cloned = stats.clone();
417        assert_eq!(cloned.bytes_scanned, stats.bytes_scanned);
418        assert_eq!(cloned.objects_found, stats.objects_found);
419        assert_eq!(cloned.valid_objects, stats.valid_objects);
420        assert_eq!(cloned.pages_found, stats.pages_found);
421        assert_eq!(cloned.duration_ms, stats.duration_ms);
422    }
423
424    #[test]
425    fn test_scan_result_creation() {
426        let objects = vec![
427            ScannedObject {
428                id: 1,
429                generation: 0,
430                offset: 100,
431                object_type: Some(ObjectType::Catalog),
432                is_valid: true,
433            },
434            ScannedObject {
435                id: 2,
436                generation: 0,
437                offset: 200,
438                object_type: Some(ObjectType::Pages),
439                is_valid: true,
440            },
441        ];
442
443        let result = ScanResult {
444            objects: objects,
445            total_objects: 2,
446            valid_objects: 2,
447            estimated_pages: 1,
448            stats: ScanStats::default(),
449        };
450
451        assert_eq!(result.objects.len(), 2);
452        assert_eq!(result.total_objects, 2);
453        assert_eq!(result.valid_objects, 2);
454        assert_eq!(result.estimated_pages, 1);
455    }
456
457    #[test]
458    fn test_scan_result_debug() {
459        let result = ScanResult {
460            objects: vec![],
461            total_objects: 0,
462            valid_objects: 0,
463            estimated_pages: 0,
464            stats: ScanStats::default(),
465        };
466
467        let debug_str = format!("{result:?}");
468        assert!(debug_str.contains("ScanResult"));
469    }
470
471    #[test]
472    fn test_scanner_reset() {
473        let mut scanner = ObjectScanner::new();
474
475        // Add some data
476        scanner.objects.insert(
477            1,
478            ScannedObject {
479                id: 1,
480                generation: 0,
481                offset: 100,
482                object_type: Some(ObjectType::Page),
483                is_valid: true,
484            },
485        );
486        scanner.stats.objects_found = 5;
487        scanner.stats.valid_objects = 4;
488        scanner.stats.bytes_scanned = 1000;
489
490        assert!(!scanner.objects.is_empty());
491        assert_eq!(scanner.stats.objects_found, 5);
492
493        // Reset
494        scanner.reset();
495
496        assert!(scanner.objects.is_empty());
497        assert_eq!(scanner.stats.objects_found, 0);
498        assert_eq!(scanner.stats.valid_objects, 0);
499        assert_eq!(scanner.stats.bytes_scanned, 0);
500    }
501
502    #[test]
503    fn test_detect_object_type_all_types() {
504        // Test Pages type (must check before Page)
505        let pages = b"<< /Type /Pages /Kids [1 0 R 2 0 R] >>";
506        assert_eq!(detect_object_type(pages), Some(ObjectType::Pages));
507
508        // Test Font
509        let font = b"<< /Type /Font /Subtype /Type1 >>";
510        assert_eq!(detect_object_type(font), Some(ObjectType::Font));
511
512        // Test Image
513        let image = b"<< /Subtype /Image /Width 100 >>";
514        assert_eq!(detect_object_type(image), Some(ObjectType::Image));
515
516        // Test Dictionary
517        let dict = b"<< /Key /Value >>";
518        assert_eq!(detect_object_type(dict), Some(ObjectType::Dictionary));
519
520        // Test Array
521        let array = b"[1 2 3 4]";
522        assert_eq!(detect_object_type(array), Some(ObjectType::Array));
523
524        // Test short content
525        let short = b"abc";
526        assert_eq!(detect_object_type(short), None);
527
528        // Test empty
529        let empty = b"";
530        assert_eq!(detect_object_type(empty), None);
531    }
532
533    #[test]
534    fn test_validate_object_various_cases() {
535        // Valid with spaces
536        assert!(validate_object(b"<< /Type /Page >>   endobj   "));
537
538        // Valid with newlines
539        assert!(validate_object(b"<< /Type /Page >>\nendobj"));
540
541        // Valid in middle of content
542        assert!(validate_object(
543            b"<< /Type /Page >> stuff endobj more stuff"
544        ));
545
546        // Invalid - too short
547        assert!(!validate_object(b"short"));
548
549        // Invalid - no endobj
550        assert!(!validate_object(b"<< /Type /Page >> no end marker"));
551
552        // Invalid - partial endobj
553        assert!(!validate_object(b"<< /Type /Page >> endob"));
554    }
555
556    #[test]
557    fn test_parse_object_header_edge_cases() {
558        // Normal case
559        assert_eq!(parse_object_header(b"42 0"), Some((42, 0)));
560
561        // With extra whitespace
562        assert_eq!(parse_object_header(b"  42   0  "), Some((42, 0)));
563
564        // With newlines
565        assert_eq!(parse_object_header(b"42\n0"), Some((42, 0)));
566
567        // Large numbers
568        assert_eq!(parse_object_header(b"999999 65535"), Some((999999, 65535)));
569
570        // Invalid - not numbers
571        assert_eq!(parse_object_header(b"abc def"), None);
572
573        // Invalid - only one number
574        assert_eq!(parse_object_header(b"42"), None);
575
576        // Invalid - empty
577        assert_eq!(parse_object_header(b""), None);
578
579        // Invalid - non-UTF8
580        assert_eq!(parse_object_header(&[0xFF, 0xFE, 0xFD]), None);
581    }
582
583    #[test]
584    fn test_find_object_start_multiple() {
585        let buffer = b"first obj at 9 obj and another obj";
586        assert_eq!(find_object_start(buffer), Some(6)); // First " obj"
587
588        // Find next occurrence
589        let next_search = &buffer[7..];
590        assert_eq!(find_object_start(next_search), Some(8));
591
592        // No obj pattern
593        assert_eq!(find_object_start(b"no_object_pattern_here"), None);
594
595        // obj at the very end
596        assert_eq!(find_object_start(b"ends with obj"), Some(10));
597
598        // obj at the beginning
599        assert_eq!(find_object_start(b" obj starts here"), Some(1));
600    }
601
602    #[test]
603    fn test_scanner_stats_access() {
604        let scanner = ObjectScanner::new();
605        let stats = scanner.stats();
606        assert_eq!(stats.objects_found, 0);
607        assert_eq!(stats.valid_objects, 0);
608        assert_eq!(stats.bytes_scanned, 0);
609        assert_eq!(stats.pages_found, 0);
610        assert_eq!(stats.duration_ms, 0);
611    }
612
613    #[test]
614    fn test_quick_scan_nonexistent_file() {
615        let temp_dir = std::env::temp_dir();
616        let temp_path = temp_dir.join("nonexistent_scanner_test.pdf");
617
618        let result = quick_scan(&temp_path);
619        assert!(result.is_err());
620    }
621
622    #[test]
623    fn test_scan_file_empty() {
624        use std::fs::File;
625
626        let temp_dir = std::env::temp_dir();
627        let temp_path = temp_dir.join("empty_scan_test.pdf");
628        let _file = File::create(&temp_path).unwrap();
629
630        let mut scanner = ObjectScanner::new();
631        let result = scanner.scan_file(&temp_path).unwrap();
632
633        assert_eq!(result.total_objects, 0);
634        assert_eq!(result.valid_objects, 0);
635        assert_eq!(result.estimated_pages, 0);
636        assert!(result.objects.is_empty());
637
638        // Cleanup
639        let _ = std::fs::remove_file(temp_path);
640    }
641
642    #[test]
643    fn test_scan_file_with_objects() {
644        use std::fs::File;
645        use std::io::Write;
646
647        let temp_dir = std::env::temp_dir();
648        let temp_path = temp_dir.join("objects_scan_test.pdf");
649        let mut file = File::create(&temp_path).unwrap();
650
651        // Write some PDF-like content
652        file.write_all(b"%PDF-1.7\n").unwrap();
653        file.write_all(b"1 0 obj\n<< /Type /Catalog >>\nendobj\n")
654            .unwrap();
655        file.write_all(b"2 0 obj\n<< /Type /Pages >>\nendobj\n")
656            .unwrap();
657        file.write_all(b"3 0 obj\n<< /Type /Page >>\nendobj\n")
658            .unwrap();
659        file.write_all(b"%%EOF").unwrap();
660
661        let result = quick_scan(&temp_path).unwrap();
662
663        assert!(result.total_objects > 0);
664        assert_eq!(result.estimated_pages, 1); // One page found
665        assert!(!result.objects.is_empty());
666
667        // Cleanup
668        let _ = std::fs::remove_file(temp_path);
669    }
670
671    #[test]
672    fn test_scan_buffer_various_objects() {
673        let mut scanner = ObjectScanner::new();
674
675        let buffer = b"1 0 obj\n<< /Type /Page >>\nendobj\n2 0 obj\n<< /Type /Font >>\nendobj";
676        scanner.scan_buffer(buffer, 0).unwrap();
677
678        assert_eq!(scanner.stats.objects_found, 2);
679        assert_eq!(scanner.stats.valid_objects, 2);
680        assert_eq!(scanner.stats.pages_found, 1);
681        assert_eq!(scanner.objects.len(), 2);
682    }
683
684    #[test]
685    fn test_scan_buffer_invalid_objects() {
686        let mut scanner = ObjectScanner::new();
687
688        // Objects without endobj
689        let buffer = b"1 0 obj\n<< /Type /Page >>\n2 0 obj\n<< /Type /Font >>";
690        scanner.scan_buffer(buffer, 100).unwrap();
691
692        assert_eq!(scanner.stats.objects_found, 2);
693        assert_eq!(scanner.stats.valid_objects, 0); // No valid objects (missing endobj)
694        assert_eq!(scanner.stats.pages_found, 1); // Page type still detected
695    }
696
697    #[test]
698    fn test_scanned_object_with_offset() {
699        let obj1 = ScannedObject {
700            id: 1,
701            generation: 0,
702            offset: 0,
703            object_type: Some(ObjectType::Catalog),
704            is_valid: true,
705        };
706
707        let obj2 = ScannedObject {
708            id: 2,
709            generation: 0,
710            offset: 1024,
711            object_type: Some(ObjectType::Page),
712            is_valid: true,
713        };
714
715        assert!(obj1.offset < obj2.offset);
716        assert_ne!(obj1.object_type, obj2.object_type);
717    }
718
719    #[test]
720    fn test_object_type_other_variant() {
721        let other1 = ObjectType::Other("CustomType".to_string());
722        let other2 = ObjectType::Other("CustomType".to_string());
723        let other3 = ObjectType::Other("DifferentType".to_string());
724
725        assert_eq!(other1, other2);
726        assert_ne!(other1, other3);
727
728        match &other1 {
729            ObjectType::Other(name) => assert_eq!(name, "CustomType"),
730            _ => panic!("Expected Other variant"),
731        }
732    }
733
734    #[test]
735    fn test_scan_result_sorting() {
736        let mut objects = [
737            ScannedObject {
738                id: 3,
739                generation: 0,
740                offset: 300,
741                object_type: Some(ObjectType::Page),
742                is_valid: true,
743            },
744            ScannedObject {
745                id: 1,
746                generation: 0,
747                offset: 100,
748                object_type: Some(ObjectType::Catalog),
749                is_valid: true,
750            },
751            ScannedObject {
752                id: 2,
753                generation: 0,
754                offset: 200,
755                object_type: Some(ObjectType::Pages),
756                is_valid: true,
757            },
758        ];
759
760        // Simulate what scan_file does
761        objects.sort_by_key(|obj| obj.id);
762
763        assert_eq!(objects[0].id, 1);
764        assert_eq!(objects[1].id, 2);
765        assert_eq!(objects[2].id, 3);
766    }
767}