use crate::bytes::Reader;
use crate::error::{Result, StryptError, UnsupportedKind};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[non_exhaustive]
pub enum Format {
Jpeg,
Png,
Webp,
Pdf,
}
impl Format {
#[must_use]
pub const fn id(self) -> &'static str {
match self {
Self::Jpeg => "jpeg",
Self::Png => "png",
Self::Webp => "webp",
Self::Pdf => "pdf",
}
}
#[must_use]
pub const fn conventional_extension(self) -> &'static str {
match self {
Self::Jpeg => "jpg",
Self::Png => "png",
Self::Webp => "webp",
Self::Pdf => "pdf",
}
}
}
impl std::fmt::Display for Format {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(match self {
Self::Jpeg => "JPEG",
Self::Png => "PNG",
Self::Webp => "WebP",
Self::Pdf => "PDF",
})
}
}
const PDF_HEADER_SEARCH_WINDOW: usize = 1024;
pub fn detect(data: &[u8]) -> Result<Format> {
if let Some(format) = detect_supported(data) {
return Ok(format);
}
if let Some(kind) = detect_unsupported(data) {
return Err(StryptError::UnsupportedFormat { format: kind });
}
Err(StryptError::UnrecognisedFormat)
}
fn detect_supported(data: &[u8]) -> Option<Format> {
if starts_with(data, &[0xFF, 0xD8, 0xFF]) {
return Some(Format::Jpeg);
}
if starts_with(data, &[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A]) {
return Some(Format::Png);
}
if is_riff_with_form(data, *b"WEBP") {
return Some(Format::Webp);
}
if find_pdf_header(data).is_some() {
return Some(Format::Pdf);
}
None
}
fn detect_unsupported(data: &[u8]) -> Option<UnsupportedKind> {
if starts_with(data, b"PK\x03\x04")
|| starts_with(data, b"PK\x05\x06")
|| starts_with(data, b"PK\x07\x08")
{
return Some(UnsupportedKind::ZipContainer);
}
if starts_with(data, b"GIF87a") || starts_with(data, b"GIF89a") {
return Some(UnsupportedKind::Gif);
}
if starts_with(data, &[b'I', b'I', 0x2A, 0x00]) || starts_with(data, &[b'M', b'M', 0x00, 0x2A])
{
return Some(UnsupportedKind::Tiff);
}
if data.get(4..8) == Some(b"ftyp") {
return Some(UnsupportedKind::IsoBaseMedia);
}
if starts_with(data, b"OggS") {
return Some(UnsupportedKind::Ogg);
}
if starts_with(data, b"fLaC") {
return Some(UnsupportedKind::Flac);
}
if starts_with(data, b"ID3") {
return Some(UnsupportedKind::Mp3);
}
if let (Some(&0xFF), Some(&second)) = (data.first(), data.get(1))
&& (second & 0xE0) == 0xE0
{
return Some(UnsupportedKind::Mp3);
}
if starts_with(data, b"RIFF") {
return Some(UnsupportedKind::OtherRiff);
}
if looks_like_xml(data) {
return Some(UnsupportedKind::Xml);
}
None
}
fn starts_with(data: &[u8], prefix: &[u8]) -> bool {
data.get(0..prefix.len()) == Some(prefix)
}
fn is_riff_with_form(data: &[u8], form: [u8; 4]) -> bool {
let mut r = Reader::new(data);
if r.peek(4) != Some(b"RIFF") {
return false;
}
if r.skip(8).is_none() {
return false;
}
r.peek(4) == Some(form.as_slice())
}
fn find_pdf_header(data: &[u8]) -> Option<usize> {
const HEADER: &[u8] = b"%PDF-";
let window = data.get(0..PDF_HEADER_SEARCH_WINDOW).unwrap_or(data);
window
.windows(HEADER.len())
.position(|candidate| candidate == HEADER)
}
fn looks_like_xml(data: &[u8]) -> bool {
let mut r = Reader::new(data);
if r.peek(3) == Some(&[0xEF, 0xBB, 0xBF]) && r.skip(3).is_none() {
return false;
}
for _ in 0..64 {
match r.peek(1) {
Some([b' ' | b'\t' | b'\r' | b'\n']) => {
if r.skip(1).is_none() {
return false;
}
}
_ => break,
}
}
r.peek(5) == Some(b"<?xml") || r.peek(4) == Some(b"<svg") || r.peek(9) == Some(b"<!DOCTYPE")
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use super::*;
#[test]
fn the_four_supported_formats_are_recognised() {
assert_eq!(detect(&[0xFF, 0xD8, 0xFF, 0xE0]).unwrap(), Format::Jpeg);
assert_eq!(
detect(&[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A]).unwrap(),
Format::Png
);
assert_eq!(
detect(b"RIFF\x00\x00\x00\x00WEBPVP8 ").unwrap(),
Format::Webp
);
assert_eq!(detect(b"%PDF-1.7\n").unwrap(), Format::Pdf);
}
#[test]
fn a_pdf_named_jpg_is_still_a_pdf() {
assert_eq!(
detect(b"%PDF-1.4\n%\xe2\xe3\xcf\xd3\n").unwrap(),
Format::Pdf
);
}
#[test]
fn a_pdf_header_inside_a_jpeg_does_not_win() {
let mut data = vec![0xFF, 0xD8, 0xFF, 0xE1];
data.extend_from_slice(b"junk %PDF-1.7 junk");
assert_eq!(detect(&data).unwrap(), Format::Jpeg);
}
#[test]
fn a_pdf_with_a_leading_preamble_is_recognised() {
let mut data = b"\r\n<!-- inserted by a broken proxy -->\r\n".to_vec();
data.extend_from_slice(b"%PDF-1.5\n");
assert_eq!(detect(&data).unwrap(), Format::Pdf);
}
#[test]
fn a_pdf_header_beyond_the_search_window_is_not_scanned_for() {
let mut data = vec![b'x'; PDF_HEADER_SEARCH_WINDOW];
data.extend_from_slice(b"%PDF-1.7\n");
assert!(matches!(
detect(&data),
Err(StryptError::UnrecognisedFormat)
));
}
#[test]
fn a_non_webp_riff_is_named_rather_than_mishandled() {
let e = detect(b"RIFF\x00\x00\x00\x00WAVEfmt ").unwrap_err();
assert!(matches!(
e,
StryptError::UnsupportedFormat {
format: UnsupportedKind::OtherRiff
}
));
}
#[test]
fn phase_two_formats_are_named_in_the_refusal() {
for (bytes, expected) in [
(&b"PK\x03\x04"[..], UnsupportedKind::ZipContainer),
(&b"GIF89a"[..], UnsupportedKind::Gif),
(&b"II\x2A\x00"[..], UnsupportedKind::Tiff),
(&b"OggS"[..], UnsupportedKind::Ogg),
(&b"fLaC"[..], UnsupportedKind::Flac),
(&b"ID3\x04"[..], UnsupportedKind::Mp3),
(
&b"\x00\x00\x00\x18ftypavif"[..],
UnsupportedKind::IsoBaseMedia,
),
(&b"<?xml version=\"1.0\"?><svg/>"[..], UnsupportedKind::Xml),
] {
let got = detect(bytes).unwrap_err();
assert!(
matches!(got, StryptError::UnsupportedFormat { format } if format == expected),
"detecting {expected:?} gave {got:?}"
);
}
}
#[test]
fn nothing_recognisable_is_an_error_never_a_silent_pass() {
assert!(matches!(detect(b""), Err(StryptError::UnrecognisedFormat)));
assert!(matches!(
detect(b"hello world"),
Err(StryptError::UnrecognisedFormat)
));
}
#[test]
fn truncated_magic_numbers_do_not_panic() {
let signatures: [&[u8]; 4] = [
&[0xFF, 0xD8, 0xFF],
&[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A],
b"RIFF\x00\x00\x00\x00WEBP",
b"%PDF-1.7",
];
for sig in signatures {
for n in 0..=sig.len() {
let prefix = sig.get(0..n).unwrap_or_default();
let _ = detect(prefix);
}
}
}
}