use crate::complete::base::skip;
use nom::number::complete::{be_u16 as parse_be_u16, be_u32 as parse_be_u32};
use nom::IResult as NomResult;
pub struct ImageDimensions {
pub width: usize,
pub height: usize,
}
pub const PNG_START: [u8; 8] = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a];
pub const GIF_START: [u8; 6] = [0x47, 0x49, 0x46, 0x38, 0x39, 0x61];
pub const JPEG_START: [u8; 2] = [0xff, 0xd8];
pub const ERRONEOUS_JPEG_START: [u8; 6] = [0xff, 0xd9, 0xff, 0xd8, 0xff, 0xd8];
const PNG_IHDR_CHUNK_TYPE: u32 = 0x49_48_44_52;
pub fn get_png_image_dimensions(input: &[u8]) -> Result<ImageDimensions, ()> {
let (input, ()) = skip::<_, _, ()>(12usize)(input).map_err(|_| ())?;
let (input, chunk_type) = parse_be_u32::<&[u8], ()>(input).map_err(|_| ())?;
if chunk_type != PNG_IHDR_CHUNK_TYPE {
return Err(());
}
let (input, width) = parse_be_u32::<&[u8], ()>(input).map_err(|_| ())?;
let (_, height) = parse_be_u32::<&[u8], ()>(input).map_err(|_| ())?;
Ok(ImageDimensions {
width: width as usize,
height: height as usize,
})
}
pub fn get_jpeg_image_dimensions(mut input: &[u8]) -> Result<ImageDimensions, ()> {
loop {
let (next_input, chunk) = take_next_jpeg_chunk(input).map_err(|_| ())?;
if chunk.is_some() {
debug_assert!(next_input.len() < input.len());
}
input = next_input;
if let Some(chunk) = chunk {
if chunk.len() < 9 {
continue;
}
let code: u8 = chunk[1];
if !is_jpeg_sof(code) {
continue;
}
let frame_height = u16::from_be_bytes([chunk[5], chunk[6]]);
let frame_width = u16::from_be_bytes([chunk[7], chunk[8]]);
return Ok(ImageDimensions {
width: usize::from(frame_width),
height: usize::from(frame_height),
});
} else {
return Err(());
}
}
}
fn take_next_jpeg_chunk(input: &[u8]) -> NomResult<&[u8], Option<&[u8]>> {
use nom::bytes::complete::take;
use nom::combinator::map;
let mut search: usize = 0;
while search + 1 < input.len() {
let cur_byte: u8 = input[search];
let marker_type: Option<JpegMarkerType> = if cur_byte == 0xff {
get_jpeg_marker_type(input[search + 1])
} else {
None
};
match marker_type {
Some(marker_type) => {
let input = &input[search..];
return match marker_type {
JpegMarkerType::Standalone => map(take(2usize), Some)(input),
JpegMarkerType::Sequence => {
let (_, size) = parse_be_u16(input)?;
map(take(2usize + usize::from(size)), Some)(input)
}
};
}
None => search += 1,
}
}
Ok((&[], None))
}
#[derive(Debug, Eq, PartialEq)]
enum JpegMarkerType {
Standalone,
Sequence,
}
fn get_jpeg_marker_type(code: u8) -> Option<JpegMarkerType> {
match code {
0x00 => None, 0x01 => Some(JpegMarkerType::Standalone), 0xd0..=0xd7 => Some(JpegMarkerType::Standalone), 0xd8 => Some(JpegMarkerType::Standalone), 0xd9 => Some(JpegMarkerType::Standalone), 0xff => None, _ => Some(JpegMarkerType::Sequence), }
}
fn is_jpeg_sof(code: u8) -> bool {
code & 0xfc == 0xc0
}
pub fn get_gif_image_dimensions(input: &[u8]) -> Result<ImageDimensions, ()> {
let (input, ()) = skip::<_, _, ()>(6usize)(input).map_err(|_| ())?;
let (input, width) = parse_be_u16::<_, ()>(input).map_err(|_| ())?;
let (_, height) = parse_be_u16::<_, ()>(input).map_err(|_| ())?;
Ok(ImageDimensions {
width: width as usize,
height: height as usize,
})
}
pub(crate) enum SniffedImageType {
Jpeg,
Png,
Gif,
}
pub(crate) fn sniff_image_type(image_data: &[u8], allow_erroneous_jpeg: bool) -> Result<SniffedImageType, ()> {
if is_sniffed_jpeg(image_data, allow_erroneous_jpeg) {
Ok(SniffedImageType::Jpeg)
} else if is_sniffed_png(image_data) {
Ok(SniffedImageType::Png)
} else if is_sniffed_gif(image_data) {
Ok(SniffedImageType::Gif)
} else {
Err(())
}
}
pub(crate) fn is_sniffed_jpeg(image_data: &[u8], allow_erroneous: bool) -> bool {
test_image_start(image_data, &JPEG_START) || (allow_erroneous && test_image_start(image_data, &ERRONEOUS_JPEG_START))
}
pub(crate) fn is_sniffed_png(image_data: &[u8]) -> bool {
test_image_start(image_data, &PNG_START)
}
pub(crate) fn is_sniffed_gif(image_data: &[u8]) -> bool {
test_image_start(image_data, &GIF_START)
}
fn test_image_start(image_data: &[u8], start_bytes: &[u8]) -> bool {
image_data.len() >= start_bytes.len() && image_data[..start_bytes.len()] == *start_bytes
}