Skip to content

Avoid returning to ReadChunkData state #555

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kornelski wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Avoid returning to ReadChunkData state #555

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5a57c1d
Allocate whole chunk before reading
kornelski Dec 21, 2024
0184117
Take raw_bytes
kornelski Dec 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 64 additions & 69 deletions src/decoder/stream.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::convert::TryInto;
use std::error;
use std::fmt;
use std::io;
use std::mem;
use std::{borrow::Cow, cmp::min};

use crc32fast::Hasher as Crc32;
Expand Down Expand Up @@ -63,8 +64,8 @@ enum State {
/// In this state we are reading chunk data from external input, and appending it to
/// `ChunkState::raw_bytes`.
ReadChunkData(ChunkType),
/// In this state we check if all chunk data has been already read into `ChunkState::raw_bytes`
/// and if so then we parse the chunk. Otherwise, we go back to the `ReadChunkData` state.
/// In this state all chunk data has been already read into `ChunkState::raw_bytes`
/// and we parse the chunk.
ParseChunkData(ChunkType),
/// In this state we are reading image data from external input and feeding it directly into
/// `StreamingDecoder::inflater`.
Expand Down Expand Up @@ -726,53 +727,40 @@ impl StreamingDecoder {
}
ParseChunkData(type_str) => {
debug_assert!(type_str != IDAT && type_str != chunk::fdAT);
if self.current_chunk.remaining == 0 {
// Got complete chunk.
Ok((0, self.parse_chunk(type_str)?))
} else {
// Make sure we have room to read more of the chunk.
// We need it fully before parsing.
self.reserve_current_chunk()?;

self.state = Some(ReadChunkData(type_str));
Ok((0, Decoded::PartialChunk(type_str)))
}
// Got complete chunk.
debug_assert_eq!(0, self.current_chunk.remaining);
Ok((0, self.parse_chunk(type_str)?))
}
ReadChunkData(type_str) => {
debug_assert!(type_str != IDAT && type_str != chunk::fdAT);
if self.current_chunk.remaining == 0 {
self.state = Some(State::new_u32(U32ValueKind::Crc(type_str)));
Ok((0, Decoded::Nothing))
let ChunkState {
crc,
remaining,
raw_bytes,
type_: _,
} = &mut self.current_chunk;

let n = min(*remaining as usize, buf.len());
let buf = &buf[..n];

// the buffer has been reserved ahead of time
if raw_bytes.capacity() - raw_bytes.len() >= buf.len() {
raw_bytes.extend_from_slice(buf);
} else {
let ChunkState {
crc,
remaining,
raw_bytes,
type_: _,
} = &mut self.current_chunk;

let buf_avail = raw_bytes.capacity() - raw_bytes.len();
let bytes_avail = min(buf.len(), buf_avail);
let n = min(*remaining, bytes_avail as u32);
if buf_avail == 0 {
self.state = Some(ParseChunkData(type_str));
Ok((0, Decoded::Nothing))
} else {
let buf = &buf[..n as usize];
if !self.decode_options.ignore_crc {
crc.update(buf);
}
raw_bytes.extend_from_slice(buf);

*remaining -= n;
if *remaining == 0 {
self.state = Some(ParseChunkData(type_str));
} else {
self.state = Some(ReadChunkData(type_str));
}
Ok((n as usize, Decoded::Nothing))
}
debug_assert!(false); // this won't happen, but return optimizes better
return Err(DecodingError::LimitsExceeded);
}
if !self.decode_options.ignore_crc {
crc.update(buf);
}

*remaining -= n as u32;
self.state = Some(if *remaining == 0 {
ParseChunkData(type_str)
} else {
ReadChunkData(type_str)
});
Ok((n, Decoded::Nothing))
}
ImageData(type_str) => {
debug_assert!(type_str == IDAT || type_str == chunk::fdAT);
Expand Down Expand Up @@ -848,6 +836,15 @@ impl StreamingDecoder {
});
return Ok(Decoded::ImageDataFlushed);
}

self.current_chunk.type_ = type_str;
if !self.decode_options.ignore_crc {
self.current_chunk.crc.reset();
self.current_chunk.crc.update(&type_str.0);
}
self.current_chunk.remaining = length;
self.current_chunk.raw_bytes.clear();

self.state = match type_str {
chunk::fdAT => {
if !self.ready_for_fdat_chunks {
Expand Down Expand Up @@ -877,15 +874,12 @@ impl StreamingDecoder {
self.have_idat = true;
Some(State::ImageData(type_str))
}
_ => Some(State::ReadChunkData(type_str)),
_ => {
// IDAT/fDAT won't use current_chunk.raw_bytes, so it won't need it grown
self.reserve_current_chunk()?;
Some(State::ReadChunkData(type_str))
}
};
self.current_chunk.type_ = type_str;
if !self.decode_options.ignore_crc {
self.current_chunk.crc.reset();
self.current_chunk.crc.update(&type_str.0);
}
self.current_chunk.remaining = length;
self.current_chunk.raw_bytes.clear();
Ok(Decoded::ChunkBegin(length, type_str))
}
U32ValueKind::Crc(type_str) => {
Expand Down Expand Up @@ -958,19 +952,24 @@ impl StreamingDecoder {
}

fn reserve_current_chunk(&mut self) -> Result<(), DecodingError> {
let max = self.limits.bytes;
let buffer = &mut self.current_chunk.raw_bytes;

// Double if necessary, but no more than until the limit is reached.
let reserve_size = max.saturating_sub(buffer.capacity()).min(buffer.len());
self.limits.reserve_bytes(reserve_size)?;
buffer.reserve_exact(reserve_size);
// reserve before any data is read
debug_assert_eq!(0, buffer.len());

if buffer.capacity() == buffer.len() {
Err(DecodingError::LimitsExceeded)
} else {
Ok(())
let required_capacity = self.current_chunk.remaining as usize;
let actual_capacity = buffer.capacity();

if required_capacity <= actual_capacity {
return Ok(());
}

// try_reserve takes the new size, but limits take only size increase
self.limits
.reserve_bytes(required_capacity - actual_capacity)?;
buffer
.try_reserve_exact(required_capacity)
.map_err(|_| DecodingError::LimitsExceeded)
}

fn parse_chunk(&mut self, type_str: ChunkType) -> Result<Decoded, DecodingError> {
Expand Down Expand Up @@ -1104,9 +1103,7 @@ impl StreamingDecoder {
FormatErrorInner::DuplicateChunk { kind: chunk::PLTE }.into(),
))
} else {
self.limits
.reserve_bytes(self.current_chunk.raw_bytes.len())?;
info.palette = Some(Cow::Owned(self.current_chunk.raw_bytes.clone()));
info.palette = Some(Cow::Owned(mem::take(&mut self.current_chunk.raw_bytes)));
Ok(Decoded::Nothing)
}
}
Expand Down Expand Up @@ -1191,9 +1188,7 @@ impl StreamingDecoder {
));
}
let (color_type, bit_depth) = { (info.color_type, info.bit_depth as u8) };
self.limits
.reserve_bytes(self.current_chunk.raw_bytes.len())?;
let mut vec = self.current_chunk.raw_bytes.clone();
let mut vec = mem::take(&mut self.current_chunk.raw_bytes);
let len = vec.len();
match color_type {
ColorType::Grayscale => {
Expand Down Expand Up @@ -1763,7 +1758,7 @@ impl StreamingDecoder {
ColorType::Grayscale | ColorType::GrayscaleAlpha => 2,
ColorType::Rgb | ColorType::Rgba => 6,
};
let vec = self.current_chunk.raw_bytes.clone();
let vec = mem::take(&mut self.current_chunk.raw_bytes);
let len = vec.len();
if len == expected {
info.bkgd = Some(Cow::Owned(vec));
Expand Down Expand Up @@ -1810,7 +1805,7 @@ impl Default for ChunkState {
type_: ChunkType([0; 4]),
crc: Crc32::new(),
remaining: 0,
raw_bytes: Vec::with_capacity(CHUNK_BUFFER_SIZE),
raw_bytes: Vec::new(),
}
}
}
Expand Down
Loading