From 0dd34d1a65a34819ae07041b2f8886ff27a87533 Mon Sep 17 00:00:00 2001 From: Linghong Date: Fri, 10 Jul 2026 07:56:37 +0000 Subject: [PATCH] =?UTF-8?q?=E6=8A=A0=E5=9B=BE=E7=BB=93=E6=9E=9C=E4=B8=8B?= =?UTF-8?q?=E8=BD=BD/=E8=A7=A3=E7=A0=81=E5=8A=A0=E5=A4=A7=E5=B0=8F?= =?UTF-8?q?=E4=B8=8E=E5=B0=BA=E5=AF=B8=E4=B8=8A=E9=99=90=EF=BC=8C=E9=98=B2?= =?UTF-8?q?=E6=AD=A2=E4=B8=8A=E6=B8=B8=E5=93=8D=E5=BA=94=E6=92=91=E7=88=86?= =?UTF-8?q?=E5=86=85=E5=AD=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit response.bytes() 直接分配任意大小响应，作为 BgFilter 降级路径且逐帧动画可并发多次 调用，异常/恶意上游会造成 API 进程内存压力。复用 BgFilter 路径的做法： - 下载：先按 Content-Length 拒绝，再流式 chunk 累加，超 32MB 立即中断 - 解码：ImageReader + Limits（边长 ≤8192、max_alloc 上限） Co-Authored-By: Claude Opus 4.8 --- server-rs/crates/platform-matting/src/lib.rs | 67 ++++++++++++++++---- 1 file changed, 54 insertions(+), 13 deletions(-) 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,