1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

use image::codecs::jpeg::JpegDecoder;
use image::ImageDecoder;
use sketching::*;

use super::ImageValidationError;

const JPEG_MAGIC: [u8; 2] = [0xff, 0xd8];
const EOI_MAGIC: [u8; 2] = [0xff, 0xd9];
const SOS_MARKER: [u8; 2] = [0xff, 0xda];

/// Checks to see if it has a valid JPEG magic bytes header
pub fn check_jpg_header(contents: &[u8]) -> Result<(), ImageValidationError> {
    if !contents.starts_with(&JPEG_MAGIC) {
        return Err(ImageValidationError::InvalidImage(
            "Failed to parse JPEG file, invalid magic bytes".to_string(),
        ));
    }
    Ok(())
}

// It's public so we can use it in benchmarking
/// Check to see if JPG is affected by acropalypse issues, returns `Ok(true)` if it is
/// based on <https://github.com/lordofpipes/acropadetect/blob/main/src/detect.ts>
pub fn has_trailer(contents: &Vec<u8>) -> Result<bool, ImageValidationError> {
    let buf = contents.as_slice();

    let mut pos = JPEG_MAGIC.len();

    while pos < buf.len() {
        let marker = &buf[pos..pos + 2];
        pos += 2;

        let segment_size_bytes: &[u8] = &buf[pos..pos + 2];
        let segment_size = u16::from_be_bytes(segment_size_bytes.try_into().map_err(|_| {
            ImageValidationError::InvalidImage("JPEG segment size bytes were invalid!".to_string())
        })?);
        // we do not add 2 because the size prefix includes the size of the size prefix
        pos += segment_size as usize;

        if marker == SOS_MARKER {
            break;
        }
    }

    // setting this to a big value so we can see if we don't find the EOI marker
    let mut eoi_index = buf.len() * 2;
    trace!("buffer length: {}", buf.len());

    // iterate through the file looking for the EOI_MAGIC bytes
    for i in pos..=(buf.len() - EOI_MAGIC.len()) {
        if buf[i..(i + EOI_MAGIC.len())] == EOI_MAGIC {
            eoi_index = i;
            break;
        }
    }

    if eoi_index > buf.len() {
        Err(ImageValidationError::InvalidImage(
            "End of image magic bytes not found in JPEG".to_string(),
        ))
    } else if (eoi_index + 2) < buf.len() {
        // there's still bytes in the buffer after the EOI magic bytes
        #[cfg(any(test, debug_assertions))]
        println!(
            "we're at pos: {} and buf len is {}, is not OK",
            eoi_index,
            buf.len()
        );
        Ok(true)
    } else {
        #[cfg(any(test, debug_assertions))]
        println!(
            "we're at pos: {} and buf len is {}, is OK",
            eoi_index,
            buf.len()
        );
        Ok(false)
    }
}

pub fn validate_decoding(
    filename: &str,
    contents: &[u8],
    limits: image::io::Limits,
) -> Result<(), ImageValidationError> {
    let mut decoder = match JpegDecoder::new(contents) {
        Ok(val) => val,
        Err(err) => {
            return Err(ImageValidationError::InvalidImage(format!(
                "Failed to parse {} as JPG: {:?}",
                filename, err
            )))
        }
    };

    match decoder.set_limits(limits) {
        Err(err) => {
            sketching::admin_warn!(
                "Image validation failed while validating {}: {:?}",
                filename,
                err
            );
            Err(ImageValidationError::ExceedsMaxDimensions)
        }
        Ok(_) => Ok(()),
    }
}

#[test]
fn test_jpg_has_trailer() {
    let file_contents = std::fs::read(format!(
        "{}/src/valueset/image/test_images/oversize_dimensions.jpg",
        env!("CARGO_MANIFEST_DIR")
    ))
    .expect("Failed to read file");
    assert!(!has_trailer(&file_contents).expect("Failed to check for JPEG trailer"));

    // checking a known bad image
    let file_contents = std::fs::read(format!(
        "{}/src/valueset/image/test_images/windows11_3_cropped.jpg",
        env!("CARGO_MANIFEST_DIR")
    ))
    .expect("Failed to read file");
    // let test_bytes = vec![0xff, 0xd8, 0xff, 0xda, 0xff, 0xd9];

    assert!(has_trailer(&file_contents).expect("Failed to check for JPEG trailer"));
}