1use borrow::Cow;
2use io::{Read, Write};
3use ops::{Deref, DerefMut};
4use std::{borrow, error, fmt, io, mem, ops, result};
5
6use crc32fast::Hasher as Crc32;
7use flate2::write::ZlibEncoder;
8
9use crate::chunk::{self, ChunkType};
10use crate::common::{
11 AnimationControl, BitDepth, BlendOp, BytesPerPixel, ColorType, Compression, DisposeOp,
12 FrameControl, Info, ParameterError, ParameterErrorKind, PixelDimensions, ScaledFloat,
13};
14use crate::filter::{filter, AdaptiveFilterType, FilterType};
15use crate::text_metadata::{
16 EncodableTextChunk, ITXtChunk, TEXtChunk, TextEncodingError, ZTXtChunk,
17};
18use crate::traits::WriteBytesExt;
19
20pub type Result<T> = result::Result<T, EncodingError>;
21
22#[derive(Debug)]
23pub enum EncodingError {
24 IoError(io::Error),
25 Format(FormatError),
26 Parameter(ParameterError),
27 LimitsExceeded,
28}
29
30#[derive(Debug)]
31pub struct FormatError {
32 inner: FormatErrorKind,
33}
34
35#[derive(Debug)]
36enum FormatErrorKind {
37 ZeroWidth,
38 ZeroHeight,
39 InvalidColorCombination(BitDepth, ColorType),
40 NoPalette,
41 WrittenTooMuch(usize),
43 NotAnimated,
44 OutOfBounds,
45 EndReached,
46 ZeroFrames,
47 MissingFrames,
48 MissingData(usize),
49 Unrecoverable,
50 BadTextEncoding(TextEncodingError),
51}
52
53impl error::Error for EncodingError {
54 fn cause(&self) -> Option<&(dyn error::Error + 'static)> {
55 match self {
56 EncodingError::IoError(err) => Some(err),
57 _ => None,
58 }
59 }
60}
61
62impl fmt::Display for EncodingError {
63 fn fmt(&self, fmt: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
64 use self::EncodingError::*;
65 match self {
66 IoError(err) => write!(fmt, "{}", err),
67 Format(desc) => write!(fmt, "{}", desc),
68 Parameter(desc) => write!(fmt, "{}", desc),
69 LimitsExceeded => write!(fmt, "Limits are exceeded."),
70 }
71 }
72}
73
74impl fmt::Display for FormatError {
75 fn fmt(&self, fmt: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
76 use FormatErrorKind::*;
77 match self.inner {
78 ZeroWidth => write!(fmt, "Zero width not allowed"),
79 ZeroHeight => write!(fmt, "Zero height not allowed"),
80 ZeroFrames => write!(fmt, "Zero frames not allowed"),
81 InvalidColorCombination(depth, color) => write!(
82 fmt,
83 "Invalid combination of bit-depth '{:?}' and color-type '{:?}'",
84 depth, color
85 ),
86 NoPalette => write!(fmt, "can't write indexed image without palette"),
87 WrittenTooMuch(index) => write!(fmt, "wrong data size, got {} bytes too many", index),
88 NotAnimated => write!(fmt, "not an animation"),
89 OutOfBounds => write!(
90 fmt,
91 "the dimension and position go over the frame boundaries"
92 ),
93 EndReached => write!(fmt, "all the frames have been already written"),
94 MissingFrames => write!(fmt, "there are still frames to be written"),
95 MissingData(n) => write!(fmt, "there are still {} bytes to be written", n),
96 Unrecoverable => write!(
97 fmt,
98 "a previous error put the writer into an unrecoverable state"
99 ),
100 BadTextEncoding(tee) => match tee {
101 TextEncodingError::Unrepresentable => write!(
102 fmt,
103 "The text metadata cannot be encoded into valid ISO 8859-1"
104 ),
105 TextEncodingError::InvalidKeywordSize => write!(fmt, "Invalid keyword size"),
106 TextEncodingError::CompressionError => {
107 write!(fmt, "Unable to compress text metadata")
108 }
109 },
110 }
111 }
112}
113
114impl From<io::Error> for EncodingError {
115 fn from(err: io::Error) -> EncodingError {
116 EncodingError::IoError(err)
117 }
118}
119
120impl From<EncodingError> for io::Error {
121 fn from(err: EncodingError) -> io::Error {
122 io::Error::new(io::ErrorKind::Other, err.to_string())
123 }
124}
125
126impl From<FormatErrorKind> for FormatError {
128 fn from(kind: FormatErrorKind) -> Self {
129 FormatError { inner: kind }
130 }
131}
132
133impl From<TextEncodingError> for EncodingError {
134 fn from(tee: TextEncodingError) -> Self {
135 EncodingError::Format(FormatError {
136 inner: FormatErrorKind::BadTextEncoding(tee),
137 })
138 }
139}
140
141pub struct Encoder<'a, W: Write> {
149 w: W,
150 info: Info<'a>,
151 options: Options,
152}
153
154#[derive(Default)]
156struct Options {
157 filter: FilterType,
158 adaptive_filter: AdaptiveFilterType,
159 sep_def_img: bool,
160 validate_sequence: bool,
161}
162
163impl<'a, W: Write> Encoder<'a, W> {
164 pub fn new(w: W, width: u32, height: u32) -> Encoder<'static, W> {
165 Encoder {
166 w,
167 info: Info::with_size(width, height),
168 options: Options::default(),
169 }
170 }
171
172 pub fn with_info(w: W, info: Info<'a>) -> Result<Encoder<'a, W>> {
173 if info.animation_control.is_some() != info.frame_control.is_some() {
174 return Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()));
175 }
176
177 if let Some(actl) = info.animation_control {
178 if actl.num_frames == 0 {
179 return Err(EncodingError::Format(FormatErrorKind::ZeroFrames.into()));
180 }
181 }
182
183 Ok(Encoder {
184 w,
185 info,
186 options: Options::default(),
187 })
188 }
189
190 pub fn set_animated(&mut self, num_frames: u32, num_plays: u32) -> Result<()> {
203 if num_frames == 0 {
204 return Err(EncodingError::Format(FormatErrorKind::ZeroFrames.into()));
205 }
206
207 let actl = AnimationControl {
208 num_frames,
209 num_plays,
210 };
211
212 let fctl = FrameControl {
213 sequence_number: 0,
214 width: self.info.width,
215 height: self.info.height,
216 ..Default::default()
217 };
218
219 self.info.animation_control = Some(actl);
220 self.info.frame_control = Some(fctl);
221 Ok(())
222 }
223
224 pub fn set_sep_def_img(&mut self, sep_def_img: bool) -> Result<()> {
234 if self.info.animation_control.is_some() {
235 self.options.sep_def_img = sep_def_img;
236 Ok(())
237 } else {
238 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
239 }
240 }
241
242 pub fn set_palette<T: Into<Cow<'a, [u8]>>>(&mut self, palette: T) {
245 self.info.palette = Some(palette.into());
246 }
247
248 pub fn set_trns<T: Into<Cow<'a, [u8]>>>(&mut self, trns: T) {
251 self.info.trns = Some(trns.into());
252 }
253
254 pub fn set_source_gamma(&mut self, source_gamma: ScaledFloat) {
256 self.info.source_gamma = Some(source_gamma);
257 }
258
259 pub fn set_source_chromaticities(
262 &mut self,
263 source_chromaticities: super::SourceChromaticities,
264 ) {
265 self.info.source_chromaticities = Some(source_chromaticities);
266 }
267
268 pub fn set_srgb(&mut self, rendering_intent: super::SrgbRenderingIntent) {
273 self.info.srgb = Some(rendering_intent);
274 }
275
276 pub fn write_header(self) -> Result<Writer<W>> {
280 Writer::new(self.w, PartialInfo::new(&self.info), self.options).init(&self.info)
281 }
282
283 pub fn set_color(&mut self, color: ColorType) {
289 self.info.color_type = color;
290 }
291
292 pub fn set_depth(&mut self, depth: BitDepth) {
294 self.info.bit_depth = depth;
295 }
296
297 pub fn set_compression(&mut self, compression: Compression) {
302 self.info.compression = compression;
303 }
304
305 pub fn set_filter(&mut self, filter: FilterType) {
311 self.options.filter = filter;
312 }
313
314 pub fn set_adaptive_filter(&mut self, adaptive_filter: AdaptiveFilterType) {
322 self.options.adaptive_filter = adaptive_filter;
323 }
324
325 pub fn set_frame_delay(&mut self, numerator: u16, denominator: u16) -> Result<()> {
343 if let Some(ref mut fctl) = self.info.frame_control {
344 fctl.delay_den = denominator;
345 fctl.delay_num = numerator;
346 Ok(())
347 } else {
348 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
349 }
350 }
351
352 pub fn set_blend_op(&mut self, op: BlendOp) -> Result<()> {
375 if let Some(ref mut fctl) = self.info.frame_control {
376 fctl.blend_op = op;
377 Ok(())
378 } else {
379 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
380 }
381 }
382
383 pub fn set_dispose_op(&mut self, op: DisposeOp) -> Result<()> {
404 if let Some(ref mut fctl) = self.info.frame_control {
405 fctl.dispose_op = op;
406 Ok(())
407 } else {
408 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
409 }
410 }
411 pub fn set_pixel_dims(&mut self, pixel_dims: Option<PixelDimensions>) {
412 self.info.pixel_dims = pixel_dims
413 }
414 pub fn add_text_chunk(&mut self, keyword: String, text: String) -> Result<()> {
416 let text_chunk = TEXtChunk::new(keyword, text);
417 self.info.uncompressed_latin1_text.push(text_chunk);
418 Ok(())
419 }
420
421 pub fn add_ztxt_chunk(&mut self, keyword: String, text: String) -> Result<()> {
423 let text_chunk = ZTXtChunk::new(keyword, text);
424 self.info.compressed_latin1_text.push(text_chunk);
425 Ok(())
426 }
427
428 pub fn add_itxt_chunk(&mut self, keyword: String, text: String) -> Result<()> {
433 let text_chunk = ITXtChunk::new(keyword, text);
434 self.info.utf8_text.push(text_chunk);
435 Ok(())
436 }
437
438 pub fn validate_sequence(&mut self, validate: bool) {
449 self.options.validate_sequence = validate;
450 }
451}
452
453pub struct Writer<W: Write> {
462 w: W,
464 info: PartialInfo,
466 options: Options,
468 images_written: u64,
470 animation_written: u32,
472 iend_written: bool,
475}
476
477struct PartialInfo {
479 width: u32,
480 height: u32,
481 bit_depth: BitDepth,
482 color_type: ColorType,
483 frame_control: Option<FrameControl>,
484 animation_control: Option<AnimationControl>,
485 compression: Compression,
486 has_palette: bool,
487}
488
489impl PartialInfo {
490 fn new(info: &Info) -> Self {
491 PartialInfo {
492 width: info.width,
493 height: info.height,
494 bit_depth: info.bit_depth,
495 color_type: info.color_type,
496 frame_control: info.frame_control,
497 animation_control: info.animation_control,
498 compression: info.compression,
499 has_palette: info.palette.is_some(),
500 }
501 }
502
503 fn bpp_in_prediction(&self) -> BytesPerPixel {
504 self.to_info().bpp_in_prediction()
506 }
507
508 fn raw_row_length(&self) -> usize {
509 self.to_info().raw_row_length()
511 }
512
513 fn raw_row_length_from_width(&self, width: u32) -> usize {
514 self.to_info().raw_row_length_from_width(width)
516 }
517
518 fn to_info(&self) -> Info<'static> {
521 Info {
522 width: self.width,
523 height: self.height,
524 bit_depth: self.bit_depth,
525 color_type: self.color_type,
526 frame_control: self.frame_control,
527 animation_control: self.animation_control,
528 compression: self.compression,
529 ..Default::default()
530 }
531 }
532}
533
534const DEFAULT_BUFFER_LENGTH: usize = 4 * 1024;
535
536pub(crate) fn write_chunk<W: Write>(mut w: W, name: chunk::ChunkType, data: &[u8]) -> Result<()> {
537 w.write_be(data.len() as u32)?;
538 w.write_all(&name.0)?;
539 w.write_all(data)?;
540 let mut crc = Crc32::new();
541 crc.update(&name.0);
542 crc.update(data);
543 w.write_be(crc.finalize())?;
544 Ok(())
545}
546
547impl<W: Write> Writer<W> {
548 fn new(w: W, info: PartialInfo, options: Options) -> Writer<W> {
549 Writer {
550 w,
551 info,
552 options,
553 images_written: 0,
554 animation_written: 0,
555 iend_written: false,
556 }
557 }
558
559 fn init(mut self, info: &Info<'_>) -> Result<Self> {
560 if self.info.width == 0 {
561 return Err(EncodingError::Format(FormatErrorKind::ZeroWidth.into()));
562 }
563
564 if self.info.height == 0 {
565 return Err(EncodingError::Format(FormatErrorKind::ZeroHeight.into()));
566 }
567
568 if self
569 .info
570 .color_type
571 .is_combination_invalid(self.info.bit_depth)
572 {
573 return Err(EncodingError::Format(
574 FormatErrorKind::InvalidColorCombination(self.info.bit_depth, self.info.color_type)
575 .into(),
576 ));
577 }
578
579 self.w.write_all(&[137, 80, 78, 71, 13, 10, 26, 10])?; info.encode(&mut self.w)?;
581
582 Ok(self)
583 }
584
585 pub fn write_chunk(&mut self, name: ChunkType, data: &[u8]) -> Result<()> {
590 use std::convert::TryFrom;
591
592 if u32::try_from(data.len()).map_or(true, |length| length > i32::MAX as u32) {
593 let kind = FormatErrorKind::WrittenTooMuch(data.len() - i32::MAX as usize);
594 return Err(EncodingError::Format(kind.into()));
595 }
596
597 write_chunk(&mut self.w, name, data)
598 }
599
600 pub fn write_text_chunk<T: EncodableTextChunk>(&mut self, text_chunk: &T) -> Result<()> {
601 text_chunk.encode(&mut self.w)
602 }
603
604 fn validate_new_image(&self) -> Result<()> {
606 if !self.options.validate_sequence {
607 return Ok(());
608 }
609
610 match self.info.animation_control {
611 None => {
612 if self.images_written == 0 {
613 Ok(())
614 } else {
615 Err(EncodingError::Format(FormatErrorKind::EndReached.into()))
616 }
617 }
618 Some(_) => {
619 if self.info.frame_control.is_some() {
620 Ok(())
621 } else {
622 Err(EncodingError::Format(FormatErrorKind::EndReached.into()))
623 }
624 }
625 }
626 }
627
628 fn validate_sequence_done(&self) -> Result<()> {
629 if !self.options.validate_sequence {
630 return Ok(());
631 }
632
633 if (self.info.animation_control.is_some() && self.info.frame_control.is_some())
634 || self.images_written == 0
635 {
636 Err(EncodingError::Format(FormatErrorKind::MissingFrames.into()))
637 } else {
638 Ok(())
639 }
640 }
641
642 const MAX_IDAT_CHUNK_LEN: u32 = u32::MAX >> 1;
643 #[allow(non_upper_case_globals)]
644 const MAX_fdAT_CHUNK_LEN: u32 = (u32::MAX >> 1) - 4;
645
646 pub fn write_image_data(&mut self, data: &[u8]) -> Result<()> {
648 if self.info.color_type == ColorType::Indexed && !self.info.has_palette {
649 return Err(EncodingError::Format(FormatErrorKind::NoPalette.into()));
650 }
651
652 self.validate_new_image()?;
653
654 let width: usize;
655 let height: usize;
656 if let Some(ref mut fctl) = self.info.frame_control {
657 width = fctl.width as usize;
658 height = fctl.height as usize;
659 } else {
660 width = self.info.width as usize;
661 height = self.info.height as usize;
662 }
663
664 let in_len = self.info.raw_row_length_from_width(width as u32) - 1;
665 let data_size = in_len * height;
666 if data_size != data.len() {
667 return Err(EncodingError::Parameter(
668 ParameterErrorKind::ImageBufferSize {
669 expected: data_size,
670 actual: data.len(),
671 }
672 .into(),
673 ));
674 }
675
676 let prev = vec![0; in_len];
677 let mut prev = prev.as_slice();
678
679 let bpp = self.info.bpp_in_prediction();
680 let filter_method = self.options.filter;
681 let adaptive_method = self.options.adaptive_filter;
682
683 let zlib_encoded = match self.info.compression {
684 Compression::Fast => {
685 let mut compressor = fdeflate::Compressor::new(std::io::Cursor::new(Vec::new()))?;
686
687 let mut current = vec![0; in_len + 1];
688 for line in data.chunks(in_len) {
689 let filter_type = filter(
690 filter_method,
691 adaptive_method,
692 bpp,
693 prev,
694 line,
695 &mut current[1..],
696 );
697
698 current[0] = filter_type as u8;
699 compressor.write_data(¤t)?;
700 prev = line;
701 }
702
703 let compressed = compressor.finish()?.into_inner();
704 if compressed.len()
705 > fdeflate::StoredOnlyCompressor::<()>::compressed_size((in_len + 1) * height)
706 {
707 let mut compressor =
715 fdeflate::StoredOnlyCompressor::new(std::io::Cursor::new(Vec::new()))?;
716 for line in data.chunks(in_len) {
717 compressor.write_data(&[0])?;
718 compressor.write_data(line)?;
719 }
720 compressor.finish()?.into_inner()
721 } else {
722 compressed
723 }
724 }
725 _ => {
726 let mut current = vec![0; in_len];
727
728 let mut zlib = ZlibEncoder::new(Vec::new(), self.info.compression.to_options());
729 for line in data.chunks(in_len) {
730 let filter_type = filter(
731 filter_method,
732 adaptive_method,
733 bpp,
734 prev,
735 line,
736 &mut current,
737 );
738
739 zlib.write_all(&[filter_type as u8])?;
740 zlib.write_all(¤t)?;
741 prev = line;
742 }
743 zlib.finish()?
744 }
745 };
746
747 match self.info.frame_control {
748 None => {
749 self.write_zlib_encoded_idat(&zlib_encoded)?;
750 }
751 Some(_) if self.should_skip_frame_control_on_default_image() => {
752 self.write_zlib_encoded_idat(&zlib_encoded)?;
753 }
754 Some(ref mut fctl) => {
755 fctl.encode(&mut self.w)?;
756 fctl.sequence_number = fctl.sequence_number.wrapping_add(1);
757 self.animation_written += 1;
758
759 if self.images_written == 0 {
761 self.write_zlib_encoded_idat(&zlib_encoded)?;
762 } else {
763 let buff_size = zlib_encoded.len().min(Self::MAX_fdAT_CHUNK_LEN as usize);
764 let mut alldata = vec![0u8; 4 + buff_size];
765 for chunk in zlib_encoded.chunks(Self::MAX_fdAT_CHUNK_LEN as usize) {
766 alldata[..4].copy_from_slice(&fctl.sequence_number.to_be_bytes());
767 alldata[4..][..chunk.len()].copy_from_slice(chunk);
768 write_chunk(&mut self.w, chunk::fdAT, &alldata[..4 + chunk.len()])?;
769 fctl.sequence_number = fctl.sequence_number.wrapping_add(1);
770 }
771 }
772 }
773 }
774
775 self.increment_images_written();
776
777 Ok(())
778 }
779
780 fn increment_images_written(&mut self) {
781 self.images_written = self.images_written.saturating_add(1);
782
783 if let Some(actl) = self.info.animation_control {
784 if actl.num_frames <= self.animation_written {
785 self.info.frame_control = None;
787 }
788 }
789 }
790
791 fn write_iend(&mut self) -> Result<()> {
792 self.iend_written = true;
793 self.write_chunk(chunk::IEND, &[])
794 }
795
796 fn should_skip_frame_control_on_default_image(&self) -> bool {
797 self.options.sep_def_img && self.images_written == 0
798 }
799
800 fn write_zlib_encoded_idat(&mut self, zlib_encoded: &[u8]) -> Result<()> {
801 for chunk in zlib_encoded.chunks(Self::MAX_IDAT_CHUNK_LEN as usize) {
802 self.write_chunk(chunk::IDAT, chunk)?;
803 }
804 Ok(())
805 }
806
807 pub fn set_filter(&mut self, filter: FilterType) {
813 self.options.filter = filter;
814 }
815
816 pub fn set_adaptive_filter(&mut self, adaptive_filter: AdaptiveFilterType) {
823 self.options.adaptive_filter = adaptive_filter;
824 }
825
826 pub fn set_frame_delay(&mut self, numerator: u16, denominator: u16) -> Result<()> {
836 if let Some(ref mut fctl) = self.info.frame_control {
837 fctl.delay_den = denominator;
838 fctl.delay_num = numerator;
839 Ok(())
840 } else {
841 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
842 }
843 }
844
845 pub fn set_frame_dimension(&mut self, width: u32, height: u32) -> Result<()> {
858 if let Some(ref mut fctl) = self.info.frame_control {
859 if Some(width) > self.info.width.checked_sub(fctl.x_offset)
860 || Some(height) > self.info.height.checked_sub(fctl.y_offset)
861 {
862 return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into()));
863 } else if width == 0 {
864 return Err(EncodingError::Format(FormatErrorKind::ZeroWidth.into()));
865 } else if height == 0 {
866 return Err(EncodingError::Format(FormatErrorKind::ZeroHeight.into()));
867 }
868 fctl.width = width;
869 fctl.height = height;
870 Ok(())
871 } else {
872 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
873 }
874 }
875
876 pub fn set_frame_position(&mut self, x: u32, y: u32) -> Result<()> {
887 if let Some(ref mut fctl) = self.info.frame_control {
888 if Some(x) > self.info.width.checked_sub(fctl.width)
889 || Some(y) > self.info.height.checked_sub(fctl.height)
890 {
891 return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into()));
892 }
893 fctl.x_offset = x;
894 fctl.y_offset = y;
895 Ok(())
896 } else {
897 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
898 }
899 }
900
901 pub fn reset_frame_dimension(&mut self) -> Result<()> {
911 if let Some(ref mut fctl) = self.info.frame_control {
912 fctl.width = self.info.width - fctl.x_offset;
913 fctl.height = self.info.height - fctl.y_offset;
914 Ok(())
915 } else {
916 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
917 }
918 }
919
920 pub fn reset_frame_position(&mut self) -> Result<()> {
928 if let Some(ref mut fctl) = self.info.frame_control {
929 fctl.x_offset = 0;
930 fctl.y_offset = 0;
931 Ok(())
932 } else {
933 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
934 }
935 }
936
937 pub fn set_blend_op(&mut self, op: BlendOp) -> Result<()> {
951 if let Some(ref mut fctl) = self.info.frame_control {
952 fctl.blend_op = op;
953 Ok(())
954 } else {
955 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
956 }
957 }
958
959 pub fn set_dispose_op(&mut self, op: DisposeOp) -> Result<()> {
971 if let Some(ref mut fctl) = self.info.frame_control {
972 fctl.dispose_op = op;
973 Ok(())
974 } else {
975 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
976 }
977 }
978
979 pub fn stream_writer(&mut self) -> Result<StreamWriter<W>> {
988 self.stream_writer_with_size(DEFAULT_BUFFER_LENGTH)
989 }
990
991 pub fn stream_writer_with_size(&mut self, size: usize) -> Result<StreamWriter<W>> {
997 StreamWriter::new(ChunkOutput::Borrowed(self), size)
998 }
999
1000 pub fn into_stream_writer(self) -> Result<StreamWriter<'static, W>> {
1008 self.into_stream_writer_with_size(DEFAULT_BUFFER_LENGTH)
1009 }
1010
1011 pub fn into_stream_writer_with_size(self, size: usize) -> Result<StreamWriter<'static, W>> {
1017 StreamWriter::new(ChunkOutput::Owned(self), size)
1018 }
1019
1020 pub fn finish(mut self) -> Result<()> {
1026 self.validate_sequence_done()?;
1027 self.write_iend()?;
1028 self.w.flush()?;
1029
1030 drop(self);
1032 Ok(())
1033 }
1034}
1035
1036impl<W: Write> Drop for Writer<W> {
1037 fn drop(&mut self) {
1038 if !self.iend_written {
1039 let _ = self.write_iend();
1040 }
1041 }
1042}
1043
1044enum ChunkOutput<'a, W: Write> {
1045 Borrowed(&'a mut Writer<W>),
1046 Owned(Writer<W>),
1047}
1048
1049impl<'a, W: Write> Deref for ChunkOutput<'a, W> {
1051 type Target = Writer<W>;
1052
1053 fn deref(&self) -> &Self::Target {
1054 match self {
1055 ChunkOutput::Borrowed(writer) => writer,
1056 ChunkOutput::Owned(writer) => writer,
1057 }
1058 }
1059}
1060
1061impl<'a, W: Write> DerefMut for ChunkOutput<'a, W> {
1062 fn deref_mut(&mut self) -> &mut Self::Target {
1063 match self {
1064 ChunkOutput::Borrowed(writer) => writer,
1065 ChunkOutput::Owned(writer) => writer,
1066 }
1067 }
1068}
1069
1070struct ChunkWriter<'a, W: Write> {
1085 writer: ChunkOutput<'a, W>,
1086 buffer: Vec<u8>,
1087 index: usize,
1089 curr_chunk: ChunkType,
1090}
1091
1092impl<'a, W: Write> ChunkWriter<'a, W> {
1093 fn new(writer: ChunkOutput<'a, W>, buf_len: usize) -> ChunkWriter<'a, W> {
1094 const CAP: usize = u32::MAX as usize >> 1;
1101 let curr_chunk = if writer.images_written == 0 {
1102 chunk::IDAT
1103 } else {
1104 chunk::fdAT
1105 };
1106 ChunkWriter {
1107 writer,
1108 buffer: vec![0; CAP.min(buf_len)],
1109 index: 0,
1110 curr_chunk,
1111 }
1112 }
1113
1114 fn next_frame_info(&self) -> (usize, usize) {
1121 let wrt = self.writer.deref();
1122
1123 let width: usize;
1124 let height: usize;
1125 if let Some(fctl) = wrt.info.frame_control {
1126 width = fctl.width as usize;
1127 height = fctl.height as usize;
1128 } else {
1129 width = wrt.info.width as usize;
1130 height = wrt.info.height as usize;
1131 }
1132
1133 let in_len = wrt.info.raw_row_length_from_width(width as u32) - 1;
1134 let data_size = in_len * height;
1135
1136 (in_len, data_size)
1137 }
1138
1139 fn write_header(&mut self) -> Result<()> {
1142 assert_eq!(self.index, 0, "Called when not flushed");
1143 let wrt = self.writer.deref_mut();
1144
1145 self.curr_chunk = if wrt.images_written == 0 {
1146 chunk::IDAT
1147 } else {
1148 chunk::fdAT
1149 };
1150
1151 match wrt.info.frame_control {
1152 Some(_) if wrt.should_skip_frame_control_on_default_image() => {}
1153 Some(ref mut fctl) => {
1154 fctl.encode(&mut wrt.w)?;
1155 fctl.sequence_number += 1;
1156 }
1157 _ => {}
1158 }
1159
1160 Ok(())
1161 }
1162
1163 fn set_fctl(&mut self, f: FrameControl) {
1168 if let Some(ref mut fctl) = self.writer.info.frame_control {
1169 *fctl = FrameControl {
1171 sequence_number: fctl.sequence_number,
1172 ..f
1173 };
1174 } else {
1175 panic!("This function must be called on an animated PNG")
1176 }
1177 }
1178
1179 fn flush_inner(&mut self) -> io::Result<()> {
1181 if self.index > 0 {
1182 write_chunk(
1184 &mut self.writer.w,
1185 self.curr_chunk,
1186 &self.buffer[..self.index],
1187 )?;
1188
1189 self.index = 0;
1190 }
1191 Ok(())
1192 }
1193}
1194
1195impl<'a, W: Write> Write for ChunkWriter<'a, W> {
1196 fn write(&mut self, mut data: &[u8]) -> io::Result<usize> {
1197 if data.is_empty() {
1198 return Ok(0);
1199 }
1200
1201 if self.index == 0 {
1203 let wrt = self.writer.deref_mut();
1204
1205 let no_fctl = wrt.should_skip_frame_control_on_default_image();
1207 if wrt.info.frame_control.is_some() && !no_fctl {
1208 let fctl = wrt.info.frame_control.as_mut().unwrap();
1209 self.buffer[0..4].copy_from_slice(&fctl.sequence_number.to_be_bytes());
1210 fctl.sequence_number += 1;
1211 self.index = 4;
1212 }
1213 }
1214
1215 let written = data.len().min(self.buffer.len() - self.index);
1218 data = &data[..written];
1219
1220 self.buffer[self.index..][..written].copy_from_slice(data);
1221 self.index += written;
1222
1223 if self.index == self.buffer.len() {
1225 self.flush_inner()?;
1226 }
1227
1228 Ok(written)
1229 }
1230
1231 fn flush(&mut self) -> io::Result<()> {
1232 self.flush_inner()
1233 }
1234}
1235
1236impl<W: Write> Drop for ChunkWriter<'_, W> {
1237 fn drop(&mut self) {
1238 let _ = self.flush();
1239 }
1240}
1241
1242enum Wrapper<'a, W: Write> {
1260 Chunk(ChunkWriter<'a, W>),
1261 Zlib(ZlibEncoder<ChunkWriter<'a, W>>),
1262 Unrecoverable,
1263 None,
1265}
1266
1267impl<'a, W: Write> Wrapper<'a, W> {
1268 fn take(&mut self) -> Wrapper<'a, W> {
1271 let mut swap = Wrapper::None;
1272 mem::swap(self, &mut swap);
1273 swap
1274 }
1275}
1276
1277pub struct StreamWriter<'a, W: Write> {
1285 writer: Wrapper<'a, W>,
1288 prev_buf: Vec<u8>,
1289 curr_buf: Vec<u8>,
1290 index: usize,
1292 line_len: usize,
1294 to_write: usize,
1296
1297 width: u32,
1298 height: u32,
1299
1300 bpp: BytesPerPixel,
1301 filter: FilterType,
1302 adaptive_filter: AdaptiveFilterType,
1303 fctl: Option<FrameControl>,
1304 compression: Compression,
1305}
1306
1307impl<'a, W: Write> StreamWriter<'a, W> {
1308 fn new(writer: ChunkOutput<'a, W>, buf_len: usize) -> Result<StreamWriter<'a, W>> {
1309 let PartialInfo {
1310 width,
1311 height,
1312 frame_control: fctl,
1313 compression,
1314 ..
1315 } = writer.info;
1316
1317 let bpp = writer.info.bpp_in_prediction();
1318 let in_len = writer.info.raw_row_length() - 1;
1319 let filter = writer.options.filter;
1320 let adaptive_filter = writer.options.adaptive_filter;
1321 let prev_buf = vec![0; in_len];
1322 let curr_buf = vec![0; in_len];
1323
1324 let mut chunk_writer = ChunkWriter::new(writer, buf_len);
1325 let (line_len, to_write) = chunk_writer.next_frame_info();
1326 chunk_writer.write_header()?;
1327 let zlib = ZlibEncoder::new(chunk_writer, compression.to_options());
1328
1329 Ok(StreamWriter {
1330 writer: Wrapper::Zlib(zlib),
1331 index: 0,
1332 prev_buf,
1333 curr_buf,
1334 bpp,
1335 filter,
1336 width,
1337 height,
1338 adaptive_filter,
1339 line_len,
1340 to_write,
1341 fctl,
1342 compression,
1343 })
1344 }
1345
1346 pub fn set_filter(&mut self, filter: FilterType) {
1355 self.filter = filter;
1356 }
1357
1358 pub fn set_adaptive_filter(&mut self, adaptive_filter: AdaptiveFilterType) {
1366 self.adaptive_filter = adaptive_filter;
1367 }
1368
1369 pub fn set_frame_delay(&mut self, numerator: u16, denominator: u16) -> Result<()> {
1379 if let Some(ref mut fctl) = self.fctl {
1380 fctl.delay_den = denominator;
1381 fctl.delay_num = numerator;
1382 Ok(())
1383 } else {
1384 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1385 }
1386 }
1387
1388 pub fn set_frame_dimension(&mut self, width: u32, height: u32) -> Result<()> {
1399 if let Some(ref mut fctl) = self.fctl {
1400 if Some(width) > self.width.checked_sub(fctl.x_offset)
1401 || Some(height) > self.height.checked_sub(fctl.y_offset)
1402 {
1403 return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into()));
1404 } else if width == 0 {
1405 return Err(EncodingError::Format(FormatErrorKind::ZeroWidth.into()));
1406 } else if height == 0 {
1407 return Err(EncodingError::Format(FormatErrorKind::ZeroHeight.into()));
1408 }
1409 fctl.width = width;
1410 fctl.height = height;
1411 Ok(())
1412 } else {
1413 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1414 }
1415 }
1416
1417 pub fn set_frame_position(&mut self, x: u32, y: u32) -> Result<()> {
1426 if let Some(ref mut fctl) = self.fctl {
1427 if Some(x) > self.width.checked_sub(fctl.width)
1428 || Some(y) > self.height.checked_sub(fctl.height)
1429 {
1430 return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into()));
1431 }
1432 fctl.x_offset = x;
1433 fctl.y_offset = y;
1434 Ok(())
1435 } else {
1436 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1437 }
1438 }
1439
1440 pub fn reset_frame_dimension(&mut self) -> Result<()> {
1450 if let Some(ref mut fctl) = self.fctl {
1451 fctl.width = self.width - fctl.x_offset;
1452 fctl.height = self.height - fctl.y_offset;
1453 Ok(())
1454 } else {
1455 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1456 }
1457 }
1458
1459 pub fn reset_frame_position(&mut self) -> Result<()> {
1467 if let Some(ref mut fctl) = self.fctl {
1468 fctl.x_offset = 0;
1469 fctl.y_offset = 0;
1470 Ok(())
1471 } else {
1472 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1473 }
1474 }
1475
1476 pub fn set_blend_op(&mut self, op: BlendOp) -> Result<()> {
1490 if let Some(ref mut fctl) = self.fctl {
1491 fctl.blend_op = op;
1492 Ok(())
1493 } else {
1494 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1495 }
1496 }
1497
1498 pub fn set_dispose_op(&mut self, op: DisposeOp) -> Result<()> {
1510 if let Some(ref mut fctl) = self.fctl {
1511 fctl.dispose_op = op;
1512 Ok(())
1513 } else {
1514 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1515 }
1516 }
1517
1518 pub fn finish(mut self) -> Result<()> {
1525 if self.to_write > 0 {
1526 let err = FormatErrorKind::MissingData(self.to_write).into();
1527 return Err(EncodingError::Format(err));
1528 }
1529
1530 self.flush()?;
1532
1533 if let Wrapper::Chunk(wrt) = self.writer.take() {
1534 wrt.writer.validate_sequence_done()?;
1535 }
1536
1537 Ok(())
1538 }
1539
1540 fn new_frame(&mut self) -> Result<()> {
1545 let wrt = match &mut self.writer {
1546 Wrapper::Chunk(wrt) => wrt,
1547 Wrapper::Unrecoverable => {
1548 let err = FormatErrorKind::Unrecoverable.into();
1549 return Err(EncodingError::Format(err));
1550 }
1551 Wrapper::Zlib(_) => unreachable!("never called on a half-finished frame"),
1552 Wrapper::None => unreachable!(),
1553 };
1554 wrt.flush()?;
1555 wrt.writer.validate_new_image()?;
1556
1557 if let Some(fctl) = self.fctl {
1558 wrt.set_fctl(fctl);
1559 }
1560 let (scansize, size) = wrt.next_frame_info();
1561 self.line_len = scansize;
1562 self.to_write = size;
1563
1564 wrt.write_header()?;
1565 wrt.writer.increment_images_written();
1566
1567 match self.writer.take() {
1569 Wrapper::Chunk(wrt) => {
1570 let encoder = ZlibEncoder::new(wrt, self.compression.to_options());
1571 self.writer = Wrapper::Zlib(encoder);
1572 }
1573 _ => unreachable!(),
1574 };
1575
1576 Ok(())
1577 }
1578}
1579
1580impl<'a, W: Write> Write for StreamWriter<'a, W> {
1581 fn write(&mut self, mut data: &[u8]) -> io::Result<usize> {
1582 if let Wrapper::Unrecoverable = self.writer {
1583 let err = FormatErrorKind::Unrecoverable.into();
1584 return Err(EncodingError::Format(err).into());
1585 }
1586
1587 if data.is_empty() {
1588 return Ok(0);
1589 }
1590
1591 if self.to_write == 0 {
1592 match self.writer.take() {
1593 Wrapper::Zlib(wrt) => match wrt.finish() {
1594 Ok(chunk) => self.writer = Wrapper::Chunk(chunk),
1595 Err(err) => {
1596 self.writer = Wrapper::Unrecoverable;
1597 return Err(err);
1598 }
1599 },
1600 chunk @ Wrapper::Chunk(_) => self.writer = chunk,
1601 Wrapper::Unrecoverable => unreachable!(),
1602 Wrapper::None => unreachable!(),
1603 };
1604
1605 self.new_frame()?;
1607 }
1608
1609 let written = data.read(&mut self.curr_buf[..self.line_len][self.index..])?;
1610 self.index += written;
1611 self.to_write -= written;
1612
1613 if self.index == self.line_len {
1614 let mut filtered = vec![0; self.curr_buf.len()];
1616 let filter_type = filter(
1617 self.filter,
1618 self.adaptive_filter,
1619 self.bpp,
1620 &self.prev_buf,
1621 &self.curr_buf,
1622 &mut filtered,
1623 );
1624 let wrt = match &mut self.writer {
1626 Wrapper::Zlib(wrt) => wrt,
1627 _ => unreachable!(),
1628 };
1629
1630 wrt.write_all(&[filter_type as u8])?;
1631 wrt.write_all(&filtered)?;
1632 mem::swap(&mut self.prev_buf, &mut self.curr_buf);
1633 self.index = 0;
1634 }
1635
1636 Ok(written)
1637 }
1638
1639 fn flush(&mut self) -> io::Result<()> {
1640 match &mut self.writer {
1641 Wrapper::Zlib(wrt) => wrt.flush()?,
1642 Wrapper::Chunk(wrt) => wrt.flush()?,
1643 Wrapper::Unrecoverable | Wrapper::None => {
1646 let err = FormatErrorKind::Unrecoverable.into();
1647 return Err(EncodingError::Format(err).into());
1648 }
1649 }
1650
1651 if self.index > 0 {
1652 let err = FormatErrorKind::WrittenTooMuch(self.index).into();
1653 return Err(EncodingError::Format(err).into());
1654 }
1655
1656 Ok(())
1657 }
1658}
1659
1660impl<W: Write> Drop for StreamWriter<'_, W> {
1661 fn drop(&mut self) {
1662 let _ = self.flush();
1663 }
1664}
1665
1666impl Compression {
1671 fn to_options(self) -> flate2::Compression {
1672 #[allow(deprecated)]
1673 match self {
1674 Compression::Default => flate2::Compression::default(),
1675 Compression::Fast => flate2::Compression::fast(),
1676 Compression::Best => flate2::Compression::best(),
1677 #[allow(deprecated)]
1678 Compression::Huffman => flate2::Compression::none(),
1679 #[allow(deprecated)]
1680 Compression::Rle => flate2::Compression::none(),
1681 }
1682 }
1683}
1684
1685#[cfg(test)]
1686mod tests {
1687 use super::*;
1688 use crate::Decoder;
1689
1690 use rand::{thread_rng, Rng};
1691 use std::cmp;
1692 use std::fs::File;
1693 use std::io::Cursor;
1694
1695 #[test]
1696 fn roundtrip() {
1697 for _ in 0..10 {
1699 for path in glob::glob("tests/pngsuite/*.png")
1700 .unwrap()
1701 .map(|r| r.unwrap())
1702 {
1703 if path.file_name().unwrap().to_str().unwrap().starts_with('x') {
1704 continue;
1706 }
1707 eprintln!("{}", path.display());
1708 let decoder = Decoder::new(File::open(path).unwrap());
1710 let mut reader = decoder.read_info().unwrap();
1711 let mut buf = vec![0; reader.output_buffer_size()];
1712 let info = reader.next_frame(&mut buf).unwrap();
1713 let mut out = Vec::new();
1715 {
1716 let mut wrapper = RandomChunkWriter {
1717 rng: thread_rng(),
1718 w: &mut out,
1719 };
1720
1721 let mut encoder = Encoder::new(&mut wrapper, info.width, info.height);
1722 encoder.set_color(info.color_type);
1723 encoder.set_depth(info.bit_depth);
1724 if let Some(palette) = &reader.info().palette {
1725 encoder.set_palette(palette.clone());
1726 }
1727 let mut encoder = encoder.write_header().unwrap();
1728 encoder.write_image_data(&buf).unwrap();
1729 }
1730 let decoder = Decoder::new(&*out);
1732 let mut reader = decoder.read_info().unwrap();
1733 let mut buf2 = vec![0; reader.output_buffer_size()];
1734 reader.next_frame(&mut buf2).unwrap();
1735 assert_eq!(buf, buf2);
1737 }
1738 }
1739 }
1740
1741 #[test]
1742 fn roundtrip_stream() {
1743 for _ in 0..10 {
1745 for path in glob::glob("tests/pngsuite/*.png")
1746 .unwrap()
1747 .map(|r| r.unwrap())
1748 {
1749 if path.file_name().unwrap().to_str().unwrap().starts_with('x') {
1750 continue;
1752 }
1753 let decoder = Decoder::new(File::open(path).unwrap());
1755 let mut reader = decoder.read_info().unwrap();
1756 let mut buf = vec![0; reader.output_buffer_size()];
1757 let info = reader.next_frame(&mut buf).unwrap();
1758 let mut out = Vec::new();
1760 {
1761 let mut wrapper = RandomChunkWriter {
1762 rng: thread_rng(),
1763 w: &mut out,
1764 };
1765
1766 let mut encoder = Encoder::new(&mut wrapper, info.width, info.height);
1767 encoder.set_color(info.color_type);
1768 encoder.set_depth(info.bit_depth);
1769 if let Some(palette) = &reader.info().palette {
1770 encoder.set_palette(palette.clone());
1771 }
1772 let mut encoder = encoder.write_header().unwrap();
1773 let mut stream_writer = encoder.stream_writer().unwrap();
1774
1775 let mut outer_wrapper = RandomChunkWriter {
1776 rng: thread_rng(),
1777 w: &mut stream_writer,
1778 };
1779
1780 outer_wrapper.write_all(&buf).unwrap();
1781 }
1782 let decoder = Decoder::new(&*out);
1784 let mut reader = decoder.read_info().unwrap();
1785 let mut buf2 = vec![0; reader.output_buffer_size()];
1786 reader.next_frame(&mut buf2).unwrap();
1787 assert_eq!(buf, buf2);
1789 }
1790 }
1791 }
1792
1793 #[test]
1794 fn image_palette() -> Result<()> {
1795 for &bit_depth in &[1u8, 2, 4, 8] {
1796 let path = format!("tests/pngsuite/basn3p0{}.png", bit_depth);
1798 let decoder = Decoder::new(File::open(&path).unwrap());
1799 let mut reader = decoder.read_info().unwrap();
1800
1801 let mut decoded_pixels = vec![0; reader.output_buffer_size()];
1802 let info = reader.info();
1803 assert_eq!(
1804 info.width as usize * info.height as usize * usize::from(bit_depth),
1805 decoded_pixels.len() * 8
1806 );
1807 let info = reader.next_frame(&mut decoded_pixels).unwrap();
1808 let indexed_data = decoded_pixels;
1809
1810 let palette = reader.info().palette.as_ref().unwrap();
1811 let mut out = Vec::new();
1812 {
1813 let mut encoder = Encoder::new(&mut out, info.width, info.height);
1814 encoder.set_depth(BitDepth::from_u8(bit_depth).unwrap());
1815 encoder.set_color(ColorType::Indexed);
1816 encoder.set_palette(palette.as_ref());
1817
1818 let mut writer = encoder.write_header().unwrap();
1819 writer.write_image_data(&indexed_data).unwrap();
1820 }
1821
1822 let decoder = Decoder::new(&*out);
1824 let mut reader = decoder.read_info().unwrap();
1825 let mut redecoded = vec![0; reader.output_buffer_size()];
1826 reader.next_frame(&mut redecoded).unwrap();
1827 assert_eq!(indexed_data, redecoded);
1829 }
1830 Ok(())
1831 }
1832
1833 #[test]
1834 fn expect_error_on_wrong_image_len() -> Result<()> {
1835 let width = 10;
1836 let height = 10;
1837
1838 let output = vec![0u8; 1024];
1839 let writer = Cursor::new(output);
1840 let mut encoder = Encoder::new(writer, width as u32, height as u32);
1841 encoder.set_depth(BitDepth::Eight);
1842 encoder.set_color(ColorType::Rgb);
1843 let mut png_writer = encoder.write_header()?;
1844
1845 let correct_image_size = width * height * 3;
1846 let image = vec![0u8; correct_image_size + 1];
1847 let result = png_writer.write_image_data(image.as_ref());
1848 assert!(result.is_err());
1849
1850 Ok(())
1851 }
1852
1853 #[test]
1854 fn expect_error_on_empty_image() -> Result<()> {
1855 let output = vec![0u8; 1024];
1856 let mut writer = Cursor::new(output);
1857
1858 let encoder = Encoder::new(&mut writer, 0, 0);
1859 assert!(encoder.write_header().is_err());
1860
1861 let encoder = Encoder::new(&mut writer, 100, 0);
1862 assert!(encoder.write_header().is_err());
1863
1864 let encoder = Encoder::new(&mut writer, 0, 100);
1865 assert!(encoder.write_header().is_err());
1866
1867 Ok(())
1868 }
1869
1870 #[test]
1871 fn expect_error_on_invalid_bit_depth_color_type_combination() -> Result<()> {
1872 let output = vec![0u8; 1024];
1873 let mut writer = Cursor::new(output);
1874
1875 let mut encoder = Encoder::new(&mut writer, 1, 1);
1876 encoder.set_depth(BitDepth::One);
1877 encoder.set_color(ColorType::Rgb);
1878 assert!(encoder.write_header().is_err());
1879
1880 let mut encoder = Encoder::new(&mut writer, 1, 1);
1881 encoder.set_depth(BitDepth::One);
1882 encoder.set_color(ColorType::GrayscaleAlpha);
1883 assert!(encoder.write_header().is_err());
1884
1885 let mut encoder = Encoder::new(&mut writer, 1, 1);
1886 encoder.set_depth(BitDepth::One);
1887 encoder.set_color(ColorType::Rgba);
1888 assert!(encoder.write_header().is_err());
1889
1890 let mut encoder = Encoder::new(&mut writer, 1, 1);
1891 encoder.set_depth(BitDepth::Two);
1892 encoder.set_color(ColorType::Rgb);
1893 assert!(encoder.write_header().is_err());
1894
1895 let mut encoder = Encoder::new(&mut writer, 1, 1);
1896 encoder.set_depth(BitDepth::Two);
1897 encoder.set_color(ColorType::GrayscaleAlpha);
1898 assert!(encoder.write_header().is_err());
1899
1900 let mut encoder = Encoder::new(&mut writer, 1, 1);
1901 encoder.set_depth(BitDepth::Two);
1902 encoder.set_color(ColorType::Rgba);
1903 assert!(encoder.write_header().is_err());
1904
1905 let mut encoder = Encoder::new(&mut writer, 1, 1);
1906 encoder.set_depth(BitDepth::Four);
1907 encoder.set_color(ColorType::Rgb);
1908 assert!(encoder.write_header().is_err());
1909
1910 let mut encoder = Encoder::new(&mut writer, 1, 1);
1911 encoder.set_depth(BitDepth::Four);
1912 encoder.set_color(ColorType::GrayscaleAlpha);
1913 assert!(encoder.write_header().is_err());
1914
1915 let mut encoder = Encoder::new(&mut writer, 1, 1);
1916 encoder.set_depth(BitDepth::Four);
1917 encoder.set_color(ColorType::Rgba);
1918 assert!(encoder.write_header().is_err());
1919
1920 let mut encoder = Encoder::new(&mut writer, 1, 1);
1921 encoder.set_depth(BitDepth::Sixteen);
1922 encoder.set_color(ColorType::Indexed);
1923 assert!(encoder.write_header().is_err());
1924
1925 Ok(())
1926 }
1927
1928 #[test]
1929 fn can_write_header_with_valid_bit_depth_color_type_combination() -> Result<()> {
1930 let output = vec![0u8; 1024];
1931 let mut writer = Cursor::new(output);
1932
1933 let mut encoder = Encoder::new(&mut writer, 1, 1);
1934 encoder.set_depth(BitDepth::One);
1935 encoder.set_color(ColorType::Grayscale);
1936 assert!(encoder.write_header().is_ok());
1937
1938 let mut encoder = Encoder::new(&mut writer, 1, 1);
1939 encoder.set_depth(BitDepth::One);
1940 encoder.set_color(ColorType::Indexed);
1941 assert!(encoder.write_header().is_ok());
1942
1943 let mut encoder = Encoder::new(&mut writer, 1, 1);
1944 encoder.set_depth(BitDepth::Two);
1945 encoder.set_color(ColorType::Grayscale);
1946 assert!(encoder.write_header().is_ok());
1947
1948 let mut encoder = Encoder::new(&mut writer, 1, 1);
1949 encoder.set_depth(BitDepth::Two);
1950 encoder.set_color(ColorType::Indexed);
1951 assert!(encoder.write_header().is_ok());
1952
1953 let mut encoder = Encoder::new(&mut writer, 1, 1);
1954 encoder.set_depth(BitDepth::Four);
1955 encoder.set_color(ColorType::Grayscale);
1956 assert!(encoder.write_header().is_ok());
1957
1958 let mut encoder = Encoder::new(&mut writer, 1, 1);
1959 encoder.set_depth(BitDepth::Four);
1960 encoder.set_color(ColorType::Indexed);
1961 assert!(encoder.write_header().is_ok());
1962
1963 let mut encoder = Encoder::new(&mut writer, 1, 1);
1964 encoder.set_depth(BitDepth::Eight);
1965 encoder.set_color(ColorType::Grayscale);
1966 assert!(encoder.write_header().is_ok());
1967
1968 let mut encoder = Encoder::new(&mut writer, 1, 1);
1969 encoder.set_depth(BitDepth::Eight);
1970 encoder.set_color(ColorType::Rgb);
1971 assert!(encoder.write_header().is_ok());
1972
1973 let mut encoder = Encoder::new(&mut writer, 1, 1);
1974 encoder.set_depth(BitDepth::Eight);
1975 encoder.set_color(ColorType::Indexed);
1976 assert!(encoder.write_header().is_ok());
1977
1978 let mut encoder = Encoder::new(&mut writer, 1, 1);
1979 encoder.set_depth(BitDepth::Eight);
1980 encoder.set_color(ColorType::GrayscaleAlpha);
1981 assert!(encoder.write_header().is_ok());
1982
1983 let mut encoder = Encoder::new(&mut writer, 1, 1);
1984 encoder.set_depth(BitDepth::Eight);
1985 encoder.set_color(ColorType::Rgba);
1986 assert!(encoder.write_header().is_ok());
1987
1988 let mut encoder = Encoder::new(&mut writer, 1, 1);
1989 encoder.set_depth(BitDepth::Sixteen);
1990 encoder.set_color(ColorType::Grayscale);
1991 assert!(encoder.write_header().is_ok());
1992
1993 let mut encoder = Encoder::new(&mut writer, 1, 1);
1994 encoder.set_depth(BitDepth::Sixteen);
1995 encoder.set_color(ColorType::Rgb);
1996 assert!(encoder.write_header().is_ok());
1997
1998 let mut encoder = Encoder::new(&mut writer, 1, 1);
1999 encoder.set_depth(BitDepth::Sixteen);
2000 encoder.set_color(ColorType::GrayscaleAlpha);
2001 assert!(encoder.write_header().is_ok());
2002
2003 let mut encoder = Encoder::new(&mut writer, 1, 1);
2004 encoder.set_depth(BitDepth::Sixteen);
2005 encoder.set_color(ColorType::Rgba);
2006 assert!(encoder.write_header().is_ok());
2007
2008 Ok(())
2009 }
2010
2011 #[test]
2012 fn all_filters_roundtrip() -> io::Result<()> {
2013 let pixel: Vec<_> = (0..48).collect();
2014
2015 let roundtrip = |filter: FilterType| -> io::Result<()> {
2016 let mut buffer = vec![];
2017 let mut encoder = Encoder::new(&mut buffer, 4, 4);
2018 encoder.set_depth(BitDepth::Eight);
2019 encoder.set_color(ColorType::Rgb);
2020 encoder.set_filter(filter);
2021 encoder.write_header()?.write_image_data(&pixel)?;
2022
2023 let decoder = crate::Decoder::new(Cursor::new(buffer));
2024 let mut reader = decoder.read_info()?;
2025 let info = reader.info();
2026 assert_eq!(info.width, 4);
2027 assert_eq!(info.height, 4);
2028 let mut dest = vec![0; pixel.len()];
2029 reader.next_frame(&mut dest)?;
2030 assert_eq!(dest, pixel, "Deviation with filter type {:?}", filter);
2031
2032 Ok(())
2033 };
2034
2035 roundtrip(FilterType::NoFilter)?;
2036 roundtrip(FilterType::Sub)?;
2037 roundtrip(FilterType::Up)?;
2038 roundtrip(FilterType::Avg)?;
2039 roundtrip(FilterType::Paeth)?;
2040
2041 Ok(())
2042 }
2043
2044 #[test]
2045 fn some_gamma_roundtrip() -> io::Result<()> {
2046 let pixel: Vec<_> = (0..48).collect();
2047
2048 let roundtrip = |gamma: Option<ScaledFloat>| -> io::Result<()> {
2049 let mut buffer = vec![];
2050 let mut encoder = Encoder::new(&mut buffer, 4, 4);
2051 encoder.set_depth(BitDepth::Eight);
2052 encoder.set_color(ColorType::Rgb);
2053 encoder.set_filter(FilterType::Avg);
2054 if let Some(gamma) = gamma {
2055 encoder.set_source_gamma(gamma);
2056 }
2057 encoder.write_header()?.write_image_data(&pixel)?;
2058
2059 let decoder = crate::Decoder::new(Cursor::new(buffer));
2060 let mut reader = decoder.read_info()?;
2061 assert_eq!(
2062 reader.info().source_gamma,
2063 gamma,
2064 "Deviation with gamma {:?}",
2065 gamma
2066 );
2067 let mut dest = vec![0; pixel.len()];
2068 let info = reader.next_frame(&mut dest)?;
2069 assert_eq!(info.width, 4);
2070 assert_eq!(info.height, 4);
2071
2072 Ok(())
2073 };
2074
2075 roundtrip(None)?;
2076 roundtrip(Some(ScaledFloat::new(0.35)))?;
2077 roundtrip(Some(ScaledFloat::new(0.45)))?;
2078 roundtrip(Some(ScaledFloat::new(0.55)))?;
2079 roundtrip(Some(ScaledFloat::new(0.7)))?;
2080 roundtrip(Some(ScaledFloat::new(1.0)))?;
2081 roundtrip(Some(ScaledFloat::new(2.5)))?;
2082
2083 Ok(())
2084 }
2085
2086 #[test]
2087 fn write_image_chunks_beyond_first() -> Result<()> {
2088 let width = 10;
2089 let height = 10;
2090
2091 let output = vec![0u8; 1024];
2092 let writer = Cursor::new(output);
2093
2094 let mut encoder = Encoder::new(writer, width, height);
2098 encoder.set_depth(BitDepth::Eight);
2099 encoder.set_color(ColorType::Grayscale);
2100 let mut png_writer = encoder.write_header()?;
2101
2102 for _ in 0..3 {
2103 let correct_image_size = (width * height) as usize;
2104 let image = vec![0u8; correct_image_size];
2105 png_writer.write_image_data(image.as_ref())?;
2106 }
2107
2108 Ok(())
2109 }
2110
2111 #[test]
2112 fn image_validate_sequence_without_animation() -> Result<()> {
2113 let width = 10;
2114 let height = 10;
2115
2116 let output = vec![0u8; 1024];
2117 let writer = Cursor::new(output);
2118
2119 let mut encoder = Encoder::new(writer, width, height);
2120 encoder.set_depth(BitDepth::Eight);
2121 encoder.set_color(ColorType::Grayscale);
2122 encoder.validate_sequence(true);
2123 let mut png_writer = encoder.write_header()?;
2124
2125 let correct_image_size = (width * height) as usize;
2126 let image = vec![0u8; correct_image_size];
2127 png_writer.write_image_data(image.as_ref())?;
2128
2129 assert!(png_writer.write_image_data(image.as_ref()).is_err());
2130 Ok(())
2131 }
2132
2133 #[test]
2134 fn image_validate_animation() -> Result<()> {
2135 let width = 10;
2136 let height = 10;
2137
2138 let output = vec![0u8; 1024];
2139 let writer = Cursor::new(output);
2140 let correct_image_size = (width * height) as usize;
2141 let image = vec![0u8; correct_image_size];
2142
2143 let mut encoder = Encoder::new(writer, width, height);
2144 encoder.set_depth(BitDepth::Eight);
2145 encoder.set_color(ColorType::Grayscale);
2146 encoder.set_animated(1, 0)?;
2147 encoder.validate_sequence(true);
2148 let mut png_writer = encoder.write_header()?;
2149
2150 png_writer.write_image_data(image.as_ref())?;
2151
2152 Ok(())
2153 }
2154
2155 #[test]
2156 fn image_validate_animation2() -> Result<()> {
2157 let width = 10;
2158 let height = 10;
2159
2160 let output = vec![0u8; 1024];
2161 let writer = Cursor::new(output);
2162 let correct_image_size = (width * height) as usize;
2163 let image = vec![0u8; correct_image_size];
2164
2165 let mut encoder = Encoder::new(writer, width, height);
2166 encoder.set_depth(BitDepth::Eight);
2167 encoder.set_color(ColorType::Grayscale);
2168 encoder.set_animated(2, 0)?;
2169 encoder.validate_sequence(true);
2170 let mut png_writer = encoder.write_header()?;
2171
2172 png_writer.write_image_data(image.as_ref())?;
2173 png_writer.write_image_data(image.as_ref())?;
2174 png_writer.finish()?;
2175
2176 Ok(())
2177 }
2178
2179 #[test]
2180 fn image_validate_animation_sep_def_image() -> Result<()> {
2181 let width = 10;
2182 let height = 10;
2183
2184 let output = vec![0u8; 1024];
2185 let writer = Cursor::new(output);
2186 let correct_image_size = (width * height) as usize;
2187 let image = vec![0u8; correct_image_size];
2188
2189 let mut encoder = Encoder::new(writer, width, height);
2190 encoder.set_depth(BitDepth::Eight);
2191 encoder.set_color(ColorType::Grayscale);
2192 encoder.set_animated(1, 0)?;
2193 encoder.set_sep_def_img(true)?;
2194 encoder.validate_sequence(true);
2195 let mut png_writer = encoder.write_header()?;
2196
2197 png_writer.write_image_data(image.as_ref())?;
2198 png_writer.write_image_data(image.as_ref())?;
2199 png_writer.finish()?;
2200
2201 Ok(())
2202 }
2203
2204 #[test]
2205 fn image_validate_missing_image() -> Result<()> {
2206 let width = 10;
2207 let height = 10;
2208
2209 let output = vec![0u8; 1024];
2210 let writer = Cursor::new(output);
2211
2212 let mut encoder = Encoder::new(writer, width, height);
2213 encoder.set_depth(BitDepth::Eight);
2214 encoder.set_color(ColorType::Grayscale);
2215 encoder.validate_sequence(true);
2216 let png_writer = encoder.write_header()?;
2217
2218 assert!(png_writer.finish().is_err());
2219 Ok(())
2220 }
2221
2222 #[test]
2223 fn image_validate_missing_animated_frame() -> Result<()> {
2224 let width = 10;
2225 let height = 10;
2226
2227 let output = vec![0u8; 1024];
2228 let writer = Cursor::new(output);
2229 let correct_image_size = (width * height) as usize;
2230 let image = vec![0u8; correct_image_size];
2231
2232 let mut encoder = Encoder::new(writer, width, height);
2233 encoder.set_depth(BitDepth::Eight);
2234 encoder.set_color(ColorType::Grayscale);
2235 encoder.set_animated(2, 0)?;
2236 encoder.validate_sequence(true);
2237 let mut png_writer = encoder.write_header()?;
2238
2239 png_writer.write_image_data(image.as_ref())?;
2240 assert!(png_writer.finish().is_err());
2241
2242 Ok(())
2243 }
2244
2245 #[test]
2246 fn issue_307_stream_validation() -> Result<()> {
2247 let output = vec![0u8; 1024];
2248 let mut cursor = Cursor::new(output);
2249
2250 let encoder = Encoder::new(&mut cursor, 1, 1); let mut writer = encoder.write_header()?;
2252 let mut stream = writer.stream_writer()?;
2253
2254 let written = stream.write(&[1, 2, 3, 4])?;
2255 assert_eq!(written, 1);
2256 stream.finish()?;
2257 drop(writer);
2258
2259 {
2260 cursor.set_position(0);
2261 let mut decoder = Decoder::new(cursor).read_info().expect("A valid image");
2262 let mut buffer = [0u8; 1];
2263 decoder.next_frame(&mut buffer[..]).expect("Valid read");
2264 assert_eq!(buffer, [1]);
2265 }
2266
2267 Ok(())
2268 }
2269
2270 #[test]
2271 fn stream_filtering() -> Result<()> {
2272 let output = vec![0u8; 1024];
2273 let mut cursor = Cursor::new(output);
2274
2275 let mut encoder = Encoder::new(&mut cursor, 8, 8);
2276 encoder.set_color(ColorType::Rgba);
2277 encoder.set_filter(FilterType::Paeth);
2278 let mut writer = encoder.write_header()?;
2279 let mut stream = writer.stream_writer()?;
2280
2281 for _ in 0..8 {
2282 let written = stream.write(&[1; 32])?;
2283 assert_eq!(written, 32);
2284 }
2285 stream.finish()?;
2286 drop(writer);
2287
2288 {
2289 cursor.set_position(0);
2290 let mut decoder = Decoder::new(cursor).read_info().expect("A valid image");
2291 let mut buffer = [0u8; 256];
2292 decoder.next_frame(&mut buffer[..]).expect("Valid read");
2293 assert_eq!(buffer, [1; 256]);
2294 }
2295
2296 Ok(())
2297 }
2298
2299 #[test]
2300 #[cfg(all(unix, not(target_pointer_width = "32")))]
2301 fn exper_error_on_huge_chunk() -> Result<()> {
2302 let empty = vec![0; 1usize << 31];
2305 let writer = Cursor::new(vec![0u8; 1024]);
2306
2307 let mut encoder = Encoder::new(writer, 10, 10);
2308 encoder.set_depth(BitDepth::Eight);
2309 encoder.set_color(ColorType::Grayscale);
2310 let mut png_writer = encoder.write_header()?;
2311
2312 assert!(png_writer.write_chunk(chunk::fdAT, &empty).is_err());
2313 Ok(())
2314 }
2315
2316 #[test]
2317 #[cfg(all(unix, not(target_pointer_width = "32")))]
2318 fn exper_error_on_non_u32_chunk() -> Result<()> {
2319 let empty = vec![0; 1usize << 32];
2322 let writer = Cursor::new(vec![0u8; 1024]);
2323
2324 let mut encoder = Encoder::new(writer, 10, 10);
2325 encoder.set_depth(BitDepth::Eight);
2326 encoder.set_color(ColorType::Grayscale);
2327 let mut png_writer = encoder.write_header()?;
2328
2329 assert!(png_writer.write_chunk(chunk::fdAT, &empty).is_err());
2330 Ok(())
2331 }
2332
2333 #[test]
2334 fn finish_drops_inner_writer() -> Result<()> {
2335 struct NoWriter<'flag>(&'flag mut bool);
2336
2337 impl Write for NoWriter<'_> {
2338 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
2339 Ok(buf.len())
2340 }
2341 fn flush(&mut self) -> io::Result<()> {
2342 Ok(())
2343 }
2344 }
2345 impl Drop for NoWriter<'_> {
2346 fn drop(&mut self) {
2347 *self.0 = true;
2348 }
2349 }
2350
2351 let mut flag = false;
2352
2353 {
2354 let mut encoder = Encoder::new(NoWriter(&mut flag), 10, 10);
2355 encoder.set_depth(BitDepth::Eight);
2356 encoder.set_color(ColorType::Grayscale);
2357
2358 let mut writer = encoder.write_header()?;
2359 writer.write_image_data(&[0; 100])?;
2360 writer.finish()?;
2361 }
2362
2363 assert!(flag, "PNG finished but writer was not dropped");
2364 Ok(())
2365 }
2366
2367 struct RandomChunkWriter<R: Rng, W: Write> {
2369 rng: R,
2370 w: W,
2371 }
2372
2373 impl<R: Rng, W: Write> Write for RandomChunkWriter<R, W> {
2374 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
2375 let len = cmp::min(self.rng.gen_range(1..50), buf.len());
2377
2378 self.w.write(&buf[0..len])
2379 }
2380
2381 fn flush(&mut self) -> io::Result<()> {
2382 self.w.flush()
2383 }
2384 }
2385}