jxl_oxide/
aux_box.rs

1use std::io::Write;
2
3use brotli_decompressor::DecompressorWriter;
4use jxl_bitstream::container::box_header::ContainerBoxType;
5use jxl_bitstream::ParseEvent;
6
7use crate::Result;
8
9mod exif;
10mod jbrd;
11
12pub use exif::*;
13pub use jbrd::*;
14
15#[derive(Debug, Default)]
16pub struct AuxBoxReader {
17    data: DataKind,
18    done: bool,
19}
20
21#[derive(Default)]
22enum DataKind {
23    #[default]
24    Init,
25    NoData,
26    Raw(Vec<u8>),
27    Brotli(Box<DecompressorWriter<Vec<u8>>>),
28}
29
30impl std::fmt::Debug for DataKind {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        match self {
33            Self::Init => write!(f, "Init"),
34            Self::NoData => write!(f, "NoData"),
35            Self::Raw(buf) => f
36                .debug_tuple("Raw")
37                .field(&format_args!("{} byte(s)", buf.len()))
38                .finish(),
39            Self::Brotli(_) => f.debug_tuple("Brotli").finish(),
40        }
41    }
42}
43
44impl AuxBoxReader {
45    pub(super) fn new() -> Self {
46        Self::default()
47    }
48
49    pub(super) fn ensure_raw(&mut self) {
50        if self.done {
51            return;
52        }
53
54        match self.data {
55            DataKind::Init => {
56                self.data = DataKind::Raw(Vec::new());
57            }
58            DataKind::NoData | DataKind::Brotli(_) => {
59                panic!();
60            }
61            DataKind::Raw(_) => {}
62        }
63    }
64
65    pub(super) fn ensure_brotli(&mut self) -> Result<()> {
66        if self.done {
67            return Ok(());
68        }
69
70        match self.data {
71            DataKind::Init => {
72                let writer = DecompressorWriter::new(Vec::<u8>::new(), 4096);
73                self.data = DataKind::Brotli(Box::new(writer));
74            }
75            DataKind::NoData | DataKind::Raw(_) => {
76                panic!();
77            }
78            DataKind::Brotli(_) => {}
79        }
80        Ok(())
81    }
82}
83
84impl AuxBoxReader {
85    pub fn feed_data(&mut self, data: &[u8]) -> Result<()> {
86        if self.done {
87            return Err(std::io::Error::new(
88                std::io::ErrorKind::InvalidData,
89                "cannot feed into finalized box",
90            )
91            .into());
92        }
93
94        match self.data {
95            DataKind::Init => {
96                self.data = DataKind::Raw(data.to_vec());
97            }
98            DataKind::NoData => {
99                unreachable!();
100            }
101            DataKind::Raw(ref mut buf) => {
102                buf.extend_from_slice(data);
103            }
104            DataKind::Brotli(ref mut writer) => {
105                writer.write_all(data)?;
106            }
107        }
108        Ok(())
109    }
110
111    pub fn finalize(&mut self) -> Result<()> {
112        if self.done {
113            return Ok(());
114        }
115
116        if let DataKind::Brotli(ref mut writer) = self.data {
117            writer.flush()?;
118            writer.close()?;
119        }
120
121        match std::mem::replace(&mut self.data, DataKind::NoData) {
122            DataKind::Init | DataKind::NoData => {}
123            DataKind::Raw(buf) => self.data = DataKind::Raw(buf),
124            DataKind::Brotli(writer) => {
125                let inner = writer.into_inner().inspect_err(|_| {
126                    tracing::warn!("Brotli decompressor reported an error");
127                });
128                let buf = inner.unwrap_or_else(|buf| buf);
129                self.data = DataKind::Raw(buf);
130            }
131        }
132
133        self.done = true;
134        Ok(())
135    }
136}
137
138impl AuxBoxReader {
139    pub fn is_done(&self) -> bool {
140        self.done
141    }
142
143    pub fn data(&self) -> AuxBoxData<&[u8]> {
144        if !self.is_done() {
145            return AuxBoxData::Decoding;
146        }
147
148        match &self.data {
149            DataKind::Init | DataKind::Brotli(_) => AuxBoxData::Decoding,
150            DataKind::NoData => AuxBoxData::NotFound,
151            DataKind::Raw(buf) => AuxBoxData::Data(buf),
152        }
153    }
154}
155
156/// Auxiliary box data.
157pub enum AuxBoxData<T> {
158    /// The box has data.
159    Data(T),
160    /// The box has not been decoded yet.
161    Decoding,
162    /// The box was not found.
163    NotFound,
164}
165
166impl<T> std::fmt::Debug for AuxBoxData<T> {
167    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
168        match self {
169            Self::Data(_) => write!(f, "Data(_)"),
170            Self::Decoding => write!(f, "Decoding"),
171            Self::NotFound => write!(f, "NotFound"),
172        }
173    }
174}
175
176impl<T> AuxBoxData<T> {
177    pub fn has_data(&self) -> bool {
178        matches!(self, Self::Data(_))
179    }
180
181    pub fn is_decoding(&self) -> bool {
182        matches!(self, Self::Decoding)
183    }
184
185    pub fn is_not_found(&self) -> bool {
186        matches!(self, Self::NotFound)
187    }
188
189    pub fn unwrap(self) -> T {
190        let Self::Data(x) = self else {
191            panic!("cannot unwrap `AuxBoxData` which doesn't have any data");
192        };
193        x
194    }
195
196    pub fn unwrap_or(self, or: T) -> T {
197        match self {
198            Self::Data(x) => x,
199            _ => or,
200        }
201    }
202
203    pub fn map<U>(self, f: impl FnOnce(T) -> U) -> AuxBoxData<U> {
204        match self {
205            Self::Data(x) => AuxBoxData::Data(f(x)),
206            Self::Decoding => AuxBoxData::Decoding,
207            Self::NotFound => AuxBoxData::NotFound,
208        }
209    }
210
211    pub fn as_ref(&self) -> AuxBoxData<&T> {
212        match self {
213            Self::Data(x) => AuxBoxData::Data(x),
214            Self::Decoding => AuxBoxData::Decoding,
215            Self::NotFound => AuxBoxData::NotFound,
216        }
217    }
218}
219
220impl<T, E> AuxBoxData<std::result::Result<T, E>> {
221    pub fn transpose(self) -> std::result::Result<AuxBoxData<T>, E> {
222        match self {
223            Self::Data(Ok(x)) => Ok(AuxBoxData::Data(x)),
224            Self::Data(Err(e)) => Err(e),
225            Self::Decoding => Ok(AuxBoxData::Decoding),
226            Self::NotFound => Ok(AuxBoxData::NotFound),
227        }
228    }
229}
230
231/// Auxiliary box list of a JPEG XL container, which may contain Exif and/or XMP metadata.
232#[derive(Debug)]
233pub struct AuxBoxList {
234    boxes: Vec<(ContainerBoxType, AuxBoxReader)>,
235    jbrd: Jbrd,
236    current_box_ty: Option<ContainerBoxType>,
237    current_box: AuxBoxReader,
238    last_box: bool,
239}
240
241impl AuxBoxList {
242    pub(super) fn new() -> Self {
243        Self {
244            boxes: Vec::new(),
245            jbrd: Jbrd::new(),
246            current_box_ty: None,
247            current_box: AuxBoxReader::new(),
248            last_box: false,
249        }
250    }
251
252    pub(super) fn handle_event(&mut self, event: ParseEvent) -> Result<()> {
253        match event {
254            ParseEvent::BitstreamKind(_) => {}
255            ParseEvent::Codestream(_) => {}
256            ParseEvent::NoMoreAuxBox => {
257                self.current_box_ty = None;
258                self.last_box = true;
259            }
260            ParseEvent::AuxBoxStart {
261                ty,
262                brotli_compressed,
263                last_box,
264            } => {
265                self.current_box_ty = Some(ty);
266                if ty != ContainerBoxType::JPEG_RECONSTRUCTION {
267                    if brotli_compressed {
268                        self.current_box.ensure_brotli()?;
269                    } else {
270                        self.current_box.ensure_raw();
271                    }
272                }
273                self.last_box = last_box;
274            }
275            ParseEvent::AuxBoxData(ty, buf) => {
276                self.current_box_ty = Some(ty);
277                if ty == ContainerBoxType::JPEG_RECONSTRUCTION {
278                    self.jbrd.feed_bytes(buf)?;
279                } else {
280                    self.current_box.feed_data(buf)?;
281                }
282            }
283            ParseEvent::AuxBoxEnd(ty) => {
284                self.current_box_ty = Some(ty);
285                self.finalize()?;
286            }
287        }
288        Ok(())
289    }
290
291    fn finalize(&mut self) -> Result<()> {
292        match self.current_box_ty {
293            Some(ContainerBoxType::JPEG_RECONSTRUCTION) => {
294                self.jbrd.finalize()?;
295            }
296            Some(ty) => {
297                self.current_box.finalize()?;
298                let finished_box = std::mem::replace(&mut self.current_box, AuxBoxReader::new());
299                self.boxes.push((ty, finished_box));
300            }
301            None => {
302                return Ok(());
303            }
304        }
305
306        self.current_box_ty = None;
307        Ok(())
308    }
309
310    pub(super) fn eof(&mut self) -> Result<()> {
311        self.finalize()?;
312        self.last_box = true;
313        Ok(())
314    }
315}
316
317impl AuxBoxList {
318    pub(crate) fn jbrd(&self) -> AuxBoxData<&jxl_jbr::JpegBitstreamData> {
319        if let Some(data) = self.jbrd.data() {
320            AuxBoxData::Data(data)
321        } else if self.last_box
322            && self.current_box_ty != Some(ContainerBoxType::JPEG_RECONSTRUCTION)
323        {
324            AuxBoxData::NotFound
325        } else {
326            AuxBoxData::Decoding
327        }
328    }
329
330    fn first_of_type(&self, ty: ContainerBoxType) -> AuxBoxData<&[u8]> {
331        let maybe = self.boxes.iter().find(|&&(ty_to_test, _)| ty_to_test == ty);
332        let data = maybe.map(|(_, b)| b.data());
333
334        if let Some(data) = data {
335            data
336        } else if self.last_box && self.current_box_ty != Some(ty) {
337            AuxBoxData::NotFound
338        } else {
339            AuxBoxData::Decoding
340        }
341    }
342
343    /// Returns the first Exif metadata, if any.
344    pub fn first_exif(&self) -> Result<AuxBoxData<RawExif>> {
345        let exif = self.first_of_type(ContainerBoxType::EXIF);
346        exif.map(RawExif::new).transpose()
347    }
348
349    /// Returns the first XML metadata, if any.
350    pub fn first_xml(&self) -> AuxBoxData<&[u8]> {
351        self.first_of_type(ContainerBoxType::XML)
352    }
353}
jxl_oxide/aux_box.rs

jxl_oxide/
aux_box.rs
