use super::decode::{sniff_format, ImageFormat};
pub fn probe_dimensions(bytes: &[u8]) -> Option<(u32, u32)> {
match sniff_format(bytes)? {
ImageFormat::Png => png_dimensions(bytes),
ImageFormat::Jpeg => jpeg_dimensions(bytes),
ImageFormat::Gif => gif_dimensions(bytes),
}
}
fn gif_dimensions(bytes: &[u8]) -> Option<(u32, u32)> {
let d = bytes.get(6..10)?;
let w = u16::from_le_bytes([d[0], d[1]]) as u32;
let h = u16::from_le_bytes([d[2], d[3]]) as u32;
(w != 0 && h != 0).then_some((w, h))
}
fn png_dimensions(bytes: &[u8]) -> Option<(u32, u32)> {
if bytes.len() < 24 || &bytes[12..16] != b"IHDR" {
return None;
}
let len = u32::from_be_bytes(bytes[8..12].try_into().ok()?);
if len < 13 {
return None; }
let w = u32::from_be_bytes(bytes[16..20].try_into().ok()?);
let h = u32::from_be_bytes(bytes[20..24].try_into().ok()?);
if w == 0 || h == 0 {
None
} else {
Some((w, h))
}
}
fn jpeg_dimensions(bytes: &[u8]) -> Option<(u32, u32)> {
let mut i = 2; loop {
while *bytes.get(i)? != 0xFF {
i += 1;
}
while *bytes.get(i)? == 0xFF {
i += 1;
}
let marker = *bytes.get(i)?; i += 1;
match marker {
0x01 | 0xD0..=0xD7 => continue,
0xD9 | 0xDA => return None,
0xC0..=0xC3 | 0xC5..=0xC7 | 0xC9..=0xCB | 0xCD..=0xCF => {
let seg = bytes.get(i..i + 7)?;
let len = u16::from_be_bytes([seg[0], seg[1]]);
if len < 9 {
return None; }
let h = u16::from_be_bytes([seg[3], seg[4]]) as u32;
let w = u16::from_be_bytes([seg[5], seg[6]]) as u32;
return if w == 0 || h == 0 { None } else { Some((w, h)) };
}
_ => {
let seg = bytes.get(i..i + 2)?;
let len = u16::from_be_bytes([seg[0], seg[1]]) as usize;
if len < 2 {
return None; }
i += len;
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::base::Rgba;
use crate::gfx::decode::decode_image;
use crate::gfx::{jpeg_fixtures, png_test_encoder};
use crate::testing::hostile_corpus;
#[test]
fn probe_matches_decoder_dimensions() {
for (w, h) in [(1u32, 1u32), (3, 2), (17, 9), (64, 48)] {
let bmp = crate::gfx::Bitmap::from_fn(w, h, |x, y| {
Rgba::rgb((x * 31) as u8, (y * 17) as u8, 99)
});
let png = png_test_encoder::encode_rgba(&bmp);
assert_eq!(probe_dimensions(&png), Some((w, h)), "png {w}x{h}");
}
for jpg in [jpeg_fixtures::GRAD444, jpeg_fixtures::GRAD420] {
let decoded = decode_image(jpg).expect("fixture decodes");
assert_eq!(
probe_dimensions(jpg),
Some((decoded.width(), decoded.height())),
"jpeg probe must match the decoder"
);
}
}
#[test]
fn rejects_non_images_and_zero_dimensions() {
assert_eq!(probe_dimensions(b""), None);
assert_eq!(probe_dimensions(b"RIFF\x24\x00\x00\x00WEBPVP8 "), None);
assert_eq!(probe_dimensions(b"GIF89a\x10\x00\x08\x00"), Some((16, 8)));
assert_eq!(
probe_dimensions(b"GIF89a\x00\x00\x00\x00"),
None,
"zero canvas"
);
assert_eq!(probe_dimensions(b"\x89PNG\r\n\x1a\n"), None, "magic only");
let mut zero = Vec::new();
zero.extend_from_slice(&[0x89, b'P', b'N', b'G', b'\r', b'\n', 0x1A, b'\n']);
zero.extend_from_slice(&13u32.to_be_bytes());
zero.extend_from_slice(b"IHDR");
zero.extend_from_slice(&0u32.to_be_bytes());
zero.extend_from_slice(&7u32.to_be_bytes());
assert_eq!(probe_dimensions(&zero), None);
}
#[test]
fn probe_survives_hostile_input_and_truncation() {
for chunk in hostile_corpus(0x0144, 300) {
let _ = probe_dimensions(&chunk);
}
let bmp = crate::gfx::Bitmap::from_fn(9, 7, |x, y| {
Rgba::rgb((x * 29) as u8, (y * 37) as u8, 128)
});
let png = png_test_encoder::encode_rgba(&bmp);
for src in [&png[..], jpeg_fixtures::GRAD444] {
for cut in 0..src.len() {
let probed = probe_dimensions(&src[..cut]);
if let (Some(p), Ok(d)) = (probed, decode_image(&src[..cut])) {
assert_eq!(p, (d.width(), d.height()), "truncation at {cut}");
}
}
}
let mut rng = crate::testing::Rng::new(0x0144_0144);
for case in 0..300 {
let len = 8 + rng.below(200);
let mut bytes: Vec<u8> = (0..len).map(|_| rng.byte()).collect();
if case % 2 == 0 {
bytes[..8].copy_from_slice(&[0x89, b'P', b'N', b'G', b'\r', b'\n', 0x1A, b'\n']);
} else {
bytes[0] = 0xFF;
bytes[1] = 0xD8;
}
let _ = probe_dimensions(&bytes);
}
}
}