/src/image/src/hooks.rs

Source
//! This module provides a way to register decoding hooks for image formats not directly supported
//! by this crate.

use std::{
    collections::HashMap,
    ffi::{OsStr, OsString},
    io::{BufRead, BufReader, Read, Seek},
    sync::{Arc, RwLock},
};

use crate::{ImageDecoder, ImageResult};

trait ReadSeek: Read + Seek {}
impl<T: Read + Seek> ReadSeek for T {}

/// Stores ascii lowercase extension to hook mapping
static DECODING_HOOKS: RwLock<Option<HashMap<OsString, Arc<DecodingHookFn>>>> = RwLock::new(None);

type DetectionHook = (&'static [u8], &'static [u8], OsString);
static GUESS_FORMAT_HOOKS: RwLock<Vec<DetectionHook>> = RwLock::new(Vec::new());

/// A wrapper around a type-erased trait object that implements `Read` and `Seek`.
pub struct GenericReader<'a>(BufReader<Box<dyn ReadSeek + 'a>>);
impl<'a> GenericReader<'a> {
    /// Creates a new `GenericReader` with a given underlying reader.
    pub(crate) fn new<R: Read + Seek + 'a>(reader: R) -> Self {
        Self(BufReader::new(Box::new(reader)))
    }
}
impl Read for GenericReader<'_> {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.0.read(buf)
    }
    fn read_vectored(&mut self, bufs: &mut [std::io::IoSliceMut<'_>]) -> std::io::Result<usize> {
        self.0.read_vectored(bufs)
    }
    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> std::io::Result<usize> {
        self.0.read_to_end(buf)
    }
    fn read_to_string(&mut self, buf: &mut String) -> std::io::Result<usize> {
        self.0.read_to_string(buf)
    }
    fn read_exact(&mut self, buf: &mut [u8]) -> std::io::Result<()> {
        self.0.read_exact(buf)
    }
}
impl BufRead for GenericReader<'_> {
    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
        self.0.fill_buf()
    }
    fn consume(&mut self, amt: usize) {
        self.0.consume(amt)
    }
    fn read_until(&mut self, byte: u8, buf: &mut Vec<u8>) -> std::io::Result<usize> {
        self.0.read_until(byte, buf)
    }
    fn read_line(&mut self, buf: &mut String) -> std::io::Result<usize> {
        self.0.read_line(buf)
    }
}
impl Seek for GenericReader<'_> {
    fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result<u64> {
        self.0.seek(pos)
    }
    fn rewind(&mut self) -> std::io::Result<()> {
        self.0.rewind()
    }
    fn stream_position(&mut self) -> std::io::Result<u64> {
        self.0.stream_position()
    }
    fn seek_relative(&mut self, offset: i64) -> std::io::Result<()> {
        self.0.seek_relative(offset)
    }
}

/// A function to produce an [`ImageDecoder`] for a given image format.
pub type DecodingHook = Box<DecodingHookFn>;
pub(crate) type DecodingHookFn =
    dyn for<'a> Fn(GenericReader<'a>) -> ImageResult<Box<dyn ImageDecoder + 'a>> + Send + Sync;

/// Register a new decoding hook or returns false if one already exists for the given format.
pub fn register_decoding_hook(mut extension: OsString, hook: DecodingHook) -> bool {
    extension.make_ascii_lowercase();
    let mut hooks = DECODING_HOOKS.write().unwrap();
    if hooks.is_none() {
        *hooks = Some(HashMap::new());
    }
    match hooks.as_mut().unwrap().entry(extension) {
        std::collections::hash_map::Entry::Vacant(entry) => {
            entry.insert(Arc::new(hook));
            true
        }
        std::collections::hash_map::Entry::Occupied(_) => false,
    }
}

/// Returns whether a decoding hook has been registered for the given format.
pub fn decoding_hook_registered(extension: &OsStr) -> bool {
    let extension = extension.to_ascii_lowercase();
    DECODING_HOOKS
        .read()
        .unwrap()
        .as_ref()
        .map(|hooks| hooks.contains_key(&extension))
        .unwrap_or(false)
}

/// Returns the decoding hook for the given format, if one exists.
pub(crate) fn get_decoding_hook(extension: &OsStr) -> Option<Arc<DecodingHookFn>> {
    let extension = extension.to_ascii_lowercase();
    let hooks = DECODING_HOOKS.read().unwrap();
    if let Some(hooks) = hooks.as_ref() {
        if let Some(hook) = hooks.get(&extension) {
            return Some(hook.clone());
        }
    }
    None
}

