scrubkit_core/
png.rs

1// File: crates/scrubkit-core/src/png.rs
2
3use crate::{MetadataEntry, ScrubError, ScrubResult, Scrubber};
4use std::io::Cursor;
5
6/// A Scrubber implementation for PNG files.
7#[derive(Debug, Clone)]
8pub struct PngScrubber {
9    file_bytes: Vec<u8>,
10}
11
12impl Scrubber for PngScrubber {
13    fn new(file_bytes: Vec<u8>) -> Result<Self, ScrubError> {
14        // The png::Decoder will fail if it's not a valid PNG, which is a robust check.
15        let decoder = png::Decoder::new(Cursor::new(&file_bytes));
16        if decoder.read_info().is_err() {
17            return Err(ScrubError::UnsupportedFileType(
18                "Not a valid PNG file.".to_string(),
19            ));
20        }
21        Ok(Self { file_bytes })
22    }
23
24    fn view_metadata(&self) -> Result<Vec<MetadataEntry>, ScrubError> {
25        let decoder = png::Decoder::new(Cursor::new(&self.file_bytes));
26        let reader = decoder
27            .read_info()
28            .map_err(|e| ScrubError::ParsingError(e.to_string()))?;
29        let mut metadata = Vec::new();
30
31        // Correctly iterate over the decoded text chunks.
32        for text_chunk in &reader.info().uncompressed_latin1_text {
33            metadata.push(MetadataEntry {
34                category: "tEXt/zTXt/iTXt".to_string(),
35                key: text_chunk.keyword.clone(),
36                value: text_chunk.text.clone(),
37            });
38        }
39
40        Ok(metadata)
41    }
42
43    fn scrub(&self) -> Result<ScrubResult, ScrubError> {
44        let metadata_removed = self.view_metadata()?;
45        if metadata_removed.is_empty() {
46            return Ok(ScrubResult {
47                cleaned_file_bytes: self.file_bytes.clone(),
48                metadata_removed: vec![],
49            });
50        }
51
52        // To scrub, we must re-encode the image while skipping the metadata chunks.
53        let decoder = png::Decoder::new(Cursor::new(&self.file_bytes));
54        let mut reader = decoder
55            .read_info()
56            .map_err(|e| ScrubError::ParsingError(e.to_string()))?;
57
58        // Read the image data itself.
59        let mut img_data = vec![0; reader.output_buffer_size()];
60        let info = reader
61            .next_frame(&mut img_data)
62            .map_err(|e| ScrubError::ParsingError(e.to_string()))?;
63
64        // Create a new PNG in memory
65        let mut cleaned_bytes = Vec::new();
66        {
67            // Create a new scope for the encoder and writer to ensure they are dropped
68            // and release their borrow on `cleaned_bytes` before we return it.
69            let mut encoder =
70                png::Encoder::new(Cursor::new(&mut cleaned_bytes), info.width, info.height);
71            encoder.set_color(info.color_type);
72            encoder.set_depth(info.bit_depth);
73
74            // Crucially, we do *not* write any of the textual metadata chunks to the new encoder.
75
76            let mut writer = encoder
77                .write_header()
78                .map_err(|e| ScrubError::ParsingError(e.to_string()))?;
79
80            writer
81                .write_image_data(&img_data)
82                .map_err(|e| ScrubError::ParsingError(e.to_string()))?;
83        } // encoder and writer are dropped here
84
85        // The `cleaned_bytes` vec now holds the scrubbed PNG.
86        Ok(ScrubResult {
87            cleaned_file_bytes: cleaned_bytes,
88            metadata_removed,
89        })
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    // A simple 1x1 pixel PNG with a tEXt chunk for metadata.
98    // Keyword: "Author", Text: "ScrubKit Tester"
99    const TEST_PNG_WITH_METADATA: &[u8] = &[
100        137, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13, 73, 72, 68, 82, 0, 0, 0, 1, 0, 0, 0, 1, 8, 6,
101        0, 0, 0, 31, 21, 196, 137, 0, 0, 0, 28, 116, 69, 88, 116, 65, 117, 116, 104, 111, 114, 0,
102        83, 99, 114, 117, 98, 75, 105, 116, 32, 84, 101, 115, 116, 101, 114, 215, 122, 61, 248, 0,
103        0, 0, 12, 73, 68, 65, 84, 8, 215, 99, 96, 96, 96, 248, 207, 192, 4, 0, 1, 10, 0, 255, 170,
104        222, 158, 221, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130,
105    ];
106
107    #[test]
108    fn view_metadata_finds_png_text_chunk() {
109        let scrubber = PngScrubber::new(TEST_PNG_WITH_METADATA.to_vec()).unwrap();
110        let metadata = scrubber.view_metadata().unwrap();
111        assert!(!metadata.is_empty());
112        assert_eq!(metadata[0].key, "Author");
113        assert_eq!(metadata[0].value, "ScrubKit Tester");
114    }
115
116    #[test]
117    fn scrub_removes_png_text_chunk() {
118        let scrubber = PngScrubber::new(TEST_PNG_WITH_METADATA.to_vec()).unwrap();
119        let result = scrubber.scrub().unwrap();
120
121        // The most important test is to verify that the *new* file has no metadata.
122        assert!(!result.metadata_removed.is_empty());
123
124        let new_scrubber = PngScrubber::new(result.cleaned_file_bytes).unwrap();
125        let new_metadata = new_scrubber.view_metadata().unwrap();
126        assert!(
127            new_metadata.is_empty(),
128            "Scrubbed file should have no metadata"
129        );
130    }
131}
scrubkit_core/png.rs

scrubkit_core/
png.rs
