image/
hooks.rs

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

image/
hooks.rs
