Extract (potentially Brotli-compressed) Exif metadata (#389)

* Extract (potentially Brotli-compressed) Exif metadata * Implement `ImageDecoder::exif_metadata` * Refactor a bit
tirr-c · Nov 21, 2024 · 976b8d9 · 976b8d9
1 parent 74d682b
commit 976b8d9
Show file tree

Hide file tree

Showing 11 changed files with 507 additions and 24 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/jxl-bitstream/src/container.rs b/crates/jxl-bitstream/src/container.rs
@@ -20,13 +20,14 @@ enum DetectState {
     WaitingBoxHeader,
     WaitingJxlpIndex(ContainerBoxHeader),
     InAuxBox {
-        #[allow(unused)]
         header: ContainerBoxHeader,
+        brotli_box_type: Option<ContainerBoxType>,
         bytes_left: Option<usize>,
     },
     InCodestream {
         kind: BitstreamKind,
         bytes_left: Option<usize>,
+        pending_no_more_aux_box: bool,
     },
 }
 

diff --git a/crates/jxl-bitstream/src/container/box_header.rs b/crates/jxl-bitstream/src/container/box_header.rs
@@ -3,7 +3,7 @@ use crate::Error;
 /// Box header used in JPEG XL containers.
 #[derive(Debug, Clone)]
 pub struct ContainerBoxHeader {
-    ty: ContainerBoxType,
+    pub(super) ty: ContainerBoxType,
     box_size: Option<u64>,
     is_last: bool,
 }

diff --git a/crates/jxl-bitstream/src/container/parse.rs b/crates/jxl-bitstream/src/container/parse.rs
@@ -40,6 +40,7 @@ impl<'inner, 'buf> ParseEvents<'inner, 'buf> {
                         *state = DetectState::InCodestream {
                             kind: BitstreamKind::BareCodestream,
                             bytes_left: None,
+                            pending_no_more_aux_box: true,
                         };
                         return Ok(Some(ParseEvent::BitstreamKind(
                             BitstreamKind::BareCodestream,
@@ -56,12 +57,14 @@ impl<'inner, 'buf> ParseEvents<'inner, 'buf> {
                         *state = DetectState::InCodestream {
                             kind: BitstreamKind::Invalid,
                             bytes_left: None,
+                            pending_no_more_aux_box: true,
                         };
                         return Ok(Some(ParseEvent::BitstreamKind(BitstreamKind::Invalid)));
                     } else {
                         return Ok(None);
                     }
                 }
+
                 DetectState::WaitingBoxHeader => match ContainerBoxHeader::parse(buf)? {
                     HeaderParseResult::Done {
                         header,
@@ -84,9 +87,11 @@ impl<'inner, 'buf> ParseEvents<'inner, 'buf> {
                                 }
                             }
 
+                            let bytes_left = header.box_size().map(|x| x as usize);
                             *state = DetectState::InCodestream {
                                 kind: BitstreamKind::Container,
-                                bytes_left: header.box_size().map(|x| x as usize),
+                                bytes_left,
+                                pending_no_more_aux_box: bytes_left.is_none(),
                             };
                         } else if tbox == ContainerBoxType::PARTIAL_CODESTREAM {
                             if let Some(box_size) = header.box_size() {
@@ -115,11 +120,37 @@ impl<'inner, 'buf> ParseEvents<'inner, 'buf> {
                             *state = DetectState::WaitingJxlpIndex(header);
                         } else {
                             let bytes_left = header.box_size().map(|x| x as usize);
-                            *state = DetectState::InAuxBox { header, bytes_left };
+                            let ty = header.box_type();
+                            let mut brotli_compressed = ty == ContainerBoxType::BROTLI_COMPRESSED;
+                            if brotli_compressed {
+                                if let Some(0..=3) = bytes_left {
+                                    tracing::error!(
+                                        bytes_left = bytes_left.unwrap(),
+                                        "Brotli-compressed box is too small"
+                                    );
+                                    return Err(Error::InvalidBox);
+                                }
+                                brotli_compressed = true;
+                            }
+
+                            *state = DetectState::InAuxBox {
+                                header,
+                                brotli_box_type: None,
+                                bytes_left,
+                            };
+
+                            if !brotli_compressed {
+                                return Ok(Some(ParseEvent::AuxBoxStart {
+                                    ty,
+                                    brotli_compressed: false,
+                                    last_box: bytes_left.is_none(),
+                                }));
+                            }
                         }
                     }
                     HeaderParseResult::NeedMoreData => return Ok(None),
                 },
+
                 DetectState::WaitingJxlpIndex(header) => {
                     let &[b0, b1, b2, b3, ..] = &**buf else {
                         return Ok(None);
@@ -149,11 +180,23 @@ impl<'inner, 'buf> ParseEvents<'inner, 'buf> {
                         }
                     }
 
+                    let bytes_left = header.box_size().map(|x| x as usize - 4);
                     *state = DetectState::InCodestream {
                         kind: BitstreamKind::Container,
-                        bytes_left: header.box_size().map(|x| x as usize - 4),
+                        bytes_left,
+                        pending_no_more_aux_box: bytes_left.is_none(),
                     };
                 }
+
+                // JXL codestream box is the last box; emit "no more aux box" event.
+                DetectState::InCodestream {
+                    pending_no_more_aux_box: pending @ true,
+                    ..
+                } => {
+                    *pending = false;
+                    return Ok(Some(ParseEvent::NoMoreAuxBox));
+                }
+
                 DetectState::InCodestream {
                     bytes_left: None, ..
                 } => {
@@ -178,30 +221,78 @@ impl<'inner, 'buf> ParseEvents<'inner, 'buf> {
                     };
                     return Ok(Some(ParseEvent::Codestream(payload)));
                 }
+
+                // Read brob payload box type.
                 DetectState::InAuxBox {
-                    header: _,
-                    bytes_left: None,
+                    header:
+                        ContainerBoxHeader {
+                            ty: ContainerBoxType::BROTLI_COMPRESSED,
+                            ..
+                        },
+                    brotli_box_type: brotli_box_type @ None,
+                    bytes_left,
                 } => {
-                    let _payload = *buf;
-                    *buf = &[];
-                    // FIXME: emit auxiliary box event
+                    if buf.len() < 4 {
+                        return Ok(None);
+                    }
+
+                    let (ty_slice, remaining) = buf.split_at(4);
+                    *buf = remaining;
+                    if let Some(bytes_left) = bytes_left {
+                        *bytes_left -= 4;
+                    }
+
+                    let mut ty = [0u8; 4];
+                    ty.copy_from_slice(ty_slice);
+                    let is_reserved_box_type =
+                        &ty[..3] == b"jxl" || &ty == b"brob" || &ty == b"jbrd";
+                    if is_reserved_box_type {
+                        return Err(Error::ValidationFailed(
+                            "brob box, jxl boxes and jbrd box cannot be Brotli-compressed",
+                        ));
+                    }
+
+                    let ty = ContainerBoxType(ty);
+                    *brotli_box_type = Some(ty);
+
+                    return Ok(Some(ParseEvent::AuxBoxStart {
+                        ty,
+                        brotli_compressed: true,
+                        last_box: bytes_left.is_none(),
+                    }));
                 }
+
                 DetectState::InAuxBox {
-                    header: _,
-                    bytes_left: Some(bytes_left),
+                    header: ContainerBoxHeader { ty, .. },
+                    brotli_box_type,
+                    bytes_left,
                 } => {
-                    let _payload = if buf.len() >= *bytes_left {
-                        let (payload, remaining) = buf.split_at(*bytes_left);
-                        *state = DetectState::WaitingBoxHeader;
-                        *buf = remaining;
-                        payload
+                    let ty = if let Some(ty) = brotli_box_type {
+                        *ty
                     } else {
-                        let payload = *buf;
-                        *bytes_left -= buf.len();
-                        *buf = &[];
-                        payload
+                        *ty
                     };
-                    // FIXME: emit auxiliary box event
+
+                    let payload = match bytes_left {
+                        Some(0) => {
+                            *state = DetectState::WaitingBoxHeader;
+                            return Ok(Some(ParseEvent::AuxBoxEnd(ty)));
+                        }
+                        Some(bytes_left) => {
+                            let num_bytes_to_read = (*bytes_left).min(buf.len());
+                            let (payload, remaining) = buf.split_at(num_bytes_to_read);
+                            *bytes_left -= num_bytes_to_read;
+                            *buf = remaining;
+                            payload
+                        }
+                        None => {
+                            let payload = *buf;
+                            *buf = &[];
+                            payload
+                        }
+                    };
+
+                    return Ok(Some(ParseEvent::AuxBoxData(ty, payload)));
                 }
             }
         }
@@ -250,6 +341,14 @@ pub enum ParseEvent<'buf> {
     /// Returned data may be partial. Complete codestream can be obtained by concatenating all data
     /// of `Codestream` events.
     Codestream(&'buf [u8]),
+    NoMoreAuxBox,
+    AuxBoxStart {
+        ty: ContainerBoxType,
+        brotli_compressed: bool,
+        last_box: bool,
+    },
+    AuxBoxData(ContainerBoxType, &'buf [u8]),
+    AuxBoxEnd(ContainerBoxType),
 }
 
 impl std::fmt::Debug for ParseEvent<'_> {
@@ -260,6 +359,23 @@ impl std::fmt::Debug for ParseEvent<'_> {
                 .debug_tuple("Codestream")
                 .field(&format_args!("{} byte(s)", buf.len()))
                 .finish(),
+            Self::NoMoreAuxBox => write!(f, "NoMoreAuxBox"),
+            Self::AuxBoxStart {
+                ty,
+                brotli_compressed,
+                last_box,
+            } => f
+                .debug_struct("AuxBoxStart")
+                .field("ty", ty)
+                .field("brotli_compressed", brotli_compressed)
+                .field("last_box", last_box)
+                .finish(),
+            Self::AuxBoxData(ty, buf) => f
+                .debug_tuple("AuxBoxData")
+                .field(ty)
+                .field(&format_args!("{} byte(s)", buf.len()))
+                .finish(),
+            Self::AuxBoxEnd(ty) => f.debug_tuple("AuxBoxEnd").field(&ty).finish(),
         }
     }
 }
diff --git a/crates/jxl-oxide-cli/src/info.rs b/crates/jxl-oxide-cli/src/info.rs
@@ -63,6 +63,19 @@ pub fn handle_info(args: InfoArgs) -> Result<()> {
         }
     }
 
