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 encode_iso_8859_1, 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 #[doc(hidden)]
273 #[deprecated(note = "use set_source_srgb")]
274 pub fn set_srgb(&mut self, rendering_intent: super::SrgbRenderingIntent) {
275 self.info.set_source_srgb(rendering_intent);
276 self.info.source_gamma = Some(crate::srgb::substitute_gamma());
277 self.info.source_chromaticities = Some(crate::srgb::substitute_chromaticities());
278 }
279
280 pub fn set_source_srgb(&mut self, rendering_intent: super::SrgbRenderingIntent) {
287 self.info.set_source_srgb(rendering_intent);
288 }
289
290 pub fn write_header(self) -> Result<Writer<W>> {
294 Writer::new(self.w, PartialInfo::new(&self.info), self.options).init(&self.info)
295 }
296
297 pub fn set_color(&mut self, color: ColorType) {
303 self.info.color_type = color;
304 }
305
306 pub fn set_depth(&mut self, depth: BitDepth) {
308 self.info.bit_depth = depth;
309 }
310
311 pub fn set_compression(&mut self, compression: Compression) {
316 self.info.compression = compression;
317 }
318
319 pub fn set_filter(&mut self, filter: FilterType) {
325 self.options.filter = filter;
326 }
327
328 pub fn set_adaptive_filter(&mut self, adaptive_filter: AdaptiveFilterType) {
335 self.options.adaptive_filter = adaptive_filter;
336 }
337
338 pub fn set_frame_delay(&mut self, numerator: u16, denominator: u16) -> Result<()> {
356 if let Some(ref mut fctl) = self.info.frame_control {
357 fctl.delay_den = denominator;
358 fctl.delay_num = numerator;
359 Ok(())
360 } else {
361 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
362 }
363 }
364
365 pub fn set_blend_op(&mut self, op: BlendOp) -> Result<()> {
388 if let Some(ref mut fctl) = self.info.frame_control {
389 fctl.blend_op = op;
390 Ok(())
391 } else {
392 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
393 }
394 }
395
396 pub fn set_dispose_op(&mut self, op: DisposeOp) -> Result<()> {
417 if let Some(ref mut fctl) = self.info.frame_control {
418 fctl.dispose_op = op;
419 Ok(())
420 } else {
421 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
422 }
423 }
424 pub fn set_pixel_dims(&mut self, pixel_dims: Option<PixelDimensions>) {
425 self.info.pixel_dims = pixel_dims
426 }
427 pub fn add_text_chunk(&mut self, keyword: String, text: String) -> Result<()> {
429 let text_chunk = TEXtChunk::new(keyword, text);
430 self.info.uncompressed_latin1_text.push(text_chunk);
431 Ok(())
432 }
433
434 pub fn add_ztxt_chunk(&mut self, keyword: String, text: String) -> Result<()> {
436 let text_chunk = ZTXtChunk::new(keyword, text);
437 self.info.compressed_latin1_text.push(text_chunk);
438 Ok(())
439 }
440
441 pub fn add_itxt_chunk(&mut self, keyword: String, text: String) -> Result<()> {
446 let text_chunk = ITXtChunk::new(keyword, text);
447 self.info.utf8_text.push(text_chunk);
448 Ok(())
449 }
450
451 pub fn validate_sequence(&mut self, validate: bool) {
462 self.options.validate_sequence = validate;
463 }
464}
465
466pub struct Writer<W: Write> {
475 w: W,
477 info: PartialInfo,
479 options: Options,
481 images_written: u64,
483 animation_written: u32,
485 iend_written: bool,
488}
489
490struct PartialInfo {
492 width: u32,
493 height: u32,
494 bit_depth: BitDepth,
495 color_type: ColorType,
496 frame_control: Option<FrameControl>,
497 animation_control: Option<AnimationControl>,
498 compression: Compression,
499 has_palette: bool,
500}
501
502impl PartialInfo {
503 fn new(info: &Info) -> Self {
504 PartialInfo {
505 width: info.width,
506 height: info.height,
507 bit_depth: info.bit_depth,
508 color_type: info.color_type,
509 frame_control: info.frame_control,
510 animation_control: info.animation_control,
511 compression: info.compression,
512 has_palette: info.palette.is_some(),
513 }
514 }
515
516 fn bpp_in_prediction(&self) -> BytesPerPixel {
517 self.to_info().bpp_in_prediction()
519 }
520
521 fn raw_row_length(&self) -> usize {
522 self.to_info().raw_row_length()
524 }
525
526 fn raw_row_length_from_width(&self, width: u32) -> usize {
527 self.to_info().raw_row_length_from_width(width)
529 }
530
531 fn to_info(&self) -> Info<'static> {
534 Info {
535 width: self.width,
536 height: self.height,
537 bit_depth: self.bit_depth,
538 color_type: self.color_type,
539 frame_control: self.frame_control,
540 animation_control: self.animation_control,
541 compression: self.compression,
542 ..Default::default()
543 }
544 }
545}
546
547const DEFAULT_BUFFER_LENGTH: usize = 4 * 1024;
548
549pub(crate) fn write_chunk<W: Write>(mut w: W, name: chunk::ChunkType, data: &[u8]) -> Result<()> {
550 w.write_be(data.len() as u32)?;
551 w.write_all(&name.0)?;
552 w.write_all(data)?;
553 let mut crc = Crc32::new();
554 crc.update(&name.0);
555 crc.update(data);
556 w.write_be(crc.finalize())?;
557 Ok(())
558}
559
560impl<W: Write> Writer<W> {
561 fn new(w: W, info: PartialInfo, options: Options) -> Writer<W> {
562 Writer {
563 w,
564 info,
565 options,
566 images_written: 0,
567 animation_written: 0,
568 iend_written: false,
569 }
570 }
571
572 fn init(mut self, info: &Info<'_>) -> Result<Self> {
573 if self.info.width == 0 {
574 return Err(EncodingError::Format(FormatErrorKind::ZeroWidth.into()));
575 }
576
577 if self.info.height == 0 {
578 return Err(EncodingError::Format(FormatErrorKind::ZeroHeight.into()));
579 }
580
581 if self
582 .info
583 .color_type
584 .is_combination_invalid(self.info.bit_depth)
585 {
586 return Err(EncodingError::Format(
587 FormatErrorKind::InvalidColorCombination(self.info.bit_depth, self.info.color_type)
588 .into(),
589 ));
590 }
591
592 self.w.write_all(&[137, 80, 78, 71, 13, 10, 26, 10])?; #[allow(deprecated)]
594 info.encode(&mut self.w)?;
595
596 Ok(self)
597 }
598
599 pub fn write_chunk(&mut self, name: ChunkType, data: &[u8]) -> Result<()> {
604 use std::convert::TryFrom;
605
606 if u32::try_from(data.len()).map_or(true, |length| length > i32::MAX as u32) {
607 let kind = FormatErrorKind::WrittenTooMuch(data.len() - i32::MAX as usize);
608 return Err(EncodingError::Format(kind.into()));
609 }
610
611 write_chunk(&mut self.w, name, data)
612 }
613
614 pub fn write_text_chunk<T: EncodableTextChunk>(&mut self, text_chunk: &T) -> Result<()> {
615 text_chunk.encode(&mut self.w)
616 }
617
618 fn validate_new_image(&self) -> Result<()> {
620 if !self.options.validate_sequence {
621 return Ok(());
622 }
623
624 match self.info.animation_control {
625 None => {
626 if self.images_written == 0 {
627 Ok(())
628 } else {
629 Err(EncodingError::Format(FormatErrorKind::EndReached.into()))
630 }
631 }
632 Some(_) => {
633 if self.info.frame_control.is_some() {
634 Ok(())
635 } else {
636 Err(EncodingError::Format(FormatErrorKind::EndReached.into()))
637 }
638 }
639 }
640 }
641
642 fn validate_sequence_done(&self) -> Result<()> {
643 if !self.options.validate_sequence {
644 return Ok(());
645 }
646
647 if (self.info.animation_control.is_some() && self.info.frame_control.is_some())
648 || self.images_written == 0
649 {
650 Err(EncodingError::Format(FormatErrorKind::MissingFrames.into()))
651 } else {
652 Ok(())
653 }
654 }
655
656 const MAX_IDAT_CHUNK_LEN: u32 = u32::MAX >> 1;
657 #[allow(non_upper_case_globals)]
658 const MAX_fdAT_CHUNK_LEN: u32 = (u32::MAX >> 1) - 4;
659
660 pub fn write_image_data(&mut self, data: &[u8]) -> Result<()> {
662 if self.info.color_type == ColorType::Indexed && !self.info.has_palette {
663 return Err(EncodingError::Format(FormatErrorKind::NoPalette.into()));
664 }
665
666 self.validate_new_image()?;
667
668 let width: usize;
669 let height: usize;
670 if let Some(ref mut fctl) = self.info.frame_control {
671 width = fctl.width as usize;
672 height = fctl.height as usize;
673 } else {
674 width = self.info.width as usize;
675 height = self.info.height as usize;
676 }
677
678 let in_len = self.info.raw_row_length_from_width(width as u32) - 1;
679 let data_size = in_len * height;
680 if data_size != data.len() {
681 return Err(EncodingError::Parameter(
682 ParameterErrorKind::ImageBufferSize {
683 expected: data_size,
684 actual: data.len(),
685 }
686 .into(),
687 ));
688 }
689
690 let prev = vec![0; in_len];
691 let mut prev = prev.as_slice();
692
693 let bpp = self.info.bpp_in_prediction();
694 let filter_method = self.options.filter;
695 let adaptive_method = self.options.adaptive_filter;
696
697 let zlib_encoded = match self.info.compression {
698 Compression::Fast => {
699 let mut compressor = fdeflate::Compressor::new(std::io::Cursor::new(Vec::new()))?;
700
701 let mut current = vec![0; in_len + 1];
702 for line in data.chunks(in_len) {
703 let filter_type = filter(
704 filter_method,
705 adaptive_method,
706 bpp,
707 prev,
708 line,
709 &mut current[1..],
710 );
711
712 current[0] = filter_type as u8;
713 compressor.write_data(¤t)?;
714 prev = line;
715 }
716
717 let compressed = compressor.finish()?.into_inner();
718 if compressed.len()
719 > fdeflate::StoredOnlyCompressor::<()>::compressed_size((in_len + 1) * height)
720 {
721 let mut compressor =
729 fdeflate::StoredOnlyCompressor::new(std::io::Cursor::new(Vec::new()))?;
730 for line in data.chunks(in_len) {
731 compressor.write_data(&[0])?;
732 compressor.write_data(line)?;
733 }
734 compressor.finish()?.into_inner()
735 } else {
736 compressed
737 }
738 }
739 _ => {
740 let mut current = vec![0; in_len];
741
742 let mut zlib = ZlibEncoder::new(Vec::new(), self.info.compression.to_options());
743 for line in data.chunks(in_len) {
744 let filter_type = filter(
745 filter_method,
746 adaptive_method,
747 bpp,
748 prev,
749 line,
750 &mut current,
751 );
752
753 zlib.write_all(&[filter_type as u8])?;
754 zlib.write_all(¤t)?;
755 prev = line;
756 }
757 zlib.finish()?
758 }
759 };
760
761 match self.info.frame_control {
762 None => {
763 self.write_zlib_encoded_idat(&zlib_encoded)?;
764 }
765 Some(_) if self.should_skip_frame_control_on_default_image() => {
766 self.write_zlib_encoded_idat(&zlib_encoded)?;
767 }
768 Some(ref mut fctl) => {
769 fctl.encode(&mut self.w)?;
770 fctl.sequence_number = fctl.sequence_number.wrapping_add(1);
771 self.animation_written += 1;
772
773 if self.images_written == 0 {
775 self.write_zlib_encoded_idat(&zlib_encoded)?;
776 } else {
777 let buff_size = zlib_encoded.len().min(Self::MAX_fdAT_CHUNK_LEN as usize);
778 let mut alldata = vec![0u8; 4 + buff_size];
779 for chunk in zlib_encoded.chunks(Self::MAX_fdAT_CHUNK_LEN as usize) {
780 alldata[..4].copy_from_slice(&fctl.sequence_number.to_be_bytes());
781 alldata[4..][..chunk.len()].copy_from_slice(chunk);
782 write_chunk(&mut self.w, chunk::fdAT, &alldata[..4 + chunk.len()])?;
783 fctl.sequence_number = fctl.sequence_number.wrapping_add(1);
784 }
785 }
786 }
787 }
788
789 self.increment_images_written();
790
791 Ok(())
792 }
793
794 fn increment_images_written(&mut self) {
795 self.images_written = self.images_written.saturating_add(1);
796
797 if let Some(actl) = self.info.animation_control {
798 if actl.num_frames <= self.animation_written {
799 self.info.frame_control = None;
801 }
802 }
803 }
804
805 fn write_iend(&mut self) -> Result<()> {
806 self.iend_written = true;
807 self.write_chunk(chunk::IEND, &[])
808 }
809
810 fn should_skip_frame_control_on_default_image(&self) -> bool {
811 self.options.sep_def_img && self.images_written == 0
812 }
813
814 fn write_zlib_encoded_idat(&mut self, zlib_encoded: &[u8]) -> Result<()> {
815 for chunk in zlib_encoded.chunks(Self::MAX_IDAT_CHUNK_LEN as usize) {
816 self.write_chunk(chunk::IDAT, chunk)?;
817 }
818 Ok(())
819 }
820
821 pub fn set_filter(&mut self, filter: FilterType) {
827 self.options.filter = filter;
828 }
829
830 pub fn set_adaptive_filter(&mut self, adaptive_filter: AdaptiveFilterType) {
837 self.options.adaptive_filter = adaptive_filter;
838 }
839
840 pub fn set_frame_delay(&mut self, numerator: u16, denominator: u16) -> Result<()> {
850 if let Some(ref mut fctl) = self.info.frame_control {
851 fctl.delay_den = denominator;
852 fctl.delay_num = numerator;
853 Ok(())
854 } else {
855 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
856 }
857 }
858
859 pub fn set_frame_dimension(&mut self, width: u32, height: u32) -> Result<()> {
872 if let Some(ref mut fctl) = self.info.frame_control {
873 if Some(width) > self.info.width.checked_sub(fctl.x_offset)
874 || Some(height) > self.info.height.checked_sub(fctl.y_offset)
875 {
876 return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into()));
877 } else if width == 0 {
878 return Err(EncodingError::Format(FormatErrorKind::ZeroWidth.into()));
879 } else if height == 0 {
880 return Err(EncodingError::Format(FormatErrorKind::ZeroHeight.into()));
881 }
882 fctl.width = width;
883 fctl.height = height;
884 Ok(())
885 } else {
886 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
887 }
888 }
889
890 pub fn set_frame_position(&mut self, x: u32, y: u32) -> Result<()> {
901 if let Some(ref mut fctl) = self.info.frame_control {
902 if Some(x) > self.info.width.checked_sub(fctl.width)
903 || Some(y) > self.info.height.checked_sub(fctl.height)
904 {
905 return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into()));
906 }
907 fctl.x_offset = x;
908 fctl.y_offset = y;
909 Ok(())
910 } else {
911 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
912 }
913 }
914
915 pub fn reset_frame_dimension(&mut self) -> Result<()> {
925 if let Some(ref mut fctl) = self.info.frame_control {
926 fctl.width = self.info.width - fctl.x_offset;
927 fctl.height = self.info.height - fctl.y_offset;
928 Ok(())
929 } else {
930 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
931 }
932 }
933
934 pub fn reset_frame_position(&mut self) -> Result<()> {
942 if let Some(ref mut fctl) = self.info.frame_control {
943 fctl.x_offset = 0;
944 fctl.y_offset = 0;
945 Ok(())
946 } else {
947 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
948 }
949 }
950
951 pub fn set_blend_op(&mut self, op: BlendOp) -> Result<()> {
965 if let Some(ref mut fctl) = self.info.frame_control {
966 fctl.blend_op = op;
967 Ok(())
968 } else {
969 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
970 }
971 }
972
973 pub fn set_dispose_op(&mut self, op: DisposeOp) -> Result<()> {
985 if let Some(ref mut fctl) = self.info.frame_control {
986 fctl.dispose_op = op;
987 Ok(())
988 } else {
989 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
990 }
991 }
992
993 pub fn stream_writer(&mut self) -> Result<StreamWriter<W>> {
1002 self.stream_writer_with_size(DEFAULT_BUFFER_LENGTH)
1003 }
1004
1005 pub fn stream_writer_with_size(&mut self, size: usize) -> Result<StreamWriter<W>> {
1011 StreamWriter::new(ChunkOutput::Borrowed(self), size)
1012 }
1013
1014 pub fn into_stream_writer(self) -> Result<StreamWriter<'static, W>> {
1022 self.into_stream_writer_with_size(DEFAULT_BUFFER_LENGTH)
1023 }
1024
1025 pub fn into_stream_writer_with_size(self, size: usize) -> Result<StreamWriter<'static, W>> {
1031 StreamWriter::new(ChunkOutput::Owned(self), size)
1032 }
1033
1034 pub fn finish(mut self) -> Result<()> {
1040 self.validate_sequence_done()?;
1041 self.write_iend()?;
1042 self.w.flush()?;
1043
1044 drop(self);
1046 Ok(())
1047 }
1048}
1049
1050impl<W: Write> Drop for Writer<W> {
1051 fn drop(&mut self) {
1052 if !self.iend_written {
1053 let _ = self.write_iend();
1054 }
1055 }
1056}
1057
1058pub(crate) fn write_iccp_chunk<W: Write>(
1060 w: &mut W,
1061 profile_name: &str,
1062 icc_profile: &[u8],
1063) -> Result<()> {
1064 let profile_name = encode_iso_8859_1(profile_name)?;
1065 if profile_name.len() < 1 || profile_name.len() > 79 {
1066 return Err(TextEncodingError::InvalidKeywordSize.into());
1067 }
1068
1069 let estimated_compressed_size = icc_profile.len() * 3 / 4;
1070 let chunk_size = profile_name
1071 .len()
1072 .checked_add(2) .and_then(|s| s.checked_add(estimated_compressed_size))
1074 .ok_or(EncodingError::LimitsExceeded)?;
1075
1076 let mut data = Vec::new();
1077 data.try_reserve_exact(chunk_size)
1078 .map_err(|_| EncodingError::LimitsExceeded)?;
1079
1080 data.extend(profile_name.into_iter().chain([0, 0]));
1081
1082 let mut encoder = ZlibEncoder::new(data, flate2::Compression::default());
1083 encoder.write_all(icc_profile)?;
1084
1085 write_chunk(w, chunk::iCCP, &encoder.finish()?)
1086}
1087
1088enum ChunkOutput<'a, W: Write> {
1089 Borrowed(&'a mut Writer<W>),
1090 Owned(Writer<W>),
1091}
1092
1093impl<'a, W: Write> Deref for ChunkOutput<'a, W> {
1095 type Target = Writer<W>;
1096
1097 fn deref(&self) -> &Self::Target {
1098 match self {
1099 ChunkOutput::Borrowed(writer) => writer,
1100 ChunkOutput::Owned(writer) => writer,
1101 }
1102 }
1103}
1104
1105impl<'a, W: Write> DerefMut for ChunkOutput<'a, W> {
1106 fn deref_mut(&mut self) -> &mut Self::Target {
1107 match self {
1108 ChunkOutput::Borrowed(writer) => writer,
1109 ChunkOutput::Owned(writer) => writer,
1110 }
1111 }
1112}
1113
1114struct ChunkWriter<'a, W: Write> {
1129 writer: ChunkOutput<'a, W>,
1130 buffer: Vec<u8>,
1131 index: usize,
1133 curr_chunk: ChunkType,
1134}
1135
1136impl<'a, W: Write> ChunkWriter<'a, W> {
1137 fn new(writer: ChunkOutput<'a, W>, buf_len: usize) -> ChunkWriter<'a, W> {
1138 const CAP: usize = u32::MAX as usize >> 1;
1145 let curr_chunk = if writer.images_written == 0 {
1146 chunk::IDAT
1147 } else {
1148 chunk::fdAT
1149 };
1150 ChunkWriter {
1151 writer,
1152 buffer: vec![0; CAP.min(buf_len)],
1153 index: 0,
1154 curr_chunk,
1155 }
1156 }
1157
1158 fn next_frame_info(&self) -> (usize, usize) {
1165 let wrt = self.writer.deref();
1166
1167 let width: usize;
1168 let height: usize;
1169 if let Some(fctl) = wrt.info.frame_control {
1170 width = fctl.width as usize;
1171 height = fctl.height as usize;
1172 } else {
1173 width = wrt.info.width as usize;
1174 height = wrt.info.height as usize;
1175 }
1176
1177 let in_len = wrt.info.raw_row_length_from_width(width as u32) - 1;
1178 let data_size = in_len * height;
1179
1180 (in_len, data_size)
1181 }
1182
1183 fn write_header(&mut self) -> Result<()> {
1186 assert_eq!(self.index, 0, "Called when not flushed");
1187 let wrt = self.writer.deref_mut();
1188
1189 self.curr_chunk = if wrt.images_written == 0 {
1190 chunk::IDAT
1191 } else {
1192 chunk::fdAT
1193 };
1194
1195 match wrt.info.frame_control {
1196 Some(_) if wrt.should_skip_frame_control_on_default_image() => {}
1197 Some(ref mut fctl) => {
1198 fctl.encode(&mut wrt.w)?;
1199 fctl.sequence_number += 1;
1200 }
1201 _ => {}
1202 }
1203
1204 Ok(())
1205 }
1206
1207 fn set_fctl(&mut self, f: FrameControl) {
1212 if let Some(ref mut fctl) = self.writer.info.frame_control {
1213 *fctl = FrameControl {
1215 sequence_number: fctl.sequence_number,
1216 ..f
1217 };
1218 } else {
1219 panic!("This function must be called on an animated PNG")
1220 }
1221 }
1222
1223 fn flush_inner(&mut self) -> io::Result<()> {
1225 if self.index > 0 {
1226 write_chunk(
1228 &mut self.writer.w,
1229 self.curr_chunk,
1230 &self.buffer[..self.index],
1231 )?;
1232
1233 self.index = 0;
1234 }
1235 Ok(())
1236 }
1237}
1238
1239impl<'a, W: Write> Write for ChunkWriter<'a, W> {
1240 fn write(&mut self, mut data: &[u8]) -> io::Result<usize> {
1241 if data.is_empty() {
1242 return Ok(0);
1243 }
1244
1245 if self.index == 0 {
1247 let wrt = self.writer.deref_mut();
1248
1249 let no_fctl = wrt.should_skip_frame_control_on_default_image();
1251 if wrt.info.frame_control.is_some() && !no_fctl {
1252 let fctl = wrt.info.frame_control.as_mut().unwrap();
1253 self.buffer[0..4].copy_from_slice(&fctl.sequence_number.to_be_bytes());
1254 fctl.sequence_number += 1;
1255 self.index = 4;
1256 }
1257 }
1258
1259 let written = data.len().min(self.buffer.len() - self.index);
1262 data = &data[..written];
1263
1264 self.buffer[self.index..][..written].copy_from_slice(data);
1265 self.index += written;
1266
1267 if self.index == self.buffer.len() {
1269 self.flush_inner()?;
1270 }
1271
1272 Ok(written)
1273 }
1274
1275 fn flush(&mut self) -> io::Result<()> {
1276 self.flush_inner()
1277 }
1278}
1279
1280impl<W: Write> Drop for ChunkWriter<'_, W> {
1281 fn drop(&mut self) {
1282 let _ = self.flush();
1283 }
1284}
1285
1286enum Wrapper<'a, W: Write> {
1304 Chunk(ChunkWriter<'a, W>),
1305 Zlib(ZlibEncoder<ChunkWriter<'a, W>>),
1306 Unrecoverable,
1307 None,
1309}
1310
1311impl<'a, W: Write> Wrapper<'a, W> {
1312 fn take(&mut self) -> Wrapper<'a, W> {
1315 let mut swap = Wrapper::None;
1316 mem::swap(self, &mut swap);
1317 swap
1318 }
1319}
1320
1321pub struct StreamWriter<'a, W: Write> {
1329 writer: Wrapper<'a, W>,
1332 prev_buf: Vec<u8>,
1333 curr_buf: Vec<u8>,
1334 index: usize,
1336 line_len: usize,
1338 to_write: usize,
1340
1341 width: u32,
1342 height: u32,
1343
1344 bpp: BytesPerPixel,
1345 filter: FilterType,
1346 adaptive_filter: AdaptiveFilterType,
1347 fctl: Option<FrameControl>,
1348 compression: Compression,
1349}
1350
1351impl<'a, W: Write> StreamWriter<'a, W> {
1352 fn new(writer: ChunkOutput<'a, W>, buf_len: usize) -> Result<StreamWriter<'a, W>> {
1353 let PartialInfo {
1354 width,
1355 height,
1356 frame_control: fctl,
1357 compression,
1358 ..
1359 } = writer.info;
1360
1361 let bpp = writer.info.bpp_in_prediction();
1362 let in_len = writer.info.raw_row_length() - 1;
1363 let filter = writer.options.filter;
1364 let adaptive_filter = writer.options.adaptive_filter;
1365 let prev_buf = vec![0; in_len];
1366 let curr_buf = vec![0; in_len];
1367
1368 let mut chunk_writer = ChunkWriter::new(writer, buf_len);
1369 let (line_len, to_write) = chunk_writer.next_frame_info();
1370 chunk_writer.write_header()?;
1371 let zlib = ZlibEncoder::new(chunk_writer, compression.to_options());
1372
1373 Ok(StreamWriter {
1374 writer: Wrapper::Zlib(zlib),
1375 index: 0,
1376 prev_buf,
1377 curr_buf,
1378 bpp,
1379 filter,
1380 width,
1381 height,
1382 adaptive_filter,
1383 line_len,
1384 to_write,
1385 fctl,
1386 compression,
1387 })
1388 }
1389
1390 pub fn set_filter(&mut self, filter: FilterType) {
1399 self.filter = filter;
1400 }
1401
1402 pub fn set_adaptive_filter(&mut self, adaptive_filter: AdaptiveFilterType) {
1410 self.adaptive_filter = adaptive_filter;
1411 }
1412
1413 pub fn set_frame_delay(&mut self, numerator: u16, denominator: u16) -> Result<()> {
1423 if let Some(ref mut fctl) = self.fctl {
1424 fctl.delay_den = denominator;
1425 fctl.delay_num = numerator;
1426 Ok(())
1427 } else {
1428 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1429 }
1430 }
1431
1432 pub fn set_frame_dimension(&mut self, width: u32, height: u32) -> Result<()> {
1443 if let Some(ref mut fctl) = self.fctl {
1444 if Some(width) > self.width.checked_sub(fctl.x_offset)
1445 || Some(height) > self.height.checked_sub(fctl.y_offset)
1446 {
1447 return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into()));
1448 } else if width == 0 {
1449 return Err(EncodingError::Format(FormatErrorKind::ZeroWidth.into()));
1450 } else if height == 0 {
1451 return Err(EncodingError::Format(FormatErrorKind::ZeroHeight.into()));
1452 }
1453 fctl.width = width;
1454 fctl.height = height;
1455 Ok(())
1456 } else {
1457 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1458 }
1459 }
1460
1461 pub fn set_frame_position(&mut self, x: u32, y: u32) -> Result<()> {
1470 if let Some(ref mut fctl) = self.fctl {
1471 if Some(x) > self.width.checked_sub(fctl.width)
1472 || Some(y) > self.height.checked_sub(fctl.height)
1473 {
1474 return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into()));
1475 }
1476 fctl.x_offset = x;
1477 fctl.y_offset = y;
1478 Ok(())
1479 } else {
1480 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1481 }
1482 }
1483
1484 pub fn reset_frame_dimension(&mut self) -> Result<()> {
1494 if let Some(ref mut fctl) = self.fctl {
1495 fctl.width = self.width - fctl.x_offset;
1496 fctl.height = self.height - fctl.y_offset;
1497 Ok(())
1498 } else {
1499 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1500 }
1501 }
1502
1503 pub fn reset_frame_position(&mut self) -> Result<()> {
1511 if let Some(ref mut fctl) = self.fctl {
1512 fctl.x_offset = 0;
1513 fctl.y_offset = 0;
1514 Ok(())
1515 } else {
1516 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1517 }
1518 }
1519
1520 pub fn set_blend_op(&mut self, op: BlendOp) -> Result<()> {
1534 if let Some(ref mut fctl) = self.fctl {
1535 fctl.blend_op = op;
1536 Ok(())
1537 } else {
1538 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1539 }
1540 }
1541
1542 pub fn set_dispose_op(&mut self, op: DisposeOp) -> Result<()> {
1554 if let Some(ref mut fctl) = self.fctl {
1555 fctl.dispose_op = op;
1556 Ok(())
1557 } else {
1558 Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1559 }
1560 }
1561
1562 pub fn finish(mut self) -> Result<()> {
1569 if self.to_write > 0 {
1570 let err = FormatErrorKind::MissingData(self.to_write).into();
1571 return Err(EncodingError::Format(err));
1572 }
1573
1574 self.flush()?;
1576
1577 if let Wrapper::Chunk(wrt) = self.writer.take() {
1578 wrt.writer.validate_sequence_done()?;
1579 }
1580
1581 Ok(())
1582 }
1583
1584 fn new_frame(&mut self) -> Result<()> {
1589 let wrt = match &mut self.writer {
1590 Wrapper::Chunk(wrt) => wrt,
1591 Wrapper::Unrecoverable => {
1592 let err = FormatErrorKind::Unrecoverable.into();
1593 return Err(EncodingError::Format(err));
1594 }
1595 Wrapper::Zlib(_) => unreachable!("never called on a half-finished frame"),
1596 Wrapper::None => unreachable!(),
1597 };
1598 wrt.flush()?;
1599 wrt.writer.validate_new_image()?;
1600
1601 if let Some(fctl) = self.fctl {
1602 wrt.set_fctl(fctl);
1603 }
1604 let (scansize, size) = wrt.next_frame_info();
1605 self.line_len = scansize;
1606 self.to_write = size;
1607
1608 wrt.write_header()?;
1609 wrt.writer.increment_images_written();
1610
1611 match self.writer.take() {
1613 Wrapper::Chunk(wrt) => {
1614 let encoder = ZlibEncoder::new(wrt, self.compression.to_options());
1615 self.writer = Wrapper::Zlib(encoder);
1616 }
1617 _ => unreachable!(),
1618 };
1619
1620 Ok(())
1621 }
1622}
1623
1624impl<'a, W: Write> Write for StreamWriter<'a, W> {
1625 fn write(&mut self, mut data: &[u8]) -> io::Result<usize> {
1626 if let Wrapper::Unrecoverable = self.writer {
1627 let err = FormatErrorKind::Unrecoverable.into();
1628 return Err(EncodingError::Format(err).into());
1629 }
1630
1631 if data.is_empty() {
1632 return Ok(0);
1633 }
1634
1635 if self.to_write == 0 {
1636 match self.writer.take() {
1637 Wrapper::Zlib(wrt) => match wrt.finish() {
1638 Ok(chunk) => self.writer = Wrapper::Chunk(chunk),
1639 Err(err) => {
1640 self.writer = Wrapper::Unrecoverable;
1641 return Err(err);
1642 }
1643 },
1644 chunk @ Wrapper::Chunk(_) => self.writer = chunk,
1645 Wrapper::Unrecoverable => unreachable!(),
1646 Wrapper::None => unreachable!(),
1647 };
1648
1649 self.new_frame()?;
1651 }
1652
1653 let written = data.read(&mut self.curr_buf[..self.line_len][self.index..])?;
1654 self.index += written;
1655 self.to_write -= written;
1656
1657 if self.index == self.line_len {
1658 let mut filtered = vec![0; self.curr_buf.len()];
1660 let filter_type = filter(
1661 self.filter,
1662 self.adaptive_filter,
1663 self.bpp,
1664 &self.prev_buf,
1665 &self.curr_buf,
1666 &mut filtered,
1667 );
1668 let wrt = match &mut self.writer {
1670 Wrapper::Zlib(wrt) => wrt,
1671 _ => unreachable!(),
1672 };
1673
1674 wrt.write_all(&[filter_type as u8])?;
1675 wrt.write_all(&filtered)?;
1676 mem::swap(&mut self.prev_buf, &mut self.curr_buf);
1677 self.index = 0;
1678 }
1679
1680 Ok(written)
1681 }
1682
1683 fn flush(&mut self) -> io::Result<()> {
1684 match &mut self.writer {
1685 Wrapper::Zlib(wrt) => wrt.flush()?,
1686 Wrapper::Chunk(wrt) => wrt.flush()?,
1687 Wrapper::Unrecoverable | Wrapper::None => {
1690 let err = FormatErrorKind::Unrecoverable.into();
1691 return Err(EncodingError::Format(err).into());
1692 }
1693 }
1694
1695 if self.index > 0 {
1696 let err = FormatErrorKind::WrittenTooMuch(self.index).into();
1697 return Err(EncodingError::Format(err).into());
1698 }
1699
1700 Ok(())
1701 }
1702}
1703
1704impl<W: Write> Drop for StreamWriter<'_, W> {
1705 fn drop(&mut self) {
1706 let _ = self.flush();
1707 }
1708}
1709
1710impl Compression {
1715 fn to_options(self) -> flate2::Compression {
1716 #[allow(deprecated)]
1717 match self {
1718 Compression::Default => flate2::Compression::default(),
1719 Compression::Fast => flate2::Compression::fast(),
1720 Compression::Best => flate2::Compression::best(),
1721 #[allow(deprecated)]
1722 Compression::Huffman => flate2::Compression::none(),
1723 #[allow(deprecated)]
1724 Compression::Rle => flate2::Compression::none(),
1725 }
1726 }
1727}
1728
1729#[cfg(test)]
1730mod tests {
1731 use super::*;
1732 use crate::Decoder;
1733
1734 use rand::{thread_rng, Rng};
1735 use std::cmp;
1736 use std::fs::File;
1737 use std::io::Cursor;
1738
1739 #[test]
1740 fn roundtrip() {
1741 for _ in 0..10 {
1743 for path in glob::glob("tests/pngsuite/*.png")
1744 .unwrap()
1745 .map(|r| r.unwrap())
1746 {
1747 if path.file_name().unwrap().to_str().unwrap().starts_with('x') {
1748 continue;
1750 }
1751 eprintln!("{}", path.display());
1752 let decoder = Decoder::new(File::open(path).unwrap());
1754 let mut reader = decoder.read_info().unwrap();
1755 let mut buf = vec![0; reader.output_buffer_size()];
1756 let info = reader.next_frame(&mut buf).unwrap();
1757 let mut out = Vec::new();
1759 {
1760 let mut wrapper = RandomChunkWriter {
1761 rng: thread_rng(),
1762 w: &mut out,
1763 };
1764
1765 let mut encoder = Encoder::new(&mut wrapper, info.width, info.height);
1766 encoder.set_color(info.color_type);
1767 encoder.set_depth(info.bit_depth);
1768 if let Some(palette) = &reader.info().palette {
1769 encoder.set_palette(palette.clone());
1770 }
1771 let mut encoder = encoder.write_header().unwrap();
1772 encoder.write_image_data(&buf).unwrap();
1773 }
1774 let decoder = Decoder::new(&*out);
1776 let mut reader = decoder.read_info().unwrap();
1777 let mut buf2 = vec![0; reader.output_buffer_size()];
1778 reader.next_frame(&mut buf2).unwrap();
1779 assert_eq!(buf, buf2);
1781 }
1782 }
1783 }
1784
1785 #[test]
1786 fn roundtrip_stream() {
1787 for _ in 0..10 {
1789 for path in glob::glob("tests/pngsuite/*.png")
1790 .unwrap()
1791 .map(|r| r.unwrap())
1792 {
1793 if path.file_name().unwrap().to_str().unwrap().starts_with('x') {
1794 continue;
1796 }
1797 let decoder = Decoder::new(File::open(path).unwrap());
1799 let mut reader = decoder.read_info().unwrap();
1800 let mut buf = vec![0; reader.output_buffer_size()];
1801 let info = reader.next_frame(&mut buf).unwrap();
1802 let mut out = Vec::new();
1804 {
1805 let mut wrapper = RandomChunkWriter {
1806 rng: thread_rng(),
1807 w: &mut out,
1808 };
1809
1810 let mut encoder = Encoder::new(&mut wrapper, info.width, info.height);
1811 encoder.set_color(info.color_type);
1812 encoder.set_depth(info.bit_depth);
1813 if let Some(palette) = &reader.info().palette {
1814 encoder.set_palette(palette.clone());
1815 }
1816 let mut encoder = encoder.write_header().unwrap();
1817 let mut stream_writer = encoder.stream_writer().unwrap();
1818
1819 let mut outer_wrapper = RandomChunkWriter {
1820 rng: thread_rng(),
1821 w: &mut stream_writer,
1822 };
1823
1824 outer_wrapper.write_all(&buf).unwrap();
1825 }
1826 let decoder = Decoder::new(&*out);
1828 let mut reader = decoder.read_info().unwrap();
1829 let mut buf2 = vec![0; reader.output_buffer_size()];
1830 reader.next_frame(&mut buf2).unwrap();
1831 assert_eq!(buf, buf2);
1833 }
1834 }
1835 }
1836
1837 #[test]
1838 fn image_palette() -> Result<()> {
1839 for &bit_depth in &[1u8, 2, 4, 8] {
1840 let path = format!("tests/pngsuite/basn3p0{}.png", bit_depth);
1842 let decoder = Decoder::new(File::open(&path).unwrap());
1843 let mut reader = decoder.read_info().unwrap();
1844
1845 let mut decoded_pixels = vec![0; reader.output_buffer_size()];
1846 let info = reader.info();
1847 assert_eq!(
1848 info.width as usize * info.height as usize * usize::from(bit_depth),
1849 decoded_pixels.len() * 8
1850 );
1851 let info = reader.next_frame(&mut decoded_pixels).unwrap();
1852 let indexed_data = decoded_pixels;
1853
1854 let palette = reader.info().palette.as_ref().unwrap();
1855 let mut out = Vec::new();
1856 {
1857 let mut encoder = Encoder::new(&mut out, info.width, info.height);
1858 encoder.set_depth(BitDepth::from_u8(bit_depth).unwrap());
1859 encoder.set_color(ColorType::Indexed);
1860 encoder.set_palette(palette.as_ref());
1861
1862 let mut writer = encoder.write_header().unwrap();
1863 writer.write_image_data(&indexed_data).unwrap();
1864 }
1865
1866 let decoder = Decoder::new(&*out);
1868 let mut reader = decoder.read_info().unwrap();
1869 let mut redecoded = vec![0; reader.output_buffer_size()];
1870 reader.next_frame(&mut redecoded).unwrap();
1871 assert_eq!(indexed_data, redecoded);
1873 }
1874 Ok(())
1875 }
1876
1877 #[test]
1878 fn expect_error_on_wrong_image_len() -> Result<()> {
1879 let width = 10;
1880 let height = 10;
1881
1882 let output = vec![0u8; 1024];
1883 let writer = Cursor::new(output);
1884 let mut encoder = Encoder::new(writer, width as u32, height as u32);
1885 encoder.set_depth(BitDepth::Eight);
1886 encoder.set_color(ColorType::Rgb);
1887 let mut png_writer = encoder.write_header()?;
1888
1889 let correct_image_size = width * height * 3;
1890 let image = vec![0u8; correct_image_size + 1];
1891 let result = png_writer.write_image_data(image.as_ref());
1892 assert!(result.is_err());
1893
1894 Ok(())
1895 }
1896
1897 #[test]
1898 fn expect_error_on_empty_image() -> Result<()> {
1899 let output = vec![0u8; 1024];
1900 let mut writer = Cursor::new(output);
1901
1902 let encoder = Encoder::new(&mut writer, 0, 0);
1903 assert!(encoder.write_header().is_err());
1904
1905 let encoder = Encoder::new(&mut writer, 100, 0);
1906 assert!(encoder.write_header().is_err());
1907
1908 let encoder = Encoder::new(&mut writer, 0, 100);
1909 assert!(encoder.write_header().is_err());
1910
1911 Ok(())
1912 }
1913
1914 #[test]
1915 fn expect_error_on_invalid_bit_depth_color_type_combination() -> Result<()> {
1916 let output = vec![0u8; 1024];
1917 let mut writer = Cursor::new(output);
1918
1919 let mut encoder = Encoder::new(&mut writer, 1, 1);
1920 encoder.set_depth(BitDepth::One);
1921 encoder.set_color(ColorType::Rgb);
1922 assert!(encoder.write_header().is_err());
1923
1924 let mut encoder = Encoder::new(&mut writer, 1, 1);
1925 encoder.set_depth(BitDepth::One);
1926 encoder.set_color(ColorType::GrayscaleAlpha);
1927 assert!(encoder.write_header().is_err());
1928
1929 let mut encoder = Encoder::new(&mut writer, 1, 1);
1930 encoder.set_depth(BitDepth::One);
1931 encoder.set_color(ColorType::Rgba);
1932 assert!(encoder.write_header().is_err());
1933
1934 let mut encoder = Encoder::new(&mut writer, 1, 1);
1935 encoder.set_depth(BitDepth::Two);
1936 encoder.set_color(ColorType::Rgb);
1937 assert!(encoder.write_header().is_err());
1938
1939 let mut encoder = Encoder::new(&mut writer, 1, 1);
1940 encoder.set_depth(BitDepth::Two);
1941 encoder.set_color(ColorType::GrayscaleAlpha);
1942 assert!(encoder.write_header().is_err());
1943
1944 let mut encoder = Encoder::new(&mut writer, 1, 1);
1945 encoder.set_depth(BitDepth::Two);
1946 encoder.set_color(ColorType::Rgba);
1947 assert!(encoder.write_header().is_err());
1948
1949 let mut encoder = Encoder::new(&mut writer, 1, 1);
1950 encoder.set_depth(BitDepth::Four);
1951 encoder.set_color(ColorType::Rgb);
1952 assert!(encoder.write_header().is_err());
1953
1954 let mut encoder = Encoder::new(&mut writer, 1, 1);
1955 encoder.set_depth(BitDepth::Four);
1956 encoder.set_color(ColorType::GrayscaleAlpha);
1957 assert!(encoder.write_header().is_err());
1958
1959 let mut encoder = Encoder::new(&mut writer, 1, 1);
1960 encoder.set_depth(BitDepth::Four);
1961 encoder.set_color(ColorType::Rgba);
1962 assert!(encoder.write_header().is_err());
1963
1964 let mut encoder = Encoder::new(&mut writer, 1, 1);
1965 encoder.set_depth(BitDepth::Sixteen);
1966 encoder.set_color(ColorType::Indexed);
1967 assert!(encoder.write_header().is_err());
1968
1969 Ok(())
1970 }
1971
1972 #[test]
1973 fn can_write_header_with_valid_bit_depth_color_type_combination() -> Result<()> {
1974 let output = vec![0u8; 1024];
1975 let mut writer = Cursor::new(output);
1976
1977 let mut encoder = Encoder::new(&mut writer, 1, 1);
1978 encoder.set_depth(BitDepth::One);
1979 encoder.set_color(ColorType::Grayscale);
1980 assert!(encoder.write_header().is_ok());
1981
1982 let mut encoder = Encoder::new(&mut writer, 1, 1);
1983 encoder.set_depth(BitDepth::One);
1984 encoder.set_color(ColorType::Indexed);
1985 assert!(encoder.write_header().is_ok());
1986
1987 let mut encoder = Encoder::new(&mut writer, 1, 1);
1988 encoder.set_depth(BitDepth::Two);
1989 encoder.set_color(ColorType::Grayscale);
1990 assert!(encoder.write_header().is_ok());
1991
1992 let mut encoder = Encoder::new(&mut writer, 1, 1);
1993 encoder.set_depth(BitDepth::Two);
1994 encoder.set_color(ColorType::Indexed);
1995 assert!(encoder.write_header().is_ok());
1996
1997 let mut encoder = Encoder::new(&mut writer, 1, 1);
1998 encoder.set_depth(BitDepth::Four);
1999 encoder.set_color(ColorType::Grayscale);
2000 assert!(encoder.write_header().is_ok());
2001
2002 let mut encoder = Encoder::new(&mut writer, 1, 1);
2003 encoder.set_depth(BitDepth::Four);
2004 encoder.set_color(ColorType::Indexed);
2005 assert!(encoder.write_header().is_ok());
2006
2007 let mut encoder = Encoder::new(&mut writer, 1, 1);
2008 encoder.set_depth(BitDepth::Eight);
2009 encoder.set_color(ColorType::Grayscale);
2010 assert!(encoder.write_header().is_ok());
2011
2012 let mut encoder = Encoder::new(&mut writer, 1, 1);
2013 encoder.set_depth(BitDepth::Eight);
2014 encoder.set_color(ColorType::Rgb);
2015 assert!(encoder.write_header().is_ok());
2016
2017 let mut encoder = Encoder::new(&mut writer, 1, 1);
2018 encoder.set_depth(BitDepth::Eight);
2019 encoder.set_color(ColorType::Indexed);
2020 assert!(encoder.write_header().is_ok());
2021
2022 let mut encoder = Encoder::new(&mut writer, 1, 1);
2023 encoder.set_depth(BitDepth::Eight);
2024 encoder.set_color(ColorType::GrayscaleAlpha);
2025 assert!(encoder.write_header().is_ok());
2026
2027 let mut encoder = Encoder::new(&mut writer, 1, 1);
2028 encoder.set_depth(BitDepth::Eight);
2029 encoder.set_color(ColorType::Rgba);
2030 assert!(encoder.write_header().is_ok());
2031
2032 let mut encoder = Encoder::new(&mut writer, 1, 1);
2033 encoder.set_depth(BitDepth::Sixteen);
2034 encoder.set_color(ColorType::Grayscale);
2035 assert!(encoder.write_header().is_ok());
2036
2037 let mut encoder = Encoder::new(&mut writer, 1, 1);
2038 encoder.set_depth(BitDepth::Sixteen);
2039 encoder.set_color(ColorType::Rgb);
2040 assert!(encoder.write_header().is_ok());
2041
2042 let mut encoder = Encoder::new(&mut writer, 1, 1);
2043 encoder.set_depth(BitDepth::Sixteen);
2044 encoder.set_color(ColorType::GrayscaleAlpha);
2045 assert!(encoder.write_header().is_ok());
2046
2047 let mut encoder = Encoder::new(&mut writer, 1, 1);
2048 encoder.set_depth(BitDepth::Sixteen);
2049 encoder.set_color(ColorType::Rgba);
2050 assert!(encoder.write_header().is_ok());
2051
2052 Ok(())
2053 }
2054
2055 #[test]
2056 fn all_filters_roundtrip() -> io::Result<()> {
2057 let pixel: Vec<_> = (0..48).collect();
2058
2059 let roundtrip = |filter: FilterType| -> io::Result<()> {
2060 let mut buffer = vec![];
2061 let mut encoder = Encoder::new(&mut buffer, 4, 4);
2062 encoder.set_depth(BitDepth::Eight);
2063 encoder.set_color(ColorType::Rgb);
2064 encoder.set_filter(filter);
2065 encoder.write_header()?.write_image_data(&pixel)?;
2066
2067 let decoder = crate::Decoder::new(Cursor::new(buffer));
2068 let mut reader = decoder.read_info()?;
2069 let info = reader.info();
2070 assert_eq!(info.width, 4);
2071 assert_eq!(info.height, 4);
2072 let mut dest = vec![0; pixel.len()];
2073 reader.next_frame(&mut dest)?;
2074 assert_eq!(dest, pixel, "Deviation with filter type {:?}", filter);
2075
2076 Ok(())
2077 };
2078
2079 roundtrip(FilterType::NoFilter)?;
2080 roundtrip(FilterType::Sub)?;
2081 roundtrip(FilterType::Up)?;
2082 roundtrip(FilterType::Avg)?;
2083 roundtrip(FilterType::Paeth)?;
2084
2085 Ok(())
2086 }
2087
2088 #[test]
2089 fn some_gamma_roundtrip() -> io::Result<()> {
2090 let pixel: Vec<_> = (0..48).collect();
2091
2092 let roundtrip = |gamma: Option<ScaledFloat>| -> io::Result<()> {
2093 let mut buffer = vec![];
2094 let mut encoder = Encoder::new(&mut buffer, 4, 4);
2095 encoder.set_depth(BitDepth::Eight);
2096 encoder.set_color(ColorType::Rgb);
2097 encoder.set_filter(FilterType::Avg);
2098 if let Some(gamma) = gamma {
2099 encoder.set_source_gamma(gamma);
2100 }
2101 encoder.write_header()?.write_image_data(&pixel)?;
2102
2103 let decoder = crate::Decoder::new(Cursor::new(buffer));
2104 let mut reader = decoder.read_info()?;
2105 assert_eq!(
2106 reader.info().source_gamma,
2107 gamma,
2108 "Deviation with gamma {:?}",
2109 gamma
2110 );
2111 let mut dest = vec![0; pixel.len()];
2112 let info = reader.next_frame(&mut dest)?;
2113 assert_eq!(info.width, 4);
2114 assert_eq!(info.height, 4);
2115
2116 Ok(())
2117 };
2118
2119 roundtrip(None)?;
2120 roundtrip(Some(ScaledFloat::new(0.35)))?;
2121 roundtrip(Some(ScaledFloat::new(0.45)))?;
2122 roundtrip(Some(ScaledFloat::new(0.55)))?;
2123 roundtrip(Some(ScaledFloat::new(0.7)))?;
2124 roundtrip(Some(ScaledFloat::new(1.0)))?;
2125 roundtrip(Some(ScaledFloat::new(2.5)))?;
2126
2127 Ok(())
2128 }
2129
2130 #[test]
2131 fn write_image_chunks_beyond_first() -> Result<()> {
2132 let width = 10;
2133 let height = 10;
2134
2135 let output = vec![0u8; 1024];
2136 let writer = Cursor::new(output);
2137
2138 let mut encoder = Encoder::new(writer, width, height);
2142 encoder.set_depth(BitDepth::Eight);
2143 encoder.set_color(ColorType::Grayscale);
2144 let mut png_writer = encoder.write_header()?;
2145
2146 for _ in 0..3 {
2147 let correct_image_size = (width * height) as usize;
2148 let image = vec![0u8; correct_image_size];
2149 png_writer.write_image_data(image.as_ref())?;
2150 }
2151
2152 Ok(())
2153 }
2154
2155 #[test]
2156 fn image_validate_sequence_without_animation() -> Result<()> {
2157 let width = 10;
2158 let height = 10;
2159
2160 let output = vec![0u8; 1024];
2161 let writer = Cursor::new(output);
2162
2163 let mut encoder = Encoder::new(writer, width, height);
2164 encoder.set_depth(BitDepth::Eight);
2165 encoder.set_color(ColorType::Grayscale);
2166 encoder.validate_sequence(true);
2167 let mut png_writer = encoder.write_header()?;
2168
2169 let correct_image_size = (width * height) as usize;
2170 let image = vec![0u8; correct_image_size];
2171 png_writer.write_image_data(image.as_ref())?;
2172
2173 assert!(png_writer.write_image_data(image.as_ref()).is_err());
2174 Ok(())
2175 }
2176
2177 #[test]
2178 fn image_validate_animation() -> Result<()> {
2179 let width = 10;
2180 let height = 10;
2181
2182 let output = vec![0u8; 1024];
2183 let writer = Cursor::new(output);
2184 let correct_image_size = (width * height) as usize;
2185 let image = vec![0u8; correct_image_size];
2186
2187 let mut encoder = Encoder::new(writer, width, height);
2188 encoder.set_depth(BitDepth::Eight);
2189 encoder.set_color(ColorType::Grayscale);
2190 encoder.set_animated(1, 0)?;
2191 encoder.validate_sequence(true);
2192 let mut png_writer = encoder.write_header()?;
2193
2194 png_writer.write_image_data(image.as_ref())?;
2195
2196 Ok(())
2197 }
2198
2199 #[test]
2200 fn image_validate_animation2() -> Result<()> {
2201 let width = 10;
2202 let height = 10;
2203
2204 let output = vec![0u8; 1024];
2205 let writer = Cursor::new(output);
2206 let correct_image_size = (width * height) as usize;
2207 let image = vec![0u8; correct_image_size];
2208
2209 let mut encoder = Encoder::new(writer, width, height);
2210 encoder.set_depth(BitDepth::Eight);
2211 encoder.set_color(ColorType::Grayscale);
2212 encoder.set_animated(2, 0)?;
2213 encoder.validate_sequence(true);
2214 let mut png_writer = encoder.write_header()?;
2215
2216 png_writer.write_image_data(image.as_ref())?;
2217 png_writer.write_image_data(image.as_ref())?;
2218 png_writer.finish()?;
2219
2220 Ok(())
2221 }
2222
2223 #[test]
2224 fn image_validate_animation_sep_def_image() -> Result<()> {
2225 let width = 10;
2226 let height = 10;
2227
2228 let output = vec![0u8; 1024];
2229 let writer = Cursor::new(output);
2230 let correct_image_size = (width * height) as usize;
2231 let image = vec![0u8; correct_image_size];
2232
2233 let mut encoder = Encoder::new(writer, width, height);
2234 encoder.set_depth(BitDepth::Eight);
2235 encoder.set_color(ColorType::Grayscale);
2236 encoder.set_animated(1, 0)?;
2237 encoder.set_sep_def_img(true)?;
2238 encoder.validate_sequence(true);
2239 let mut png_writer = encoder.write_header()?;
2240
2241 png_writer.write_image_data(image.as_ref())?;
2242 png_writer.write_image_data(image.as_ref())?;
2243 png_writer.finish()?;
2244
2245 Ok(())
2246 }
2247
2248 #[test]
2249 fn image_validate_missing_image() -> Result<()> {
2250 let width = 10;
2251 let height = 10;
2252
2253 let output = vec![0u8; 1024];
2254 let writer = Cursor::new(output);
2255
2256 let mut encoder = Encoder::new(writer, width, height);
2257 encoder.set_depth(BitDepth::Eight);
2258 encoder.set_color(ColorType::Grayscale);
2259 encoder.validate_sequence(true);
2260 let png_writer = encoder.write_header()?;
2261
2262 assert!(png_writer.finish().is_err());
2263 Ok(())
2264 }
2265
2266 #[test]
2267 fn image_validate_missing_animated_frame() -> Result<()> {
2268 let width = 10;
2269 let height = 10;
2270
2271 let output = vec![0u8; 1024];
2272 let writer = Cursor::new(output);
2273 let correct_image_size = (width * height) as usize;
2274 let image = vec![0u8; correct_image_size];
2275
2276 let mut encoder = Encoder::new(writer, width, height);
2277 encoder.set_depth(BitDepth::Eight);
2278 encoder.set_color(ColorType::Grayscale);
2279 encoder.set_animated(2, 0)?;
2280 encoder.validate_sequence(true);
2281 let mut png_writer = encoder.write_header()?;
2282
2283 png_writer.write_image_data(image.as_ref())?;
2284 assert!(png_writer.finish().is_err());
2285
2286 Ok(())
2287 }
2288
2289 #[test]
2290 fn issue_307_stream_validation() -> Result<()> {
2291 let output = vec![0u8; 1024];
2292 let mut cursor = Cursor::new(output);
2293
2294 let encoder = Encoder::new(&mut cursor, 1, 1); let mut writer = encoder.write_header()?;
2296 let mut stream = writer.stream_writer()?;
2297
2298 let written = stream.write(&[1, 2, 3, 4])?;
2299 assert_eq!(written, 1);
2300 stream.finish()?;
2301 drop(writer);
2302
2303 {
2304 cursor.set_position(0);
2305 let mut decoder = Decoder::new(cursor).read_info().expect("A valid image");
2306 let mut buffer = [0u8; 1];
2307 decoder.next_frame(&mut buffer[..]).expect("Valid read");
2308 assert_eq!(buffer, [1]);
2309 }
2310
2311 Ok(())
2312 }
2313
2314 #[test]
2315 fn stream_filtering() -> Result<()> {
2316 let output = vec![0u8; 1024];
2317 let mut cursor = Cursor::new(output);
2318
2319 let mut encoder = Encoder::new(&mut cursor, 8, 8);
2320 encoder.set_color(ColorType::Rgba);
2321 encoder.set_filter(FilterType::Paeth);
2322 let mut writer = encoder.write_header()?;
2323 let mut stream = writer.stream_writer()?;
2324
2325 for _ in 0..8 {
2326 let written = stream.write(&[1; 32])?;
2327 assert_eq!(written, 32);
2328 }
2329 stream.finish()?;
2330 drop(writer);
2331
2332 {
2333 cursor.set_position(0);
2334 let mut decoder = Decoder::new(cursor).read_info().expect("A valid image");
2335 let mut buffer = [0u8; 256];
2336 decoder.next_frame(&mut buffer[..]).expect("Valid read");
2337 assert_eq!(buffer, [1; 256]);
2338 }
2339
2340 Ok(())
2341 }
2342
2343 #[test]
2344 #[cfg(all(unix, not(target_pointer_width = "32")))]
2345 fn exper_error_on_huge_chunk() -> Result<()> {
2346 let empty = vec![0; 1usize << 31];
2349 let writer = Cursor::new(vec![0u8; 1024]);
2350
2351 let mut encoder = Encoder::new(writer, 10, 10);
2352 encoder.set_depth(BitDepth::Eight);
2353 encoder.set_color(ColorType::Grayscale);
2354 let mut png_writer = encoder.write_header()?;
2355
2356 assert!(png_writer.write_chunk(chunk::fdAT, &empty).is_err());
2357 Ok(())
2358 }
2359
2360 #[test]
2361 #[cfg(all(unix, not(target_pointer_width = "32")))]
2362 fn exper_error_on_non_u32_chunk() -> Result<()> {
2363 let empty = vec![0; 1usize << 32];
2366 let writer = Cursor::new(vec![0u8; 1024]);
2367
2368 let mut encoder = Encoder::new(writer, 10, 10);
2369 encoder.set_depth(BitDepth::Eight);
2370 encoder.set_color(ColorType::Grayscale);
2371 let mut png_writer = encoder.write_header()?;
2372
2373 assert!(png_writer.write_chunk(chunk::fdAT, &empty).is_err());
2374 Ok(())
2375 }
2376
2377 #[test]
2378 fn finish_drops_inner_writer() -> Result<()> {
2379 struct NoWriter<'flag>(&'flag mut bool);
2380
2381 impl Write for NoWriter<'_> {
2382 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
2383 Ok(buf.len())
2384 }
2385 fn flush(&mut self) -> io::Result<()> {
2386 Ok(())
2387 }
2388 }
2389 impl Drop for NoWriter<'_> {
2390 fn drop(&mut self) {
2391 *self.0 = true;
2392 }
2393 }
2394
2395 let mut flag = false;
2396
2397 {
2398 let mut encoder = Encoder::new(NoWriter(&mut flag), 10, 10);
2399 encoder.set_depth(BitDepth::Eight);
2400 encoder.set_color(ColorType::Grayscale);
2401
2402 let mut writer = encoder.write_header()?;
2403 writer.write_image_data(&[0; 100])?;
2404 writer.finish()?;
2405 }
2406
2407 assert!(flag, "PNG finished but writer was not dropped");
2408 Ok(())
2409 }
2410
2411 struct RandomChunkWriter<R: Rng, W: Write> {
2413 rng: R,
2414 w: W,
2415 }
2416
2417 impl<R: Rng, W: Write> Write for RandomChunkWriter<R, W> {
2418 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
2419 let len = cmp::min(self.rng.gen_range(1..50), buf.len());
2421
2422 self.w.write(&buf[0..len])
2423 }
2424
2425 fn flush(&mut self) -> io::Result<()> {
2426 self.w.flush()
2427 }
2428 }
2429}