From 675dadf6b4f22f0e694859cff17cc2d28b1939db Mon Sep 17 00:00:00 2001
From: Alula <6276139+alula@users.noreply.github.com>
Date: Sat, 7 Dec 2024 00:09:02 +0100
Subject: [PATCH] Implement writing ICC profiles for JPEG and PNG images.

---
 src/codecs/jpeg/encoder.rs | 47 ++++++++++++++++++++++++++++++++++++++
 src/codecs/png.rs          | 19 ++++++++++++++-
 2 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/src/codecs/jpeg/encoder.rs b/src/codecs/jpeg/encoder.rs
index 0c1f074985..338cbdef59 100644
--- a/src/codecs/jpeg/encoder.rs
+++ b/src/codecs/jpeg/encoder.rs
@@ -30,6 +30,7 @@ static SOS: u8 = 0xDA;
 static DQT: u8 = 0xDB;
 // Application segments start and end
 static APP0: u8 = 0xE0;
+static APP2: u8 = 0xE2;
 
 // section K.1
 // table K.1
@@ -346,6 +347,8 @@ pub struct JpegEncoder<W> {
     chroma_actable: Cow<'static, [(u8, u16); 256]>,
 
     pixel_density: PixelDensity,
+
+    icc_profile: Vec<u8>,
 }
 
 impl<W: Write> JpegEncoder<W> {
@@ -415,6 +418,8 @@ impl<W: Write> JpegEncoder<W> {
             chroma_actable: Cow::Borrowed(&STD_CHROMA_AC_HUFF_LUT),
 
             pixel_density: PixelDensity::default(),
+
+            icc_profile: Vec::new(),
         }
     }
 
@@ -494,6 +499,9 @@ impl<W: Write> JpegEncoder<W> {
         build_jfif_header(&mut buf, self.pixel_density);
         self.writer.write_segment(APP0, &buf)?;
 
+        // Write ICC profile chunks if present
+        self.write_icc_profile_chunks()?;
+
         build_frame_header(
             &mut buf,
             8,
@@ -648,6 +656,40 @@ impl<W: Write> JpegEncoder<W> {
 
         Ok(())
     }
+
+    fn write_icc_profile_chunks(&mut self) -> io::Result<()> {
+        if self.icc_profile.is_empty() {
+            return Ok(());
+        }
+
+        const MAX_CHUNK_SIZE: usize = 65533 - 14;
+        const MAX_CHUNK_COUNT: usize = 255;
+        const MAX_ICC_PROFILE_SIZE: usize = MAX_CHUNK_SIZE * MAX_CHUNK_COUNT;
+
+        if self.icc_profile.len() > MAX_ICC_PROFILE_SIZE {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "ICC profile too large",
+            ));
+        }
+
+        let num_chunks = ((self.icc_profile.len() + MAX_CHUNK_SIZE - 1) / MAX_CHUNK_SIZE) as u8;
+
+        for (i, chunk) in self.icc_profile.chunks(MAX_CHUNK_SIZE).enumerate() {
+            let chunk_number = (i + 1) as u8;
+            let length = 14 + chunk.len();
+
+            let mut segment = Vec::with_capacity(length);
+            segment.extend_from_slice(b"ICC_PROFILE\0");
+            segment.push(chunk_number);
+            segment.push(num_chunks);
+            segment.extend_from_slice(chunk);
+
+            self.writer.write_segment(APP2, &segment)?;
+        }
+
+        Ok(())
+    }
 }
 
 impl<W: Write> ImageEncoder for JpegEncoder<W> {
@@ -661,6 +703,11 @@ impl<W: Write> ImageEncoder for JpegEncoder<W> {
     ) -> ImageResult<()> {
         self.encode(buf, width, height, color_type)
     }
+
+    fn set_icc_profile(&mut self, icc_profile: Vec<u8>) -> Result<(), UnsupportedError> {
+        self.icc_profile = icc_profile;
+        Ok(())
+    }
 }
 
 fn build_jfif_header(m: &mut Vec<u8>, density: PixelDensity) {
diff --git a/src/codecs/png.rs b/src/codecs/png.rs
index 8c95ba6954..8cf296541a 100644
--- a/src/codecs/png.rs
+++ b/src/codecs/png.rs
@@ -6,6 +6,7 @@
 //! * <http://www.w3.org/TR/PNG/> - The PNG Specification
 //!
 
+use std::borrow::Cow;
 use std::fmt;
 use std::io::{BufRead, Seek, Write};
 
@@ -482,6 +483,7 @@ pub struct PngEncoder<W: Write> {
     w: W,
     compression: CompressionType,
     filter: FilterType,
+    icc_profile: Vec<u8>,
 }
 
 /// Compression level of a PNG encoder. The default setting is `Fast`.
@@ -535,6 +537,7 @@ impl<W: Write> PngEncoder<W> {
             w,
             compression: CompressionType::default(),
             filter: FilterType::default(),
+            icc_profile: Vec::new(),
         }
     }
 
@@ -559,6 +562,7 @@ impl<W: Write> PngEncoder<W> {
             w,
             compression,
             filter,
+            icc_profile: Vec::new(),
         }
     }
 
@@ -604,7 +608,15 @@ impl<W: Write> PngEncoder<W> {
             FilterType::Adaptive => (png::FilterType::Sub, png::AdaptiveFilterType::Adaptive),
         };
 
-        let mut encoder = png::Encoder::new(self.w, width, height);
+        let mut info = png::Info::with_size(width, height);
+
+        if !self.icc_profile.is_empty() {
+            info.icc_profile = Some(Cow::Borrowed(&self.icc_profile));
+        }
+
+        let mut encoder =
+            png::Encoder::with_info(self.w, info).map_err(|e| ImageError::IoError(e.into()))?;
+
         encoder.set_color(ct);
         encoder.set_depth(bits);
         encoder.set_compression(comp);
@@ -669,6 +681,11 @@ impl<W: Write> ImageEncoder for PngEncoder<W> {
             ))),
         }
     }
+
+    fn set_icc_profile(&mut self, icc_profile: Vec<u8>) -> Result<(), UnsupportedError> {
+        self.icc_profile = icc_profile;
+        Ok(())
+    }
 }
 
 impl ImageError {