mod stream;
mod zlib;
use self::stream::{get_info, CHUNCK_BUFFER_SIZE};
pub use self::stream::{Decoded, DecodingError, StreamingDecoder};
use std::borrow;
use std::io::{BufRead, BufReader, Read, Write};
use std::mem;
use std::ops::Range;
use crate::chunk;
use crate::common::{BitDepth, BytesPerPixel, ColorType, Info, Transformations};
use crate::filter::{unfilter, FilterType};
use crate::utils;
#[derive(Debug, PartialEq, Eq)]
pub struct OutputInfo {
pub width: u32,
pub height: u32,
pub color_type: ColorType,
pub bit_depth: BitDepth,
pub line_size: usize,
}
impl OutputInfo {
pub fn buffer_size(&self) -> usize {
self.line_size * self.height as usize
}
}
#[derive(Clone, Copy, Debug)]
pub struct Limits {
pub bytes: usize,
}
impl Default for Limits {
fn default() -> Limits {
Limits {
bytes: 1024 * 1024 * 64,
}
}
}
pub struct Decoder<R: Read> {
r: R,
transform: Transformations,
limits: Limits,
}
struct InterlacedRow<'data> {
data: &'data [u8],
interlace: InterlaceInfo,
}
enum InterlaceInfo {
None,
Adam7 { pass: u8, line: u32, width: u32 },
}
impl<R: Read> Decoder<R> {
pub fn new(r: R) -> Decoder<R> {
Decoder::new_with_limits(r, Limits::default())
}
pub fn new_with_limits(r: R, limits: Limits) -> Decoder<R> {
Decoder {
r,
transform: crate::Transformations::EXPAND
| crate::Transformations::SCALE_16
| crate::Transformations::STRIP_16,
limits,
}
}
pub fn set_limits(&mut self, limits: Limits) {
self.limits = limits;
}
pub fn read_info(self) -> Result<(OutputInfo, Reader<R>), DecodingError> {
let mut r = Reader::new(self.r, StreamingDecoder::new(), self.transform, self.limits);
r.init()?;
let color_type = r.info().color_type;
let bit_depth = r.info().bit_depth;
if color_type.is_combination_invalid(bit_depth) {
return Err(DecodingError::Format(
format!(
"Invalid color/depth combination in header: {:?}/{:?}",
color_type, bit_depth
)
.into(),
));
}
if r.checked_output_buffer_size().is_none() {
return Err(DecodingError::LimitsExceeded);
}
let (ct, bits) = r.output_color_type();
let info = {
let info = r.info();
OutputInfo {
width: info.width,
height: info.height,
color_type: ct,
bit_depth: bits,
line_size: r.output_line_size(info.width),
}
};
Ok((info, r))
}
pub fn set_transformations(&mut self, transform: Transformations) {
self.transform = transform;
}
}
struct ReadDecoder<R: Read> {
reader: BufReader<R>,
decoder: StreamingDecoder,
at_eof: bool,
}
impl<R: Read> ReadDecoder<R> {
fn decode_next(&mut self, image_data: &mut Vec<u8>) -> Result<Option<Decoded>, DecodingError> {
while !self.at_eof {
let (consumed, result) = {
let buf = self.reader.fill_buf()?;
if buf.is_empty() {
return Err(DecodingError::Format("unexpected EOF".into()));
}
self.decoder.update(buf, image_data)?
};
self.reader.consume(consumed);
match result {
Decoded::Nothing => (),
Decoded::ImageEnd => self.at_eof = true,
result => return Ok(Some(result)),
}
}
Ok(None)
}
fn finished_decoding(&mut self) -> Result<(), DecodingError> {
while !self.at_eof {
let buf = self.reader.fill_buf()?;
if buf.is_empty() {
return Err(DecodingError::Format("unexpected EOF after image".into()));
}
let (consumed, event) = self.decoder.update(buf, &mut vec![])?;
self.reader.consume(consumed);
match event {
Decoded::Nothing => (),
Decoded::ImageEnd => self.at_eof = true,
Decoded::ChunkComplete(_, _) | Decoded::ChunkBegin(_, _) | Decoded::ImageData => {}
Decoded::ImageDataFlushed => return Ok(()),
Decoded::PartialChunk(_) => {}
new => unreachable!("{:?}", new),
}
}
Err(DecodingError::Format("unexpected EOF after image".into()))
}
fn info(&self) -> Option<&Info> {
get_info(&self.decoder)
}
}
pub struct Reader<R: Read> {
decoder: ReadDecoder<R>,
bpp: BytesPerPixel,
subframe: SubframeInfo,
fctl_read: u32,
next_frame: SubframeIdx,
prev: Vec<u8>,
current: Vec<u8>,
scan_start: usize,
transform: Transformations,
processed: Vec<u8>,
limits: Limits,
}
struct SubframeInfo {
width: u32,
rowlen: usize,
interlace: InterlaceIter,
consumed_and_flushed: bool,
}
#[derive(Clone)]
enum InterlaceIter {
None(Range<u32>),
Adam7(utils::Adam7Iterator),
}
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
enum SubframeIdx {
Uninit,
Initial,
Some(u32),
End,
}
macro_rules! get_info(
($this:expr) => {
$this.decoder.info().unwrap()
}
);
impl<R: Read> Reader<R> {
fn new(r: R, d: StreamingDecoder, t: Transformations, limits: Limits) -> Reader<R> {
Reader {
decoder: ReadDecoder {
reader: BufReader::with_capacity(CHUNCK_BUFFER_SIZE, r),
decoder: d,
at_eof: false,
},
bpp: BytesPerPixel::One,
subframe: SubframeInfo::not_yet_init(),
fctl_read: 0,
next_frame: SubframeIdx::Initial,
prev: Vec::new(),
current: Vec::new(),
scan_start: 0,
transform: t,
processed: Vec::new(),
limits,
}
}
fn init(&mut self) -> Result<(), DecodingError> {
if self.next_frame == self.subframe_idx() {
return Ok(());
} else if self.next_frame == SubframeIdx::End {
return Err(DecodingError::Other("End of image has been reached".into()));
}
loop {
match self.decoder.decode_next(&mut Vec::new())? {
Some(Decoded::ChunkBegin(_, chunk::IDAT))
| Some(Decoded::ChunkBegin(_, chunk::fdAT)) => break,
Some(Decoded::FrameControl(_)) => {
self.subframe = SubframeInfo::new(self.info());
self.next_frame = SubframeIdx::Some(self.fctl_read);
self.fctl_read += 1;
}
None => return Err(DecodingError::Format("IDAT chunk missing".into())),
Some(Decoded::Header { .. }) => {
self.validate_buffer_sizes()?;
}
_ => {}
}
}
{
let info = match self.decoder.info() {
Some(info) => info,
None => return Err(DecodingError::Format("IHDR chunk missing".into())),
};
self.bpp = info.bpp_in_prediction();
self.subframe = SubframeInfo::new(info);
}
self.allocate_out_buf()?;
self.prev = vec![0; self.subframe.rowlen];
Ok(())
}
fn reset_current(&mut self) {
self.current.clear();
self.scan_start = 0;
}
pub fn info(&self) -> &Info {
get_info!(self)
}
fn subframe_idx(&self) -> SubframeIdx {
let info = match self.decoder.info() {
None => return SubframeIdx::Uninit,
Some(info) => info,
};
match info.frame_control() {
None => SubframeIdx::Initial,
Some(_) => SubframeIdx::Some(self.fctl_read - 1),
}
}
fn finished_frame(&mut self) {
let info = self.info();
let past_end_subframe = match info.animation_control() {
None => 0,
Some(ac) => ac.num_frames,
};
self.next_frame = match self.next_frame {
SubframeIdx::Uninit => unreachable!("Next frame can never be initial"),
SubframeIdx::End => unreachable!("Next frame called when already at image end"),
SubframeIdx::Initial if past_end_subframe == 0 => SubframeIdx::End,
SubframeIdx::Initial => SubframeIdx::Some(0),
SubframeIdx::Some(idx) if past_end_subframe <= idx + 1 => SubframeIdx::End,
SubframeIdx::Some(idx) => SubframeIdx::Some(idx + 1),
}
}
pub fn next_frame(&mut self, buf: &mut [u8]) -> Result<(), DecodingError> {
self.init()?;
let (color_type, bit_depth) = self.output_color_type();
if buf.len() < self.output_buffer_size() {
return Err(DecodingError::Other(
"supplied buffer is too small to hold the image".into(),
));
}
self.reset_current();
let width = self.info().width;
if get_info!(self).interlaced {
while let Some((row, adam7)) = self.next_interlaced_row()? {
let (pass, line, _) = adam7.unwrap();
let samples = color_type.samples() as u8;
utils::expand_pass(buf, width, row, pass, line, samples * (bit_depth as u8));
}
} else {
let mut len = 0;
while let Some(row) = self.next_row()? {
len += (&mut buf[len..]).write(row)?;
}
}
if !self.subframe.consumed_and_flushed {
self.decoder.finished_decoding()?;
}
self.finished_frame();
Ok(())
}
pub fn next_row(&mut self) -> Result<Option<&[u8]>, DecodingError> {
self.next_interlaced_row().map(|v| v.map(|v| v.0))
}
pub fn next_interlaced_row(
&mut self,
) -> Result<Option<(&[u8], Option<(u8, u32, u32)>)>, DecodingError> {
match self.next_interlaced_row_impl() {
Err(err) => Err(err),
Ok(None) => Ok(None),
Ok(Some(row)) => {
let interlace = match row.interlace {
InterlaceInfo::None => None,
InterlaceInfo::Adam7 { pass, line, width } => Some((pass, line, width)),
};
Ok(Some((row.data, interlace)))
}
}
}
fn next_interlaced_row_impl(&mut self) -> Result<Option<InterlacedRow<'_>>, DecodingError> {
use crate::common::ColorType::*;
let transform = self.transform;
if transform == crate::Transformations::IDENTITY {
return self.next_raw_interlaced_row();
}
let mut buffer = mem::replace(&mut self.processed, Vec::new());
let (got_next, adam7) = if let Some(row) = self.next_raw_interlaced_row()? {
(&mut buffer[..]).write_all(row.data)?;
(true, row.interlace)
} else {
(false, InterlaceInfo::None)
};
let _ = mem::replace(&mut self.processed, buffer);
if !got_next {
return Ok(None);
}
let (color_type, bit_depth, trns) = {
let info = get_info!(self);
(info.color_type, info.bit_depth as u8, info.trns.is_some())
};
let output_buffer = if let InterlaceInfo::Adam7 { width, .. } = adam7 {
let width = self
.line_size(width)
.expect("Adam7 interlaced rows are shorter than the buffer.");
&mut self.processed[..width]
} else {
&mut *self.processed
};
let mut len = output_buffer.len();
if transform.contains(crate::Transformations::EXPAND) {
match color_type {
Indexed => expand_paletted(output_buffer, get_info!(self))?,
Grayscale | GrayscaleAlpha if bit_depth < 8 => {
expand_gray_u8(output_buffer, get_info!(self))
}
Grayscale | RGB if trns => {
let channels = color_type.samples();
let trns = get_info!(self).trns.as_ref().unwrap();
if bit_depth == 8 {
utils::expand_trns_line(output_buffer, &*trns, channels);
} else {
utils::expand_trns_line16(output_buffer, &*trns, channels);
}
}
_ => (),
}
}
if bit_depth == 16
&& transform
.intersects(crate::Transformations::SCALE_16 | crate::Transformations::STRIP_16)
{
len /= 2;
for i in 0..len {
output_buffer[i] = output_buffer[2 * i];
}
}
Ok(Some(InterlacedRow {
data: &output_buffer[..len],
interlace: adam7,
}))
}
pub fn output_color_type(&mut self) -> (ColorType, BitDepth) {
self.imm_output_color_type()
}
pub(crate) fn imm_output_color_type(&self) -> (ColorType, BitDepth) {
use crate::common::ColorType::*;
let t = self.transform;
let info = get_info!(self);
if t == crate::Transformations::IDENTITY {
(info.color_type, info.bit_depth)
} else {
let bits = match info.bit_depth as u8 {
16 if t.intersects(
crate::Transformations::SCALE_16 | crate::Transformations::STRIP_16,
) =>
{
8
}
n if n < 8 && t.contains(crate::Transformations::EXPAND) => 8,
n => n,
};
let color_type = if t.contains(crate::Transformations::EXPAND) {
let has_trns = info.trns.is_some();
match info.color_type {
Grayscale if has_trns => GrayscaleAlpha,
RGB if has_trns => RGBA,
Indexed if has_trns => RGBA,
Indexed => RGB,
ct => ct,
}
} else {
info.color_type
};
(color_type, BitDepth::from_u8(bits).unwrap())
}
}
pub fn output_buffer_size(&self) -> usize {
let (width, height) = get_info!(self).size();
let size = self.output_line_size(width);
size * height as usize
}
fn validate_buffer_sizes(&self) -> Result<(), DecodingError> {
if self.info().checked_raw_row_length().is_none() {
return Err(DecodingError::LimitsExceeded);
}
if self.checked_output_buffer_size().is_none() {
return Err(DecodingError::LimitsExceeded);
}
Ok(())
}
fn checked_output_buffer_size(&self) -> Option<usize> {
let (width, height) = get_info!(self).size();
let (color, depth) = self.imm_output_color_type();
let rowlen = color.checked_raw_row_length(depth, width)? - 1;
let height: usize = std::convert::TryFrom::try_from(height).ok()?;
rowlen.checked_mul(height)
}
pub fn output_line_size(&self, width: u32) -> usize {
let (color, depth) = self.imm_output_color_type();
color.raw_row_length_from_width(depth, width) - 1
}
fn line_size(&self, width: u32) -> Option<usize> {
use crate::common::ColorType::*;
let t = self.transform;
let info = get_info!(self);
let trns = info.trns.is_some();
let expanded = if info.bit_depth == BitDepth::Sixteen {
BitDepth::Sixteen
} else {
BitDepth::Eight
};
let (color, depth) = match info.color_type {
Indexed if trns && t.contains(Transformations::EXPAND) => (RGBA, expanded),
Indexed if t.contains(Transformations::EXPAND) => (RGB, expanded),
RGB if trns && t.contains(Transformations::EXPAND) => (RGBA, expanded),
Grayscale if trns && t.contains(Transformations::EXPAND) => (GrayscaleAlpha, expanded),
Grayscale if t.contains(Transformations::EXPAND) => (Grayscale, expanded),
GrayscaleAlpha if t.contains(Transformations::EXPAND) => (GrayscaleAlpha, expanded),
other => (other, info.bit_depth),
};
color.checked_raw_row_length(depth, width).map(|n| n - 1)
}
fn allocate_out_buf(&mut self) -> Result<(), DecodingError> {
let width = self.subframe.width;
let bytes = self.limits.bytes;
let buflen = match self.line_size(width) {
Some(buflen) if buflen <= bytes => buflen,
_ => return Err(DecodingError::LimitsExceeded),
};
self.processed.resize(buflen, 0u8);
Ok(())
}
fn next_pass(&mut self) -> Option<(usize, InterlaceInfo)> {
match self.subframe.interlace {
InterlaceIter::Adam7(ref mut adam7) => {
let last_pass = adam7.current_pass();
let (pass, line, width) = adam7.next()?;
let rowlen = get_info!(self).raw_row_length_from_width(width);
if last_pass != pass {
self.prev.clear();
self.prev.resize(rowlen, 0u8);
}
Some((rowlen, InterlaceInfo::Adam7 { pass, line, width }))
}
InterlaceIter::None(ref mut height) => {
let _ = height.next()?;
Some((self.subframe.rowlen, InterlaceInfo::None))
}
}
}
fn next_raw_interlaced_row(&mut self) -> Result<Option<InterlacedRow<'_>>, DecodingError> {
let bpp = self.bpp;
let (rowlen, passdata) = match self.next_pass() {
Some((rowlen, passdata)) => (rowlen, passdata),
None => return Ok(None),
};
loop {
if self.current.len() - self.scan_start >= rowlen {
let row = &mut self.current[self.scan_start..];
let filter = match FilterType::from_u8(row[0]) {
None => {
self.scan_start += rowlen;
return Err(DecodingError::Format(
format!("invalid filter method ({})", row[0]).into(),
));
}
Some(filter) => filter,
};
if let Err(message) =
unfilter(filter, bpp, &self.prev[1..rowlen], &mut row[1..rowlen])
{
return Err(DecodingError::Format(borrow::Cow::Borrowed(message)));
}
self.prev[..rowlen].copy_from_slice(&row[..rowlen]);
self.scan_start += rowlen;
return Ok(Some(InterlacedRow {
data: &self.prev[1..rowlen],
interlace: passdata,
}));
} else {
if self.subframe.consumed_and_flushed {
return Err(DecodingError::Format(
format!("not enough data for image").into(),
));
}
if self.scan_start > 0 {
self.current.drain(..self.scan_start).for_each(drop);
self.scan_start = 0;
}
let val = self.decoder.decode_next(&mut self.current)?;
match val {
Some(Decoded::ImageData) => {}
Some(Decoded::ImageDataFlushed) => {
self.subframe.consumed_and_flushed = true;
}
None => {
if !self.current.is_empty() {
return Err(DecodingError::Format("file truncated".into()));
} else {
return Ok(None);
}
}
_ => (),
}
}
}
}
}
impl SubframeInfo {
fn not_yet_init() -> Self {
SubframeInfo {
width: 0,
rowlen: 0,
interlace: InterlaceIter::None(0..0),
consumed_and_flushed: false,
}
}
fn new(info: &Info) -> Self {
let (width, height) = if let Some(fc) = info.frame_control {
(fc.width, fc.height)
} else {
(info.width, info.height)
};
let interlace = if info.interlaced {
InterlaceIter::Adam7(utils::Adam7Iterator::new(width, height))
} else {
InterlaceIter::None(0..height)
};
SubframeInfo {
width,
rowlen: info.raw_row_length_from_width(width),
interlace,
consumed_and_flushed: false,
}
}
}
fn expand_paletted(buffer: &mut [u8], info: &Info) -> Result<(), DecodingError> {
if let Some(palette) = info.palette.as_ref() {
if let BitDepth::Sixteen = info.bit_depth {
Err(DecodingError::Format(
"Bit depth '16' is not valid for paletted images".into(),
))
} else {
let black = [0, 0, 0];
if let Some(ref trns) = info.trns {
utils::unpack_bits(buffer, 4, info.bit_depth as u8, |i, chunk| {
let (rgb, a) = (
palette
.get(3 * i as usize..3 * i as usize + 3)
.unwrap_or(&black),
*trns.get(i as usize).unwrap_or(&0xFF),
);
chunk[0] = rgb[0];
chunk[1] = rgb[1];
chunk[2] = rgb[2];
chunk[3] = a;
});
} else {
utils::unpack_bits(buffer, 3, info.bit_depth as u8, |i, chunk| {
let rgb = palette
.get(3 * i as usize..3 * i as usize + 3)
.unwrap_or(&black);
chunk[0] = rgb[0];
chunk[1] = rgb[1];
chunk[2] = rgb[2];
})
}
Ok(())
}
} else {
Err(DecodingError::Format("missing palette".into()))
}
}
fn expand_gray_u8(buffer: &mut [u8], info: &Info) {
let rescale = true;
let scaling_factor = if rescale {
(255) / ((1u16 << info.bit_depth as u8) - 1) as u8
} else {
1
};
if let Some(ref trns) = info.trns {
utils::unpack_bits(buffer, 2, info.bit_depth as u8, |pixel, chunk| {
if pixel == trns[0] {
chunk[1] = 0
} else {
chunk[1] = 0xFF
}
chunk[0] = pixel * scaling_factor
})
} else {
utils::unpack_bits(buffer, 1, info.bit_depth as u8, |val, chunk| {
chunk[0] = val * scaling_factor
})
}
}
#[cfg(test)]
mod tests {
use super::Decoder;
use std::io::{BufRead, Read, Result};
use std::mem::discriminant;
struct SmalBuf<R: BufRead> {
inner: R,
cap: usize,
}
impl<R: BufRead> SmalBuf<R> {
fn new(inner: R, cap: usize) -> Self {
SmalBuf { inner, cap }
}
}
impl<R: BufRead> Read for SmalBuf<R> {
fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
let len = buf.len().min(self.cap);
self.inner.read(&mut buf[..len])
}
}
impl<R: BufRead> BufRead for SmalBuf<R> {
fn fill_buf(&mut self) -> Result<&[u8]> {
let buf = self.inner.fill_buf()?;
let len = buf.len().min(self.cap);
Ok(&buf[..len])
}
fn consume(&mut self, amt: usize) {
assert!(amt <= self.cap);
self.inner.consume(amt)
}
}
#[test]
fn no_data_dup_on_finish() {
const IMG: &[u8] = include_bytes!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/bugfixes/x_issue#214.png"
));
let (info, mut normal) = Decoder::new(IMG).read_info().unwrap();
let mut buffer = vec![0; info.buffer_size()];
let normal = normal.next_frame(&mut buffer).unwrap_err();
let smal = Decoder::new(SmalBuf::new(IMG, 1))
.read_info()
.unwrap()
.1
.next_frame(&mut buffer)
.unwrap_err();
assert_eq!(discriminant(&normal), discriminant(&smal));
}
}