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.
use std::io::Cursor;
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use crate::bit_reader::*;
use crate::entropy_coding::decode::Histograms;
use crate::entropy_coding::decode::SymbolReader;
use crate::error::{Error, Result};
use crate::headers::encodings::*;
use crate::util::NewWithCapacity;
use crate::util::tracing_wrappers::warn;
mod header;
mod stream;
mod tag;
use header::read_header;
use stream::{IccStream, read_varint_from_reader};
use tag::{read_single_command, read_tag_list};
const ICC_CONTEXTS: usize = 41;
const ICC_HEADER_SIZE: u64 = 128;
fn read_icc_inner(stream: &mut IccStream) -> Result<Vec<u8>, Error> {
let output_size = stream.read_varint()?;
let commands_size = stream.read_varint()?;
if stream.bytes_read().saturating_add(commands_size) > stream.len() {
return Err(Error::InvalidIccStream);
}
// Simple check to avoid allocating too large buffer.
if output_size > (1 << 28) {
return Err(Error::IccTooLarge);
}
if output_size + 65536 < stream.len() {
return Err(Error::IccTooLarge);
}
// Extract command stream first.
let commands = stream.read_to_vec_exact(commands_size as usize)?;
let mut commands_stream = Cursor::new(commands);
// `stream` contains data stream from here.
let data_stream = stream;
// Decode ICC profile header.
let mut decoded_profile = read_header(data_stream, output_size)?;
if output_size <= ICC_HEADER_SIZE {
return Ok(decoded_profile);
}
// Convert to slice writer to prevent buffer from growing.
// `read_header` above returns buffer with capacity of `output_size`, so this doesn't realloc.
debug_assert_eq!(decoded_profile.capacity(), output_size as usize);
decoded_profile.resize(output_size as usize, 0);
let mut decoded_profile_writer = Cursor::new(&mut *decoded_profile);
decoded_profile_writer.set_position(ICC_HEADER_SIZE);
// Decode tag list.
let v = read_varint_from_reader(&mut commands_stream)?;
if let Some(num_tags) = v.checked_sub(1) {
if (output_size - ICC_HEADER_SIZE) / 12 < num_tags {
warn!(output_size, num_tags, "num_tags too large");
return Err(Error::InvalidIccStream);
}
let num_tags = num_tags as u32;
decoded_profile_writer
.write_u32::<BigEndian>(num_tags)
.map_err(|_| Error::InvalidIccStream)?;
read_tag_list(
data_stream,
&mut commands_stream,
&mut decoded_profile_writer,
num_tags,
output_size,
)?;
}
// Decode tag data.
// Will not enter the loop if end of stream was reached while decoding tag list.
while let Ok(command) = commands_stream.read_u8() {
read_single_command(
data_stream,
&mut commands_stream,
&mut decoded_profile_writer,
command,
)?;
}
// Validate output size.
let actual_len = decoded_profile_writer.position();
if actual_len != output_size {
warn!(output_size, actual_len, "ICC profile size mismatch");
return Err(Error::InvalidIccStream);
}
Ok(decoded_profile)
}
/// Struct to incrementally decode an ICC profile.
pub struct IncrementalIccReader {
histograms: Histograms,
reader: SymbolReader,
out_buf: Vec<u8>,
len: usize,
// [prev, prev_prev]
prev_bytes: [u8; 2],
}
impl IncrementalIccReader {
pub fn new(br: &mut BitReader) -> Result<Self> {
let len = u64::read_unconditional(&(), br, &Empty {})?;
if len > 1u64 << 20 {
return Err(Error::IccTooLarge);
}
let len = len as usize;
let histograms = Histograms::decode(ICC_CONTEXTS, br, true)?;
let reader = SymbolReader::new(&histograms, br, None)?;
Ok(Self {
histograms,
reader,
len,
out_buf: Vec::new_with_capacity(len)?,
prev_bytes: [0, 0],
})
}
fn get_icc_ctx(&self) -> u32 {
if self.out_buf.len() <= ICC_HEADER_SIZE as usize {
return 0;
}
let [b1, b2] = self.prev_bytes;
let p1 = match b1 {
b'a'..=b'z' | b'A'..=b'Z' => 0,
b'0'..=b'9' | b'.' | b',' => 1,
0..=1 => 2 + b1 as u32,
2..=15 => 4,
241..=254 => 5,
255 => 6,
_ => 7,
};
let p2 = match b2 {
b'a'..=b'z' | b'A'..=b'Z' => 0,
b'0'..=b'9' | b'.' | b',' => 1,
0..=15 => 2,
241..=255 => 3,
_ => 4,
};
1 + p1 + 8 * p2
}
pub fn num_coded_bytes(&self) -> usize {
self.len
}
pub fn remaining(&self) -> usize {
assert!(self.num_coded_bytes() >= self.out_buf.len());
self.num_coded_bytes() - self.out_buf.len()
}
pub fn read_one(&mut self, br: &mut BitReader) -> Result<()> {
let ctx = self.get_icc_ctx() as usize;
let checkpoint = self.reader.checkpoint::<1>();
let sym = self.reader.read_unsigned(&self.histograms, br, ctx);
if let Err(err) = br.check_for_error() {
self.reader.restore(checkpoint);
return Err(err);
}
if sym >= 256 {
warn!(sym, "Invalid symbol in ICC stream");
return Err(Error::InvalidIccStream);
}
let b = sym as u8;
self.out_buf.push(b);
self.prev_bytes = [b, self.prev_bytes[0]];
Ok(())
}
pub fn read_all(&mut self, br: &mut BitReader) -> Result<()> {
for _ in self.out_buf.len()..self.num_coded_bytes() {
self.read_one(br)?;
}
Ok(())
}
pub fn finalize(self, br: &mut BitReader) -> Result<Vec<u8>> {
assert_eq!(self.num_coded_bytes(), self.out_buf.len());
self.reader.check_final_state(&self.histograms, br)?;
let mut stream = IccStream::new(self.out_buf);
let profile = read_icc_inner(&mut stream)?;
stream.finalize()?;
Ok(profile)
}
}