1 mod stream;
2
3 use self::stream::{get_info, CHUNCK_BUFFER_SIZE};
4 pub use self::stream::{Decoded, DecodingError, StreamingDecoder};
5
6 use std::borrow;
7 use std::io::{BufRead, BufReader, Read, Write};
8 use std::mem;
9
10 use crate::chunk::IDAT;
11 use crate::common::{BitDepth, ColorType, Info, Transformations};
12 use crate::filter::{unfilter, FilterType};
13 use crate::utils;
14
15 /*
16 pub enum InterlaceHandling {
17 /// Outputs the raw rows
18 RawRows,
19 /// Fill missing the pixels from the existing ones
20 Rectangle,
21 /// Only fill the needed pixels
22 Sparkle
23 }
24 */
25
26 /// Output info
27 #[derive(Debug)]
28 pub struct OutputInfo {
29 pub width: u32,
30 pub height: u32,
31 pub color_type: ColorType,
32 pub bit_depth: BitDepth,
33 pub line_size: usize,
34 }
35
36 impl OutputInfo {
37 /// Returns the size needed to hold a decoded frame
buffer_size(&self) -> usize38 pub fn buffer_size(&self) -> usize {
39 self.line_size * self.height as usize
40 }
41 }
42
43 #[derive(Clone, Copy, Debug)]
44 /// Limits on the resources the `Decoder` is allowed too use
45 pub struct Limits {
46 /// maximum number of bytes the decoder is allowed to allocate, default is 64Mib
47 pub bytes: usize,
48 }
49
50 impl Default for Limits {
default() -> Limits51 fn default() -> Limits {
52 Limits {
53 bytes: 1024 * 1024 * 64,
54 }
55 }
56 }
57
58 /// PNG Decoder
59 pub struct Decoder<R: Read> {
60 /// Reader
61 r: R,
62 /// Output transformations
63 transform: Transformations,
64 /// Limits on resources the Decoder is allowed to use
65 limits: Limits,
66 }
67
68 struct InterlacedRow<'data> {
69 data: &'data [u8],
70 interlace: InterlaceInfo,
71 }
72
73 enum InterlaceInfo {
74 None,
75 Adam7 { pass: u8, line: u32, width: u32 },
76 }
77
78 impl<R: Read> Decoder<R> {
new(r: R) -> Decoder<R>79 pub fn new(r: R) -> Decoder<R> {
80 Decoder::new_with_limits(r, Limits::default())
81 }
82
new_with_limits(r: R, limits: Limits) -> Decoder<R>83 pub fn new_with_limits(r: R, limits: Limits) -> Decoder<R> {
84 Decoder {
85 r,
86 transform: crate::Transformations::EXPAND
87 | crate::Transformations::SCALE_16
88 | crate::Transformations::STRIP_16,
89 limits,
90 }
91 }
92
93 /// Limit resource usage
94 ///
95 /// ```
96 /// use std::fs::File;
97 /// use png::{Decoder, Limits};
98 /// // This image is 32x32 pixels, so the deocder will allocate more than four bytes
99 /// let mut limits = Limits::default();
100 /// limits.bytes = 4;
101 /// let mut decoder = Decoder::new_with_limits(File::open("tests/pngsuite/basi0g01.png").unwrap(), limits);
102 /// assert!(decoder.read_info().is_err());
103 /// // This image is 32x32 pixels, so the decoder will allocate less than 10Kib
104 /// let mut limits = Limits::default();
105 /// limits.bytes = 10*1024;
106 /// let mut decoder = Decoder::new_with_limits(File::open("tests/pngsuite/basi0g01.png").unwrap(), limits);
107 /// assert!(decoder.read_info().is_ok());
108 /// ```
set_limits(&mut self, limits: Limits)109 pub fn set_limits(&mut self, limits: Limits) {
110 self.limits = limits;
111 }
112
113 /// Reads all meta data until the first IDAT chunk
read_info(self) -> Result<(OutputInfo, Reader<R>), DecodingError>114 pub fn read_info(self) -> Result<(OutputInfo, Reader<R>), DecodingError> {
115 let mut r = Reader::new(self.r, StreamingDecoder::new(), self.transform, self.limits);
116 r.init()?;
117 let (ct, bits) = r.output_color_type();
118 let info = {
119 let info = r.info();
120 OutputInfo {
121 width: info.width,
122 height: info.height,
123 color_type: ct,
124 bit_depth: bits,
125 line_size: r.output_line_size(info.width),
126 }
127 };
128 Ok((info, r))
129 }
130
131 /// Set the allowed and performed transformations.
132 ///
133 /// A transformation is a pre-processing on the raw image data modifying content or encoding.
134 /// Many options have an impact on memory or CPU usage during decoding.
set_transformations(&mut self, transform: Transformations)135 pub fn set_transformations(&mut self, transform: Transformations) {
136 self.transform = transform;
137 }
138 }
139
140 struct ReadDecoder<R: Read> {
141 reader: BufReader<R>,
142 decoder: StreamingDecoder,
143 at_eof: bool,
144 }
145
146 impl<R: Read> ReadDecoder<R> {
147 /// Returns the next decoded chunk. If the chunk is an ImageData chunk, its contents are written
148 /// into image_data.
decode_next(&mut self, image_data: &mut Vec<u8>) -> Result<Option<Decoded>, DecodingError>149 fn decode_next(&mut self, image_data: &mut Vec<u8>) -> Result<Option<Decoded>, DecodingError> {
150 while !self.at_eof {
151 let (consumed, result) = {
152 let buf = self.reader.fill_buf()?;
153 if buf.is_empty() {
154 return Err(DecodingError::Format("unexpected EOF".into()));
155 }
156 self.decoder.update(buf, image_data)?
157 };
158 self.reader.consume(consumed);
159 match result {
160 Decoded::Nothing => (),
161 Decoded::ImageEnd => self.at_eof = true,
162 result => return Ok(Some(result)),
163 }
164 }
165 Ok(None)
166 }
167
info(&self) -> Option<&Info>168 fn info(&self) -> Option<&Info> {
169 get_info(&self.decoder)
170 }
171 }
172
173 /// PNG reader (mostly high-level interface)
174 ///
175 /// Provides a high level that iterates over lines or whole images.
176 pub struct Reader<R: Read> {
177 decoder: ReadDecoder<R>,
178 bpp: usize,
179 rowlen: usize,
180 adam7: Option<utils::Adam7Iterator>,
181 /// Previous raw line
182 prev: Vec<u8>,
183 /// Current raw line
184 current: Vec<u8>,
185 /// Start index of the current scan line.
186 scan_start: usize,
187 /// Output transformations
188 transform: Transformations,
189 /// Processed line
190 processed: Vec<u8>,
191 limits: Limits,
192 }
193
194 macro_rules! get_info(
195 ($this:expr) => {
196 $this.decoder.info().unwrap()
197 }
198 );
199
200 impl<R: Read> Reader<R> {
201 /// Creates a new PNG reader
new(r: R, d: StreamingDecoder, t: Transformations, limits: Limits) -> Reader<R>202 fn new(r: R, d: StreamingDecoder, t: Transformations, limits: Limits) -> Reader<R> {
203 Reader {
204 decoder: ReadDecoder {
205 reader: BufReader::with_capacity(CHUNCK_BUFFER_SIZE, r),
206 decoder: d,
207 at_eof: false,
208 },
209 bpp: 0,
210 rowlen: 0,
211 adam7: None,
212 prev: Vec::new(),
213 current: Vec::new(),
214 scan_start: 0,
215 transform: t,
216 processed: Vec::new(),
217 limits,
218 }
219 }
220
221 /// Reads all meta data until the first IDAT chunk
init(&mut self) -> Result<(), DecodingError>222 fn init(&mut self) -> Result<(), DecodingError> {
223 use crate::Decoded::*;
224 if self.decoder.info().is_some() {
225 Ok(())
226 } else {
227 loop {
228 match self.decoder.decode_next(&mut Vec::new())? {
229 Some(ChunkBegin(_, IDAT)) => break,
230 None => return Err(DecodingError::Format("IDAT chunk missing".into())),
231 _ => (),
232 }
233 }
234 {
235 let info = match self.decoder.info() {
236 Some(info) => info,
237 None => return Err(DecodingError::Format("IHDR chunk missing".into())),
238 };
239 self.bpp = info.bytes_per_pixel();
240 self.rowlen = info.raw_row_length();
241 if info.interlaced {
242 self.adam7 = Some(utils::Adam7Iterator::new(info.width, info.height))
243 }
244 }
245 self.allocate_out_buf()?;
246 self.prev = vec![0; self.rowlen];
247 Ok(())
248 }
249 }
250
info(&self) -> &Info251 pub fn info(&self) -> &Info {
252 get_info!(self)
253 }
254
255 /// Decodes the next frame into `buf`
next_frame(&mut self, buf: &mut [u8]) -> Result<(), DecodingError>256 pub fn next_frame(&mut self, buf: &mut [u8]) -> Result<(), DecodingError> {
257 // TODO 16 bit
258 let (color_type, bit_depth) = self.output_color_type();
259 let width = get_info!(self).width;
260 if buf.len() < self.output_buffer_size() {
261 return Err(DecodingError::Other(
262 "supplied buffer is too small to hold the image".into(),
263 ));
264 }
265 if get_info!(self).interlaced {
266 while let Some((row, adam7)) = self.next_interlaced_row()? {
267 let (pass, line, _) = adam7.unwrap();
268 let samples = color_type.samples() as u8;
269 utils::expand_pass(buf, width, row, pass, line, samples * (bit_depth as u8));
270 }
271 } else {
272 let mut len = 0;
273 while let Some(row) = self.next_row()? {
274 len += (&mut buf[len..]).write(row)?;
275 }
276 }
277 Ok(())
278 }
279
280 /// Returns the next processed row of the image
next_row(&mut self) -> Result<Option<&[u8]>, DecodingError>281 pub fn next_row(&mut self) -> Result<Option<&[u8]>, DecodingError> {
282 self.next_interlaced_row().map(|v| v.map(|v| v.0))
283 }
284
285 /// Returns the next processed row of the image
next_interlaced_row( &mut self, ) -> Result<Option<(&[u8], Option<(u8, u32, u32)>)>, DecodingError>286 pub fn next_interlaced_row(
287 &mut self,
288 ) -> Result<Option<(&[u8], Option<(u8, u32, u32)>)>, DecodingError> {
289 match self.next_interlaced_row_impl() {
290 Err(err) => Err(err),
291 Ok(None) => Ok(None),
292 Ok(Some(row)) => {
293 let interlace = match row.interlace {
294 InterlaceInfo::None => None,
295 InterlaceInfo::Adam7 { pass, line, width } => Some((pass, line, width)),
296 };
297
298 Ok(Some((row.data, interlace)))
299 }
300 }
301 }
302
303 /// Fetch the next interlaced row and filter it according to our own transformations.
next_interlaced_row_impl(&mut self) -> Result<Option<InterlacedRow<'_>>, DecodingError>304 fn next_interlaced_row_impl(&mut self) -> Result<Option<InterlacedRow<'_>>, DecodingError> {
305 use crate::common::ColorType::*;
306 let transform = self.transform;
307
308 if transform == crate::Transformations::IDENTITY {
309 return self.next_raw_interlaced_row();
310 }
311
312 // swap buffer to circumvent borrow issues
313 let mut buffer = mem::replace(&mut self.processed, Vec::new());
314 let (got_next, adam7) = if let Some(row) = self.next_raw_interlaced_row()? {
315 (&mut buffer[..]).write_all(row.data)?;
316 (true, row.interlace)
317 } else {
318 (false, InterlaceInfo::None)
319 };
320 // swap back
321 let _ = mem::replace(&mut self.processed, buffer);
322
323 if !got_next {
324 return Ok(None);
325 }
326
327 let (color_type, bit_depth, trns) = {
328 let info = get_info!(self);
329 (info.color_type, info.bit_depth as u8, info.trns.is_some())
330 };
331 let output_buffer = if let InterlaceInfo::Adam7 { width, .. } = adam7 {
332 let width = self.line_size(width);
333 &mut self.processed[..width]
334 } else {
335 &mut *self.processed
336 };
337
338 let mut len = output_buffer.len();
339 if transform.contains(crate::Transformations::EXPAND) {
340 match color_type {
341 Indexed => expand_paletted(output_buffer, get_info!(self))?,
342 Grayscale | GrayscaleAlpha if bit_depth < 8 => {
343 expand_gray_u8(output_buffer, get_info!(self))
344 }
345 Grayscale | RGB if trns => {
346 let channels = color_type.samples();
347 let trns = get_info!(self).trns.as_ref().unwrap();
348 if bit_depth == 8 {
349 utils::expand_trns_line(output_buffer, &*trns, channels);
350 } else {
351 utils::expand_trns_line16(output_buffer, &*trns, channels);
352 }
353 }
354 _ => (),
355 }
356 }
357
358 if bit_depth == 16
359 && transform
360 .intersects(crate::Transformations::SCALE_16 | crate::Transformations::STRIP_16)
361 {
362 len /= 2;
363 for i in 0..len {
364 output_buffer[i] = output_buffer[2 * i];
365 }
366 }
367
368 Ok(Some(InterlacedRow {
369 data: &output_buffer[..len],
370 interlace: adam7,
371 }))
372 }
373
374 /// Returns the color type and the number of bits per sample
375 /// of the data returned by `Reader::next_row` and Reader::frames`.
output_color_type(&mut self) -> (ColorType, BitDepth)376 pub fn output_color_type(&mut self) -> (ColorType, BitDepth) {
377 use crate::common::ColorType::*;
378 let t = self.transform;
379 let info = get_info!(self);
380 if t == crate::Transformations::IDENTITY {
381 (info.color_type, info.bit_depth)
382 } else {
383 let bits = match info.bit_depth as u8 {
384 16 if t.intersects(
385 crate::Transformations::SCALE_16 | crate::Transformations::STRIP_16,
386 ) =>
387 {
388 8
389 }
390 n if n < 8 && t.contains(crate::Transformations::EXPAND) => 8,
391 n => n,
392 };
393 let color_type = if t.contains(crate::Transformations::EXPAND) {
394 let has_trns = info.trns.is_some();
395 match info.color_type {
396 Grayscale if has_trns => GrayscaleAlpha,
397 RGB if has_trns => RGBA,
398 Indexed if has_trns => RGBA,
399 Indexed => RGB,
400 ct => ct,
401 }
402 } else {
403 info.color_type
404 };
405 (color_type, BitDepth::from_u8(bits).unwrap())
406 }
407 }
408
409 /// Returns the number of bytes required to hold a deinterlaced image frame
410 /// that is decoded using the given input transformations.
output_buffer_size(&self) -> usize411 pub fn output_buffer_size(&self) -> usize {
412 let (width, height) = get_info!(self).size();
413 let size = self.output_line_size(width);
414 size * height as usize
415 }
416
417 /// Returns the number of bytes required to hold a deinterlaced row.
output_line_size(&self, width: u32) -> usize418 pub fn output_line_size(&self, width: u32) -> usize {
419 let size = self.line_size(width);
420 if get_info!(self).bit_depth as u8 == 16
421 && self
422 .transform
423 .intersects(crate::Transformations::SCALE_16 | crate::Transformations::STRIP_16)
424 {
425 size / 2
426 } else {
427 size
428 }
429 }
430
431 /// Returns the number of bytes required to decode a deinterlaced row.
line_size(&self, width: u32) -> usize432 fn line_size(&self, width: u32) -> usize {
433 use crate::common::ColorType::*;
434 let t = self.transform;
435 let info = get_info!(self);
436 let trns = info.trns.is_some();
437 // TODO 16 bit
438 let bits = match info.color_type {
439 Indexed if trns && t.contains(crate::Transformations::EXPAND) => 4 * 8,
440 Indexed if t.contains(crate::Transformations::EXPAND) => 3 * 8,
441 RGB if trns && t.contains(crate::Transformations::EXPAND) => 4 * 8,
442 Grayscale if trns && t.contains(crate::Transformations::EXPAND) => 2 * 8,
443 Grayscale if t.contains(crate::Transformations::EXPAND) => 1 * 8,
444 GrayscaleAlpha if t.contains(crate::Transformations::EXPAND) => 2 * 8,
445 // divide by 2 as it will get mutiplied by two later
446 _ if info.bit_depth as u8 == 16 => info.bits_per_pixel() / 2,
447 _ => info.bits_per_pixel(),
448 } * width as usize
449 * if info.bit_depth as u8 == 16 { 2 } else { 1 };
450 let len = bits / 8;
451 let extra = bits % 8;
452 len + match extra {
453 0 => 0,
454 _ => 1,
455 }
456 }
457
allocate_out_buf(&mut self) -> Result<(), DecodingError>458 fn allocate_out_buf(&mut self) -> Result<(), DecodingError> {
459 let width = get_info!(self).width;
460 let bytes = self.limits.bytes;
461 if bytes < self.line_size(width) {
462 return Err(DecodingError::LimitsExceeded);
463 }
464 self.processed = vec![0; self.line_size(width)];
465 Ok(())
466 }
467
468 /// Returns the next raw scanline of the image interlace pass.
469 /// The scanline is filtered against the previous scanline according to the specification.
next_raw_interlaced_row(&mut self) -> Result<Option<InterlacedRow<'_>>, DecodingError>470 fn next_raw_interlaced_row(&mut self) -> Result<Option<InterlacedRow<'_>>, DecodingError> {
471 let _ = get_info!(self);
472 let bpp = self.bpp;
473 let (rowlen, passdata) = if let Some(ref mut adam7) = self.adam7 {
474 let last_pass = adam7.current_pass();
475 if let Some((pass, line, len)) = adam7.next() {
476 let rowlen = get_info!(self).raw_row_length_from_width(len);
477 if last_pass != pass {
478 self.prev.clear();
479 self.prev.resize(rowlen, 0u8);
480 }
481 (rowlen, Some((pass, line, len)))
482 } else {
483 return Ok(None);
484 }
485 } else {
486 (self.rowlen, None)
487 };
488 loop {
489 if self.current.len() - self.scan_start >= rowlen {
490 let row = &mut self.current[self.scan_start..];
491 let filter = match FilterType::from_u8(row[0]) {
492 None => {
493 return Err(DecodingError::Format(
494 format!("invalid filter method ({})", row[0]).into(),
495 ))
496 }
497 Some(filter) => filter,
498 };
499
500 if let Err(message) =
501 unfilter(filter, bpp, &self.prev[1..rowlen], &mut row[1..rowlen])
502 {
503 return Err(DecodingError::Format(borrow::Cow::Borrowed(message)));
504 }
505
506 self.prev[..rowlen].copy_from_slice(&row[..rowlen]);
507 self.scan_start += rowlen;
508
509 return Ok(Some(InterlacedRow {
510 data: &self.prev[1..rowlen],
511 interlace: match passdata {
512 None => InterlaceInfo::None,
513 Some((pass, line, width)) => InterlaceInfo::Adam7 { pass, line, width },
514 },
515 }));
516 } else {
517 // Clear the current buffer before appending more data.
518 if self.scan_start > 0 {
519 self.current.drain(..self.scan_start).for_each(drop);
520 self.scan_start = 0;
521 }
522
523 let val = self.decoder.decode_next(&mut self.current)?;
524 match val {
525 Some(Decoded::ImageData) => {}
526 None => {
527 if !self.current.is_empty() {
528 return Err(DecodingError::Format("file truncated".into()));
529 } else {
530 return Ok(None);
531 }
532 }
533 _ => (),
534 }
535 }
536 }
537 }
538 }
539
expand_paletted(buffer: &mut [u8], info: &Info) -> Result<(), DecodingError>540 fn expand_paletted(buffer: &mut [u8], info: &Info) -> Result<(), DecodingError> {
541 if let Some(palette) = info.palette.as_ref() {
542 if let BitDepth::Sixteen = info.bit_depth {
543 Err(DecodingError::Format(
544 "Bit depth '16' is not valid for paletted images".into(),
545 ))
546 } else {
547 let black = [0, 0, 0];
548 if let Some(ref trns) = info.trns {
549 utils::unpack_bits(buffer, 4, info.bit_depth as u8, |i, chunk| {
550 let (rgb, a) = (
551 palette
552 .get(3 * i as usize..3 * i as usize + 3)
553 .unwrap_or(&black),
554 *trns.get(i as usize).unwrap_or(&0xFF),
555 );
556 chunk[0] = rgb[0];
557 chunk[1] = rgb[1];
558 chunk[2] = rgb[2];
559 chunk[3] = a;
560 });
561 } else {
562 utils::unpack_bits(buffer, 3, info.bit_depth as u8, |i, chunk| {
563 let rgb = palette
564 .get(3 * i as usize..3 * i as usize + 3)
565 .unwrap_or(&black);
566 chunk[0] = rgb[0];
567 chunk[1] = rgb[1];
568 chunk[2] = rgb[2];
569 })
570 }
571 Ok(())
572 }
573 } else {
574 Err(DecodingError::Format("missing palette".into()))
575 }
576 }
577
expand_gray_u8(buffer: &mut [u8], info: &Info)578 fn expand_gray_u8(buffer: &mut [u8], info: &Info) {
579 let rescale = true;
580 let scaling_factor = if rescale {
581 (255) / ((1u16 << info.bit_depth as u8) - 1) as u8
582 } else {
583 1
584 };
585 if let Some(ref trns) = info.trns {
586 utils::unpack_bits(buffer, 2, info.bit_depth as u8, |pixel, chunk| {
587 if pixel == trns[0] {
588 chunk[1] = 0
589 } else {
590 chunk[1] = 0xFF
591 }
592 chunk[0] = pixel * scaling_factor
593 })
594 } else {
595 utils::unpack_bits(buffer, 1, info.bit_depth as u8, |val, chunk| {
596 chunk[0] = val * scaling_factor
597 })
598 }
599 }
600