fix(media-proxy): serve large simple jxl, fit full size (#3408)

fluxer_media_proxy/src/media_process/image_probe.rs

@@ -15,9 +15,24 @@ pub(crate) fn probe_animated(
     media_limits: &MediaLimits,
     input: &[u8],
 ) -> Result<Option<AnimatedProbe>, MediaError> {
+    Ok(
+        probe_header(media_limits, input)?.map(|dims| AnimatedProbe {
+            width: dims.width as c_int,
+            height: dims.height as c_int,
+            pages: dims.pages as c_int,
+        }),
+    )
+}
+
+fn probe_header(
+    media_limits: &MediaLimits,
+    input: &[u8],
+) -> Result<Option<ImageDimensions>, MediaError> {
     let mut width = 0;
     let mut height = 0;
     let mut pages = 0;
+    let mut bands = 0;
+    let mut format = 0;
     let rc = unsafe {
         native::fluxer_vips_probe_animated(
             input.as_ptr().cast(),
@@ -25,6 +40,8 @@ pub(crate) fn probe_animated(
             &mut width,
             &mut height,
             &mut pages,
+            &mut bands,
+            &mut format,
         )
     };
     match native::NativeStatus::from_code(rc) {
@@ -39,14 +56,16 @@ pub(crate) fn probe_animated(
                 .expect("non-success native status must map to an error"));
         }
     }
-    if pages <= 0 || width <= 0 || height <= 0 {
+    if pages <= 0 || width <= 0 || height <= 0 || bands <= 0 {
         return Err(MediaError::InvalidImageDimensions);
     }
     validate_animation_geometry(media_limits, width as u32, height as u32, pages as u32)?;
-    Ok(Some(AnimatedProbe {
-        width,
-        height,
-        pages,
+    Ok(Some(ImageDimensions {
+        width: width as u32,
+        height: height as u32,
+        pages: pages as u32,
+        bands: bands as u32,
+        eight_bit: format == unsafe { native::fluxer_vips_format_uchar },
     }))
 }
 
@@ -55,12 +74,8 @@ pub(crate) fn probe_image_dims(
     input: &[u8],
 ) -> Result<ImageDimensions, MediaError> {
     ensure_vips_init()?;
-    if let Some(probe) = probe_animated(media_limits, input)? {
-        return Ok(ImageDimensions {
-            width: probe.width as u32,
-            height: probe.height as u32,
-            pages: probe.pages as u32,
-        });
+    if let Some(dims) = probe_header(media_limits, input)? {
+        return Ok(dims);
     }
     if mime::sniff(input).mime == "image/bmp"
         && let Some(image) = try_decode_bmp(input, media_limits)?
@@ -89,6 +104,10 @@ fn probe_loaded_image_dims(
         width,
         height,
         pages,
+        bands: unsafe { native::fluxer_vips_image_get_bands(image.as_ptr()) } as u32,
+        eight_bit: unsafe {
+            native::fluxer_vips_image_get_format(image.as_ptr()) == native::fluxer_vips_format_uchar
+        },
     })
 }
 

fluxer_media_proxy/src/media_process/loaded_image.rs

@@ -66,7 +66,13 @@ pub(crate) fn validate_jxl_decode_pixels(
     let pixels = (dims.width as usize)
         .saturating_mul(dims.height as usize)
         .saturating_mul(dims.pages.max(1) as usize);
-    if pixels > constants::MAX_JXL_DECODE_PIXELS {
+    let simple = dims.pages == 1 && matches!(dims.bands, 1 | 3) && dims.eight_bit;
+    let cap = if simple {
+        constants::MAX_JXL_SIMPLE_DECODE_PIXELS
+    } else {
+        constants::MAX_JXL_DECODE_PIXELS
+    };
+    if pixels > cap {
         return Err(MediaError::TooLargeToProcess);
     }
     Ok(())

fluxer_media_proxy/src/media_process/mod.rs

@@ -138,6 +138,8 @@ pub struct ImageDimensions {
     pub width: u32,
     pub height: u32,
     pub pages: u32,
+    pub bands: u32,
+    pub eight_bit: bool,
 }
 
 #[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]

fluxer_media_proxy/src/media_process/native_cost.rs

@@ -1,9 +1,9 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
-use super::MediaError;
 use super::image_probe::probe_image_dims;
 use super::loaded_image::static_output_pixels;
 use super::native_runtime::ensure_vips_init;
+use super::{ImageDimensions, MediaError};
 use crate::{
     constants, image_quality::ImageQuality, image_transform::ImageOptions,
     media_limits::MediaLimits, mime, output_format::OutputFormat,
@@ -48,7 +48,8 @@ pub fn image_transform_cost(
     let sniffed_mime = mime::sniff(input).mime;
     let dims = probe_image_dims(media_limits, input)?;
     let source_pixels = dims.width as usize * dims.height as usize;
-    let output_pixels = static_output_pixels(options, dims);
+    let options = fit_jxl_output(input.len(), sniffed_mime, options, dims, media_limits);
+    let output_pixels = static_output_pixels(&options, dims);
     let encode = encode_bytes_per_pixel(options.format, options.quality);
     let pixel_cost = if options.is_animated() && dims.pages > 1 {
         source_pixels.max(output_pixels) * RGBA_BYTES_PER_PIXEL * ANIMATED_CANVAS_BUFFERS
@@ -59,6 +60,38 @@ pub fn image_transform_cost(
     Ok(NativeCost::Bytes(input.len().saturating_add(pixel_cost)))
 }
 
+pub(super) fn fit_jxl_output(
+    input_len: usize,
+    sniffed_mime: &str,
+    options: &ImageOptions,
+    dims: ImageDimensions,
+    media_limits: &MediaLimits,
+) -> ImageOptions {
+    let mut fitted = *options;
+    if sniffed_mime != "image/jxl" || options.is_animated() {
+        return fitted;
+    }
+    let source_pixels = dims.width as usize * dims.height as usize;
+    let affordable_pixels = media_limits
+        .native_memory_budget_bytes()
+        .saturating_sub(input_len)
+        .saturating_sub(source_pixels * decode_bytes_per_pixel(sniffed_mime))
+        / encode_bytes_per_pixel(options.format, options.quality);
+    let max_pixels = affordable_pixels.min(constants::MAX_STATIC_OUTPUT_PIXELS);
+    if max_pixels == 0 || static_output_pixels(options, dims) <= max_pixels {
+        return fitted;
+    }
+    let (width, height) = match (options.wants_cover_crop(), options.width, options.height) {
+        (true, Some(width), Some(height)) => (width, height),
+        _ => (dims.width, dims.height),
+    };
+    let (width, height) = (f64::from(width), f64::from(height));
+    let shrink = (max_pixels as f64 / (width * height)).sqrt();
+    fitted.width = Some(((width * shrink).floor() as u32).max(1));
+    fitted.height = Some(((height * shrink).floor() as u32).max(1));
+    fitted
+}
+
 pub fn av_native_cost(input_len: usize) -> NativeCost {
     NativeCost::UpToBudget(
         input_len

fluxer_media_proxy/src/media_process/tests/animated_apng.rs

@@ -57,7 +57,9 @@ fn animated_apng_input_transforms_through_ffmpeg_decode_path() {
         ImageDimensions {
             width: 8,
             height: 8,
-            pages: 1
+            pages: 1,
+            bands: 3,
+            eight_bit: true,
         },
         probe_image_dims(&test_media_limits(), &static_webp.bytes).unwrap()
     );

fluxer_media_proxy/src/media_process/tests/heif.rs

@@ -37,7 +37,9 @@ fn animated_heic_sequence_decodes_every_frame_and_encodes_an_animation() {
         ImageDimensions {
             width: 64,
             height: 64,
-            pages: 4
+            pages: 4,
+            bands: 3,
+            eight_bit: true,
         },
         probe_image_dims(&test_media_limits(), &animated.bytes).expect("probes webp")
     );
@@ -85,7 +87,9 @@ fn still_heic_decodes_as_a_single_frame_when_animation_is_requested() {
         ImageDimensions {
             width: 64,
             height: 64,
-            pages: 1
+            pages: 1,
+            bands: 3,
+            eight_bit: true,
         },
         probe_image_dims(&test_media_limits(), &still.bytes).expect("probes webp")
     );
@@ -113,7 +117,9 @@ fn avif_direct_decode_applies_resize_and_crop() {
         ImageDimensions {
             width: 32,
             height: 24,
-            pages: 1
+            pages: 1,
+            bands: 3,
+            eight_bit: true,
         },
         probe_image_dims(&test_media_limits(), &resized.bytes).unwrap()
     );
@@ -135,7 +141,9 @@ fn avif_direct_decode_applies_resize_and_crop() {
         ImageDimensions {
             width: 32,
             height: 32,
-            pages: 1
+            pages: 1,
+            bands: 3,
+            eight_bit: true,
         },
         probe_image_dims(&test_media_limits(), &cropped.bytes).unwrap()
     );

fluxer_media_proxy/src/media_process/tests/metadata.rs

@@ -14,9 +14,11 @@ use crate::{
     nsfw::{NSFWClient, NSFWPolicy},
     output_format::OutputFormat,
     test_fixtures::{
-        fixture_audio_mp3_with_png_cover_art, fixture_audio_mp4_with_attached_picture,
-        fixture_audio_only_mp4, fixture_h264_mp4, fixture_mkv_with_png_video_stream,
-        fixture_mp4_with_undecodable_video, synthetic_png, synthetic_wav,
+        JXL_3000_LOSSY_2_FRAMES, JXL_4100_LOSSLESS_RGB16, JXL_4100_LOSSY_ALPHA,
+        JXL_4100_LOSSY_RGB_DEPTH, JXL_7000X4000_LOSSY_RGB, fixture_audio_mp3_with_png_cover_art,
+        fixture_audio_mp4_with_attached_picture, fixture_audio_only_mp4, fixture_h264_mp4,
+        fixture_mkv_with_png_video_stream, fixture_mp4_with_undecodable_video, synthetic_png,
+        synthetic_wav,
     },
 };
 use std::{io::Write as _, sync::Arc};
@@ -252,22 +254,38 @@ fn metadata_succeeds_without_a_placeholder_when_thumbhash_generation_fails() {
 
 #[test]
 fn metadata_refuses_jxl_above_the_decode_pixel_cap() {
-    let err = tokio::runtime::Builder::new_current_thread()
-        .build()
-        .unwrap()
-        .block_on(async {
-            metadata_json_with_options(
-                crate::test_fixtures::JXL_4100_LOSSY_ALPHA,
-                "big.jxl",
-                MetadataOptions::default(),
-                &test_media_limits(),
-                &NSFWClient::disabled(),
-                &TransformMetrics::new(),
-            )
-            .await
-            .unwrap_err()
-        });
-    assert_eq!(MediaError::TooLargeToProcess, err);
+    for jxl in [
+        JXL_4100_LOSSY_ALPHA,
+        JXL_4100_LOSSY_RGB_DEPTH,
+        JXL_4100_LOSSLESS_RGB16,
+        JXL_3000_LOSSY_2_FRAMES,
+    ] {
+        let err = tokio::runtime::Builder::new_current_thread()
+            .build()
+            .unwrap()
+            .block_on(async {
+                metadata_json_with_options(
+                    jxl,
+                    "big.jxl",
+                    MetadataOptions::default(),
+                    &test_media_limits(),
+                    &NSFWClient::disabled(),
+                    &TransformMetrics::new(),
+                )
+                .await
+                .unwrap_err()
+            });
+        assert_eq!(MediaError::TooLargeToProcess, err);
+    }
+}
+
+#[test]
+fn metadata_admits_8_bit_rgb_jxl_above_the_decode_pixel_cap() {
+    let value = metadata_value(JXL_7000X4000_LOSSY_RGB, "big.jxl");
+    assert_eq!("image/jxl", value["content_type"]);
+    assert_eq!(7000, value["width"]);
+    assert_eq!(4000, value["height"]);
+    assert!(value["placeholder"].as_str().is_some_and(|p| !p.is_empty()));
 }
 
 #[test]

fluxer_media_proxy/src/media_process/tests/transform.rs

@@ -1,16 +1,22 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
-use super::super::image_probe::load_image;
+use super::super::image_probe::{load_image, probe_image_dims};
+use super::super::loaded_image::validate_jxl_decode_pixels;
 use super::super::native_runtime::ensure_vips_init;
 use super::super::transform::source_supports_pages;
-use super::super::{ImageOptions, MediaError, NativeCost, image_transform_cost};
-use super::fixtures::{animated_mode, metadata_value, png_zlib_level, transform_image};
+use super::super::{ImageDimensions, ImageOptions, MediaError, NativeCost, image_transform_cost};
+use super::fixtures::{
+    animated_mode, metadata_value, png_zlib_level, test_media_limits, transform_image,
+};
 use crate::{
+    constants,
+    metrics::transform::TransformMetrics,
     mime, native,
     output_format::OutputFormat,
     test_fixtures::{
-        JPEG_ROTATED_MISMATCHED_SCAN, JXL_64_LOSSY_ALPHA, JXL_4100_LOSSY_ALPHA, png_dimensions,
-        synthetic_bmp, synthetic_png,
+        JPEG_ROTATED_MISMATCHED_SCAN, JXL_64_LOSSY_ALPHA, JXL_3000_LOSSY_2_FRAMES,
+        JXL_4100_LOSSLESS_RGB16, JXL_4100_LOSSY_ALPHA, JXL_4100_LOSSY_RGB_DEPTH,
+        JXL_7000X4000_LOSSY_RGB, png_dimensions, synthetic_bmp, synthetic_png,
     },
 };
 
@@ -109,20 +115,132 @@ fn source_supports_pages_matches_libvips_loader_list() {
 
 #[test]
 fn refuses_jxl_above_the_decode_pixel_cap_before_decoding() {
-    for width in [None, Some(128)] {
-        let err = transform_image(
-            JXL_4100_LOSSY_ALPHA,
-            &ImageOptions {
-                width,
-                format: OutputFormat::WebP,
-                ..Default::default()
-            },
-        )
-        .unwrap_err();
-        assert_eq!(MediaError::TooLargeToProcess, err);
+    for jxl in [
+        JXL_4100_LOSSY_ALPHA,
+        JXL_4100_LOSSY_RGB_DEPTH,
+        JXL_4100_LOSSLESS_RGB16,
+        JXL_3000_LOSSY_2_FRAMES,
+    ] {
+        for width in [None, Some(128)] {
+            let err = transform_image(
+                jxl,
+                &ImageOptions {
+                    width,
+                    format: OutputFormat::WebP,
+                    ..Default::default()
+                },
+            )
+            .unwrap_err();
+            assert_eq!(MediaError::TooLargeToProcess, err);
+        }
     }
 }
 
+#[test]
+fn jxl_decode_cap_depends_on_the_header_layout() {
+    let dims = |width, height, pages, bands, eight_bit| ImageDimensions {
+        width,
+        height,
+        pages,
+        bands,
+        eight_bit,
+    };
+    let simple_cap = (8192, 6144);
+    let complex_cap = (4096, 4096);
+    for bands in [1, 3] {
+        let (width, height) = simple_cap;
+        assert_eq!(
+            Ok(()),
+            validate_jxl_decode_pixels("image/jxl", dims(width, height, 1, bands, true))
+        );
+        assert_eq!(
+            Err(MediaError::TooLargeToProcess),
+            validate_jxl_decode_pixels("image/jxl", dims(width, height + 1, 1, bands, true))
+        );
+    }
+    for (pages, bands, eight_bit) in [(1, 2, true), (1, 4, true), (1, 3, false), (2, 3, true)] {
+        let (width, height) = complex_cap;
+        assert_eq!(
+            Ok(()),
+            validate_jxl_decode_pixels(
+                "image/jxl",
+                dims(width, height / pages, pages, bands, eight_bit)
+            )
+        );
+        assert_eq!(
+            Err(MediaError::TooLargeToProcess),
+            validate_jxl_decode_pixels(
+                "image/jxl",
+                dims(width, height / pages + 1, pages, bands, eight_bit)
+            )
+        );
+    }
+    assert_eq!(
+        Ok(()),
+        validate_jxl_decode_pixels("image/png", dims(16384, 16384, 1, 4, false))
+    );
+}
+
+#[test]
+fn thumbnails_8_bit_rgb_jxl_above_the_decode_pixel_cap() {
+    let out = transform_image(
+        JXL_7000X4000_LOSSY_RGB,
+        &ImageOptions {
+            width: Some(550),
+            format: OutputFormat::WebP,
+            ..Default::default()
+        },
+    )
+    .unwrap();
+    assert_eq!("image/webp", out.content_type);
+    let dims = probe_image_dims(&test_media_limits(), &out.bytes).unwrap();
+    assert_eq!((550, 314), (dims.width, dims.height));
+}
+
+#[test]
+fn full_size_jxl_above_the_static_output_cap_is_downscaled_to_fit() {
+    let options = ImageOptions {
+        format: OutputFormat::WebP,
+        ..Default::default()
+    };
+    let out = transform_image(JXL_7000X4000_LOSSY_RGB, &options).unwrap();
+    assert_eq!("image/webp", out.content_type);
+    let dims = probe_image_dims(&test_media_limits(), &out.bytes).unwrap();
+    assert_eq!((6636, 3792), (dims.width, dims.height));
+    assert!(dims.width as usize * dims.height as usize <= constants::MAX_STATIC_OUTPUT_PIXELS);
+
+    let cost =
+        image_transform_cost(JXL_7000X4000_LOSSY_RGB, &options, &test_media_limits()).unwrap();
+    assert_eq!(
+        NativeCost::Bytes(JXL_7000X4000_LOSSY_RGB.len() + 7000 * 4000 * 40 + 6636 * 3792 * 28),
+        cost
+    );
+}
+
+#[test]
+fn full_size_jxl_is_downscaled_to_fit_the_native_memory_budget() {
+    let options = ImageOptions {
+        format: OutputFormat::WebP,
+        ..Default::default()
+    };
+    let decode = JXL_7000X4000_LOSSY_RGB.len() + 7000 * 4000 * 40;
+    let budget = decode + 10_000_000 * 28;
+    let media_limits = test_media_limits().with_native_memory_budget(budget);
+    let out = super::super::transform_image(
+        JXL_7000X4000_LOSSY_RGB,
+        &options,
+        &media_limits,
+        &TransformMetrics::new(),
+    )
+    .unwrap();
+    let dims = probe_image_dims(&test_media_limits(), &out.bytes).unwrap();
+    assert_eq!((4183, 2390), (dims.width, dims.height));
+    assert_eq!(
+        NativeCost::Bytes(decode + 4183 * 2390 * 28),
+        image_transform_cost(JXL_7000X4000_LOSSY_RGB, &options, &media_limits).unwrap()
+    );
+}
+
 #[test]
 fn transforms_jxl_within_the_decode_pixel_cap() {
     let out = transform_image(

fluxer_media_proxy/src/media_process/transform.rs

@@ -16,6 +16,7 @@ use super::loaded_image::{
     static_output_pixels, validate_dimensions_u32, validate_jxl_decode_pixels,
     validate_static_output_pixels, validate_vips_image,
 };
+use super::native_cost::fit_jxl_output;
 use super::native_runtime::{clear_vips_error, ensure_vips_init, last_vips_error};
 use super::transform_plan::{
     AutoQualitySource, output_is_sdr, resolve_auto_quality, should_transform_animated_webp_direct,
@@ -80,17 +81,18 @@ fn heif_source_may_be_a_sequence(sniffed: mime::SniffInfo) -> bool {
     }
 }
 
-fn ensure_decode_cost_admissible(
+fn admitted_transform_options(
     input: &[u8],
     sniffed_mime: &str,
     options: &ImageOptions,
     media_limits: &MediaLimits,
-) -> Result<(), MediaError> {
+) -> Result<ImageOptions, MediaError> {
     let animated = options.is_animated();
     if animated && sniffed_mime != "image/jxl" {
-        return Ok(());
+        return Ok(*options);
     }
     let dims = probe_image_dims(media_limits, input)?;
+    let options = &fit_jxl_output(input.len(), sniffed_mime, options, dims, media_limits);
     let source_pixels = dims.width as usize * dims.height as usize;
     if source_pixels >= constants::LARGE_SOURCE_LOG_PIXELS {
         tracing::info!(
@@ -125,7 +127,7 @@ fn ensure_decode_cost_admissible(
             "native transform refused by decode cost limits"
         );
     }
-    admitted
+    admitted.map(|()| *options)
 }
 
 fn ensure_animated_output_admissible(
@@ -175,7 +177,7 @@ pub fn transform_image(
     ensure_vips_init()?;
     let sniffed = mime::sniff(input);
     let animated = options.is_animated();
-    ensure_decode_cost_admissible(input, sniffed.mime, options, media_limits)?;
+    let options = &admitted_transform_options(input, sniffed.mime, options, media_limits)?;
     let format = effective_transform_format(sniffed.mime, options.format, animated);
     let animated_avif = sniffed.mime == "image/avif" && sniffed.animated;
     let full_canvas_animation = animated