Source code
Revision control
Copy as Markdown
Other Tools
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
#[cfg(feature = "brotli")]
use std::borrow::Cow;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::io::{IoSliceMut, Read};
use crate::api::{JxlBitstreamInput, JxlSignature, ProfileLevel, SmallBuffer, check_signature};
use crate::error::{Error, Result};
#[cfg(feature = "brotli")]
use crate::util::NewWithCapacity;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ParseState {
SignatureNeeded,
// We need at least this many bytes in the buffer to determine what to do
// with the next box.
BoxNeeded(usize),
// The next bytes are raw codestream. None = no limit.
Codestream(Option<u64>),
// Skip the next bytes.
Skip(Option<u64>),
// The next bytes are auxiliary box contents. None = no limit.
Aux(Option<u64>),
// The next bytes should be buffered in a jxlp box.
OOOJxlp(u32, Option<u64>),
// After the last codestream box, no more container bytes: no further
// codestream in file.
Complete,
}
// Relevant box types.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CodestreamBoxType {
None,
Jxlc,
// (index, is_last)
Jxlp(u32, bool),
}
impl CodestreamBoxType {
fn is_last(&self) -> bool {
matches!(self, Self::Jxlc | Self::Jxlp(_, true))
}
}
struct OOOJxlpBox {
data: Vec<u8>,
consumed: usize,
is_last: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct JxlAuxBoxType(pub [u8; 4]);
impl JxlAuxBoxType {
pub const EXIF: Self = JxlAuxBoxType(*b"Exif");
pub const XMP: Self = JxlAuxBoxType(*b"xml ");
pub const JUMBF: Self = JxlAuxBoxType(*b"jumb");
pub const JPEG_RECONSTRUCTION_DATA: Self = JxlAuxBoxType(*b"jbrd");
pub const HDR_GAIN_MAP: Self = JxlAuxBoxType(*b"jhgm");
}
pub struct JxlAuxBox {
ty: JxlAuxBoxType,
data: Vec<u8>,
brotli: bool,
}
impl std::fmt::Debug for JxlAuxBox {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let len = self.data.len();
f.debug_struct("JxlAuxBox")
.field("ty", &self.ty)
.field(
"data",
&format_args!("({len} byte{})", if len == 1 { "" } else { "s" }),
)
.field("brotli", &self.brotli)
.finish()
}
}
impl JxlAuxBox {
pub fn box_type(&self) -> JxlAuxBoxType {
self.ty
}
/// Returns the raw box data, potentially compressed.
pub fn raw_data(&self) -> &[u8] {
&self.data
}
/// Returns if the raw data is Brotli-compressed.
pub fn is_compressed(&self) -> bool {
self.brotli
}
#[cfg(feature = "brotli")]
pub fn data(&self, mut trailing_data: &[u8]) -> Result<Cow<'_, [u8]>> {
use std::io::prelude::*;
let output = if self.is_compressed() {
let mut data = &self.data[..];
let r = (&mut data).chain(&mut trailing_data);
let mut output = Vec::new();
let mut brotli = brotli_decompressor::Decompressor::new(r, 1024);
let mut buf = [0u8; 1024];
loop {
let n = brotli.read(&mut buf).map_err(crate::error::Error::Brotli)?;
if n == 0 {
break;
}
output.try_reserve(n)?;
output.extend_from_slice(&buf[..n]);
}
Cow::Owned(output)
} else if trailing_data.is_empty() {
Cow::Borrowed(&self.data[..])
} else {
let mut output = Vec::new_with_capacity(self.data.len() + trailing_data.len())?;
output.extend_from_slice(&self.data);
output.extend_from_slice(trailing_data);
Cow::Owned(output)
};
Ok(output)
}
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct BoxParserCheckpoint {
box_type: CodestreamBoxType,
pub(crate) file_position: u64,
codestream_left: Option<u64>,
is_valid_checkpoint: bool,
pub(super) consumed_codestream: u64,
next_aux_box_idx: usize,
}
pub(super) struct BoxParser {
local_buffer: SmallBuffer,
version: Option<u32>,
ooo_jxlp_buffer: HashMap<u32, OOOJxlpBox>,
state: ParseState,
latest_codestream_box: CodestreamBoxType,
// Map from codestream position at the start of the box
// to box info
codestream_pos_to_box: BTreeMap<u64, BoxParserCheckpoint>,
allow_checkpoint: bool,
container_level: Option<ProfileLevel>,
aux: AuxBoxState,
input_closed: bool,
}
#[derive(Default)]
struct AuxBoxState {
boxes_to_extract: HashSet<JxlAuxBoxType>,
boxes: HashMap<JxlAuxBoxType, Vec<JxlAuxBox>>,
box_buffer: Option<JxlAuxBox>,
next_box_idx: usize,
seen_box_count: usize,
}
impl BoxParser {
pub(super) fn with_aux_boxes(box_types: impl IntoIterator<Item = JxlAuxBoxType>) -> Self {
BoxParser {
local_buffer: SmallBuffer::new(128),
state: ParseState::SignatureNeeded,
latest_codestream_box: CodestreamBoxType::None,
ooo_jxlp_buffer: HashMap::new(),
version: None,
container_level: None,
codestream_pos_to_box: BTreeMap::new(),
allow_checkpoint: true,
aux: AuxBoxState {
boxes_to_extract: box_types.into_iter().collect(),
..Default::default()
},
input_closed: false,
}
}
fn start_aux_box(&mut self, box_type: JxlAuxBoxType, brotli: bool, content_len: Option<u64>) {
let box_idx = self.aux.next_box_idx;
self.aux.next_box_idx += 1;
let is_new = box_idx >= self.aux.seen_box_count;
self.aux.seen_box_count = self.aux.seen_box_count.max(self.aux.next_box_idx);
if is_new && self.aux.boxes_to_extract.contains(&box_type) {
self.aux.box_buffer = Some(JxlAuxBox {
ty: box_type,
brotli,
data: Vec::new(),
});
self.state = ParseState::Aux(content_len);
} else {
self.state = ParseState::Skip(content_len);
}
}
pub(super) fn reset(&mut self) {
self.local_buffer.reset();
self.version = None;
self.ooo_jxlp_buffer.clear();
self.state = ParseState::SignatureNeeded;
self.latest_codestream_box = CodestreamBoxType::None;
self.codestream_pos_to_box.clear();
self.allow_checkpoint = true;
self.container_level = None;
self.aux.boxes.clear();
self.aux.box_buffer = None;
self.aux.next_box_idx = 0;
self.aux.seen_box_count = 0;
self.input_closed = false;
}
pub(super) fn reset_to_checkpoint(&mut self, box_checkpoint: BoxParserCheckpoint) {
self.local_buffer.reset();
self.ooo_jxlp_buffer.clear();
self.latest_codestream_box = box_checkpoint.box_type;
self.state = ParseState::Codestream(box_checkpoint.codestream_left);
// Do not allow creating new checkpoints after a reset.
self.allow_checkpoint = false;
if self.aux.box_buffer.take().is_some() {
self.aux.seen_box_count -= 1;
}
self.aux.next_box_idx = box_checkpoint.next_aux_box_idx;
self.input_closed = false;
}
pub(super) fn close_input(&mut self) {
self.input_closed = true;
}
pub(super) fn state_checkpoint(
&self,
codestream_pos: u64,
) -> Result<Option<BoxParserCheckpoint>> {
if !self.allow_checkpoint {
Ok(None)
} else {
// last element with key <= codestream_pos.
let (&start, &(mut b)) = self
.codestream_pos_to_box
.range(..=codestream_pos)
.last()
.unwrap();
if !b.is_valid_checkpoint {
// Disallow files that have frame starts in boxes that are not
// valid checkpoints.
return Err(Error::InvalidBox);
}
let overflow = codestream_pos - start;
if let Some(bytes) = &mut b.codestream_left {
*bytes -= overflow;
}
b.file_position += overflow;
b.consumed_codestream = codestream_pos;
Ok(Some(b))
}
}
pub(super) fn total_bytes_consumed(&self, codestream_bytes_consumed: u64) -> u64 {
let (start, b) = self.codestream_pos_to_box.last_key_value().unwrap();
b.file_position + codestream_bytes_consumed.saturating_sub(*start)
}
pub(super) fn aux_boxes(&self, box_type: JxlAuxBoxType) -> &[JxlAuxBox] {
self.aux
.boxes
.get(&box_type)
.map(|v| &**v)
.unwrap_or_default()
}
pub(super) fn container_level(&self) -> Option<ProfileLevel> {
self.container_level
}
pub(super) fn trailing_box(&self) -> Option<&JxlAuxBox> {
if self.state != ParseState::Aux(None) {
return None;
}
self.aux.box_buffer.as_ref()
}
pub(super) fn is_codestream_complete(&self) -> bool {
if !self.latest_codestream_box.is_last() {
return false;
}
match self.state {
ParseState::Codestream(Some(0)) => true,
ParseState::Codestream(_) => false,
_ => true,
}
}
fn add_checkpoint(&mut self) {
if !self.allow_checkpoint {
return;
}
// For a codestream box to be a valid checkpoint, it needs to be decoded
// in-order (i.e. all other logically-preceding jxlp boxes have been seen already)
// and we must not have seen any future jxlp boxes (which would be stored in the
// ooo jxlp buffer).
let (codestream_left, is_valid_checkpoint) = match self.state {
ParseState::Codestream(c) => (c, self.ooo_jxlp_buffer.is_empty()),
ParseState::OOOJxlp(_, c) => (c, false),
_ => (None, false),
};
let codestream_position = self
.codestream_pos_to_box
.last_key_value()
.map(|(p, c)| *p + c.codestream_left.unwrap())
.unwrap_or(0);
self.codestream_pos_to_box.insert(
codestream_position,
BoxParserCheckpoint {
file_position: self.local_buffer.consumed(),
box_type: self.latest_codestream_box,
codestream_left,
is_valid_checkpoint,
consumed_codestream: 0,
next_aux_box_idx: self.aux.next_box_idx,
},
);
}
fn injected_jxlp(&mut self) -> Option<&mut OOOJxlpBox> {
// Skip 0-sized jxlp boxes.
while matches!(self.state, ParseState::BoxNeeded(_) | ParseState::Complete)
&& let CodestreamBoxType::Jxlp(j, _) = self.latest_codestream_box
&& self
.ooo_jxlp_buffer
.get(&j)
.is_some_and(|b| b.consumed == b.data.len())
{
self.check_ooo_jxlp_done();
}
if matches!(self.state, ParseState::BoxNeeded(_) | ParseState::Complete)
&& let CodestreamBoxType::Jxlp(j, _) = self.latest_codestream_box
&& let Some(b) = self.ooo_jxlp_buffer.get_mut(&j)
{
Some(b)
} else {
None
}
}
fn check_ooo_jxlp_done(&mut self) {
let CodestreamBoxType::Jxlp(j, is_last) = self.latest_codestream_box else {
unreachable!()
};
let Some(b) = self.ooo_jxlp_buffer.get_mut(&j) else {
unreachable!()
};
if b.consumed == b.data.len() {
self.ooo_jxlp_buffer.remove(&j);
if is_last {
self.state = ParseState::Complete;
} else if let Some(b) = self.ooo_jxlp_buffer.get(&(j + 1)) {
self.latest_codestream_box = CodestreamBoxType::Jxlp(j + 1, b.is_last);
}
}
}
fn mark_codestream_read(&mut self, n: usize) {
let ParseState::Codestream(c) = &mut self.state else {
unreachable!()
};
if c.is_none() {
return;
}
let new_size = c.unwrap() - n as u64;
if new_size == 0 {
self.state = if self.latest_codestream_box.is_last() {
ParseState::Complete
} else {
if let CodestreamBoxType::Jxlp(j, _) = self.latest_codestream_box
&& let Some(b) = self.ooo_jxlp_buffer.get(&(j + 1))
{
self.latest_codestream_box = CodestreamBoxType::Jxlp(j + 1, b.is_last);
self.add_checkpoint();
}
ParseState::BoxNeeded(8)
}
} else {
*c = Some(new_size);
}
}
fn available_bytes_inner(&mut self, input: &mut dyn JxlBitstreamInput) -> Result<usize> {
Ok(self.local_buffer.len() + input.available_bytes()?)
}
fn read_inner(
&mut self,
input: &mut dyn JxlBitstreamInput,
mut buf: &mut [IoSliceMut],
) -> Result<usize> {
let mut n = 0;
if !self.local_buffer.is_empty() {
let mut ll = &self.local_buffer[..];
n = ll.read_vectored(buf)?;
IoSliceMut::advance_slices(&mut buf, n);
self.local_buffer.consume(n);
}
let read = input.read(buf)?;
self.local_buffer.mark_consumed(read as u64);
Ok(n + read)
}
fn skip_inner(
&mut self,
input: &mut dyn JxlBitstreamInput,
mut amount: usize,
) -> Result<usize> {
let mut n = 0;
if !self.local_buffer.is_empty() {
let ll = self.local_buffer.len();
n = ll.min(amount);
amount -= n;
self.local_buffer.consume(n);
}
let skip = input.skip(amount)?;
self.local_buffer.mark_consumed(skip as u64);
Ok(n + skip)
}
fn read_until_at_least(
&mut self,
input: &mut dyn JxlBitstreamInput,
min_size: usize,
) -> Result<()> {
loop {
let n = self.local_buffer.refill(|b| Ok(input.read(b)?))?;
if n == 0 || self.local_buffer.len() >= min_size {
break;
}
}
if self.local_buffer.len() < min_size {
return Err(Error::OutOfBounds(min_size - self.local_buffer.len()));
}
Ok(())
}
fn advance_to_codestream(&mut self, input: &mut dyn JxlBitstreamInput) -> Result<()> {
loop {
match self.state {
ParseState::Codestream(Some(0)) => self.state = ParseState::BoxNeeded(8),
ParseState::Complete | ParseState::Codestream(_) => return Ok(()),
ParseState::SignatureNeeded => {
let codestream_signature_len = JxlSignature::Codestream.signature_len();
self.read_until_at_least(input, codestream_signature_len)?;
match check_signature(&self.local_buffer) {
JxlSignature::NeedMoreInput { size_hint } => {
return Err(Error::OutOfBounds(size_hint));
}
JxlSignature::None => {
return Err(Error::InvalidSignature);
}
JxlSignature::Codestream => {
self.state = ParseState::Codestream(None);
self.latest_codestream_box = CodestreamBoxType::Jxlc;
self.add_checkpoint();
return Ok(());
}
JxlSignature::Container => {
let l = JxlSignature::Container.signature_len();
self.local_buffer.consume(l);
self.state = ParseState::BoxNeeded(8);
}
}
}
ParseState::Skip(count) => {
if count == Some(0) {
self.state = ParseState::BoxNeeded(8);
continue;
}
let to_skip = count.unwrap_or(u64::MAX).min(usize::MAX as u64) as usize;
let n = self.skip_inner(input, to_skip)? as u64;
if n == 0 {
return Err(Error::OutOfBounds(to_skip));
}
self.state = ParseState::Skip(count.map(|x| x - n));
}
ParseState::Aux(None) => {
let local_buffer_len = self.local_buffer.len();
let buf = self.aux.box_buffer.as_mut().unwrap();
if local_buffer_len > 0 {
buf.data.extend_from_slice(&self.local_buffer);
self.local_buffer.consume(local_buffer_len);
}
return Ok(());
}
ParseState::Aux(Some(count)) => {
if count == 0 {
let buf = self.aux.box_buffer.take().unwrap();
self.aux.boxes.entry(buf.ty).or_default().push(buf);
self.state = if self.latest_codestream_box.is_last() {
ParseState::Complete
} else {
ParseState::BoxNeeded(8)
};
continue;
}
let total = self.handle_aux_box(input, count)?;
if total == 0 {
return Err(Error::OutOfBounds(count.min(usize::MAX as u64) as usize));
}
self.state = ParseState::Aux(Some(count - total));
}
ParseState::OOOJxlp(id, count) => {
if count == Some(0) {
self.state = ParseState::BoxNeeded(8);
continue;
}
// Temporarily take the buffer out of the jxlp state to make the borrow checker happy.
let mut buf = self.ooo_jxlp_buffer.remove(&id).unwrap();
let mut total = 0;
loop {
let space = buf
.data
.len()
.max(1024)
.min(self.available_bytes_inner(input)?)
as u64;
let space = count.map(|x| (x - total).min(space)).unwrap_or(space) as usize;
if space == 0 {
break;
}
buf.data.try_reserve(space)?;
let cur = buf.data.len();
buf.data.resize(cur + space, 0);
let n =
self.read_inner(input, &mut [IoSliceMut::new(&mut buf.data[cur..])])?;
buf.data.truncate(cur + n);
if n == 0 {
break;
}
total += n as u64;
}
self.ooo_jxlp_buffer.insert(id, buf);
if total == 0 {
return Err(Error::OutOfBounds(
count.map(|x| x.min(usize::MAX as u64)).unwrap_or(1) as usize,
));
}
self.state = ParseState::OOOJxlp(id, count.map(|x| x - total));
}
ParseState::BoxNeeded(min_size) => self.parse_box(input, min_size)?,
}
}
}
fn consume_trailing_data(&mut self, input: &mut dyn JxlBitstreamInput) -> Result<()> {
if self.aux.boxes_to_extract.is_empty() {
return Ok(());
}
loop {
match self.state {
ParseState::Codestream(None) => {
return Ok(());
}
ParseState::Complete => {
if self.available_bytes_inner(input)? == 0 {
if self.input_closed {
return Ok(());
}
return Err(Error::OutOfBounds(8));
}
self.state = ParseState::BoxNeeded(8);
}
ParseState::Skip(None) => return Ok(()),
ParseState::Codestream(count) | ParseState::Skip(count) => {
if count == Some(0) {
self.state = ParseState::Complete;
continue;
}
let to_skip = count.unwrap_or(u64::MAX).min(usize::MAX as u64) as usize;
let n = self.skip_inner(input, to_skip)? as u64;
if n == 0 {
if self.input_closed {
return Err(Error::InvalidBox);
}
return Err(Error::OutOfBounds(to_skip));
}
self.state = ParseState::Skip(count.map(|x| x - n));
}
ParseState::Aux(None) => {
let local_buffer_len = self.local_buffer.len();
let buf = self.aux.box_buffer.as_mut().unwrap();
if local_buffer_len > 0 {
buf.data.extend_from_slice(&self.local_buffer);
self.local_buffer.consume(local_buffer_len);
}
return Ok(());
}
ParseState::Aux(Some(count)) => {
if count == 0 {
let buf = self.aux.box_buffer.take().unwrap();
self.aux.boxes.entry(buf.ty).or_default().push(buf);
self.state = ParseState::Complete;
continue;
}
let total = self.handle_aux_box(input, count)?;
if total == 0 {
if self.input_closed {
return Err(Error::InvalidBox);
}
return Err(Error::OutOfBounds(count.min(usize::MAX as u64) as usize));
}
self.state = ParseState::Aux(Some(count - total));
}
ParseState::BoxNeeded(min_size) => self.parse_box(input, min_size)?,
_ => {
return Err(Error::InvalidBox);
}
}
}
}
fn handle_aux_box(&mut self, input: &mut dyn JxlBitstreamInput, count: u64) -> Result<u64> {
let mut buf = self.aux.box_buffer.take().unwrap();
let mut total = 0;
loop {
let space = buf
.data
.len()
.max(1024)
.min(self.available_bytes_inner(input)?) as u64;
let space = (count - total).min(space) as usize;
if space == 0 {
break;
}
buf.data.try_reserve(space)?;
let cur = buf.data.len();
buf.data.resize(cur + space, 0);
let n = self.read_inner(input, &mut [IoSliceMut::new(&mut buf.data[cur..])])?;
buf.data.truncate(cur + n);
if n == 0 {
break;
}
total += n as u64;
}
self.aux.box_buffer = Some(buf);
Ok(total)
}
fn parse_box(&mut self, input: &mut dyn JxlBitstreamInput, required_size: usize) -> Result<()> {
match self.read_until_at_least(input, required_size) {
Err(Error::OutOfBounds(_)) if self.input_closed => return Err(Error::InvalidBox),
res => res?,
}
let min_len = match &self.local_buffer[..] {
[0, 0, 0, 1, ..] => 16,
_ => 8,
};
let ty: [_; 4] = self.local_buffer[4..8].try_into().unwrap();
if &ty != b"ftyp" && self.version.is_none() {
// ftyp should be the first box.
return Err(Error::InvalidBox);
}
let extra_len = match &ty {
b"jxlp" => 4,
b"ftyp" => 8,
b"brob" => 4,
b"jxll" => 1,
_ => 0,
};
if self.local_buffer.len() < extra_len + min_len {
if self.input_closed {
return Err(Error::InvalidBox);
}
self.state = ParseState::BoxNeeded(extra_len + min_len);
return Ok(());
}
let box_len = match &self.local_buffer[..] {
[0, 0, 0, 1, ..] => u64::from_be_bytes(self.local_buffer[8..16].try_into().unwrap()),
_ => u32::from_be_bytes(self.local_buffer[0..4].try_into().unwrap()) as u64,
};
self.local_buffer.consume(min_len);
let content_len = if box_len == 0 {
None
} else if let Some(d) = box_len.checked_sub((extra_len + min_len) as u64) {
Some(d)
} else {
return Err(Error::InvalidBox);
};
match &ty {
b"ftyp" => {
if self.version.is_some() {
return Err(Error::InvalidBox);
}
if &self.local_buffer[0..4] != b"jxl " {
return Err(Error::InvalidBox);
}
let ver = u32::from_be_bytes(self.local_buffer[4..8].try_into().unwrap());
if ver > 1 {
return Err(Error::InvalidBox);
}
self.version = Some(ver);
self.local_buffer.consume(8);
self.state = ParseState::Skip(content_len);
}
b"jxlc" => {
if matches!(self.latest_codestream_box, CodestreamBoxType::Jxlp(..)) {
return Err(Error::InvalidBox);
}
self.latest_codestream_box = CodestreamBoxType::Jxlc;
self.state = ParseState::Codestream(content_len);
self.add_checkpoint();
}
b"jxlp" => {
let index = u32::from_be_bytes(self.local_buffer[..4].try_into().unwrap());
self.local_buffer.consume(4);
let last = index & 0x80000000 != 0;
let idx = index & 0x7fffffff;
let wanted_idx = match self.latest_codestream_box {
CodestreamBoxType::Jxlc | CodestreamBoxType::Jxlp(_, true) => {
return Err(Error::InvalidBox);
}
CodestreamBoxType::None => 0,
CodestreamBoxType::Jxlp(i, _) => i + 1,
};
if idx < wanted_idx {
return Err(Error::InvalidBox);
}
if idx > wanted_idx {
// Out-of-order jxlp
if self.version != Some(1) {
return Err(Error::InvalidBox);
}
if self.ooo_jxlp_buffer.contains_key(&idx) {
return Err(Error::InvalidBox);
}
if content_len.is_none() {
return Err(Error::InvalidBox);
}
self.state = ParseState::OOOJxlp(idx, content_len);
self.ooo_jxlp_buffer.insert(
idx,
OOOJxlpBox {
data: vec![],
consumed: 0,
is_last: last,
},
);
// We create a known-invalid checkpoint here to ensure
// that the checks work properly, and to count codestream
// length correctly.
self.add_checkpoint();
return Ok(());
}
// In-order jxlp
self.latest_codestream_box = CodestreamBoxType::Jxlp(idx, last);
self.state = ParseState::Codestream(content_len);
self.add_checkpoint();
}
b"brob" => {
let inner_type = &self.local_buffer[..4];
if matches!(inner_type, b"brob" | [b'j', b'x', b'l', _] | b"jbrd") {
return Err(Error::InvalidBox);
}
let inner_type = JxlAuxBoxType(inner_type.try_into().unwrap());
self.local_buffer.consume(4);
self.start_aux_box(inner_type, true, content_len);
}
b"jxll" => {
if self.version.is_none() || self.latest_codestream_box != CodestreamBoxType::None {
return Err(Error::InvalidBox);
}
if self.container_level.is_some() {
return Err(Error::InvalidBox);
}
if content_len != Some(0) {
return Err(Error::InvalidBox);
}
let level = match self.local_buffer[0] {
5 => ProfileLevel::Main5,
10 => ProfileLevel::Main10,
_ => return Err(Error::InvalidBox),
};
self.local_buffer.consume(1);
self.container_level = Some(level);
self.state = ParseState::BoxNeeded(8);
}
code => {
let ty = JxlAuxBoxType(*code);
self.start_aux_box(ty, false, content_len);
}
}
Ok(())
}
}
pub(super) struct CodestreamInput<'a> {
box_parser: &'a mut BoxParser,
input: &'a mut dyn JxlBitstreamInput,
}
impl<'a> CodestreamInput<'a> {
pub(super) fn new(box_parser: &'a mut BoxParser, input: &'a mut dyn JxlBitstreamInput) -> Self {
Self { box_parser, input }
}
pub(super) fn box_parser(&self) -> &BoxParser {
self.box_parser
}
pub(super) fn consume_trailing_data(&mut self) -> Result<()> {
self.box_parser.consume_trailing_data(self.input)
}
// The methods below have the same semantics as the methods in JxlBitstreamInput.
pub(super) fn available_bytes(&mut self) -> Result<usize> {
if let Some(b) = self.box_parser.injected_jxlp() {
return Ok(b.data.len() - b.consumed);
}
self.box_parser.advance_to_codestream(self.input)?;
let ParseState::Codestream(r) = self.box_parser.state else {
return Ok(0);
};
let available = self.box_parser.available_bytes_inner(self.input)?;
Ok(r.map(|x| x.min(available as u64) as usize)
.unwrap_or(available))
}
pub(super) fn read(&mut self, buf: &mut [IoSliceMut]) -> Result<usize> {
if let Some(b) = self.box_parser.injected_jxlp() {
let mut r = &b.data[b.consumed..];
let n = r.read_vectored(buf)?;
b.consumed += n;
self.box_parser.check_ooo_jxlp_done();
return Ok(n);
}
self.box_parser.advance_to_codestream(self.input)?;
let ParseState::Codestream(r) = self.box_parser.state else {
return Ok(0);
};
let limit = r.unwrap_or(u64::MAX);
// If the input buffers have more space than the remaining codestream in the current
// box, split the read to ensure we don't over-read.
let n = if buf.iter().map(|x| x.len() as u64).sum::<u64>() <= limit {
self.box_parser.read_inner(self.input, buf)?
} else {
let mut cumsum = 0;
let mut upper = 0;
while cumsum + buf[upper].len() as u64 <= limit {
cumsum += buf[upper].len() as u64;
upper += 1;
}
let n = self.box_parser.read_inner(self.input, &mut buf[..upper])?;
let left = (limit - n as u64) as usize;
n + self
.box_parser
.read_inner(self.input, &mut [IoSliceMut::new(&mut buf[upper][..left])])?
};
self.box_parser.mark_codestream_read(n);
Ok(n)
}
pub(super) fn skip(&mut self, amount: usize) -> Result<usize> {
if let Some(b) = self.box_parser.injected_jxlp() {
let n = (b.data.len() - b.consumed).min(amount);
b.consumed += n;
self.box_parser.check_ooo_jxlp_done();
return Ok(n);
}
self.box_parser.advance_to_codestream(self.input)?;
let ParseState::Codestream(r) = self.box_parser.state else {
return Ok(0);
};
let n = self.box_parser.skip_inner(
self.input,
r.unwrap_or(u64::MAX).min(amount as u64) as usize,
)?;
self.box_parser.mark_codestream_read(n);
Ok(n)
}
}
#[cfg(test)]
mod tests {
use std::io::IoSliceMut;
use super::{BoxParser, CodestreamInput};
/// Regression: a zero-length skippable box must not leave the parser stuck at
/// `SkippableBox(0)` when more container input is available.
#[test]
fn zero_length_skippable_box_does_not_hang() {
let data = include_bytes!("../../tests/testdata/zero_length_skippable_box.jxl");
let mut parser = BoxParser::with_aux_boxes(None);
let mut input = data.as_slice();
{
let mut input = CodestreamInput::new(&mut parser, &mut input);
let mut buf = vec![0; 1024];
while matches!(input.read(&mut [IoSliceMut::new(&mut buf)]), Ok(x) if x > 0) {}
}
// All input should be consumed.
assert!(input.is_empty());
}
}