fix(media-proxy): budget native image work by pod memory (#3385)

fluxer_media_proxy/src/media_process/loaded_image.rs

@@ -1,8 +1,9 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
 use super::native_runtime::clear_vips_error;
-use super::{MediaError, native_status_error};
+use super::{ImageDimensions, MediaError, native_status_error};
 use crate::{
+    constants,
     image_transform::{ImageOptions, ResizeMode},
     media_limits::MediaLimits,
     native,
@@ -55,6 +56,51 @@ pub(crate) fn validate_dimensions_u32(
     Ok(())
 }
 
+pub(crate) fn validate_jxl_decode_pixels(
+    sniffed_mime: &str,
+    dims: ImageDimensions,
+) -> Result<(), MediaError> {
+    if sniffed_mime != "image/jxl" {
+        return Ok(());
+    }
+    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 {
+        return Err(MediaError::InvalidImageDimensions);
+    }
+    Ok(())
+}
+
+pub(crate) fn static_output_pixels(options: &ImageOptions, dims: ImageDimensions) -> usize {
+    let source_width = dims.width as f64;
+    let source_height = dims.height as f64;
+    let source_pixels = dims.width as usize * dims.height as usize;
+    if options.wants_cover_crop()
+        && let (Some(width), Some(height)) = (options.width, options.height)
+    {
+        return (width as usize * height as usize).min(source_pixels);
+    }
+    let width_scale = options
+        .width
+        .map_or(1.0, |width| width as f64 / source_width);
+    let height_scale = options
+        .height
+        .map_or(1.0, |height| height as f64 / source_height);
+    let scale = width_scale.min(height_scale).min(1.0);
+    ((source_width * scale).round() * (source_height * scale).round()) as usize
+}
+
+pub(crate) fn validate_static_output_pixels(
+    options: &ImageOptions,
+    dims: ImageDimensions,
+) -> Result<(), MediaError> {
+    if static_output_pixels(options, dims) > constants::MAX_STATIC_OUTPUT_PIXELS {
+        return Err(MediaError::InvalidImageDimensions);
+    }
+    Ok(())
+}
+
 pub(crate) fn validate_dimensions(
     media_limits: &MediaLimits,
     width: c_int,

fluxer_media_proxy/src/media_process/metadata.rs

@@ -3,7 +3,7 @@
 use super::MediaError;
 use super::av_metadata::{AVMetadata, AVProbe, NSFW_PREVIEW_MAX_DIMENSION, probe_av_metadata};
 use super::image_probe::probe_image_dims;
-use super::loaded_image::validate_dimensions_u32;
+use super::loaded_image::{validate_dimensions_u32, validate_jxl_decode_pixels};
 use super::nsfw_processing::{
     NSFWScanPreparation, NSFWScanSource, classify_nsfw_buffers, nsfw_scan_buffers,
 };
@@ -22,6 +22,7 @@ use sha2::{Digest, Sha256};
 pub struct MetadataOptions {
     pub placeholder: bool,
     pub nsfw: NSFWPolicy,
+    pub deadline_ms: Option<i64>,
 }
 
 impl Default for MetadataOptions {
@@ -29,6 +30,7 @@ impl Default for MetadataOptions {
         Self {
             placeholder: true,
             nsfw: NSFWPolicy::Disabled,
+            deadline_ms: None,
         }
     }
 }
@@ -50,12 +52,12 @@ struct MetadataResponse {
     nsfw_probability: f32,
 }
 
-struct MetadataBlocking {
+pub struct PreparedMetadata {
     response: MetadataResponse,
     nsfw_scan: Option<NSFWScanPreparation>,
 }
 
-impl MetadataBlocking {
+impl PreparedMetadata {
     fn take_nsfw_scan(&mut self) -> Option<NSFWScanPreparation> {
         self.nsfw_scan.take()
     }
@@ -128,7 +130,7 @@ fn probe_av_metadata_without_requiring_a_frame(
     }
 }
 
-fn metadata_blocking(request: MetadataBlockingRequest<'_>) -> Result<MetadataBlocking, MediaError> {
+fn metadata_blocking(request: MetadataBlockingRequest<'_>) -> Result<PreparedMetadata, MediaError> {
     let MetadataBlockingRequest {
         media_limits,
         metrics,
@@ -145,7 +147,9 @@ fn metadata_blocking(request: MetadataBlockingRequest<'_>) -> Result<MetadataBlo
     let initial_category = mime::category(sniffed.mime).ok_or(MediaError::UnsupportedMediaType)?;
     let is_image = initial_category == mime::Category::Image;
     let dims = if is_image {
-        Some(probe_image_dims(media_limits, input)?)
+        let dims = probe_image_dims(media_limits, input)?;
+        validate_jxl_decode_pixels(sniffed.mime, dims)?;
+        Some(dims)
     } else {
         None
     };
@@ -190,7 +194,7 @@ fn metadata_blocking(request: MetadataBlockingRequest<'_>) -> Result<MetadataBlo
     let placeholder = if options.placeholder {
         let hash = if is_image {
             optional_thumbhash(
-                encode_thumbhash(media_limits, input, None),
+                encode_thumbhash(media_limits, input, options.deadline_ms),
                 "image_metadata",
             )
         } else {
@@ -219,7 +223,7 @@ fn metadata_blocking(request: MetadataBlockingRequest<'_>) -> Result<MetadataBlo
             frame_count: frames_count.max(sniffed.frames),
             input,
             duration_seconds: av_probe.as_ref().and_then(|probe| probe.duration_seconds),
-            deadline_ms: None,
+            deadline_ms: options.deadline_ms,
         })?
     } else {
         None
@@ -239,7 +243,7 @@ fn metadata_blocking(request: MetadataBlockingRequest<'_>) -> Result<MetadataBlo
 
     let format = metadata_format(content_type);
     let content_hash = hex::encode(Sha256::digest(input));
-    Ok(MetadataBlocking {
+    Ok(PreparedMetadata {
         response: MetadataResponse {
             content_type: content_type.to_owned(),
             size: input.len(),
@@ -257,13 +261,37 @@ fn metadata_blocking(request: MetadataBlockingRequest<'_>) -> Result<MetadataBlo
     })
 }
 
-fn metadata_finalize(prepared: MetadataBlocking, verdict: NSFWClassification) -> MetadataResponse {
+fn metadata_finalize(prepared: PreparedMetadata, verdict: NSFWClassification) -> MetadataResponse {
     let mut response = prepared.response;
     response.nsfw = verdict.is_nsfw;
     response.nsfw_probability = verdict.probability;
     response
 }
 
+pub fn prepare_metadata(
+    input: &[u8],
+    options: &MetadataOptions,
+    media_limits: &MediaLimits,
+    metrics: &TransformMetrics,
+) -> Result<PreparedMetadata, MediaError> {
+    metadata_blocking(MetadataBlockingRequest {
+        media_limits,
+        metrics,
+        input,
+        options,
+    })
+}
+
+pub async fn finish_metadata(
+    mut prepared: PreparedMetadata,
+    nsfw_client: &NSFWClient,
+) -> Result<String, MediaError> {
+    let scan = prepared.take_nsfw_scan();
+    let verdict = classify_nsfw_buffers(nsfw_client, scan).await?;
+    let response = metadata_finalize(prepared, verdict);
+    serde_json::to_string(&response).map_err(|_| MediaError::MediaEncodeFailed)
+}
+
 pub async fn metadata_json_with_options(
     input: &[u8],
     _filename: &str,
@@ -272,14 +300,6 @@ pub async fn metadata_json_with_options(
     nsfw_client: &NSFWClient,
     metrics: &TransformMetrics,
 ) -> Result<String, MediaError> {
-    let mut prepared = metadata_blocking(MetadataBlockingRequest {
-        media_limits,
-        metrics,
-        input,
-        options: &options,
-    })?;
-    let scan = prepared.take_nsfw_scan();
-    let verdict = classify_nsfw_buffers(nsfw_client, scan).await?;
-    let response = metadata_finalize(prepared, verdict);
-    serde_json::to_string(&response).map_err(|_| MediaError::MediaEncodeFailed)
+    let prepared = prepare_metadata(input, &options, media_limits, metrics)?;
+    finish_metadata(prepared, nsfw_client).await
 }

fluxer_media_proxy/src/media_process/mod.rs

@@ -18,6 +18,7 @@ mod encoding;
 mod image_probe;
 mod loaded_image;
 mod metadata;
+mod native_cost;
 pub(crate) mod native_runtime;
 mod nsfw_processing;
 mod placeholder;
@@ -32,7 +33,11 @@ mod tests;
 pub use av_metadata::{
     AVMetadata, AVMetadataFrame, AVProbe, NSFW_PREVIEW_MAX_DIMENSION, probe_av_metadata,
 };
-pub use metadata::{MetadataOptions, metadata_json_with_options};
+pub use metadata::{
+    MetadataOptions, PreparedMetadata, finish_metadata, metadata_json_with_options,
+    prepare_metadata,
+};
+pub use native_cost::{AV_NATIVE_COST_BYTES, image_transform_cost, metadata_cost};
 pub use nsfw_processing::encode_static_image_for_nsfw;
 pub use transform::transform_image;
 pub use video_thumbnail::{

fluxer_media_proxy/src/media_process/native_cost.rs

@@ -0,0 +1,72 @@
+// 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 crate::{
+    image_quality::ImageQuality, image_transform::ImageOptions, media_limits::MediaLimits, mime,
+    output_format::OutputFormat,
+};
+
+pub const AV_NATIVE_COST_BYTES: usize = 256 << 20;
+const ANIMATED_CANVAS_BUFFERS: usize = 8;
+const RGBA_BYTES_PER_PIXEL: usize = 4;
+
+fn decode_bytes_per_pixel(sniffed_mime: &str) -> usize {
+    match sniffed_mime {
+        "image/jxl" => 40,
+        "image/webp" => 20,
+        "image/avif" | "image/heic" | "image/heif" => 12,
+        "image/jpeg" => 2,
+        "image/svg+xml" => 0,
+        _ => 8,
+    }
+}
+
+fn encode_bytes_per_pixel(format: OutputFormat, quality: ImageQuality) -> usize {
+    match (format, quality) {
+        (OutputFormat::WebP, ImageQuality::Lossless | ImageQuality::Auto) => 48,
+        (OutputFormat::WebP, _) => 28,
+        (OutputFormat::JPEG, _) => 4,
+        _ => 8,
+    }
+}
+
+pub fn image_transform_cost(
+    input: &[u8],
+    options: &ImageOptions,
+    media_limits: &MediaLimits,
+) -> Result<usize, MediaError> {
+    ensure_vips_init()?;
+    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 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
+            + output_pixels * encode
+    } else {
+        source_pixels * decode_bytes_per_pixel(sniffed_mime) + output_pixels * encode
+    };
+    Ok(input.len().saturating_add(pixel_cost))
+}
+
+pub fn metadata_cost(input: &[u8], media_limits: &MediaLimits) -> Result<usize, MediaError> {
+    let sniffed_mime = mime::sniff(input).mime;
+    match mime::category(sniffed_mime) {
+        Some(mime::Category::Image) => {
+            ensure_vips_init()?;
+            let dims = probe_image_dims(media_limits, input)?;
+            let source_pixels = dims.width as usize * dims.height as usize;
+            Ok(input
+                .len()
+                .saturating_add(source_pixels * decode_bytes_per_pixel(sniffed_mime)))
+        }
+        Some(mime::Category::Video | mime::Category::Audio) => {
+            Ok(input.len().saturating_add(AV_NATIVE_COST_BYTES))
+        }
+        _ => Ok(input.len()),
+    }
+}

fluxer_media_proxy/src/media_process/tests/metadata.rs

@@ -64,6 +64,7 @@ fn metadata_json_returns_unavailable_when_nsfw_service_fails() {
                 MetadataOptions {
                     placeholder: false,
                     nsfw: NSFWPolicy::enabled(0.85).expect("valid threshold"),
+                    deadline_ms: None,
                 },
                 &test_media_limits(),
                 &client,
@@ -237,3 +238,23 @@ fn metadata_succeeds_without_a_placeholder_when_thumbhash_generation_fails() {
     assert_eq!(64, value["height"]);
     assert_eq!(None, value.get("placeholder"));
 }
+
+#[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::InvalidImageDimensions, err);
+}

fluxer_media_proxy/src/media_process/tests/transform.rs

@@ -1,14 +1,16 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 
-use super::super::ImageOptions;
 use super::super::image_probe::load_image;
 use super::super::native_runtime::ensure_vips_init;
 use super::super::transform::source_supports_pages;
+use super::super::{ImageOptions, MediaError, image_transform_cost};
 use super::fixtures::{animated_mode, metadata_value, transform_image};
 use crate::{
     mime, native,
     output_format::OutputFormat,
-    test_fixtures::{synthetic_bmp, synthetic_png},
+    test_fixtures::{
+        JXL_64_LOSSY_ALPHA, JXL_4100_LOSSY_ALPHA, png_dimensions, synthetic_bmp, synthetic_png,
+    },
 };
 
 #[test]
@@ -103,3 +105,105 @@ fn source_supports_pages_matches_libvips_loader_list() {
     assert!(!source_supports_pages("image/bmp"));
     assert!(!source_supports_pages("application/octet-stream"));
 }
+
+#[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::InvalidImageDimensions, err);
+    }
+}
+
+#[test]
+fn transforms_jxl_within_the_decode_pixel_cap() {
+    let out = transform_image(
+        JXL_64_LOSSY_ALPHA,
+        &ImageOptions {
+            width: Some(32),
+            format: OutputFormat::WebP,
+            ..Default::default()
+        },
+    )
+    .unwrap();
+    assert_eq!("image/webp", out.content_type);
+}
+
+#[test]
+fn refuses_full_size_svg_raster_above_the_static_output_cap() {
+    let svg = br#"<svg xmlns="http://www.w3.org/2000/svg" width="16383" height="16383"></svg>"#;
+    let err = transform_image(
+        svg,
+        &ImageOptions {
+            format: OutputFormat::WebP,
+            ..Default::default()
+        },
+    )
+    .unwrap_err();
+    assert_eq!(MediaError::InvalidImageDimensions, err);
+    let err = transform_image(
+        svg,
+        &ImageOptions {
+            width: Some(16384),
+            format: OutputFormat::WebP,
+            ..Default::default()
+        },
+    )
+    .unwrap_err();
+    assert_eq!(MediaError::InvalidImageDimensions, err);
+}
+
+#[test]
+fn rasterizes_large_svg_into_a_bounded_box() {
+    let svg = br#"<svg xmlns="http://www.w3.org/2000/svg" width="16383" height="16383"></svg>"#;
+    let out = transform_image(
+        svg,
+        &ImageOptions {
+            width: Some(64),
+            height: Some(64),
+            format: OutputFormat::PNG,
+            ..Default::default()
+        },
+    )
+    .unwrap();
+    assert_eq!(Some((64, 64)), png_dimensions(&out.bytes));
+}
+
+#[test]
+fn transform_cost_charges_the_full_source_decode_for_jxl() {
+    let cost = image_transform_cost(
+        JXL_4100_LOSSY_ALPHA,
+        &ImageOptions {
+            width: Some(128),
+            format: OutputFormat::WebP,
+            ..Default::default()
+        },
+        &super::fixtures::test_media_limits(),
+    )
+    .unwrap();
+    assert!(cost >= 4100 * 4100 * 40, "cost {cost}");
+}
+
+#[test]
+fn transform_cost_charges_svg_at_its_output_size() {
+    let svg = br#"<svg xmlns="http://www.w3.org/2000/svg" width="16383" height="16383"></svg>"#;
+    let cost = image_transform_cost(
+        svg,
+        &ImageOptions {
+            width: Some(64),
+            height: Some(64),
+            format: OutputFormat::WebP,
+            ..Default::default()
+        },
+        &super::fixtures::test_media_limits(),
+    )
+    .unwrap();
+    assert!(cost < 1 << 20, "cost {cost}");
+}

fluxer_media_proxy/src/media_process/transform.rs

@@ -8,10 +8,13 @@ use super::encoding::{
     VipsEncodeRequest, anim_limits_from_options, encode_vips_image,
     try_transform_animated_webp_direct,
 };
-use super::image_probe::{animated_probe_from_image, load_image, probe_animated, try_decode_bmp};
+use super::image_probe::{
+    animated_probe_from_image, load_image, probe_animated, probe_image_dims, try_decode_bmp,
+};
 use super::loaded_image::{
     normalize_vips_image_to_uchar, page_height, resize_loaded_image, resize_loaded_image_by_scale,
-    validate_dimensions_u32, validate_vips_image,
+    validate_dimensions_u32, validate_jxl_decode_pixels, validate_static_output_pixels,
+    validate_vips_image,
 };
 use super::native_runtime::{clear_vips_error, ensure_vips_init, last_vips_error};
 use super::transform_plan::{
@@ -77,6 +80,54 @@ fn heif_source_may_be_a_sequence(sniffed: mime::SniffInfo) -> bool {
     }
 }
 
+fn ensure_decode_cost_admissible(
+    input: &[u8],
+    sniffed_mime: &str,
+    options: &ImageOptions,
+    media_limits: &MediaLimits,
+) -> Result<(), MediaError> {
+    let animated = options.is_animated();
+    if animated && sniffed_mime != "image/jxl" {
+        return Ok(());
+    }
+    let dims = probe_image_dims(media_limits, input)?;
+    let source_pixels = dims.width as usize * dims.height as usize;
+    if source_pixels >= constants::LARGE_SOURCE_LOG_PIXELS {
+        tracing::info!(
+            mime = sniffed_mime,
+            input_bytes = input.len(),
+            width = dims.width,
+            height = dims.height,
+            pages = dims.pages,
+            out_width = options.width.unwrap_or(0),
+            out_height = options.height.unwrap_or(0),
+            animated,
+            "large native transform starting"
+        );
+    }
+    let admitted = validate_jxl_decode_pixels(sniffed_mime, dims).and_then(|()| {
+        if animated {
+            Ok(())
+        } else {
+            validate_static_output_pixels(options, dims)
+        }
+    });
+    if admitted.is_err() {
+        tracing::warn!(
+            mime = sniffed_mime,
+            input_bytes = input.len(),
+            width = dims.width,
+            height = dims.height,
+            pages = dims.pages,
+            out_width = options.width.unwrap_or(0),
+            out_height = options.height.unwrap_or(0),
+            animated,
+            "native transform refused by decode cost limits"
+        );
+    }
+    admitted
+}
+
 pub fn transform_image(
     input: &[u8],
     options: &ImageOptions,
@@ -99,6 +150,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 format = effective_transform_format(sniffed.mime, options.format, animated);
     let animated_avif = sniffed.mime == "image/avif" && sniffed.animated;
     let full_canvas_animation = animated