1 extern crate crc32fast; 2 extern crate deflate; 3 4 use std::borrow::Cow; 5 use std::error; 6 use std::fmt; 7 use std::io::{self, Read, Write}; 8 use std::mem; 9 use std::result; 10 11 use crc32fast::Hasher as Crc32; 12 13 use crate::chunk; 14 use crate::common::{BitDepth, ColorType, Compression, Info}; 15 use crate::filter::{filter, FilterType}; 16 use crate::traits::WriteBytesExt; 17 18 pub type Result<T> = result::Result<T, EncodingError>; 19 20 #[derive(Debug)] 21 pub enum EncodingError { 22 IoError(io::Error), 23 Format(Cow<'static, str>), 24 } 25 26 impl error::Error for EncodingError { cause(&self) -> Option<&(dyn error::Error + 'static)>27 fn cause(&self) -> Option<&(dyn error::Error + 'static)> { 28 match self { 29 EncodingError::IoError(err) => Some(err), 30 _ => None, 31 } 32 } 33 } 34 35 impl fmt::Display for EncodingError { fmt(&self, fmt: &mut fmt::Formatter) -> result::Result<(), fmt::Error>36 fn fmt(&self, fmt: &mut fmt::Formatter) -> result::Result<(), fmt::Error> { 37 use self::EncodingError::*; 38 match self { 39 IoError(err) => write!(fmt, "{}", err), 40 Format(desc) => write!(fmt, "{}", desc), 41 } 42 } 43 } 44 45 impl From<io::Error> for EncodingError { from(err: io::Error) -> EncodingError46 fn from(err: io::Error) -> EncodingError { 47 EncodingError::IoError(err) 48 } 49 } 50 impl From<EncodingError> for io::Error { from(err: EncodingError) -> io::Error51 fn from(err: EncodingError) -> io::Error { 52 io::Error::new(io::ErrorKind::Other, err.to_string()) 53 } 54 } 55 56 /// PNG Encoder 57 pub struct Encoder<W: Write> { 58 w: W, 59 info: Info, 60 } 61 62 impl<W: Write> Encoder<W> { new(w: W, width: u32, height: u32) -> Encoder<W>63 pub fn new(w: W, width: u32, height: u32) -> Encoder<W> { 64 let mut info = Info::default(); 65 info.width = width; 66 info.height = height; 67 Encoder { w, info } 68 } 69 write_header(self) -> Result<Writer<W>>70 pub fn write_header(self) -> Result<Writer<W>> { 71 Writer::new(self.w, self.info).init() 72 } 73 74 /// Set the color of the encoded image. 75 /// 76 /// These correspond to the color types in the png IHDR data that will be written. The length 77 /// of the image data that is later supplied must match the color type, otherwise an error will 78 /// be emitted. set_color(&mut self, color: ColorType)79 pub fn set_color(&mut self, color: ColorType) { 80 self.info.color_type = color; 81 } 82 83 /// Set the indicated depth of the image data. set_depth(&mut self, depth: BitDepth)84 pub fn set_depth(&mut self, depth: BitDepth) { 85 self.info.bit_depth = depth; 86 } 87 88 /// Set compression parameters. 89 /// 90 /// Accepts a `Compression` or any type that can transform into a `Compression`. Notably `deflate::Compression` and 91 /// `deflate::CompressionOptions` which "just work". set_compression<C: Into<Compression>>(&mut self, compression: C)92 pub fn set_compression<C: Into<Compression>>(&mut self, compression: C) { 93 self.info.compression = compression.into(); 94 } 95 96 /// Set the used filter type. 97 /// 98 /// The default filter is [`FilterType::Sub`] which provides a basic prediction algorithm for 99 /// sample values based on the previous. For a potentially better compression ratio, at the 100 /// cost of more complex processing, try out [`FilterType::Paeth`]. 101 /// 102 /// [`FilterType::Sub`]: enum.FilterType.html#variant.Sub 103 /// [`FilterType::Paeth`]: enum.FilterType.html#variant.Paeth set_filter(&mut self, filter: FilterType)104 pub fn set_filter(&mut self, filter: FilterType) { 105 self.info.filter = filter; 106 } 107 } 108 109 /// PNG writer 110 pub struct Writer<W: Write> { 111 w: W, 112 info: Info, 113 } 114 115 const DEFAULT_BUFFER_LENGTH: usize = 4 * 1024; 116 117 impl<W: Write> Writer<W> { new(w: W, info: Info) -> Writer<W>118 fn new(w: W, info: Info) -> Writer<W> { 119 Writer { w, info } 120 } 121 init(mut self) -> Result<Self>122 fn init(mut self) -> Result<Self> { 123 if self.info.width == 0 { 124 return Err(EncodingError::Format("Zero width not allowed".into())); 125 } 126 127 if self.info.height == 0 { 128 return Err(EncodingError::Format("Zero height not allowed".into())); 129 } 130 131 self.w.write_all(&[137, 80, 78, 71, 13, 10, 26, 10])?; 132 let mut data = [0; 13]; 133 (&mut data[..]).write_be(self.info.width)?; 134 (&mut data[4..]).write_be(self.info.height)?; 135 data[8] = self.info.bit_depth as u8; 136 data[9] = self.info.color_type as u8; 137 data[12] = if self.info.interlaced { 1 } else { 0 }; 138 self.write_chunk(chunk::IHDR, &data)?; 139 Ok(self) 140 } 141 write_chunk(&mut self, name: [u8; 4], data: &[u8]) -> Result<()>142 pub fn write_chunk(&mut self, name: [u8; 4], data: &[u8]) -> Result<()> { 143 self.w.write_be(data.len() as u32)?; 144 self.w.write_all(&name)?; 145 self.w.write_all(data)?; 146 let mut crc = Crc32::new(); 147 crc.update(&name); 148 crc.update(data); 149 self.w.write_be(crc.finalize())?; 150 Ok(()) 151 } 152 153 /// Writes the image data. write_image_data(&mut self, data: &[u8]) -> Result<()>154 pub fn write_image_data(&mut self, data: &[u8]) -> Result<()> { 155 const MAX_CHUNK_LEN: u32 = (1u32 << 31) - 1; 156 let bpp = self.info.bytes_per_pixel(); 157 let in_len = self.info.raw_row_length() - 1; 158 let prev = vec![0; in_len]; 159 let mut prev = prev.as_slice(); 160 let mut current = vec![0; in_len]; 161 let data_size = in_len * self.info.height as usize; 162 if data_size != data.len() { 163 let message = format!("wrong data size, expected {} got {}", data_size, data.len()); 164 return Err(EncodingError::Format(message.into())); 165 } 166 let mut zlib = deflate::write::ZlibEncoder::new(Vec::new(), self.info.compression.clone()); 167 let filter_method = self.info.filter; 168 for line in data.chunks(in_len) { 169 current.copy_from_slice(&line); 170 zlib.write_all(&[filter_method as u8])?; 171 filter(filter_method, bpp, &prev, &mut current); 172 zlib.write_all(¤t)?; 173 prev = line; 174 } 175 let zlib_encoded = zlib.finish()?; 176 for chunk in zlib_encoded.chunks(MAX_CHUNK_LEN as usize) { 177 self.write_chunk(chunk::IDAT, &chunk)?; 178 } 179 Ok(()) 180 } 181 182 /// Create an stream writer. 183 /// 184 /// This allows you create images that do not fit in memory. The default 185 /// chunk size is 4K, use `stream_writer_with_size` to set another chuck 186 /// size. 187 /// 188 /// This borrows the writer. This preserves it which allows manually 189 /// appending additional chunks after the image data has been written stream_writer(&mut self) -> StreamWriter<W>190 pub fn stream_writer(&mut self) -> StreamWriter<W> { 191 self.stream_writer_with_size(DEFAULT_BUFFER_LENGTH) 192 } 193 194 /// Create a stream writer with custom buffer size. 195 /// 196 /// See [`stream_writer`]. 197 /// 198 /// [`stream_writer`]: #fn.stream_writer stream_writer_with_size(&mut self, size: usize) -> StreamWriter<W>199 pub fn stream_writer_with_size(&mut self, size: usize) -> StreamWriter<W> { 200 StreamWriter::new(ChunkOutput::Borrowed(self), size) 201 } 202 203 /// Turn this into a stream writer for image data. 204 /// 205 /// This allows you create images that do not fit in memory. The default 206 /// chunk size is 4K, use `stream_writer_with_size` to set another chuck 207 /// size. into_stream_writer(self) -> StreamWriter<'static, W>208 pub fn into_stream_writer(self) -> StreamWriter<'static, W> { 209 self.into_stream_writer_with_size(DEFAULT_BUFFER_LENGTH) 210 } 211 212 /// Turn this into a stream writer with custom buffer size. 213 /// 214 /// See [`into_stream_writer`]. 215 /// 216 /// [`into_stream_writer`]: #fn.into_stream_writer into_stream_writer_with_size(self, size: usize) -> StreamWriter<'static, W>217 pub fn into_stream_writer_with_size(self, size: usize) -> StreamWriter<'static, W> { 218 StreamWriter::new(ChunkOutput::Owned(self), size) 219 } 220 } 221 222 impl<W: Write> Drop for Writer<W> { drop(&mut self)223 fn drop(&mut self) { 224 let _ = self.write_chunk(chunk::IEND, &[]); 225 } 226 } 227 228 struct ChunkWriter<'a, W: Write> { 229 writer: ChunkOutput<'a, W>, 230 buffer: Vec<u8>, 231 index: usize, 232 } 233 234 enum ChunkOutput<'a, W: Write> { 235 Borrowed(&'a mut Writer<W>), 236 Owned(Writer<W>), 237 } 238 239 impl<'a, W: Write> ChunkWriter<'a, W> { new(writer: ChunkOutput<'a, W>, buf_len: usize) -> ChunkWriter<'a, W>240 fn new(writer: ChunkOutput<'a, W>, buf_len: usize) -> ChunkWriter<'a, W> { 241 ChunkWriter { 242 writer, 243 buffer: vec![0; buf_len], 244 index: 0, 245 } 246 } 247 } 248 249 impl<'a, W: Write> AsMut<Writer<W>> for ChunkOutput<'a, W> { as_mut(&mut self) -> &mut Writer<W>250 fn as_mut(&mut self) -> &mut Writer<W> { 251 match self { 252 ChunkOutput::Borrowed(writer) => writer, 253 ChunkOutput::Owned(writer) => writer, 254 } 255 } 256 } 257 258 impl<'a, W: Write> Write for ChunkWriter<'a, W> { write(&mut self, mut buf: &[u8]) -> io::Result<usize>259 fn write(&mut self, mut buf: &[u8]) -> io::Result<usize> { 260 let written = buf.read(&mut self.buffer[self.index..])?; 261 self.index += written; 262 263 if self.index + 1 >= self.buffer.len() { 264 self.writer 265 .as_mut() 266 .write_chunk(chunk::IDAT, &self.buffer)?; 267 self.index = 0; 268 } 269 270 Ok(written) 271 } 272 flush(&mut self) -> io::Result<()>273 fn flush(&mut self) -> io::Result<()> { 274 if self.index > 0 { 275 self.writer 276 .as_mut() 277 .write_chunk(chunk::IDAT, &self.buffer[..=self.index])?; 278 } 279 self.index = 0; 280 Ok(()) 281 } 282 } 283 284 impl<'a, W: Write> Drop for ChunkWriter<'a, W> { drop(&mut self)285 fn drop(&mut self) { 286 let _ = self.flush(); 287 } 288 } 289 290 /// Streaming png writer 291 /// 292 /// This may silently fail in the destructor, so it is a good idea to call 293 /// [`finish`](#method.finish) or [`flush`](https://doc.rust-lang.org/stable/std/io/trait.Write.html#tymethod.flush) before dropping. 294 pub struct StreamWriter<'a, W: Write> { 295 writer: deflate::write::ZlibEncoder<ChunkWriter<'a, W>>, 296 prev_buf: Vec<u8>, 297 curr_buf: Vec<u8>, 298 index: usize, 299 bpp: usize, 300 filter: FilterType, 301 } 302 303 impl<'a, W: Write> StreamWriter<'a, W> { new(mut writer: ChunkOutput<'a, W>, buf_len: usize) -> StreamWriter<'a, W>304 fn new(mut writer: ChunkOutput<'a, W>, buf_len: usize) -> StreamWriter<'a, W> { 305 let bpp = writer.as_mut().info.bytes_per_pixel(); 306 let in_len = writer.as_mut().info.raw_row_length() - 1; 307 let filter = writer.as_mut().info.filter; 308 let prev_buf = vec![0; in_len]; 309 let curr_buf = vec![0; in_len]; 310 311 let compression = writer.as_mut().info.compression.clone(); 312 let chunk_writer = ChunkWriter::new(writer, buf_len); 313 let zlib = deflate::write::ZlibEncoder::new(chunk_writer, compression); 314 315 StreamWriter { 316 writer: zlib, 317 index: 0, 318 prev_buf, 319 curr_buf, 320 bpp, 321 filter, 322 } 323 } 324 finish(mut self) -> Result<()>325 pub fn finish(mut self) -> Result<()> { 326 // TODO: call `writer.finish` somehow? 327 self.flush()?; 328 Ok(()) 329 } 330 } 331 332 impl<'a, W: Write> Write for StreamWriter<'a, W> { write(&mut self, mut buf: &[u8]) -> io::Result<usize>333 fn write(&mut self, mut buf: &[u8]) -> io::Result<usize> { 334 let written = buf.read(&mut self.curr_buf[self.index..])?; 335 self.index += written; 336 337 if self.index >= self.curr_buf.len() { 338 self.writer.write_all(&[self.filter as u8])?; 339 filter(self.filter, self.bpp, &self.prev_buf, &mut self.curr_buf); 340 self.writer.write_all(&self.curr_buf)?; 341 mem::swap(&mut self.prev_buf, &mut self.curr_buf); 342 self.index = 0; 343 } 344 345 Ok(written) 346 } 347 flush(&mut self) -> io::Result<()>348 fn flush(&mut self) -> io::Result<()> { 349 self.writer.flush()?; 350 if self.index > 0 { 351 let message = format!("wrong data size, got {} bytes too many", self.index); 352 return Err(EncodingError::Format(message.into()).into()); 353 } 354 Ok(()) 355 } 356 } 357 358 impl<'a, W: Write> Drop for StreamWriter<'a, W> { drop(&mut self)359 fn drop(&mut self) { 360 let _ = self.flush(); 361 } 362 } 363 364 #[cfg(test)] 365 mod tests { 366 use super::*; 367 368 extern crate glob; 369 370 use rand::{thread_rng, Rng}; 371 use std::fs::File; 372 use std::io::Write; 373 use std::{cmp, io}; 374 375 #[test] roundtrip()376 fn roundtrip() { 377 // More loops = more random testing, but also more test wait time 378 for _ in 0..10 { 379 for path in glob::glob("tests/pngsuite/*.png") 380 .unwrap() 381 .map(|r| r.unwrap()) 382 { 383 if path.file_name().unwrap().to_str().unwrap().starts_with("x") { 384 // x* files are expected to fail to decode 385 continue; 386 } 387 // Decode image 388 let decoder = crate::Decoder::new(File::open(path).unwrap()); 389 let (info, mut reader) = decoder.read_info().unwrap(); 390 if info.line_size != 32 { 391 // TODO encoding only works with line size 32? 392 continue; 393 } 394 let mut buf = vec![0; info.buffer_size()]; 395 reader.next_frame(&mut buf).unwrap(); 396 // Encode decoded image 397 let mut out = Vec::new(); 398 { 399 let mut wrapper = RandomChunkWriter { 400 rng: thread_rng(), 401 w: &mut out, 402 }; 403 404 let mut encoder = Encoder::new(&mut wrapper, info.width, info.height) 405 .write_header() 406 .unwrap(); 407 encoder.write_image_data(&buf).unwrap(); 408 } 409 // Decode encoded decoded image 410 let decoder = crate::Decoder::new(&*out); 411 let (info, mut reader) = decoder.read_info().unwrap(); 412 let mut buf2 = vec![0; info.buffer_size()]; 413 reader.next_frame(&mut buf2).unwrap(); 414 // check if the encoded image is ok: 415 assert_eq!(buf, buf2); 416 } 417 } 418 } 419 420 #[test] roundtrip_stream()421 fn roundtrip_stream() { 422 // More loops = more random testing, but also more test wait time 423 for _ in 0..10 { 424 for path in glob::glob("tests/pngsuite/*.png") 425 .unwrap() 426 .map(|r| r.unwrap()) 427 { 428 if path.file_name().unwrap().to_str().unwrap().starts_with("x") { 429 // x* files are expected to fail to decode 430 continue; 431 } 432 // Decode image 433 let decoder = crate::Decoder::new(File::open(path).unwrap()); 434 let (info, mut reader) = decoder.read_info().unwrap(); 435 if info.line_size != 32 { 436 // TODO encoding only works with line size 32? 437 continue; 438 } 439 let mut buf = vec![0; info.buffer_size()]; 440 reader.next_frame(&mut buf).unwrap(); 441 // Encode decoded image 442 let mut out = Vec::new(); 443 { 444 let mut wrapper = RandomChunkWriter { 445 rng: thread_rng(), 446 w: &mut out, 447 }; 448 449 let mut encoder = Encoder::new(&mut wrapper, info.width, info.height) 450 .write_header() 451 .unwrap(); 452 let mut stream_writer = encoder.stream_writer(); 453 454 let mut outer_wrapper = RandomChunkWriter { 455 rng: thread_rng(), 456 w: &mut stream_writer, 457 }; 458 459 outer_wrapper.write_all(&buf).unwrap(); 460 } 461 // Decode encoded decoded image 462 let decoder = crate::Decoder::new(&*out); 463 let (info, mut reader) = decoder.read_info().unwrap(); 464 let mut buf2 = vec![0; info.buffer_size()]; 465 reader.next_frame(&mut buf2).unwrap(); 466 // check if the encoded image is ok: 467 assert_eq!(buf, buf2); 468 } 469 } 470 } 471 472 #[test] expect_error_on_wrong_image_len() -> Result<()>473 fn expect_error_on_wrong_image_len() -> Result<()> { 474 use std::io::Cursor; 475 476 let width = 10; 477 let height = 10; 478 479 let output = vec![0u8; 1024]; 480 let writer = Cursor::new(output); 481 let mut encoder = Encoder::new(writer, width as u32, height as u32); 482 encoder.set_depth(BitDepth::Eight); 483 encoder.set_color(ColorType::RGB); 484 let mut png_writer = encoder.write_header()?; 485 486 let correct_image_size = width * height * 3; 487 let image = vec![0u8; correct_image_size + 1]; 488 let result = png_writer.write_image_data(image.as_ref()); 489 assert!(result.is_err()); 490 491 Ok(()) 492 } 493 494 #[test] expect_error_on_empty_image() -> Result<()>495 fn expect_error_on_empty_image() -> Result<()> { 496 use std::io::Cursor; 497 498 let output = vec![0u8; 1024]; 499 let mut writer = Cursor::new(output); 500 501 let encoder = Encoder::new(&mut writer, 0, 0); 502 assert!(encoder.write_header().is_err()); 503 504 let encoder = Encoder::new(&mut writer, 100, 0); 505 assert!(encoder.write_header().is_err()); 506 507 let encoder = Encoder::new(&mut writer, 0, 100); 508 assert!(encoder.write_header().is_err()); 509 510 Ok(()) 511 } 512 513 #[test] all_filters_roundtrip() -> io::Result<()>514 fn all_filters_roundtrip() -> io::Result<()> { 515 let pixel: Vec<_> = (0..48).collect(); 516 517 let roundtrip = |filter: FilterType| -> io::Result<()> { 518 let mut buffer = vec![]; 519 let mut encoder = Encoder::new(&mut buffer, 4, 4); 520 encoder.set_depth(BitDepth::Eight); 521 encoder.set_color(ColorType::RGB); 522 encoder.set_filter(filter); 523 encoder.write_header()?.write_image_data(&pixel)?; 524 525 let decoder = crate::Decoder::new(io::Cursor::new(buffer)); 526 let (info, mut reader) = decoder.read_info()?; 527 assert_eq!(info.width, 4); 528 assert_eq!(info.height, 4); 529 let mut dest = vec![0; pixel.len()]; 530 reader.next_frame(&mut dest)?; 531 assert_eq!(dest, pixel, "Deviation with filter type {:?}", filter); 532 533 Ok(()) 534 }; 535 536 roundtrip(FilterType::NoFilter)?; 537 roundtrip(FilterType::Sub)?; 538 roundtrip(FilterType::Up)?; 539 roundtrip(FilterType::Avg)?; 540 roundtrip(FilterType::Paeth)?; 541 542 Ok(()) 543 } 544 545 /// A Writer that only writes a few bytes at a time 546 struct RandomChunkWriter<'a, R: Rng, W: Write + 'a> { 547 rng: R, 548 w: &'a mut W, 549 } 550 551 impl<'a, R: Rng, W: Write + 'a> Write for RandomChunkWriter<'a, R, W> { write(&mut self, buf: &[u8]) -> io::Result<usize>552 fn write(&mut self, buf: &[u8]) -> io::Result<usize> { 553 // choose a random length to write 554 let len = cmp::min(self.rng.gen_range(1, 50), buf.len()); 555 556 self.w.write(&buf[0..len]) 557 } 558 flush(&mut self) -> io::Result<()>559 fn flush(&mut self) -> io::Result<()> { 560 self.w.flush() 561 } 562 } 563 } 564