kanidmd_lib/valueset/image/
png.rs

Help
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_END: &[u8; 4] = b"IEND";
5
6#[derive(Debug)]
7/// This is used as part of PNG validation to identify if we've seen the end of the file, and if it suffers from
8/// Acropalypse issues by having trailing data.
9enum PngChunkStatus {
10    SeenEnd { has_trailer: bool },
11    MoreChunks,
12}
13
14/// Loop over the PNG file contents to find out if we've got valid chunks
15fn png_consume_chunks_until_iend(
16    buf: &[u8],
17) -> Result<(PngChunkStatus, &[u8]), ImageValidationError> {
18    // length[u8;4] + chunk_type[u8;4] + checksum[u8;4] + minimum size
19    if buf.len() < 12 {
20        return Err(ImageValidationError::InvalidImage(format!(
21            "PNG file is too short to be valid, got {} bytes",
22            buf.len()
23        )));
24    } else {
25        #[cfg(any(debug_assertions, test))]
26        trace!("input buflen: {}", buf.len());
27    }
28    let (length_bytes, buf) = buf.split_at(4);
29    let (chunk_type, buf) = buf.split_at(4);
30
31    // Infallible: We've definitely consumed 4 bytes
32    let length = u32::from_be_bytes(
33        length_bytes
34            .try_into()
35            .map_err(|_| ImageValidationError::InvalidImage("PNG corrupt!".to_string()))?,
36    );
37    #[cfg(any(debug_assertions, test))]
38    trace!(
39        "length_bytes: {:?} length: {} chunk_type: {:?} buflen: {}",
40        length_bytes,
41        &length,
42        &chunk_type,
43        &buf.len()
44    );
45
46    if buf.len() < (length + 4) as usize {
47        return Err(ImageValidationError::InvalidImage(format!(
48            "PNG file is too short to be valid, failed to split at the chunk length {}, had {} bytes",
49            length,
50            buf.len(),
51        )));
52    }
53    let (_, buf) = buf.split_at(length as usize);
54    #[cfg(any(debug_assertions, test))]
55    trace!("new buflen: {}", &buf.len());
56
57    let (_checksum, buf) = buf.split_at(4);
58    #[cfg(any(debug_assertions, test))]
59    trace!("post-checksum buflen: {}", &buf.len());
60
61    if chunk_type == PNG_CHUNK_END {
62        if buf.is_empty() {
63            Ok((PngChunkStatus::SeenEnd { has_trailer: false }, buf))
64        } else {
65            Ok((PngChunkStatus::SeenEnd { has_trailer: true }, buf))
66        }
67    } else {
68        Ok((PngChunkStatus::MoreChunks, buf))
69    }
70}
71
72// needs to be pub for bench things
73pub fn png_has_trailer(contents: &Vec<u8>) -> Result<bool, ImageValidationError> {
74    let buf = contents.as_slice();
75    // let magic = buf.split_off(PNG_PRELUDE.len());
76    let (magic, buf) = buf.split_at(PNG_PRELUDE.len());
77
78    let buf = buf.to_owned();
79    let mut buf = buf.as_slice();
80
81    if magic != PNG_PRELUDE {
82        return Err(ImageValidationError::InvalidPngPrelude);
83    }
84
85    loop {
86        let (status, new_buf) = png_consume_chunks_until_iend(buf)?;
87        buf = match status {
88            PngChunkStatus::SeenEnd { has_trailer } => return Ok(has_trailer),
89            PngChunkStatus::MoreChunks => new_buf,
90        };
91    }
92}
93
94// needs to be pub for bench things
95pub fn png_lodepng_validate(
96    contents: &Vec<u8>,
97    filename: &str,
98) -> Result<(), ImageValidationError> {
99    match lodepng::decode32(contents) {
100        Ok(val) => {
101            if val.width > MAX_IMAGE_WIDTH as usize || val.height > MAX_IMAGE_HEIGHT as usize {
102                admin_debug!(
103                    "PNG validation failed for {} {}",
104                    filename,
105                    ImageValidationError::ExceedsMaxWidth
106                );
107                Err(ImageValidationError::ExceedsMaxWidth)
108            } else if val.height > MAX_IMAGE_HEIGHT as usize {
109                admin_debug!(
110                    "PNG validation failed for {} {}",
111                    filename,
112                    ImageValidationError::ExceedsMaxHeight
113                );
114                Err(ImageValidationError::ExceedsMaxHeight)
115            } else {
116                Ok(())
117            }
118        }
119        Err(err) => {
120            // admin_debug!("PNG validation failed for {} {:?}", self.filename, err);
121            Err(ImageValidationError::InvalidImage(format!("{:?}", err)))
122        }
123    }
124}
125
126#[test]
127/// this tests a variety of input options for `png_consume_chunks_until_iend`
128fn test_png_consume_chunks_until_iend() {
129    let mut testchunks = vec![0, 0, 0, 1]; // the length
130
131    testchunks.extend(PNG_CHUNK_END); // ... the type of chunk we're looking at!
132    testchunks.push(1); // the data
133    testchunks.extend([0, 0, 0, 1]); // the 4-byte checksum which we ignore
134    let expected: [u8; 0] = [];
135    let testchunks_slice = testchunks.as_slice();
136    let res = png_consume_chunks_until_iend(testchunks_slice);
137
138    // simple, valid image works
139    match res {
140        Ok((result, buf)) => {
141            if let PngChunkStatus::MoreChunks = result {
142                panic!("Shouldn't have more chunks!");
143            }
144            assert_eq!(buf, &expected);
145        }
146        Err(err) => panic!("Error: {:?}", err),
147    };
148
149    // let's make sure it works with a bunch of different length inputs
150    let mut x = 11;
151    while x > 0 {
152        let newslice = &testchunks_slice[0..=x];
153        let res = png_consume_chunks_until_iend(newslice);
154        trace!("chunkstatus at size {} {:?}", x, &res);
155        assert!(res.is_err());
156        x -= 1;
157    }
158}
kanidmd_lib/valueset/image/png.rs

kanidmd_lib/valueset/image/
png.rs
