png/decoder/
stream.rs

Help
1use std::borrow::Cow;
2use std::default::Default;
3use std::error;
4use std::fmt;
5use std::io;
6use std::cmp::min;
7use std::convert::From;
8
9extern crate inflate;
10
11use self::inflate::InflateStream;
12use crc::Crc32;
13use traits::ReadBytesExt;
14use common::{ColorType, BitDepth, Info, Unit, PixelDimensions, AnimationControl, FrameControl};
15use chunk::{self, ChunkType, IHDR, IDAT, IEND};
16
17/// TODO check if these size are reasonable
18pub const CHUNCK_BUFFER_SIZE: usize = 32*1024;
19
20#[derive(Debug)]
21enum U32Value {
22    // CHUNKS
23    Length,
24    Type(u32),
25    Crc(ChunkType)
26}
27
28#[derive(Debug)]
29enum State {
30    Signature(u8, [u8; 7]),
31    U32Byte3(U32Value, u32),
32    U32Byte2(U32Value, u32),
33    U32Byte1(U32Value, u32),
34    U32(U32Value),
35    ReadChunk(ChunkType, bool),
36    PartialChunk(ChunkType),
37    DecodeData(ChunkType, usize),
38}
39
40#[derive(Debug)]
41/// Result of the decoding process
42pub enum Decoded {
43    /// Nothing decoded yet
44    Nothing,
45    Header(u32, u32, BitDepth, ColorType, bool),
46    ChunkBegin(u32, ChunkType),
47    ChunkComplete(u32, ChunkType),
48    PixelDimensions(PixelDimensions),
49    AnimationControl(AnimationControl),
50    FrameControl(FrameControl),
51    /// Decoded raw image data.
52    ImageData,
53    PartialChunk(ChunkType),
54    ImageEnd,
55}
56
57#[derive(Debug)]
58pub enum DecodingError {
59    IoError(io::Error),
60    Format(Cow<'static, str>),
61    InvalidSignature,
62    CrcMismatch {
63        /// bytes to skip to try to recover from this error
64        recover: usize,
65        /// Stored CRC32 value
66        crc_val: u32,
67        /// Calculated CRC32 sum
68        crc_sum: u32,
69        chunk: ChunkType
70    },
71    Other(Cow<'static, str>),
72    CorruptFlateStream
73}
74
75impl error::Error for DecodingError {
76    fn description(&self) -> &str {
77        use self::DecodingError::*;
78        match *self {
79            IoError(ref err) => err.description(),
80            Format(ref desc) | Other(ref desc) => &desc,
81            InvalidSignature => "invalid signature",
82            CrcMismatch { .. } => "CRC error",
83            CorruptFlateStream => "compressed data stream corrupted"
84        }
85    }
86}
87
88impl fmt::Display for DecodingError {
89    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
90        write!(fmt, "{}", (self as &error::Error).description())
91    }
92}
93
94impl From<io::Error> for DecodingError {
95    fn from(err: io::Error) -> DecodingError {
96        DecodingError::IoError(err)
97    }
98}
99
100impl From<String> for DecodingError {
101    fn from(err: String) -> DecodingError {
102        DecodingError::Other(err.into())
103    }
104}
105
106impl From<DecodingError> for io::Error {
107    fn from(err: DecodingError) -> io::Error {
108        use std::error::Error;
109        match err {
110            DecodingError::IoError(err) => err,
111            err => io::Error::new(
112                io::ErrorKind::Other,
113                err.description()
114            )
115        }
116    }
117}
118
119/// PNG StreamingDecoder (low-level interface)
120pub struct StreamingDecoder {
121    state: Option<State>,
122    current_chunk: (Crc32, u32, Vec<u8>),
123    inflater: InflateStream,
124    info: Option<Info>,
125    current_seq_no: Option<u32>,
126    have_idat: bool,
127}
128
129impl StreamingDecoder {
130    /// Creates a new StreamingDecoder
131    ///
132    /// Allocates the internal buffers.
133    pub fn new() -> StreamingDecoder {
134        StreamingDecoder {
135            state: Some(State::Signature(0, [0; 7])),
136            current_chunk: (Crc32::new(), 0, Vec::with_capacity(CHUNCK_BUFFER_SIZE)),
137            inflater: if cfg!(fuzzing) {InflateStream::from_zlib_no_checksum()} else {InflateStream::from_zlib()},
138            info: None,
139            current_seq_no: None,
140            have_idat: false
141        }
142    }
143
144    /// Resets the StreamingDecoder
145    pub fn reset(&mut self) {
146        self.state = Some(State::Signature(0, [0; 7]));
147        self.current_chunk.0 = Crc32::new();
148        self.current_chunk.1 = 0;
149        self.current_chunk.2.clear();
150        self.inflater = if cfg!(fuzzing) {InflateStream::from_zlib_no_checksum()} else {InflateStream::from_zlib()};
151        self.info = None;
152        self.current_seq_no = None;
153        self.have_idat = false;
154    }
155
156    /// Low level StreamingDecoder interface.
157    ///
158    /// Allows to stream partial data to the encoder. Returns a tuple containing the bytes that have
159    /// been consumed from the input buffer and the current decoding result. If the decoded chunk
160    /// was an image data chunk, it also appends the read data to `image_data`.
161    pub fn update(&mut self, mut buf: &[u8], image_data: &mut Vec<u8>)
162    -> Result<(usize, Decoded), DecodingError> {
163        let len = buf.len();
164        while buf.len() > 0 && self.state.is_some() {
165            match self.next_state(buf, image_data) {
166                Ok((bytes, Decoded::Nothing)) => {
167                    buf = &buf[bytes..]
168                }
169                Ok((bytes, result)) => {
170                    buf = &buf[bytes..];
171                    return Ok((len-buf.len(), result));
172                }
173                Err(err) => return Err(err)
174            }
175        }
176        Ok((len-buf.len(), Decoded::Nothing))
177    }
178
179    fn next_state<'a>(&'a mut self, buf: &[u8], image_data: &mut Vec<u8>)
180    -> Result<(usize, Decoded), DecodingError> {
181        use self::State::*;
182
183        macro_rules! goto (
184            ($n:expr, $state:expr) => ({
185                self.state = Some($state);
186                Ok(($n, Decoded::Nothing))
187            });
188            ($state:expr) => ({
189                self.state = Some($state);
190                Ok((1, Decoded::Nothing))
191            });
192            ($n:expr, $state:expr, emit $res:expr) => ({
193                self.state = Some($state);
194                Ok(($n, $res))
195            });
196            ($state:expr, emit $res:expr) => ({
197                self.state = Some($state);
198                Ok((1, $res))
199            })
200        );
201
202        let current_byte = buf[0];
203
204        // Driver should ensure that state is never None
205        let state = self.state.take().unwrap();
206        //println!("state: {:?}", state);
207
208        match state {
209            Signature(i, mut signature) => if i < 7 {
210                signature[i as usize] = current_byte;
211                goto!(Signature(i+1, signature))
212            } else {
213                if signature == [137, 80, 78, 71, 13, 10, 26] && current_byte == 10 {
214                    goto!(U32(U32Value::Length))
215                } else {
216                    Err(DecodingError::InvalidSignature)
217                }
218            },
219            U32Byte3(type_, mut val) => {
220                use self::U32Value::*;
221                val |= current_byte as u32;
222                match type_ {
223                    Length => goto!(U32(Type(val))),
224                    Type(length) => {
225                        let type_str = [
226                            (val >> 24) as u8,
227                            (val >> 16) as u8,
228                            (val >> 8) as u8,
229                            val as u8
230                        ];
231                        self.current_chunk.0.reset();
232                        self.current_chunk.0.update(&type_str);
233                        self.current_chunk.1 = length;
234                        goto!(
235                            ReadChunk(type_str, true),
236                            emit Decoded::ChunkBegin(length, type_str)
237                        )
238                    },
239                    Crc(type_str) => {
240                        if cfg!(fuzzing) || val == self.current_chunk.0.checksum() {
241                            goto!(
242                                State::U32(U32Value::Length),
243                                emit if type_str == IEND {
244                                    Decoded::ImageEnd
245                                } else {
246                                    Decoded::ChunkComplete(val, type_str)
247                                }
248                            )
249                        } else {
250                            Err(DecodingError::CrcMismatch {
251                                recover: 1,
252                                crc_val: val,
253                                crc_sum: self.current_chunk.0.checksum(),
254                                chunk: type_str
255                            })
256                        }
257                    },
258                }
259            },
260            U32Byte2(type_, val) => {
261                goto!(U32Byte3(type_, val | (current_byte as u32) << 8))
262            },
263            U32Byte1(type_, val) => {
264                goto!(U32Byte2(type_, val | (current_byte as u32) << 16))
265            },
266            U32(type_) => {
267                goto!(U32Byte1(type_,       (current_byte as u32) << 24))
268            },
269            PartialChunk(type_str) => {
270                match type_str {
271                    IDAT => {
272                        self.have_idat = true;
273                        goto!(
274                            0,
275                            DecodeData(type_str, 0),
276                            emit Decoded::PartialChunk(type_str)
277                        )
278                    },
279                    chunk::fdAT => {
280                        if let Some(seq_no) = self.current_seq_no {
281                            let mut buf = &self.current_chunk.2[..];
282                            let next_seq_no = try!(buf.read_be());
283                            if next_seq_no != seq_no + 1 {
284                                return Err(DecodingError::Format(format!(
285                                    "Sequence is not in order, expected #{} got #{}.",
286                                    seq_no + 1,
287                                    next_seq_no
288                                ).into()))
289                            }
290                            self.current_seq_no = Some(next_seq_no);
291                        } else {
292                            return Err(DecodingError::Format("fcTL chunk missing before fdAT chunk.".into()))
293                        }
294                        goto!(
295                            0,
296                            DecodeData(type_str, 4),
297                            emit Decoded::PartialChunk(type_str)
298                        )
299                    },
300                    // Handle other chunks
301                    _ => {
302                        if self.current_chunk.1 == 0 { // complete chunk
303                            Ok((0, try!(self.parse_chunk(type_str))))
304                        } else {
305                            goto!(
306                                0, ReadChunk(type_str, true),
307                                emit Decoded::PartialChunk(type_str)
308                            )
309                        }
310                    }
311                }
312
313            },
314            ReadChunk(type_str, clear) => {
315                if clear {
316                    self.current_chunk.2.clear();
317                }
318                if self.current_chunk.1 > 0 {
319                    let (ref mut crc, ref mut remaining, ref mut c_buf) = self.current_chunk;
320                    let buf_avail = c_buf.capacity() - c_buf.len();
321                    let bytes_avail = min(buf.len(), buf_avail);
322                    let n = min(*remaining, bytes_avail as u32);
323                    if buf_avail == 0 {
324                        goto!(0, PartialChunk(type_str))
325                    } else {
326                        let buf = &buf[..n as usize];
327                        crc.update(buf);
328                        c_buf.extend(buf.iter().map(|&v| v));
329                        *remaining -= n;
330                        if *remaining == 0 {
331                            goto!(n as usize, PartialChunk(type_str
332                            ))
333                        } else {
334                            goto!(n as usize, ReadChunk(type_str, false))
335                        }
336
337                    }
338                } else {
339                    goto!(0, U32(U32Value::Crc(type_str)))
340                }
341            }
342            DecodeData(type_str, mut n) => {
343                let chunk_len = self.current_chunk.2.len();
344                let (c, data) = try!(self.inflater.update(&self.current_chunk.2[n..]));
345                image_data.extend_from_slice(data);
346                n += c;
347                if n == chunk_len && data.len() == 0 && c == 0 {
348                    goto!(
349                        0,
350                        ReadChunk(type_str, true),
351                        emit Decoded::ImageData
352                    )
353                } else {
354                    goto!(
355                        0,
356                        DecodeData(type_str, n),
357                        emit Decoded::ImageData
358                    )
359                }
360            }
361        }
362    }
363
364    fn parse_chunk(&mut self, type_str: [u8; 4])
365    -> Result<Decoded, DecodingError> {
366        self.state = Some(State::U32(U32Value::Crc(type_str)));
367        if self.info.is_none() && type_str != IHDR {
368            return Err(DecodingError::Format(format!(
369                "{} chunk appeared before IHDR chunk", String::from_utf8_lossy(&type_str)
370            ).into()))
371        }
372        match match type_str {
373            IHDR => {
374                self.parse_ihdr()
375            }
376            chunk::PLTE => {
377                self.parse_plte()
378            }
379            chunk::tRNS => {
380                self.parse_trns()
381            }
382            chunk::pHYs => {
383                self.parse_phys()
384            }
385            chunk::acTL => {
386                self.parse_actl()
387            }
388            chunk::fcTL => {
389                self.parse_fctl()
390            }
391            _ => Ok(Decoded::PartialChunk(type_str))
392        } {
393            Err(err) =>{
394                // Borrow of self ends here, because Decoding error does not borrow self.
395                self.state = None;
396                Err(err)
397            },
398            ok => ok
399        }
400    }
401
402    fn get_info_or_err(&self) -> Result<&Info, DecodingError> {
403        self.info.as_ref().ok_or(DecodingError::Format(
404            "IHDR chunk missing".into()
405        ))
406    }
407
408    fn parse_fctl(&mut self)
409    -> Result<Decoded, DecodingError> {
410        let mut buf = &self.current_chunk.2[..];
411        let next_seq_no = try!(buf.read_be());
412
413        // Asuming that fcTL is required before *every* fdAT-sequence
414        self.current_seq_no = Some(if let Some(seq_no) = self.current_seq_no {
415            if next_seq_no != seq_no + 1 {
416                return Err(DecodingError::Format(format!(
417                    "Sequence is not in order, expected #{} got #{}.",
418                    seq_no + 1,
419                    next_seq_no
420                ).into()))
421            }
422            next_seq_no
423        } else {
424            if next_seq_no != 0 {
425                return Err(DecodingError::Format(format!(
426                    "Sequence is not in order, expected #{} got #{}.",
427                    0,
428                    next_seq_no
429                ).into()))
430            }
431            0
432        });
433        self.inflater = if cfg!(fuzzing) {InflateStream::from_zlib_no_checksum()} else {InflateStream::from_zlib()};
434        let fc = FrameControl {
435            sequence_number: next_seq_no,
436            width: try!(buf.read_be()),
437            height: try!(buf.read_be()),
438            x_offset: try!(buf.read_be()),
439            y_offset: try!(buf.read_be()),
440            delay_num: try!(buf.read_be()),
441            delay_den: try!(buf.read_be()),
442            dispose_op: try!(buf.read_be()),
443            blend_op : try!(buf.read_be()),
444        };
445        self.info.as_mut().unwrap().frame_control = Some(fc.clone());
446        Ok(Decoded::FrameControl(fc))
447    }
448
449    fn parse_actl(&mut self)
450    -> Result<Decoded, DecodingError> {
451        if self.have_idat {
452            Err(DecodingError::Format(
453                "acTL chunk appeared after first IDAT chunk".into()
454            ))
455        } else {
456            let mut buf = &self.current_chunk.2[..];
457            let actl = AnimationControl {
458                num_frames: try!(buf.read_be()),
459                num_plays: try!(buf.read_be())
460            };
461            self.info.as_mut().unwrap().animation_control = Some(actl);
462            Ok(Decoded::AnimationControl(actl))
463        }
464    }
465
466    fn parse_plte(&mut self)
467    -> Result<Decoded, DecodingError> {
468        let mut vec = Vec::new();
469        vec.extend(self.current_chunk.2.iter().map(|&v| v));
470        self.info.as_mut().map(
471            |info| info.palette = Some(vec)
472        );
473        Ok(Decoded::Nothing)
474    }
475
476    fn parse_trns(&mut self)
477    -> Result<Decoded, DecodingError> {
478        use common::ColorType::*;
479        let (color_type, bit_depth) = {
480            let info = try!(self.get_info_or_err());
481            (info.color_type, info.bit_depth as u8)
482        };
483        let mut vec = Vec::new();
484        vec.extend(self.current_chunk.2.iter().map(|&v| v));
485        let len = vec.len();
486        let info = match self.info {
487            Some(ref mut info) => info,
488            None => return Err(DecodingError::Format(
489              "tRNS chunk occured before IHDR chunk".into()
490            ))
491        };
492        info.trns = Some(vec);
493        let vec = info.trns.as_mut().unwrap();
494        match color_type {
495            Grayscale => {
496                if len < 2 {
497                    return Err(DecodingError::Format(
498                        "not enough palette entries".into()
499                    ))
500                }
501                if bit_depth < 16 {
502                    vec[0] = vec[1];
503                    vec.truncate(1);
504                }
505                Ok(Decoded::Nothing)
506            },
507            RGB => {
508                if len < 6 {
509                    return Err(DecodingError::Format(
510                        "not enough palette entries".into()
511                    ))
512                }
513                if bit_depth < 16 {
514                    vec[0] = vec[1];
515                    vec[1] = vec[3];
516                    vec[2] = vec[5];
517                    vec.truncate(3);
518                }
519                Ok(Decoded::Nothing)
520            },
521            Indexed => {
522                let _ = info.palette.as_ref().ok_or(DecodingError::Format(
523                    "tRNS chunk occured before PLTE chunk".into()
524                ));
525                Ok(Decoded::Nothing)
526            },
527            c => Err(DecodingError::Format(
528                format!("tRNS chunk found for color type ({})", c as u8).into()
529            ))
530        }
531
532    }
533
534    fn parse_phys(&mut self)
535    -> Result<Decoded, DecodingError> {
536        if self.have_idat {
537            Err(DecodingError::Format(
538                "pHYs chunk appeared after first IDAT chunk".into()
539            ))
540        } else {
541            let mut buf = &self.current_chunk.2[..];
542            let xppu = try!(buf.read_be());
543            let yppu = try!(buf.read_be());
544            let unit = try!(buf.read_be());
545            let unit = match Unit::from_u8(unit) {
546                Some(unit) => unit,
547                None => return Err(DecodingError::Format(
548                    format!("invalid unit ({})", unit).into()
549                ))
550            };
551            let pixel_dims = PixelDimensions {
552                xppu: xppu,
553                yppu: yppu,
554                unit: unit,
555            };
556            self.info.as_mut().unwrap().pixel_dims = Some(pixel_dims);
557            Ok(Decoded::PixelDimensions(pixel_dims))
558        }
559    }
560
561    fn parse_ihdr(&mut self)
562    -> Result<Decoded, DecodingError> {
563        // TODO: check if color/bit depths combination is valid
564        let mut buf = &self.current_chunk.2[..];
565        let width = try!(buf.read_be());
566        let height = try!(buf.read_be());
567        let bit_depth = try!(buf.read_be());
568        let bit_depth = match BitDepth::from_u8(bit_depth) {
569            Some(bits) => bits,
570            None => return Err(DecodingError::Format(
571                format!("invalid bit depth ({})", bit_depth).into()
572            ))
573        };
574        let color_type = try!(buf.read_be());
575        let color_type = match ColorType::from_u8(color_type) {
576            Some(color_type) => color_type,
577            None => return Err(DecodingError::Format(
578                format!("invalid color type ({})", color_type).into()
579            ))
580        };
581        match try!(buf.read_be()) { // compression method
582            0u8 => (),
583            n => return Err(DecodingError::Format(
584                format!("unknown compression method ({})", n).into()
585            ))
586        }
587        match try!(buf.read_be()) { // filter method
588            0u8 => (),
589            n => return Err(DecodingError::Format(
590                format!("unknown filter method ({})", n).into()
591            ))
592        }
593        let interlaced = match try!(buf.read_be()) {
594            0u8 => false,
595            1 => {
596                true
597            },
598            n => return Err(DecodingError::Format(
599                format!("unknown interlace method ({})", n).into()
600            ))
601        };
602        let mut info = Info::default();
603
604        info.width = width;
605        info.height = height;
606        info.bit_depth = bit_depth;
607        info.color_type = color_type;
608        info.interlaced = interlaced;
609        self.info = Some(info);
610        Ok(Decoded::Header(
611            width,
612            height,
613            bit_depth,
614            color_type,
615            interlaced
616        ))
617    }
618}
619
620#[inline(always)]
621pub fn get_info(d: &StreamingDecoder) -> Option<&Info> {
622    d.info.as_ref()
623}
png/decoder/stream.rs

png/decoder/
stream.rs
