diff --git a/Cargo.lock b/Cargo.lock index 671d6a9d..b55dab06 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3068,7 +3068,28 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" dependencies = [ - "foreign-types-shared", + "foreign-types-shared 0.1.1", +] + +[[package]] +name = "foreign-types" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965" +dependencies = [ + "foreign-types-macros", + "foreign-types-shared 0.3.1", +] + +[[package]] +name = "foreign-types-macros" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a5c6c585bc94aaf2c7b51dd4c2ba22680844aba4c687be581871a6f518c5742" +dependencies = [ + "proc-macro2", + "quote 1.0.45", + "syn 2.0.117", ] [[package]] @@ -3077,6 +3098,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" +[[package]] +name = "foreign-types-shared" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -4847,6 +4874,29 @@ dependencies = [ "spin 0.9.8", ] +[[package]] +name = "lcms2" +version = "6.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b75877b724685dd49310bdbadbf973fc69b1d01992a6d4a861b928fc3943f87b" +dependencies = [ + "bytemuck", + "foreign-types 0.5.0", + "lcms2-sys", +] + +[[package]] +name = "lcms2-sys" +version = "4.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c2604b23848ca80b2add60f0fb2270fd980e622c25029b6597fa01cfd5f8d5f" +dependencies = [ + "cc", + "dunce", + "libc", + "pkg-config", +] + [[package]] name = "leb128fmt" version = "0.1.0" @@ -5894,7 +5944,7 @@ checksum = "951c002c75e16ea2c65b8c7e4d3d51d5530d8dfa7d060b4776828c88cfb18ecf" dependencies = [ "bitflags 2.11.0", "cfg-if", - "foreign-types", + "foreign-types 0.3.2", "libc", "once_cell", "openssl-macros", @@ -7466,6 +7516,7 @@ dependencies = [ "jxl-oxide", "kamadak-exif", "lazy_static", + "lcms2", "moka", "nanoid", "revolt-config", diff --git a/Cargo.toml b/Cargo.toml index 6ce826ef..f7a3e83a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -134,6 +134,7 @@ kamadak-exif = "0.5.4" webp = "0.3.0" image = "0.25.2" # avif encode requires dav1d system library: features = ["avif-native"] thumbhash = "0.1.0" +lcms2 = "6.1.1" # for color profile processing # File processing revolt_clamav-client = "0.1.5" diff --git a/crates/services/autumn/Cargo.toml b/crates/services/autumn/Cargo.toml index 1636a8ec..62b0bbcf 100644 --- a/crates/services/autumn/Cargo.toml +++ b/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 } diff --git a/crates/services/autumn/src/exif.rs b/crates/services/autumn/src/exif.rs index af02953e..be8e6150 100644 --- a/crates/services/autumn/src/exif.rs +++ b/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, 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, }, diff --git a/crates/services/autumn/src/main.rs b/crates/services/autumn/src/main.rs index cb936e0c..e6234318 100644 --- a/crates/services/autumn/src/main.rs +++ b/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 { diff --git a/crates/services/autumn/src/metadata.rs b/crates/services/autumn/src/metadata.rs index 8e0c89cd..c476eac4 100644 --- a/crates/services/autumn/src/metadata.rs +++ b/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() diff --git a/crates/services/autumn/src/utils.rs b/crates/services/autumn/src/utils.rs new file mode 100644 index 00000000..5e2fbfb8 --- /dev/null +++ b/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) + } +}