fix: don't strip ICC from exif (#735)

* fix: don't strip ICC from exif

Signed-off-by: ispik <ispik@ispik.dev>

* fix: refactor ICC profile handling in image processing

Signed-off-by: ispik <ispik@ispik.dev>

---------

Signed-off-by: ispik <ispik@ispik.dev>

crates/services/autumn/Cargo.toml

@@ -18,6 +18,7 @@ kamadak-exif = { workspace = true }
 # revolt_little_exif = "0.5.1"
 image = { workspace = true }
 thumbhash = { workspace = true }
+lcms2 = { workspace = true }
 
 # File processing
 revolt_clamav-client = { workspace = true }

crates/services/autumn/src/exif.rs

@@ -1,13 +1,24 @@
 use std::io::{Cursor, Read};
 
+use crate::utils::apply_icc_profile;
 use exif::Reader;
-use image::{ImageFormat, ImageReader};
+use image::{ImageEncoder, ImageReader};
 use revolt_config::report_internal_error;
 use revolt_database::Metadata;
 use revolt_result::{create_error, Result};
 use tempfile::NamedTempFile;
 use tokio::process::Command;
 
+macro_rules! encode_with_icc {
+    ($encoder:expr, $icc:expr, $image:expr, $width:expr, $height:expr, $color:expr) => {{
+        let mut encoder = $encoder;
+        if let Some(icc) = $icc {
+            let _ = encoder.set_icc_profile(icc.clone());
+        }
+        encoder.write_image($image, $width, $height, $color)
+    }};
+}
+
 /// Strip EXIF data from given file and produce new file and metadata
 pub async fn strip_metadata(
     file: NamedTempFile,
@@ -17,8 +28,8 @@ pub async fn strip_metadata(
 ) -> Result<(Vec<u8>, Metadata)> {
     match &metadata {
         Metadata::Image {
-            width,
-            height,
+            width: _,
+            height: _,
             thumbhash,
             animated,
         } => match mime {
@@ -46,11 +57,12 @@ pub async fn strip_metadata(
                 let mut cursor = Cursor::new(buf);
 
                 // Decode the image
-                let image = report_internal_error!(report_internal_error!(ImageReader::new(
-                    &mut cursor
-                )
-                .with_guessed_format())?
-                .decode());
+                let reader =
+                    report_internal_error!(ImageReader::new(&mut cursor).with_guessed_format())?;
+                let mut decoder = report_internal_error!(reader.into_decoder())?;
+                let mut icc_profile =
+                    report_internal_error!(image::ImageDecoder::icc_profile(&mut decoder))?;
+                let mut image = report_internal_error!(image::DynamicImage::from_decoder(decoder))?;
 
                 // Reset read position
                 cursor.set_position(0);
@@ -71,38 +83,68 @@ pub async fn strip_metadata(
 
                 // Apply the EXIF rotation
                 // See https://jdhao.github.io/2019/07/31/image_rotation_exif_info/
-                report_internal_error!(match &rotation {
-                    2 => image?.fliph(),
-                    3 => image?.rotate180(),
-                    4 => image?.rotate180().fliph(),
-                    5 => image?.rotate90().fliph(),
-                    6 => image?.rotate90(),
-                    7 => image?.rotate270().fliph(),
-                    8 => image?.rotate270(),
-                    _ => image?,
-                }
-                .write_to(
-                    &mut writer,
-                    match mime {
-                        "image/jpeg" => ImageFormat::Jpeg,
-                        "image/png" => ImageFormat::Png,
-                        "image/avif" => ImageFormat::Avif,
-                        "image/tiff" => ImageFormat::Tiff,
-                        _ => todo!(),
-                    },
-                ))?;
-
-                // Calculate dimensions after rotation.
-                let (width, height) = match &rotation {
-                    2 | 4 | 5 | 7 => (*height, *width),
-                    _ => (*width, *height),
+                image = match &rotation {
+                    2 => image.fliph(),
+                    3 => image.rotate180(),
+                    4 => image.rotate180().fliph(),
+                    5 => image.rotate90().fliph(),
+                    6 => image.rotate90(),
+                    7 => image.rotate270().fliph(),
+                    8 => image.rotate270(),
+                    _ => image,
                 };
 
+                if let Some(icc) = &icc_profile {
+                    image = apply_icc_profile(image, icc);
+                    icc_profile = None;
+                }
+
+                let color_type = image.color();
+                let width = image.width();
+                let height = image.height();
+
+                report_internal_error!(match mime {
+                    "image/jpeg" => encode_with_icc!(
+                        image::codecs::jpeg::JpegEncoder::new(&mut writer),
+                        &icc_profile,
+                        image.as_bytes(),
+                        width,
+                        height,
+                        color_type.into()
+                    ),
+                    "image/png" => encode_with_icc!(
+                        image::codecs::png::PngEncoder::new(&mut writer),
+                        &icc_profile,
+                        image.as_bytes(),
+                        width,
+                        height,
+                        color_type.into()
+                    ),
+                    "image/avif" => {
+                        // avif encoder doesn't implement set_icc_profile currently
+                        image::codecs::avif::AvifEncoder::new(&mut writer).write_image(
+                            image.as_bytes(),
+                            width,
+                            height,
+                            color_type.into(),
+                        )
+                    }
+                    "image/tiff" => encode_with_icc!(
+                        image::codecs::tiff::TiffEncoder::new(&mut writer),
+                        &icc_profile,
+                        image.as_bytes(),
+                        width,
+                        height,
+                        color_type.into()
+                    ),
+                    _ => unreachable!(),
+                })?;
+
                 Ok((
                     bytes,
                     Metadata::Image {
-                        width,
-                        height,
+                        width: width as isize,
+                        height: height as isize,
                         thumbhash: thumbhash.clone(),
                         animated: *animated,
                     },

crates/services/autumn/src/main.rs

@@ -18,6 +18,7 @@ pub mod exif;
 pub mod metadata;
 pub mod mime_type;
 mod ratelimits;
+mod utils;
 
 #[derive(FromRef, Clone)]
 struct AppState {

crates/services/autumn/src/metadata.rs

@@ -1,5 +1,6 @@
 use std::io::Cursor;
 
+use crate::utils::apply_icc_profile;
 use image::{GenericImageView, ImageError, ImageReader};
 use revolt_database::Metadata;
 use revolt_files::{image_size, is_animated, video_size};
@@ -27,16 +28,26 @@ pub fn generate_metadata(f: &NamedTempFile, mime_type: &str) -> Metadata {
             .map(|(width, height)| Metadata::Image {
                 width: width as isize,
                 height: height as isize,
-                thumbhash: ImageReader::open(f)
-                    .and_then(|r| r.with_guessed_format())
-                    .map_err(ImageError::from)
-                    .and_then(|r| r.decode())
-                    .map(|img| img.thumbnail(100, 100))
-                    .map(|img| (img.dimensions(), img.to_rgba8().into_raw()))
-                    .map(|((width, height), rgba)| {
-                        thumbhash::rgba_to_thumb_hash(width as usize, height as usize, &rgba)
-                    })
-                    .ok(),
+                thumbhash: (|| {
+                    let reader = ImageReader::open(f).ok()?.with_guessed_format().ok()?;
+                    let mut decoder = reader.into_decoder().ok()?;
+                    let icc_profile = image::ImageDecoder::icc_profile(&mut decoder)
+                        .ok()
+                        .flatten();
+                    let mut img = image::DynamicImage::from_decoder(decoder).ok()?;
+
+                    if let Some(icc) = icc_profile {
+                        img = apply_icc_profile(img, &icc);
+                    }
+
+                    let img = img.thumbnail(100, 100);
+                    let (width, height) = img.dimensions();
+                    Some(thumbhash::rgba_to_thumb_hash(
+                        width as usize,
+                        height as usize,
+                        &img.into_rgba8().into_raw(),
+                    ))
+                })(),
                 animated: is_animated(f, mime_type).or(Some(false)),
             })
             .unwrap_or_default()

crates/services/autumn/src/utils.rs

@@ -0,0 +1,31 @@
+/// Convert image to sRGB using the provided ICC profile.
+/// Returns the converted image, or the original if conversion fails.
+pub fn apply_icc_profile(image: image::DynamicImage, icc: &[u8]) -> image::DynamicImage {
+    let Ok(src_profile) = lcms2::Profile::new_icc(icc) else {
+        return image;
+    };
+    let dst_profile = lcms2::Profile::new_srgb();
+    let format = if image.color().has_alpha() {
+        lcms2::PixelFormat::RGBA_8
+    } else {
+        lcms2::PixelFormat::RGB_8
+    };
+    let Ok(t) = lcms2::Transform::new(
+        &src_profile,
+        format,
+        &dst_profile,
+        format,
+        lcms2::Intent::Perceptual,
+    ) else {
+        return image;
+    };
+    if image.color().has_alpha() {
+        let mut rgba_image = image.into_rgba8();
+        t.transform_in_place(rgba_image.as_mut());
+        image::DynamicImage::ImageRgba8(rgba_image)
+    } else {
+        let mut rgb_image = image.into_rgb8();
+        t.transform_in_place(rgb_image.as_mut());
+        image::DynamicImage::ImageRgb8(rgb_image)
+    }
+}