+    match image.raw_exif_data() {
+        Ok(None) => {}
+        Ok(Some(exif)) => {
+            if let Some(data) = exif.payload() {
+                let size = data.len();
+                println!("Exif metadata: {size} byte(s)");
+            }
+        }
+        Err(e) => {
+            println!("Invalid Exif metadata: {e}");
+        }
+    }
+
     if let Some(animation) = &image_meta.animation {
         println!(
             "  Animated ({}/{} ticks per second)",

diff --git a/crates/jxl-oxide/Cargo.toml b/crates/jxl-oxide/Cargo.toml
@@ -12,6 +12,7 @@ version = "0.11.0"
 edition = "2021"
 
 [dependencies]
+brotli-decompressor = "4.0.1"
 tracing.workspace = true
 
 [dependencies.bytemuck]

diff --git a/crates/jxl-oxide/examples/image-integration.rs b/crates/jxl-oxide/examples/image-integration.rs
@@ -16,6 +16,13 @@ fn main() {
 
     #[allow(unused)]
     let icc = decoder.icc_profile().unwrap();
+    let exif = decoder
+        .exif_metadata()
+        .expect("cannot decode Exif metadata");
+    if let Some(exif) = exif {
+        println!("Exif metadata found ({} byte(s))", exif.len());
+    }
+
     let image = DynamicImage::from_decoder(decoder).expect("cannot decode image");
 
     let output_file = std::fs::File::create(output_path).expect("cannot open output file");