/// Registers a format detection hook.
///
/// The signature field holds the magic bytes from the start of the file that must be matched to
/// detect the format. The mask field is optional and can be used to specify which bytes in the
/// signature should be ignored during the detection.
///
/// # Examples
///
/// ## Using the mask to ignore some bytes
///
/// ```
/// # use image::hooks::register_format_detection_hook;
/// // WebP signature is 'riff' followed by 4 bytes of length and then by 'webp'.
/// // This requires a mask to ignore the length.
/// register_format_detection_hook("webp".into(),
///      &[b'r', b'i', b'f', b'f', 0, 0, 0, 0, b'w', b'e', b'b', b'p'],
/// Some(&[0xff, 0xff, 0xff, 0xff, 0, 0, 0, 0, 0xff, 0xff, 0xff, 0xff]),
/// );
/// ```
///
/// ## Multiple signatures
///
/// ```
/// # use image::hooks::register_format_detection_hook;
/// // JPEG XL has two different signatures: https://en.wikipedia.org/wiki/JPEG_XL
/// // This function should be called twice to register them both.
/// register_format_detection_hook("jxl".into(), &[0xff, 0x0a], None);
/// register_format_detection_hook("jxl".into(),
///      &[0x00, 0x00, 0x00, 0x0c, 0x4a, 0x58, 0x4c, 0x20, 0x0d, 0x0a, 0x87, 0x0a], None,
/// );
/// ```
///
pub fn register_format_detection_hook(
    mut extension: OsString,
    signature: &'static [u8],
    mask: Option<&'static [u8]>,
) {
    extension.make_ascii_lowercase();
    GUESS_FORMAT_HOOKS
        .write()
        .unwrap()
        .push((signature, mask.unwrap_or(&[]), extension));
}

