kanidmd_lib/valueset/image/
png.rs

1use super::{ImageValidationError, MAX_IMAGE_HEIGHT, MAX_IMAGE_WIDTH};
2use crate::prelude::*;
3static PNG_PRELUDE: &[u8] = &[0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a];
4static PNG_CHUNK_IHDR: &[u8; 4] = b"IHDR";
5static PNG_CHUNK_END: &[u8; 4] = b"IEND";
6
7#[derive(Debug)]
8/// This is used as part of PNG validation to identify if we've seen the end of the file, and if it suffers from
9/// Acropalypse issues by having trailing data.
10enum PngChunkStatus {
11    SeenEnd { has_trailer: bool },
12    MoreChunks,
13}
14
15fn png_split_chunk(buf: &[u8]) -> Result<(&[u8], &[u8], &[u8]), ImageValidationError> {
16    // length[u8;4] + chunk_type[u8;4] + checksum[u8;4] + minimum size
17    if buf.len() < 12 {
18        return Err(ImageValidationError::InvalidImage(format!(
19            "PNG file is too short to be valid, got {} bytes",
20            buf.len()
21        )));
22    } else {
23        #[cfg(any(debug_assertions, test))]
24        trace!("input buflen: {}", buf.len());
25    }
26    let (length_bytes, buf) = buf.split_at(4);
27    let (chunk_type, buf) = buf.split_at(4);
28
29    // Infallible: We've definitely consumed 4 bytes
30    let length = u32::from_be_bytes(
31        length_bytes
32            .try_into()
33            .map_err(|_| ImageValidationError::InvalidImage("PNG corrupt!".to_string()))?,
34    );
35    #[cfg(any(debug_assertions, test))]
36    trace!(
37        "length_bytes: {:?} length: {} chunk_type: {:?} buflen: {}",
38        length_bytes,
39        &length,
40        &chunk_type,
41        &buf.len()
42    );
43
44    if buf.len() < (length as usize).saturating_add(4) {
45        return Err(ImageValidationError::InvalidImage(format!(
46            "PNG file is too short to be valid, failed to split at the chunk length {}, had {} bytes",
47            length,
48            buf.len(),
49        )));
50    }
51    let (chunk_data, buf) = buf.split_at(length as usize);
52    #[cfg(any(debug_assertions, test))]
53    trace!("new buflen: {}", &buf.len());
54
55    let (_checksum, buf) = buf.split_at(4);
56    #[cfg(any(debug_assertions, test))]
57    trace!("post-checksum buflen: {}", &buf.len());
58
59    Ok((chunk_type, chunk_data, buf))
60}
61
62fn png_validate_ihdr(chunk_type: &[u8], chunk_data: &[u8]) -> Result<(), ImageValidationError> {
63    if chunk_type != PNG_CHUNK_IHDR {
64        return Err(ImageValidationError::InvalidImage(
65            "PNG first chunk must be IHDR".to_string(),
66        ));
67    }
68
69    if chunk_data.len() != 13 {
70        return Err(ImageValidationError::InvalidImage(format!(
71            "PNG IHDR chunk must be 13 bytes, got {} bytes",
72            chunk_data.len()
73        )));
74    }
75
76    let (width_chunk, chunk_data) = chunk_data.split_at(4);
77    let (height_chunk, _rest) = chunk_data.split_at(4);
78    let width = u32::from_be_bytes(
79        width_chunk
80            .try_into()
81            .map_err(|_| ImageValidationError::InvalidImage("PNG corrupt!".to_string()))?,
82    );
83    let height = u32::from_be_bytes(
84        height_chunk
85            .try_into()
86            .map_err(|_| ImageValidationError::InvalidImage("PNG corrupt!".to_string()))?,
87    );
88
89    if width == 0 || height == 0 {
90        return Err(ImageValidationError::InvalidImage(
91            "PNG IHDR dimensions must be non-zero".to_string(),
92        ));
93    }
94
95    if width > MAX_IMAGE_WIDTH || height > MAX_IMAGE_HEIGHT {
96        return Err(ImageValidationError::ExceedsMaxDimensions);
97    }
98
99    Ok(())
100}
101
102/// Loop over the PNG file contents to find out if we've got valid chunks
103fn png_consume_chunks_until_iend(
104    buf: &[u8],
105) -> Result<(PngChunkStatus, &[u8]), ImageValidationError> {
106    let (chunk_type, chunk_data, buf) = png_split_chunk(buf)?;
107
108    if chunk_type == PNG_CHUNK_END {
109        if !chunk_data.is_empty() {
110            return Err(ImageValidationError::InvalidImage(
111                "PNG IEND chunk must be empty".to_string(),
112            ));
113        }
114
115        if buf.is_empty() {
116            Ok((PngChunkStatus::SeenEnd { has_trailer: false }, buf))
117        } else {
118            Ok((PngChunkStatus::SeenEnd { has_trailer: true }, buf))
119        }
120    } else {
121        Ok((PngChunkStatus::MoreChunks, buf))
122    }
123}
124
125// needs to be pub for bench things
126pub fn png_has_trailer(contents: &Vec<u8>) -> Result<bool, ImageValidationError> {
127    if contents.len() < PNG_PRELUDE.len() {
128        return Err(ImageValidationError::InvalidImage(format!(
129            "PNG file is too short to be valid, got {} bytes",
130            contents.len()
131        )));
132    }
133
134    let (magic, mut buf) = contents.as_slice().split_at(PNG_PRELUDE.len());
135    if magic != PNG_PRELUDE {
136        return Err(ImageValidationError::InvalidPngPrelude);
137    }
138
139    let (chunk_type, chunk_data, new_buf) = png_split_chunk(buf)?;
140    png_validate_ihdr(chunk_type, chunk_data)?;
141    buf = new_buf;
142
143    loop {
144        let (status, new_buf) = png_consume_chunks_until_iend(buf)?;
145        buf = match status {
146            PngChunkStatus::SeenEnd { has_trailer } => return Ok(has_trailer),
147            PngChunkStatus::MoreChunks => new_buf,
148        };
149    }
150}
151
152// needs to be pub for bench things
153pub fn png_lodepng_validate(
154    contents: &Vec<u8>,
155    filename: &str,
156) -> Result<(), ImageValidationError> {
157    match lodepng::decode32(contents) {
158        Ok(val) => {
159            if val.width > MAX_IMAGE_WIDTH as usize {
160                admin_debug!(
161                    "PNG validation failed for {} {}",
162                    filename,
163                    ImageValidationError::ExceedsMaxWidth
164                );
165                Err(ImageValidationError::ExceedsMaxWidth)
166            } else if val.height > MAX_IMAGE_HEIGHT as usize {
167                admin_debug!(
168                    "PNG validation failed for {} {}",
169                    filename,
170                    ImageValidationError::ExceedsMaxHeight
171                );
172                Err(ImageValidationError::ExceedsMaxHeight)
173            } else {
174                Ok(())
175            }
176        }
177        Err(err) => {
178            // admin_debug!("PNG validation failed for {} {:?}", self.filename, err);
179            Err(ImageValidationError::InvalidImage(format!("{err:?}")))
180        }
181    }
182}
183
184#[test]
185/// this tests a variety of input options for `png_consume_chunks_until_iend`
186fn test_png_consume_chunks_until_iend() {
187    let mut testchunks = vec![0, 0, 0, 0]; // the length
188
189    testchunks.extend(PNG_CHUNK_END); // ... the type of chunk we're looking at!
190    testchunks.extend([0, 0, 0, 1]); // the 4-byte checksum which we ignore
191    let expected: [u8; 0] = [];
192    let testchunks_slice = testchunks.as_slice();
193    let res = png_consume_chunks_until_iend(testchunks_slice);
194
195    // simple, valid image works
196    match res {
197        Ok((result, buf)) => {
198            if let PngChunkStatus::MoreChunks = result {
199                panic!("Shouldn't have more chunks!");
200            }
201            assert_eq!(buf, &expected);
202        }
203        Err(err) => panic!("Error: {err:?}"),
204    };
205
206    // let's make sure it works with a bunch of different length inputs
207    let mut x = 10;
208    while x > 0 {
209        let newslice = &testchunks_slice[0..=x];
210        let res = png_consume_chunks_until_iend(newslice);
211        trace!("chunkstatus at size {} {:?}", x, &res);
212        assert!(res.is_err());
213        x -= 1;
214    }
215}
216
217#[test]
218fn test_png_too_short() {
219    let too_short = vec![0, 0, 0, 1, 2, 3];
220    let res = png_consume_chunks_until_iend(too_short.as_slice());
221    assert!(res.is_err());
222}
223
224#[test]
225fn audit_png_short_input_err() {
226    let short = vec![0x89u8, 0x50, 0x4e, 0x47];
227    assert!(png_has_trailer(&short).is_err());
228}
229
230#[test]
231fn audit_png_chunk_length_overflow_errs() {
232    let mut data = vec![0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a];
233    data.extend_from_slice(&[0xFF, 0xFF, 0xFF, 0xFD]);
234    data.extend_from_slice(PNG_CHUNK_IHDR);
235    data.extend_from_slice(&[0u8; 8]);
236    assert!(png_has_trailer(&data).is_err());
237}
238
239#[cfg(test)]
240fn png_test_chunk(chunk_type: [u8; 4], data: &[u8]) -> Vec<u8> {
241    let mut chunk = Vec::new();
242    chunk.extend_from_slice(&(data.len() as u32).to_be_bytes());
243    chunk.extend_from_slice(&chunk_type);
244    chunk.extend_from_slice(data);
245    chunk.extend_from_slice(&[0u8; 4]);
246    chunk
247}
248
249#[cfg(test)]
250fn png_test_image_with_ihdr(width: u32, height: u32) -> Vec<u8> {
251    let mut ihdr = Vec::new();
252    ihdr.extend_from_slice(&width.to_be_bytes());
253    ihdr.extend_from_slice(&height.to_be_bytes());
254    ihdr.extend_from_slice(&[8, 6, 0, 0, 0]);
255
256    let mut data = PNG_PRELUDE.to_vec();
257    data.extend_from_slice(&png_test_chunk(*PNG_CHUNK_IHDR, &ihdr));
258    data.extend_from_slice(&png_test_chunk(*PNG_CHUNK_END, &[]));
259    data
260}
261
262#[test]
263fn audit_png_signature_without_chunks_errs() {
264    let data = PNG_PRELUDE.to_vec();
265    assert!(png_has_trailer(&data).is_err());
266}
267
268#[test]
269fn audit_png_valid_file_with_trailing_bytes_has_trailer() {
270    let mut data = include_bytes!("test_images/ok.png").to_vec();
271    data.push(0);
272    assert!(png_has_trailer(&data).expect("valid PNG with trailer should parse"));
273}
274
275#[test]
276fn audit_png_nonzero_length_iend_errs() {
277    let mut data = png_test_image_with_ihdr(1, 1);
278    data.truncate(data.len() - 12);
279    data.extend_from_slice(&png_test_chunk(*PNG_CHUNK_END, &[1]));
280    assert!(png_has_trailer(&data).is_err());
281}
282
283#[test]
284fn audit_png_first_chunk_must_be_ihdr() {
285    let mut data = PNG_PRELUDE.to_vec();
286    data.extend_from_slice(&png_test_chunk(*b"tEXt", &[]));
287    data.extend_from_slice(&png_test_chunk(*PNG_CHUNK_END, &[]));
288    assert!(png_has_trailer(&data).is_err());
289}
290
291#[test]
292fn audit_png_ihdr_length_must_be_13() {
293    let mut data = PNG_PRELUDE.to_vec();
294    data.extend_from_slice(&png_test_chunk(*PNG_CHUNK_IHDR, &[]));
295    data.extend_from_slice(&png_test_chunk(*PNG_CHUNK_END, &[]));
296    assert!(png_has_trailer(&data).is_err());
297}
298
299#[test]
300fn audit_png_ihdr_zero_width_or_height_errs() {
301    let zero_width = png_test_image_with_ihdr(0, 1);
302    let zero_height = png_test_image_with_ihdr(1, 0);
303    assert!(png_has_trailer(&zero_width).is_err());
304    assert!(png_has_trailer(&zero_height).is_err());
305}
306
307#[test]
308fn audit_png_ihdr_over_limit_dimensions_err() {
309    let too_wide = png_test_image_with_ihdr(MAX_IMAGE_WIDTH + 1, 1);
310    let too_tall = png_test_image_with_ihdr(1, MAX_IMAGE_HEIGHT + 1);
311    assert!(png_has_trailer(&too_wide).is_err());
312    assert!(png_has_trailer(&too_tall).is_err());
313}
kanidmd_lib/valueset/image/png.rs

kanidmd_lib/valueset/image/
png.rs
