kanidmd_lib/valueset/image/
jpg.rs1use std::io::Cursor;
2
3use image::codecs::jpeg::JpegDecoder;
4use image::ImageDecoder;
5use sketching::*;
6
7use super::ImageValidationError;
8
9const JPEG_MAGIC: [u8; 2] = [0xff, 0xd8];
10const EOI_MAGIC: [u8; 2] = [0xff, 0xd9];
11const SOS_MARKER: [u8; 2] = [0xff, 0xda];
12
13pub fn check_jpg_header(contents: &[u8]) -> Result<(), ImageValidationError> {
15    if !contents.starts_with(&JPEG_MAGIC) {
16        return Err(ImageValidationError::InvalidImage(
17            "Failed to parse JPEG file, invalid magic bytes".to_string(),
18        ));
19    }
20    Ok(())
21}
22
23pub fn has_trailer(contents: &Vec<u8>) -> Result<bool, ImageValidationError> {
27    let buf = contents.as_slice();
28
29    let mut pos = JPEG_MAGIC.len();
30
31    while pos < buf.len() {
32        let marker = &buf[pos..pos + 2];
33        pos += 2;
34
35        let segment_size_bytes: &[u8] = &buf[pos..pos + 2];
36        let segment_size = u16::from_be_bytes(segment_size_bytes.try_into().map_err(|_| {
37            ImageValidationError::InvalidImage("JPEG segment size bytes were invalid!".to_string())
38        })?);
39        pos += segment_size as usize;
41
42        if marker == SOS_MARKER {
43            break;
44        }
45    }
46
47    let mut eoi_index = buf.len() * 2;
49    trace!("buffer length: {}", buf.len());
50
51    for i in pos..=(buf.len() - EOI_MAGIC.len()) {
53        if buf[i..(i + EOI_MAGIC.len())] == EOI_MAGIC {
54            eoi_index = i;
55            break;
56        }
57    }
58
59    if eoi_index > buf.len() {
60        Err(ImageValidationError::InvalidImage(
61            "End of image magic bytes not found in JPEG".to_string(),
62        ))
63    } else if (eoi_index + 2) < buf.len() {
64        debug!(
66            "we're at pos: {} and buf len is {}, is not OK",
67            eoi_index,
68            buf.len()
69        );
70        Ok(true)
71    } else {
72        debug!(
73            "we're at pos: {} and buf len is {}, is OK",
74            eoi_index,
75            buf.len()
76        );
77        Ok(false)
78    }
79}
80
81pub fn validate_decoding(
82    filename: &str,
83    contents: &[u8],
84    limits: image::Limits,
85) -> Result<(), ImageValidationError> {
86    let mut decoder = match JpegDecoder::new(Cursor::new(contents)) {
87        Ok(val) => val,
88        Err(err) => {
89            return Err(ImageValidationError::InvalidImage(format!(
90                "Failed to parse {filename} as JPG: {err:?}"
91            )))
92        }
93    };
94
95    match decoder.set_limits(limits) {
96        Err(err) => {
97            sketching::admin_warn!(
98                "Image validation failed while validating {}: {:?}",
99                filename,
100                err
101            );
102            Err(ImageValidationError::ExceedsMaxDimensions)
103        }
104        Ok(_) => Ok(()),
105    }
106}
107
108#[test]
109fn test_jpg_has_trailer() {
110    let file_contents = std::fs::read(format!(
111        "{}/src/valueset/image/test_images/oversize_dimensions.jpg",
112        env!("CARGO_MANIFEST_DIR")
113    ))
114    .expect("Failed to read file");
115    assert!(!has_trailer(&file_contents).expect("Failed to check for JPEG trailer"));
116
117    let file_contents = std::fs::read(format!(
119        "{}/src/valueset/image/test_images/windows11_3_cropped.jpg",
120        env!("CARGO_MANIFEST_DIR")
121    ))
122    .expect("Failed to read file");
123    assert!(has_trailer(&file_contents).expect("Failed to check for JPEG trailer"));
126}