Skip to main content

png/decoder/
zlib.rs

1use super::{stream::FormatErrorInner, unfiltering_buffer::UnfilteringBuffer, DecodingError};
2
3use fdeflate::Decompressor;
4
5/// [PNG spec](https://www.w3.org/TR/2003/REC-PNG-20031110/#10Compression) says that
6/// "deflate/inflate compression with a sliding window (which is an upper bound on the
7/// distances appearing in the deflate stream) of at most 32768 bytes".
8///
9/// `fdeflate` requires that we keep this many most recently decompressed bytes in the
10/// `out_buffer` - this allows referring back to them when handling "length and distance
11/// codes" in the deflate stream).
12const LOOKBACK_SIZE: usize = 32768;
13
14/// A buffer for decompression and in-place filtering of PNG rowlines.
15///
16/// The underlying data structure is a vector, with additional markers dividing
17/// the vector into specific regions of bytes - see [`UnfilterRegion`] for more
18/// details.
19pub struct UnfilterBuf<'data> {
20    /// The data container.
21    pub(crate) buffer: &'data mut [u8],
22    /// The past-the-end index of the region that is allowed to be modified.
23    pub(crate) available: &'data mut usize,
24    /// The past-the-end index of the region with decompressed bytes.
25    pub(crate) filled: &'data mut usize,
26}
27
28/// `UnfilterRegion` divides a `Vec<u8>` buffer into three consecutive regions:
29///
30/// * `vector[0..available]` - bytes that may be mutated (this typically means
31///   bytes that were decompressed earlier, but user of the buffer may also use
32///   this region for storing other data)
33/// * `vector[available..filled]` - already decompressed bytes that need to be
34///   preserved. (Future decompressor calls may reference and copy bytes from
35///   this region.  The maximum `filled - available` "look back" distance for
36///   [PNG compression method 0](https://www.w3.org/TR/png-3/#10CompressionCM0)
37///   is 32768 bytes)
38/// * `vector[filled..]` - buffer where future decompressor calls can write
39///   additional decompressed bytes
40///
41/// Even though only `vector[0..available]` bytes can be mutated, it is allowed
42/// to "shift" or "move" the contents of vector, as long as the:
43///
44/// * `vector[available..filled]` bytes are preserved
45/// * `available` and `filled` offsets are updated
46///
47/// Violating the invariants described above (e.g. mutating the bytes in the
48/// `vector[available..filled]` region) may result in absurdly wacky
49/// decompression output or panics, but not undefined behavior.
50#[derive(Default, Clone, Copy)]
51pub struct UnfilterRegion {
52    /// The past-the-end index of the region that is allowed to be modified.
53    pub available: usize,
54    /// The past-the-end index of the region with decompressed bytes.
55    pub filled: usize,
56}
57
58/// Ergonomics wrapper around `miniz_oxide::inflate::stream` for zlib compressed data.
59pub(super) struct ZlibStream {
60    /// Current decoding state.
61    state: Box<fdeflate::Decompressor>,
62    /// If there has been a call to decompress already.
63    started: bool,
64    /// Ignore and do not calculate the Adler-32 checksum. Defaults to `true`.
65    ///
66    /// This flag overrides `TINFL_FLAG_COMPUTE_ADLER32`.
67    ///
68    /// This flag should not be modified after decompression has started.
69    ignore_adler32: bool,
70}
71
72impl ZlibStream {
73    pub(crate) fn new() -> Self {
74        ZlibStream {
75            state: Box::new(Decompressor::new()),
76            started: false,
77            ignore_adler32: true,
78        }
79    }
80
81    pub(crate) fn reset(&mut self) {
82        self.started = false;
83        *self.state = Decompressor::new();
84    }
85
86    /// Set the `ignore_adler32` flag and return `true` if the flag was
87    /// successfully set.
88    ///
89    /// The default is `true`.
90    ///
91    /// This flag cannot be modified after decompression has started until the
92    /// [ZlibStream] is reset.
93    pub(crate) fn set_ignore_adler32(&mut self, flag: bool) -> bool {
94        if !self.started {
95            self.ignore_adler32 = flag;
96            true
97        } else {
98            false
99        }
100    }
101
102    /// Return the `ignore_adler32` flag.
103    pub(crate) fn ignore_adler32(&self) -> bool {
104        self.ignore_adler32
105    }
106
107    /// Fill the decoded buffer as far as possible from `data`.
108    /// On success returns the number of consumed input bytes.
109    pub(crate) fn decompress(
110        &mut self,
111        data: &[u8],
112        image_data: &mut UnfilterBuf<'_>,
113    ) -> Result<usize, DecodingError> {
114        // There may be more data past the adler32 checksum at the end of the deflate stream. We
115        // match libpng's default behavior and ignore any trailing data. In the future we may want
116        // to add a flag to control this behavior.
117        if self.state.is_done() {
118            return Ok(data.len());
119        }
120
121        if !self.started && self.ignore_adler32 {
122            self.state.ignore_adler32();
123        }
124
125        let in_consumed = image_data.decompress(&mut self.state, data)?;
126        self.started = true;
127
128        Ok(in_consumed)
129    }
130
131    /// Output any remaining buffered data within the decompressor.
132    ///
133    /// Returns `Ok(true)` if all data has been decompressed and there is no
134    /// more data that will be produced, or `Ok(false)` if there's potentially
135    /// more output.
136    ///
137    /// Returns `Err` if the zlib stream is corrupt or truncated too early.
138    pub(crate) fn finish(
139        &mut self,
140        image_data: &mut UnfilterBuf<'_>,
141    ) -> Result<bool, DecodingError> {
142        if !self.started || self.state.is_done() {
143            return Ok(true);
144        }
145
146        // If the zlib stream isn't done but we've already output all the pixel
147        // data needed, then either there's too much compressed data or the
148        // checksum is missing. Those aren't allowed by the spec, but libpng
149        // generally doesn't treat them as fatal.
150        if *image_data.filled == image_data.buffer.len() {
151            return Ok(true);
152        }
153
154        let (_, out_consumed) = self
155            .state
156            .read(
157                &[],
158                &mut image_data.buffer[*image_data.available..],
159                *image_data.filled - *image_data.available,
160                false,
161            )
162            .map_err(|err| {
163                DecodingError::Format(FormatErrorInner::CorruptFlateStream { err }.into())
164            })?;
165        *image_data.filled += out_consumed;
166
167        if self.state.is_done() {
168            *image_data.available = *image_data.filled;
169            return Ok(true);
170        }
171
172        // More output is only possible if zlib stream hasn't finished and the
173        // output buffer *is* full. (Empty space in the output buffer tells us
174        // there wasn't more data to write into it.)
175        if *image_data.filled == image_data.buffer.len() {
176            *image_data.available =
177                (*image_data.available).max(image_data.filled.saturating_sub(LOOKBACK_SIZE));
178            return Ok(false);
179        }
180
181        // The zlib stream was truncated before the end of the pixel data. This
182        // would ordinarily be caught within fdeflate if we'd passed
183        // end_of_input=true. But we intentionally don't pass that flag so that
184        // we're able to drain all available pixel data first.
185        Err(DecodingError::Format(
186            FormatErrorInner::CorruptFlateStream {
187                err: fdeflate::DecompressionError::InsufficientInput,
188            }
189            .into(),
190        ))
191    }
192}
193
194impl UnfilterRegion {
195    /// Use this region to decompress new filtered rowline data.
196    ///
197    /// Pass the wrapped buffer to
198    /// [`StreamingDecoder::update`][`super::stream::StreamingDecoder::update`] to fill it with
199    /// data and update the region indices.
200    ///
201    /// May panic if invariants of [`UnfilterRegion`] are violated.
202    pub fn as_buf<'data>(&'data mut self, buffer: &'data mut Vec<u8>) -> UnfilterBuf<'data> {
203        assert!(self.available <= self.filled);
204        assert!(self.filled <= buffer.len());
205        UnfilterBuf {
206            buffer,
207            filled: &mut self.filled,
208            available: &mut self.available,
209        }
210    }
211}
212
213impl UnfilterBuf<'_> {
214    /// Pushes `input` into `fdeflate` crate and appends decompressed bytes to `self.buffer`
215    /// (adjusting `self.filled` and `self.available` depending on how many bytes have been
216    /// decompressed).
217    ///
218    /// Returns how many bytes of `input` have been consumed.
219    #[inline]
220    fn decompress(
221        &mut self,
222        decompressor: &mut fdeflate::Decompressor,
223        input: &[u8],
224    ) -> Result<usize, DecodingError> {
225        let output_limit = (*self.filled + UnfilteringBuffer::GROWTH_BYTES).min(self.buffer.len());
226        let (in_consumed, out_consumed) = decompressor
227            .read(
228                input,
229                &mut self.buffer[*self.available..output_limit],
230                *self.filled - *self.available,
231                false,
232            )
233            .map_err(|err| {
234                DecodingError::Format(FormatErrorInner::CorruptFlateStream { err }.into())
235            })?;
236
237        *self.filled += out_consumed;
238        if decompressor.is_done() {
239            *self.available = *self.filled;
240        } else if let Some(new_available) = self.filled.checked_sub(LOOKBACK_SIZE) {
241            // The decompressed data may have started in the middle of the buffer,
242            // so ensure that `self.available` never goes backward.  This is needed
243            // to avoid miscommunicating the size of the "look-back" window when calling
244            // `fdeflate::Decompressor::read` a bit earlier and passing
245            // `&mut self.buffer[*self.available..output_limit]`.
246            if new_available > *self.available {
247                *self.available = new_available;
248            }
249        }
250
251        Ok(in_consumed)
252    }
253}