/// Guesses the format extension from the start of the file using the registered detection hooks.
pub(crate) fn guess_format_extension(start: &[u8]) -> Option<OsString> {
    let hooks = GUESS_FORMAT_HOOKS.read().unwrap();
    for &(signature, mask, ref extension) in &*hooks {
        if mask.is_empty() {
            if start.starts_with(signature) {
                return Some(extension.clone());
            }
        } else if start.len() >= signature.len()
            && start
                .iter()
                .zip(signature.iter())
                .zip(mask.iter().chain(std::iter::repeat(&0xFF)))
                .all(|((&byte, &sig), &mask)| byte & mask == sig)
        {
            return Some(extension.clone());
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::io::{DecodedImageAttributes, DecoderPreparedImage};
    use crate::{load_from_memory, ColorType, DynamicImage, ImageReaderOptions};
    use std::io::Cursor;

    const MOCK_HOOK_EXTENSION: &str = "MOCKHOOK";

    const MOCK_IMAGE_OUTPUT: [u8; 9] = [255, 0, 0, 0, 255, 0, 0, 0, 255];
    struct MockDecoder {}
    impl ImageDecoder for MockDecoder {
        fn prepare_image(&mut self) -> ImageResult<DecoderPreparedImage> {
            Ok(DecoderPreparedImage::new(
                (MOCK_IMAGE_OUTPUT.len() / 3) as u32,
                1,
                ColorType::Rgb8,
            ))
        }

        fn read_image(&mut self, buf: &mut [u8]) -> ImageResult<DecodedImageAttributes> {
            buf[..MOCK_IMAGE_OUTPUT.len()].copy_from_slice(&MOCK_IMAGE_OUTPUT);
            Ok(DecodedImageAttributes::default())
        }
    }
    fn is_mock_decoder_output(image: DynamicImage) -> bool {
        image.as_rgb8().unwrap().as_raw() == &MOCK_IMAGE_OUTPUT
    }

    #[test]
    fn decoding_hook() {
        register_decoding_hook(
            MOCK_HOOK_EXTENSION.into(),
            Box::new(|_| Ok(Box::new(MockDecoder {}))),
        );

        assert!(decoding_hook_registered(OsStr::new(MOCK_HOOK_EXTENSION)));
        assert!(get_decoding_hook(OsStr::new(MOCK_HOOK_EXTENSION)).is_some());

        let image = ImageReaderOptions::open("tests/assets/hook/extension.MoCkHoOk")
            .unwrap()
            .decode()
            .unwrap();

        assert!(is_mock_decoder_output(image));
    }

    #[test]
    fn detection_hook() {
        register_decoding_hook(
            MOCK_HOOK_EXTENSION.into(),
            Box::new(|_| Ok(Box::new(MockDecoder {}))),
        );

        register_format_detection_hook(
            MOCK_HOOK_EXTENSION.into(),
            &[b'H', b'E', b'A', b'D', 0, 0, 0, 0, b'M', b'O', b'C', b'K'],
            Some(&[0xff, 0xff, 0xff, 0xff, 0, 0, 0, 0, 0xff, 0xff, 0xff, 0xff]),
        );

        const TEST_INPUT_IMAGE: [u8; 16] = *b"HEADJUNKMOCKmore";
        assert_eq!(
            guess_format_extension(&TEST_INPUT_IMAGE),
            Some(OsStr::new(MOCK_HOOK_EXTENSION).to_ascii_lowercase())
        );

        let image = ImageReaderOptions::new(Cursor::new(TEST_INPUT_IMAGE))
            .with_guessed_format()
            .unwrap()
            .decode()
            .unwrap();

        assert!(is_mock_decoder_output(image));

        let image_via_free_function = load_from_memory(&TEST_INPUT_IMAGE).unwrap();
        assert!(is_mock_decoder_output(image_via_free_function));
    }
}

Coverage Report

Created: 2026-08-13 08:17

Line	Count	Source
1		//! This module provides a way to register decoding hooks for image formats not directly supported
2		//! by this crate.
3
4		use std::{
5		collections::HashMap,
6		ffi::{OsStr, OsString},
7		io::{BufRead, BufReader, Read, Seek},
8		sync::{Arc, RwLock},
9		};
10
11		use crate::{ImageDecoder, ImageResult};
12
13		trait ReadSeek: Read + Seek {}
14		impl<T: Read + Seek> ReadSeek for T {}
15
16		/// Stores ascii lowercase extension to hook mapping
17		static DECODING_HOOKS: RwLock<Option<HashMap<OsString, Arc<DecodingHookFn>>>> = RwLock::new(None);
18
19		type DetectionHook = (&'static [u8], &'static [u8], OsString);
20		static GUESS_FORMAT_HOOKS: RwLock<Vec<DetectionHook>> = RwLock::new(Vec::new());
21
22		/// A wrapper around a type-erased trait object that implements `Read` and `Seek`.
23		pub struct GenericReader<'a>(BufReader<Box<dyn ReadSeek + 'a>>);
24		impl<'a> GenericReader<'a> {
25		/// Creates a new `GenericReader` with a given underlying reader.
26	0	pub(crate) fn new<R: Read + Seek + 'a>(reader: R) -> Self {
27	0	Self(BufReader::new(Box::new(reader)))
28	0	}
29		}
30		impl Read for GenericReader<'_> {
31	0	fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
32	0	self.0.read(buf)
33	0	}
34	0	fn read_vectored(&mut self, bufs: &mut [std::io::IoSliceMut<'_>]) -> std::io::Result<usize> {
35	0	self.0.read_vectored(bufs)
36	0	}
37	0	fn read_to_end(&mut self, buf: &mut Vec<u8>) -> std::io::Result<usize> {
38	0	self.0.read_to_end(buf)
39	0	}
40	0	fn read_to_string(&mut self, buf: &mut String) -> std::io::Result<usize> {
41	0	self.0.read_to_string(buf)
42	0	}
43	0	fn read_exact(&mut self, buf: &mut [u8]) -> std::io::Result<()> {
44	0	self.0.read_exact(buf)
45	0	}
46		}
47		impl BufRead for GenericReader<'_> {
48	0	fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
49	0	self.0.fill_buf()
50	0	}
51	0	fn consume(&mut self, amt: usize) {
52	0	self.0.consume(amt)
53	0	}
54	0	fn read_until(&mut self, byte: u8, buf: &mut Vec<u8>) -> std::io::Result<usize> {
55	0	self.0.read_until(byte, buf)
56	0	}
57	0	fn read_line(&mut self, buf: &mut String) -> std::io::Result<usize> {
58	0	self.0.read_line(buf)
59	0	}
60		}
61		impl Seek for GenericReader<'_> {
62	0	fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result<u64> {
63	0	self.0.seek(pos)
64	0	}
65	0	fn rewind(&mut self) -> std::io::Result<()> {
66	0	self.0.rewind()
67	0	}
68	0	fn stream_position(&mut self) -> std::io::Result<u64> {
69	0	self.0.stream_position()
70	0	}
71	0	fn seek_relative(&mut self, offset: i64) -> std::io::Result<()> {
72	0	self.0.seek_relative(offset)
73	0	}
74		}
75
76		/// A function to produce an [`ImageDecoder`] for a given image format.
77		pub type DecodingHook = Box<DecodingHookFn>;
78		pub(crate) type DecodingHookFn =
79		dyn for<'a> Fn(GenericReader<'a>) -> ImageResult<Box<dyn ImageDecoder + 'a>> + Send + Sync;
80
81		/// Register a new decoding hook or returns false if one already exists for the given format.
82	0	pub fn register_decoding_hook(mut extension: OsString, hook: DecodingHook) -> bool {
83	0	extension.make_ascii_lowercase();
84	0	let mut hooks = DECODING_HOOKS.write().unwrap();
85	0	if hooks.is_none() {
86	0	*hooks = Some(HashMap::new());
87	0	}
88	0	match hooks.as_mut().unwrap().entry(extension) {
89	0	std::collections::hash_map::Entry::Vacant(entry) => {
90	0	entry.insert(Arc::new(hook));
91	0	true
92		}
93	0	std::collections::hash_map::Entry::Occupied(_) => false,
94		}
95	0	}
96
97		/// Returns whether a decoding hook has been registered for the given format.
98	0	pub fn decoding_hook_registered(extension: &OsStr) -> bool {
99	0	let extension = extension.to_ascii_lowercase();
100	0	DECODING_HOOKS
101	0	.read()
102	0	.unwrap()
103	0	.as_ref()
104	0	.map(\|hooks\| hooks.contains_key(&extension))
105	0	.unwrap_or(false)
106	0	}
107
108		/// Returns the decoding hook for the given format, if one exists.
109	0	pub(crate) fn get_decoding_hook(extension: &OsStr) -> Option<Arc<DecodingHookFn>> {
110	0	let extension = extension.to_ascii_lowercase();
111	0	let hooks = DECODING_HOOKS.read().unwrap();
112	0	if let Some(hooks) = hooks.as_ref() {
113	0	if let Some(hook) = hooks.get(&extension) {
114	0	return Some(hook.clone());
115	0	}
116	0	}
117	0	None
118	0	}
119
120		/// Registers a format detection hook.
121		///
122		/// The signature field holds the magic bytes from the start of the file that must be matched to
123		/// detect the format. The mask field is optional and can be used to specify which bytes in the
124		/// signature should be ignored during the detection.
125		///
126		/// # Examples
127		///
128		/// ## Using the mask to ignore some bytes
129		///
130		/// ```
131		/// # use image::hooks::register_format_detection_hook;
132		/// // WebP signature is 'riff' followed by 4 bytes of length and then by 'webp'.
133		/// // This requires a mask to ignore the length.
134		/// register_format_detection_hook("webp".into(),
135		/// &[b'r', b'i', b'f', b'f', 0, 0, 0, 0, b'w', b'e', b'b', b'p'],
136		/// Some(&[0xff, 0xff, 0xff, 0xff, 0, 0, 0, 0, 0xff, 0xff, 0xff, 0xff]),
137		/// );
138		/// ```
139		///
140		/// ## Multiple signatures
141		///
142		/// ```
143		/// # use image::hooks::register_format_detection_hook;
144		/// // JPEG XL has two different signatures: https://en.wikipedia.org/wiki/JPEG_XL
145		/// // This function should be called twice to register them both.
146		/// register_format_detection_hook("jxl".into(), &[0xff, 0x0a], None);
147		/// register_format_detection_hook("jxl".into(),
148		/// &[0x00, 0x00, 0x00, 0x0c, 0x4a, 0x58, 0x4c, 0x20, 0x0d, 0x0a, 0x87, 0x0a], None,
149		/// );
150		/// ```
151		///
152	0	pub fn register_format_detection_hook(
153	0	mut extension: OsString,
154	0	signature: &'static [u8],
155	0	mask: Option<&'static [u8]>,
156	0	) {
157	0	extension.make_ascii_lowercase();
158	0	GUESS_FORMAT_HOOKS
159	0	.write()
160	0	.unwrap()
161	0	.push((signature, mask.unwrap_or(&[]), extension));
162	0	}
163
164		/// Guesses the format extension from the start of the file using the registered detection hooks.
165	0	pub(crate) fn guess_format_extension(start: &[u8]) -> Option<OsString> {
166	0	let hooks = GUESS_FORMAT_HOOKS.read().unwrap();
167	0	for &(signature, mask, ref extension) in &*hooks {
168	0	if mask.is_empty() {
169	0	if start.starts_with(signature) {
170	0	return Some(extension.clone());
171	0	}
172	0	} else if start.len() >= signature.len()
173	0	&& start
174	0	.iter()
175	0	.zip(signature.iter())
176	0	.zip(mask.iter().chain(std::iter::repeat(&0xFF)))
177	0	.all(\|((&byte, &sig), &mask)\| byte & mask == sig)
178		{
179	0	return Some(extension.clone());
180	0	}
181		}
182	0	None
183	0	}
184
185		#[cfg(test)]
186		mod tests {
187		use super::*;
188		use crate::io::{DecodedImageAttributes, DecoderPreparedImage};
189		use crate::{load_from_memory, ColorType, DynamicImage, ImageReaderOptions};
190		use std::io::Cursor;
191
192		const MOCK_HOOK_EXTENSION: &str = "MOCKHOOK";
193
194		const MOCK_IMAGE_OUTPUT: [u8; 9] = [255, 0, 0, 0, 255, 0, 0, 0, 255];
195		struct MockDecoder {}
196		impl ImageDecoder for MockDecoder {
197		fn prepare_image(&mut self) -> ImageResult<DecoderPreparedImage> {
198		Ok(DecoderPreparedImage::new(
199		(MOCK_IMAGE_OUTPUT.len() / 3) as u32,
200		1,
201		ColorType::Rgb8,
202		))
203		}
204
205		fn read_image(&mut self, buf: &mut [u8]) -> ImageResult<DecodedImageAttributes> {
206		buf[..MOCK_IMAGE_OUTPUT.len()].copy_from_slice(&MOCK_IMAGE_OUTPUT);
207		Ok(DecodedImageAttributes::default())
208		}
209		}
210		fn is_mock_decoder_output(image: DynamicImage) -> bool {
211		image.as_rgb8().unwrap().as_raw() == &MOCK_IMAGE_OUTPUT
212		}
213
214		#[test]
215		fn decoding_hook() {
216		register_decoding_hook(
217		MOCK_HOOK_EXTENSION.into(),
218		Box::new(\|_\| Ok(Box::new(MockDecoder {}))),
219		);
220
221		assert!(decoding_hook_registered(OsStr::new(MOCK_HOOK_EXTENSION)));
222		assert!(get_decoding_hook(OsStr::new(MOCK_HOOK_EXTENSION)).is_some());
223
224		let image = ImageReaderOptions::open("tests/assets/hook/extension.MoCkHoOk")
225		.unwrap()
226		.decode()
227		.unwrap();
228
229		assert!(is_mock_decoder_output(image));
230		}
231
232		#[test]
233		fn detection_hook() {
234		register_decoding_hook(
235		MOCK_HOOK_EXTENSION.into(),
236		Box::new(\|_\| Ok(Box::new(MockDecoder {}))),
237		);
238
239		register_format_detection_hook(
240		MOCK_HOOK_EXTENSION.into(),
241		&[b'H', b'E', b'A', b'D', 0, 0, 0, 0, b'M', b'O', b'C', b'K'],
242		Some(&[0xff, 0xff, 0xff, 0xff, 0, 0, 0, 0, 0xff, 0xff, 0xff, 0xff]),
243		);
244
245		const TEST_INPUT_IMAGE: [u8; 16] = *b"HEADJUNKMOCKmore";
246		assert_eq!(
247		guess_format_extension(&TEST_INPUT_IMAGE),
248		Some(OsStr::new(MOCK_HOOK_EXTENSION).to_ascii_lowercase())
249		);
250
251		let image = ImageReaderOptions::new(Cursor::new(TEST_INPUT_IMAGE))
252		.with_guessed_format()
253		.unwrap()
254		.decode()
255		.unwrap();
256
257		assert!(is_mock_decoder_output(image));
258
259		let image_via_free_function = load_from_memory(&TEST_INPUT_IMAGE).unwrap();
260		assert!(is_mock_decoder_output(image_via_free_function));
261		}
262		}