diff --git a/server-rs/crates/platform-matting/src/lib.rs b/server-rs/crates/platform-matting/src/lib.rs index 82f5b0535..80bdaf407 100644 --- a/server-rs/crates/platform-matting/src/lib.rs +++ b/server-rs/crates/platform-matting/src/lib.rs @@ -38,6 +38,15 @@ const MAX_INPUT_BYTES: usize = 3 * 1024 * 1024; /// SegmentCommonImage 要求每条边大于 32 像素。 const MIN_INPUT_EDGE: u32 = 32; +/// 抠图结果下载的最大字节数。上游是阿里云返回的原尺寸 RGBA PNG,正常远小于此值; +/// 设上限是为了防止异常 / 恶意上游用超大响应撑爆 API 进程内存(逐帧动画可并发多次调用)。 +/// 与 api-server BgFilter 降级路径的 `EDITOR_BACKGROUND_REMOVAL_MAX_RESPONSE_BYTES` 对齐。 +const MAX_RESULT_RESPONSE_BYTES: usize = 32 * 1024 * 1024; + +/// 抠图结果解码的最大边长像素。防止解码阶段按声明的巨幅尺寸分配像素缓冲。 +/// 与 api-server BgFilter 路径的 `EDITOR_BACKGROUND_REMOVAL_MAX_IMAGE_DIMENSION` 对齐。 +const MAX_RESULT_IMAGE_DIMENSION: u32 = 8192; + #[derive(Clone, Debug)] pub struct MattingConfig { pub endpoint: String, @@ -326,11 +335,7 @@ impl MattingClient { }) .await?; let result_bytes = self.download_result_image(&result.image_url).await?; - let result_image = image::load_from_memory(&result_bytes) - .map_err(|error| { - MattingError::Upstream(format!("解析抠图结果图片失败:{error}")) - })? - .to_rgba8(); + let result_image = decode_result_image(&result_bytes)?.to_rgba8(); if !downscaled { if result_image.dimensions() != (source_width, source_height) { @@ -360,7 +365,7 @@ impl MattingClient { } async fn download_result_image(&self, url: &str) -> Result, MattingError> { - let response = self + let mut response = self .client .get(url) .send() @@ -375,13 +380,24 @@ impl MattingClient { status.as_u16() ))); } - response - .bytes() - .await - .map(|bytes| bytes.to_vec()) - .map_err(|error| { - MattingError::Upstream(describe_result_download_body_error(&error)) - }) + // 先按 Content-Length 快速拒绝,再流式累加做兜底:不信任上游声明的长度, + // 逐块累计超阈值立即中断,避免 response.bytes() 一次性分配任意大小响应撑爆内存。 + if response + .content_length() + .is_some_and(|length| length > MAX_RESULT_RESPONSE_BYTES as u64) + { + return Err(result_response_too_large_error()); + } + let mut bytes = Vec::new(); + while let Some(chunk) = response.chunk().await.map_err(|error| { + MattingError::Upstream(describe_result_download_body_error(&error)) + })? { + if bytes.len().saturating_add(chunk.len()) > MAX_RESULT_RESPONSE_BYTES { + return Err(result_response_too_large_error()); + } + bytes.extend_from_slice(chunk.as_ref()); + } + Ok(bytes) } /// 把本地图片字节上传到 VIAPI 官方临时桶,返回可直接作为 ImageURL 的公网地址。 @@ -631,6 +647,31 @@ fn encode_rgba_png(image: &image::RgbaImage) -> Result, MattingError> { Ok(encoded) } +fn result_response_too_large_error() -> MattingError { + MattingError::Upstream(format!( + "抠图结果响应过大,超过 {MAX_RESULT_RESPONSE_BYTES} 字节上限" + )) +} + +/// 解码抠图结果时套上尺寸 / 分配上限,防止上游用巨幅尺寸声明在解码阶段撑爆内存。 +fn decode_result_image(bytes: &[u8]) -> Result { + use std::io::Cursor; + + let mut reader = image::ImageReader::new(Cursor::new(bytes)) + .with_guessed_format() + .map_err(|error| { + MattingError::Upstream(format!("识别抠图结果格式失败:{error}")) + })?; + let mut limits = image::Limits::default(); + limits.max_image_width = Some(MAX_RESULT_IMAGE_DIMENSION); + limits.max_image_height = Some(MAX_RESULT_IMAGE_DIMENSION); + limits.max_alloc = Some(MAX_RESULT_RESPONSE_BYTES as u64 * 4); + reader.limits(limits); + reader.decode().map_err(|error| { + MattingError::Upstream(format!("解析抠图结果图片失败:{error}")) + }) +} + struct ViapiStsToken { access_key_id: String, access_key_secret: String,