From 3ef4af6abf25ee06aadb896d24a5161c8968bae5 Mon Sep 17 00:00:00 2001
From: Kornel <kornel@geekhood.net>
Date: Mon, 9 Dec 2024 00:17:34 +0000
Subject: [PATCH] Fix iCCP chunk encoding (#548)

---
 src/common.rs         |  2 +-
 src/decoder/stream.rs | 23 +++++++++++++++++++++--
 src/encoder.rs        | 32 +++++++++++++++++++++++++++++++-
 src/text_metadata.rs  |  2 +-
 4 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/src/common.rs b/src/common.rs
index 9a4e1876..ce694497 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -766,7 +766,7 @@ impl Info<'_> {
                 chrms.encode(&mut w)?;
             }
             if let Some(iccp) = &self.icc_profile {
-                encoder::write_chunk(&mut w, chunk::iCCP, iccp)?;
+                encoder::write_iccp_chunk(&mut w, "_", iccp)?
             }
         }
 
diff --git a/src/decoder/stream.rs b/src/decoder/stream.rs
index 7d5de33f..88044b74 100644
--- a/src/decoder/stream.rs
+++ b/src/decoder/stream.rs
@@ -1528,9 +1528,11 @@ impl StreamingDecoder {
         let mut buf = &self.current_chunk.raw_bytes[..];
 
         // read profile name
-        let _: u8 = buf.read_be()?;
-        for _ in 1..80 {
+        for len in 0..=80 {
             let raw: u8 = buf.read_be()?;
+            if (raw == 0 && len == 0) || (raw != 0 && len == 80) {
+                return Err(DecodingError::from(TextDecodingError::InvalidKeywordSize));
+            }
             if raw == 0 {
                 break;
             }
@@ -2108,6 +2110,23 @@ mod tests {
         assert_eq!(4070462061, crc32fast::hash(&icc_profile));
     }
 
+    #[test]
+    fn test_iccp_roundtrip() {
+        let dummy_icc = b"I'm a profile";
+
+        let mut info = crate::Info::with_size(1, 1);
+        info.icc_profile = Some(dummy_icc.into());
+        let mut encoded_image = Vec::new();
+        let enc = crate::Encoder::with_info(&mut encoded_image, info).unwrap();
+        let mut enc = enc.write_header().unwrap();
+        enc.write_image_data(&[0]).unwrap();
+        enc.finish().unwrap();
+
+        let dec = crate::Decoder::new(encoded_image.as_slice());
+        let dec = dec.read_info().unwrap();
+        assert_eq!(dummy_icc, &**dec.info().icc_profile.as_ref().unwrap());
+    }
+
     #[test]
     fn test_png_with_broken_iccp() {
         let decoder = crate::Decoder::new(File::open("tests/iccp/broken_iccp.png").unwrap());
diff --git a/src/encoder.rs b/src/encoder.rs
index d694f760..7a71b468 100644
--- a/src/encoder.rs
+++ b/src/encoder.rs
@@ -13,7 +13,7 @@ use crate::common::{
 };
 use crate::filter::{filter, AdaptiveFilterType, FilterType};
 use crate::text_metadata::{
-    EncodableTextChunk, ITXtChunk, TEXtChunk, TextEncodingError, ZTXtChunk,
+    encode_iso_8859_1, EncodableTextChunk, ITXtChunk, TEXtChunk, TextEncodingError, ZTXtChunk,
 };
 use crate::traits::WriteBytesExt;
 
@@ -1040,6 +1040,36 @@ impl<W: Write> Drop for Writer<W> {
     }
 }
 
+// This should be moved to Writer after `Info::encoding` is gone
+pub(crate) fn write_iccp_chunk<W: Write>(
+    w: &mut W,
+    profile_name: &str,
+    icc_profile: &[u8],
+) -> Result<()> {
+    let profile_name = encode_iso_8859_1(profile_name)?;
+    if profile_name.len() < 1 || profile_name.len() > 79 {
+        return Err(TextEncodingError::InvalidKeywordSize.into());
+    }
+
+    let estimated_compressed_size = icc_profile.len() * 3 / 4;
+    let chunk_size = profile_name
+        .len()
+        .checked_add(2) // string NUL + compression type. Checked add optimizes out later Vec reallocations.
+        .and_then(|s| s.checked_add(estimated_compressed_size))
+        .ok_or(EncodingError::LimitsExceeded)?;
+
+    let mut data = Vec::new();
+    data.try_reserve_exact(chunk_size)
+        .map_err(|_| EncodingError::LimitsExceeded)?;
+
+    data.extend(profile_name.into_iter().chain([0, 0]));
+
+    let mut encoder = ZlibEncoder::new(data, flate2::Compression::default());
+    encoder.write_all(icc_profile)?;
+
+    write_chunk(w, chunk::iCCP, &encoder.finish()?)
+}
+
 enum ChunkOutput<'a, W: Write> {
     Borrowed(&'a mut Writer<W>),
     Owned(Writer<W>),
diff --git a/src/text_metadata.rs b/src/text_metadata.rs
index 7fc730fa..3bb491fe 100644
--- a/src/text_metadata.rs
+++ b/src/text_metadata.rs
@@ -162,7 +162,7 @@ fn decode_iso_8859_1(text: &[u8]) -> String {
     text.iter().map(|&b| b as char).collect()
 }
 
-fn encode_iso_8859_1(text: &str) -> Result<Vec<u8>, TextEncodingError> {
+pub(crate) fn encode_iso_8859_1(text: &str) -> Result<Vec<u8>, TextEncodingError> {
     encode_iso_8859_1_iter(text).collect()
 }