Skip to main content

otf_pixels_codec_jxl/
decoder.rs

1//! The JPEG XL decoder: a bare codestream or the ISOBMFF container, decoded
2//! to its first displayed frame.
3
4#![allow(
5    clippy::indexing_slicing,
6    reason = "planes are the canvas size, so every pixel index is in range"
7)]
8
9use otf_pixels_core::{
10    Animation, Codec, Decoder, Format, ImageDescriptor, Limits, Orientation, PixelFormat,
11    PixelsError, Result, Source, SourceInfo,
12};
13
14use crate::decode::{container, header::ExtraKind, render};
15
16/// The most compressed bytes read before a file is called hostile.
17const MAX_COMPRESSED: usize = 256 * 1024 * 1024;
18
19/// Decodes a JPEG XL stream.
20#[derive(Debug)]
21pub struct JxlDecoder {
22    descriptor: ImageDescriptor,
23    /// The codestream, until the first row is asked for.
24    codestream: Vec<u8>,
25    /// The decoded image, interleaved in the output format.
26    pixels: Option<Vec<u8>>,
27    row: u32,
28    orientation: Orientation,
29    icc: Option<Vec<u8>>,
30    exif: Option<Vec<u8>>,
31    xmp: Option<Vec<u8>>,
32    info: SourceInfo,
33    animation: Option<Animation>,
34}
35
36impl JxlDecoder {
37    /// Read the file and parse its headers; the pixels decode when the first
38    /// row is read.
39    ///
40    /// # Errors
41    ///
42    /// Returns [`PixelsError::Malformed`] for headers it rejects, or
43    /// [`PixelsError::LimitExceeded`] if the image exceeds `limits`.
44    pub fn new<S: Source>(mut source: S, limits: Limits) -> Result<Self> {
45        let mut bytes = Vec::new();
46        let mut chunk = [0_u8; 64 * 1024];
47        loop {
48            if bytes.len() > MAX_COMPRESSED {
49                return Err(PixelsError::malformed(
50                    "jxl",
51                    format!("stream exceeds {MAX_COMPRESSED} bytes"),
52                ));
53            }
54            match source.read(&mut chunk)? {
55                0 => break,
56                read => bytes.extend_from_slice(chunk.get(..read).unwrap_or(&[])),
57            }
58        }
59        let contents = container::parse(&bytes)?;
60        let exif = contents.exif;
61        let xmp = contents.xmp;
62        let codestream = contents.codestream;
63        // Limits first: the size is the attacker's to state, and walking the
64        // frame headers sizes their tables of contents by it.
65        let header = crate::decode::header::ImageHeader::parse(
66            &mut crate::decode::reader::BitReader::new(&codestream),
67        )?;
68        ImageDescriptor::with_limits(header.width, header.height, PixelFormat::Rgba8, &limits)?;
69        let (header, icc, durations) = render::probe(&codestream)?;
70        let (width, height) = (header.width, header.height);
71        let alpha = header.alpha().is_some();
72        let wide = header.bit_depth.bits > 8;
73        let pixel = match (header.grey(), alpha, wide, header.bit_depth.float) {
74            (_, false, _, true) => PixelFormat::RgbF32,
75            (_, true, _, true) => PixelFormat::RgbaF32,
76            (true, false, false, _) => PixelFormat::Gray8,
77            (true, false, true, _) => PixelFormat::Gray16,
78            (true, true, false, _) => PixelFormat::GrayA8,
79            (false, false, false, _) => PixelFormat::Rgb8,
80            (false, false, true, _) => PixelFormat::Rgb16,
81            (false, true, false, _) => PixelFormat::Rgba8,
82            (_, true, true, _) => PixelFormat::Rgba16,
83        };
84        let descriptor = ImageDescriptor::with_limits(width, height, pixel, &limits)?;
85        let animation = header.animation.and_then(|a| {
86            if durations.len() < 2 || a.tps_numerator == 0 {
87                return None;
88            }
89            let ms = durations
90                .iter()
91                .map(|&t| {
92                    (u64::from(t) * 1000 * u64::from(a.tps_denominator)
93                        / u64::from(a.tps_numerator)) as u32
94                })
95                .collect();
96            Animation::new(ms, a.num_loops)
97        });
98        // The codestream's own orientation; an EXIF block repeats it.
99        let orientation = Orientation::from_exif(header.orientation as u16).unwrap_or_default();
100        // XYB images are delivered as sRGB, so their profile does not apply.
101        let icc = if header.xyb { None } else { icc };
102        let mut info = SourceInfo::default();
103        info.bits_per_sample = u8::try_from(header.bit_depth.bits).ok();
104        // XYB is how libjxl codes lossy images; lossless keeps the colour
105        // space it was given.
106        info.compression = Some(if header.xyb { "lossy" } else { "lossless" });
107        Ok(Self {
108            descriptor,
109            codestream,
110            pixels: None,
111            row: 0,
112            orientation,
113            icc,
114            exif,
115            xmp,
116            info,
117            animation,
118        })
119    }
120
121    fn decode_pixels(&self) -> Result<Vec<u8>> {
122        let image = render::decode(&self.codestream)?;
123        let canvas = &image.canvas;
124        let alpha = image
125            .header
126            .extra
127            .iter()
128            .position(|e| matches!(e.kind, ExtraKind::Alpha { .. }))
129            .and_then(|i| canvas.extra.get(i));
130        let n = canvas.width * canvas.height;
131        let pixel = self.descriptor.pixel;
132        let channels: Vec<&[f32]> = if pixel.channels() <= 2 {
133            vec![&canvas.colour[0][..n]]
134        } else {
135            canvas.colour.iter().map(|p| &p[..n]).collect()
136        };
137        let mut planes = channels;
138        if pixel.has_alpha() {
139            planes.push(alpha.map_or(&[][..], |a| &a[..n]));
140        }
141        let mut out = Vec::with_capacity(self.descriptor.row_bytes() * canvas.height);
142        for i in 0..n {
143            for p in &planes {
144                let v = p.get(i).copied().unwrap_or(1.0);
145                match pixel {
146                    PixelFormat::RgbF32 | PixelFormat::RgbaF32 => {
147                        out.extend_from_slice(&v.to_ne_bytes())
148                    }
149                    PixelFormat::Gray16 | PixelFormat::Rgb16 | PixelFormat::Rgba16 => {
150                        out.extend_from_slice(
151                            &((v.clamp(0.0, 1.0) * 65535.0 + 0.5) as u16).to_ne_bytes(),
152                        );
153                    }
154                    _ => out.push((v.clamp(0.0, 1.0) * 255.0 + 0.5) as u8),
155                }
156            }
157        }
158        Ok(out)
159    }
160}
161
162impl Decoder for JxlDecoder {
163    fn descriptor(&self) -> ImageDescriptor {
164        self.descriptor
165    }
166
167    fn orientation(&self) -> Orientation {
168        self.orientation
169    }
170
171    fn icc_profile(&self) -> Option<&[u8]> {
172        self.icc.as_deref()
173    }
174
175    fn exif(&self) -> Option<&[u8]> {
176        self.exif.as_deref()
177    }
178
179    /// The container's `xml ` box.
180    fn xmp(&self) -> Option<&[u8]> {
181        self.xmp.as_deref()
182    }
183
184    /// Bit depth, and lossy (XYB) or lossless.
185    fn source_info(&self) -> SourceInfo {
186        self.info.clone()
187    }
188
189    fn animation(&self) -> Option<Animation> {
190        self.animation.clone()
191    }
192
193    fn read_row(&mut self, out: &mut [u8]) -> Result<()> {
194        if self.row >= self.descriptor.height {
195            return Err(PixelsError::invalid_argument(
196                "out",
197                format!("all {} rows have already been read", self.descriptor.height),
198            ));
199        }
200        let row_bytes = self.descriptor.row_bytes();
201        if out.len() != row_bytes {
202            return Err(PixelsError::invalid_argument(
203                "out",
204                format!("row buffer is {} bytes, expected {row_bytes}", out.len()),
205            ));
206        }
207        if self.pixels.is_none() {
208            self.pixels = Some(self.decode_pixels()?);
209            self.codestream = Vec::new();
210        }
211        let start = self.row as usize * row_bytes;
212        let row = self
213            .pixels
214            .as_deref()
215            .and_then(|p| p.get(start..start + row_bytes))
216            .ok_or_else(|| PixelsError::malformed("jxl", "decoded image is short"))?;
217        out.copy_from_slice(row);
218        self.row += 1;
219        Ok(())
220    }
221}
222
223/// Whether `prefix` starts with a JPEG XL signature, bare or contained.
224#[must_use]
225pub fn probe(prefix: &[u8]) -> bool {
226    container::probe(prefix)
227}
228
229/// The JPEG XL entry in a sniffing registry.
230#[derive(Debug, Clone, Copy, Default)]
231pub struct JxlCodec;
232
233impl Codec for JxlCodec {
234    fn format(&self) -> Format {
235        Format::Jxl
236    }
237
238    fn magic_len(&self) -> usize {
239        12
240    }
241
242    fn probe(&self, prefix: &[u8]) -> bool {
243        probe(prefix)
244    }
245}