//! generate_image AI 工具 — 调 provider 的 OpenAI 兼容图像生成端点生成图片并下载落地。 //! //! 设计目标:让 LLM 调用图像生成模型(默认 SenseNova U1 Fast,也支持 DALL-E 等 OpenAI 兼容 //! 图像端点)生成图片并落到工作区磁盘。图像生成模型不是 chat 模型,走 `/v1/images/generations` //! 端点(非 chat completions),故不能经主对话链路调用,必须用本工具显式触发。 //! //! ## 与 http_request / fetch_url / download_file 的分工 //! //! - `http_request`:结构化 API 调用,返回 body(截断 50KB),不落盘 —— 通用 HTTP,不感知图像端点。 //! - `fetch_url`:网页 → markdown 文档嗅探,不落盘。 //! - `download_file`:URL → 落盘文件(已知 URL)。 //! - `generate_image`:**调 provider 图像端点 + 自动下载落地**(封装图像生成的 provider 凭证 //! 解析 + 端点拼接 + 响应解析 + 图片下载全链路,LLM 只需给 prompt/model/output_path)。 //! //! ## 风险:High //! //! ① 付费 API(图像生成按张计费);② 写磁盘(落盘 output_path);③ 外发网络(provider 端点 //! + 图片 URL 下载)。统一 High 须人工批准。 //! //! ## 安全 //! //! 1. **provider 凭证**:从 db.ai_providers 筛 `enabled && provider_type=="openai_compat"` 的 //! provider,经 `df_storage::secret::resolve_provider_secret` 解析 key(DB 优先 fallback keyring)。 //! anthropic provider 不参与(图像端点是 OpenAI 风格)。 //! 2. **端点拼接**:`build_images_url` 智能 base_url(`/v1`/`/v4` 后缀直接补 `/images/generations`, //! 否则补 `/v1/images/generations`),对齐 `model_fetch_helpers::build_models_url` 思路。 //! 3. **provider 端点域名**:base_url 来自 DB 用户配置(设置页可填任意 URL),provider 攻陷 / //! 配置错误 / 恶意 base_url 即可打内网(127.0.0.1 / 169.254.169.254 元数据)。故 provider //! POST 端点同样走 SSRF 防护(validate_url + resolve_and_check_host + build_client), //! 与图片 URL 下载同源,口径一致(2026-08-02 走查修复)。 //! 4. **图片 URL 下载**:图片 URL 来自 provider 响应,**可能被恶意 provider 篡改指向内网** //! (provider 域名虽可信但响应内容不可信),故下载图片 URL 复用 SSRF 防护 //! (validate_url + resolve_and_check_host + build_client,与 download_file 同源)。 //! 5. **路径校验**:output_path 走 `resolve_workspace_path_with_allowed`(白名单 + symlink 逃逸 //! 拦截 + 路径遍历拒),与 write_file/download_file 同源;默认路径锚定首个持久授权目录。 //! 6. **大小上限**:50MB(图片够用),流式边下边累计,超限中止 + 删半成品(对齐 download_file)。 use std::collections::HashMap; use std::sync::Arc; use std::time::{Duration, Instant}; use base64::{engine::general_purpose::STANDARD, Engine as _}; use serde_json::{json, Value}; use tokio::io::AsyncWriteExt; use tokio::sync::RwLock; use df_storage::crud::AiProviderRepo; use df_storage::db::Database; use df_storage::models::AiProviderRecord; use df_storage::secret::resolve_provider_secret; use crate::commands::ai::audit::truncate_chars_ellipsis; use crate::commands::ai::http::{ build_client, execute_with_redirects, resolve_and_check_host, validate_url, MAX_REDIRECTS, }; use crate::commands::ai::tool_registry::resolve_workspace_path_with_allowed; use crate::state::AllowedDirs; /// 下载图片大小硬上限(字节)。50MB —— 覆盖高分辨率图片(4K壁纸级 < 20MB),超限中止。 /// 图片 URL 来自 provider 响应可能被篡改指向超大文件,流式边下边累计防磁盘炸。 const MAX_IMAGE_BYTES: u64 = 50 * 1024 * 1024; /// 生成图片 POST 请求超时(秒)。图像生成模型耗时较高(高分辨率 10-30s),给 120s 余量。 const GENERATE_TIMEOUT_SECS: u64 = 120; /// 按 base64 长度估算解码后字节数(base64 每 4 字符编码 3 字节)。 /// /// 用于 b64_json 解码前的 OOM 防护:恶意 provider 返超长 b64_json(如 200MB base64 → ~150MB /// 解码字节),若先 `STANDARD.decode` 全载入内存再检查上限,瞬时 OOM。故解码前先用本函数 /// 估算,超 MAX_IMAGE_BYTES 直接 bail 不解码。 /// /// 公式 `len * 3 / 4`:忽略 padding(`=`)和 whitespace 误差,估值**略小于实际** b64 长度对应 /// 的理论解码字节(实际含 padding 会更少)。安全方向:估算偏保守(偏小),配合解码后兜底校验 /// (`decoded.len() > MAX_IMAGE_BYTES`)双保险。`saturating_mul` 防 usize 溢出(64 位平台 /// b64.len() 不可能逼近 usize::MAX,但防御性编程)。 /// /// 抽为独立纯函数供单测覆盖边界(2026-08-02 走查修复,对齐 test-strategy 防回归)。 fn estimate_decoded_bytes(b64_len: usize) -> u64 { (b64_len as u64).saturating_mul(3) / 4 } /// generate_image 工具 handler 入口(供 tools/generate_image.rs register 调用)。 /// /// 参数: /// - prompt: 必填,图片描述(中英文均可,空串报错) /// - model: 默认 `sensenova-u1-fast`(SenseNova U1 信息图生成,走 /v1/images/generations) /// - size: 可选,图片尺寸字符串(SenseNova U1 常用 `2752x1536`/`1536x2752`,DALL-E 常用 `1024x1024`) /// - n: 默认 1,生成数量,clamp [1, 4](多张图取首张 url/b64 落地,余下仅返回 url 列表) /// - provider_id: 可选,指定 provider(多 provider 时);不传则按 model 名匹配厂商 host /// - output_path: 可选,落盘路径;不传则默认 `{首个持久授权目录}/generated_images/{model}-{时间戳}.png` /// /// 返回 {path, url(或 null 若仅 b64), model, provider_id, bytes_written, elapsed_ms} /// /// 注:handler 接收 db + allowed_dirs(经 declare_tool! 闭包捕获 tuple Arc 传入)。 /// 路径校验在本函数内完成(需 allowed_dirs 快照,与 http_request output_file 同款)。 pub(crate) async fn execute_generate_image( args: Value, db: &Arc, allowed_dirs: &Arc>, ) -> anyhow::Result { let started = Instant::now(); // ── 参数解析 ── let prompt = args.get("prompt").and_then(|v| v.as_str()) .ok_or_else(|| anyhow::anyhow!("缺少 prompt 参数"))? .trim() .to_string(); if prompt.is_empty() { anyhow::bail!("prompt 不能为空"); } let model = args.get("model").and_then(|v| v.as_str()) .unwrap_or("sensenova-u1-fast") .trim() .to_string(); if model.is_empty() { anyhow::bail!("model 不能为空"); } let size = args.get("size").and_then(|v| v.as_str()) .map(|s| s.trim().to_string()) .filter(|s| !s.is_empty()); let n = args.get("n").and_then(|v| v.as_u64()) .unwrap_or(1) .clamp(1, 4) as u32; let provider_id = args.get("provider_id").and_then(|v| v.as_str()) .map(|s| s.trim().to_string()) .filter(|s| !s.is_empty()); let output_path_raw = args.get("output_path").and_then(|v| v.as_str()) .map(|s| s.trim().to_string()) .filter(|s| !s.is_empty()); // ── 选 provider(enabled && openai_compat,按 model 名匹配厂商 host 兜底首个) ── let providers = AiProviderRepo::new(db).list_all().await?; let provider = select_image_provider(&providers, &model, provider_id.as_deref())?; let provider_id_resolved = provider.id.clone(); // ── 解析 api_key(DB 优先 fallback keyring,空则报错) ── let api_key = resolve_provider_secret(&provider); let api_key = api_key.trim(); if api_key.is_empty() { anyhow::bail!( "provider「{}」未配置 api_key,请在设置填入并保存", provider.name ); } // ── 拼端点 URL ── let endpoint = build_images_url(&provider.base_url); // ── POST 请求 ── // base_url 来自 DB 用户配置(可填任意 URL),须走 SSRF 防护(协议白名单 + 私网 IP + // DNS resolve + 重定向每跳校验),与图片 URL 下载同源(2026-08-02 走查修复)。 let body = { let mut m = serde_json::Map::new(); m.insert("model".into(), json!(model)); m.insert("prompt".into(), json!(prompt)); m.insert("n".into(), json!(n)); if let Some(sz) = &size { m.insert("size".into(), json!(sz)); } Value::Object(m) }; let body_str = serde_json::to_string(&body) .map_err(|e| anyhow::anyhow!("序列化请求 body 失败: {}", e))?; // SSRF 三层校验:词法 URL → DNS resolve IP(后续 build_client + execute_with_redirects // 对每跳重定向重复校验,防 302 绕过到内网) let (_scheme, host, port) = validate_url(&endpoint)?; resolve_and_check_host(&host, port).await?; // build_client 已含 timeout/connect_timeout + 关闭自动重定向(每跳手动校验) let client = build_client(Duration::from_secs(GENERATE_TIMEOUT_SECS))?; let mut headers = HashMap::new(); headers.insert("Authorization".into(), format!("Bearer {}", api_key)); headers.insert("Content-Type".into(), "application/json".into()); let resp = execute_with_redirects( &client, reqwest::Method::POST, endpoint.clone(), &headers, &Some(body_str), MAX_REDIRECTS, ) .await .map_err(|e| anyhow::anyhow!("调用图像生成端点失败 ({}): {}", endpoint, e))?; let status = resp.status(); if !status.is_success() { // 错误响应读 body 摘要(截断 500 chars,截断时尾部补 …)帮助定位 let err_text = resp.text().await.unwrap_or_default(); let snippet: String = truncate_chars_ellipsis(&err_text, 500); anyhow::bail!( "图像生成端点返回 HTTP {}: {}", status.as_u16(), snippet ); } let resp_body: Value = resp.json().await .map_err(|e| anyhow::anyhow!("解析图像生成响应 JSON 失败: {}", e))?; // ── 解析响应:优先 data[0].url,次选 data[0].b64_json ── let data0 = resp_body.get("data").and_then(|d| d.get(0)) .ok_or_else(|| { let snippet = body_snippet(&resp_body); anyhow::anyhow!("图像生成响应缺 data[0](响应摘要: {})", snippet) })?; let image_url_opt = data0.get("url").and_then(|u| u.as_str()) .filter(|s| !s.trim().is_empty()) .map(|s| s.trim().to_string()); let b64_opt = data0.get("b64_json").and_then(|b| b.as_str()) .filter(|s| !s.trim().is_empty()) .map(|s| s.trim().to_string()); if image_url_opt.is_none() && b64_opt.is_none() { let snippet = body_snippet(&resp_body); anyhow::bail!( "图像生成响应 data[0] 既无 url 也无 b64_json(响应摘要: {})", snippet ); } // ── 解析 output_path(LLM 传则走白名单校验;不传则默认 generated_images/ 时间戳名) ── let snap = allowed_dirs.read().await.clone(); let final_path = if let Some(raw) = &output_path_raw { resolve_workspace_path_with_allowed(raw, &snap)? .to_str() .ok_or_else(|| anyhow::anyhow!("output_path 含非法字符"))? .to_string() } else { // 默认:首个持久授权目录/generated_images/{model}-{YYYYMMDD-HHMMSS}.png let root = snap.first_persistent_dir() .ok_or_else(|| anyhow::anyhow!( "未传 output_path 且无持久授权目录,请先绑定项目目录或显式传 output_path" ))?; let ts = chrono::Local::now().format("%Y%m%d-%H%M%S"); // model 含 / 等路径分隔符风险:替换为 - 防穿透到子目录 let safe_model = model.replace(['/', '\\', ':', '*', '?', '"', '<', '>', '|'], "-"); let fname = format!("{}-{}.png", safe_model, ts); root.join("generated_images").join(fname) .to_str() .ok_or_else(|| anyhow::anyhow!("默认 output_path 含非法字符"))? .to_string() }; // ── 落盘:url 走 SSRF 防护下载流式写;b64 直接写解码后字节 ── let written_bytes = if let Some(image_url) = &image_url_opt { download_image_with_ssrf(image_url, &final_path).await? } else { // b64 路径:解码后直接写 let b64 = b64_opt.as_ref().unwrap(); // OOM 防护(2026-08-02 走查修复):解码前先按 base64 长度估算解码后字节数,超 MAX_IMAGE_BYTES // 直接 bail 不解码。恶意 provider 返超长 b64_json(如 200MB base64 → ~150MB 解码字节), // 若先 STANDARD.decode 全载入内存再检查,瞬时 OOM。详见 `estimate_decoded_bytes` 文档。 let estimated_decoded = estimate_decoded_bytes(b64.len()); if estimated_decoded > MAX_IMAGE_BYTES { anyhow::bail!( "b64_json 估算解码后约 {} 字节超过 {} 上限(原始 base64 长度 {})", estimated_decoded, MAX_IMAGE_BYTES, b64.len() ); } let decoded = STANDARD.decode(b64) .map_err(|e| anyhow::anyhow!("b64_json 解码失败: {}", e))?; // 解码后再用实际长度兜底校验(防估算偏差,如 base64 含大量空白/padding) if (decoded.len() as u64) > MAX_IMAGE_BYTES { anyhow::bail!( "b64_json 解码后 {} 字节超过 {} 上限", decoded.len(), MAX_IMAGE_BYTES ); } write_atomic(&final_path, &decoded).await?; decoded.len() as u64 }; let elapsed_ms = started.elapsed().as_millis() as u64; Ok(json!({ "path": final_path, "url": image_url_opt, "model": model, "provider_id": provider_id_resolved, "bytes_written": written_bytes, "elapsed_ms": elapsed_ms, })) } /// 从 provider 列表筛选图像生成可用 provider。 /// /// 筛选条件:`enabled && provider_type == "openai_compat"`(图像端点是 OpenAI 风格, /// anthropic provider 排除)。 /// /// 选择策略: /// 1. LLM 传 `provider_id` → 按 id 精确匹配(找不到报错)。 /// 2. 否则按 model 名匹配厂商 host 提示: /// - model 含 `sensenova` → 找 base_url 含 `sensenova` /// - model 含 `dall-e` → 找 base_url 含 `openai` /// - model 含 `google`/`imagen` → 找 base_url 含 `googleapis` /// 3. host 命中 → 用首个匹配;host 未命中 → 兜底取筛选结果首个。 /// 4. 筛选后空 → 报错(无可用 openai_compat provider)。 fn select_image_provider( providers: &[AiProviderRecord], model: &str, provider_id: Option<&str>, ) -> anyhow::Result { // 筛 enabled + openai_compat let candidates: Vec<&AiProviderRecord> = providers.iter() .filter(|p| p.enabled && p.provider_type == "openai_compat") .collect(); // 路径1:LLM 指定 provider_id,精确匹配 if let Some(pid) = provider_id { let found = candidates.iter().find(|p| p.id == pid) .ok_or_else(|| anyhow::anyhow!( "未找到指定 provider(id={},且需 enabled + openai_compat),请检查 provider_id 参数", pid ))?; return Ok((*found).clone()); } if candidates.is_empty() { anyhow::bail!("无可用 openai_compat provider,请先在设置配置支持图像生成的 provider"); } // 路径2:按 model 名匹配厂商 host let model_lower = model.to_lowercase(); let host_hint: Option<&str> = if model_lower.contains("sensenova") { Some("sensenova") } else if model_lower.contains("dall-e") { Some("openai") } else if model_lower.contains("google") || model_lower.contains("imagen") { Some("googleapis") } else { None }; if let Some(hint) = host_hint { let hit = candidates.iter() .find(|p| p.base_url.to_lowercase().contains(hint)); if let Some(p) = hit { return Ok((*p).clone()); } } // 路径3:兜底首个候选 Ok(candidates[0].clone()) } /// 拼接图像生成端点 URL。 /// /// 规则(base_url 先去尾 `/`,对齐 `model_fetch_helpers::build_models_url` 思路, /// 但目标是 `/images/generations`): /// - 以 `/v1` 或 `/v4` 结尾 → 直接补 `/images/generations`(避免 `/v1/v1/images/...`) /// - 否则 → 补 `/v1/images/generations`(最常见:用户填 `https://api.sensenova.cn`) fn build_images_url(base_url: &str) -> String { let url = base_url.trim_end_matches('/'); if url.ends_with("/v1") || url.ends_with("/v4") { format!("{url}/images/generations") } else { format!("{url}/v1/images/generations") } } /// 流式下载图片 URL 到 output_path,带 SSRF 防护 + 大小上限 + 原子写。 /// /// 图片 URL 来自 provider 响应可能被恶意 provider 篡改指向内网,故复用 SSRF 防护 /// (validate_url + resolve_and_check_host + build_client + execute_with_redirects)。 async fn download_image_with_ssrf(image_url: &str, output_path: &str) -> anyhow::Result { // SSRF 三层校验:词法 URL → DNS resolve IP → 每跳重定向重复 let (_scheme, host, port) = validate_url(image_url)?; resolve_and_check_host(&host, port).await?; let client = build_client(Duration::from_secs(GENERATE_TIMEOUT_SECS))?; let resp = execute_with_redirects( &client, reqwest::Method::GET, image_url.to_string(), &HashMap::new(), // 图片下载不带自定义头 &None, MAX_REDIRECTS, ) .await?; if !resp.status().is_success() { anyhow::bail!( "下载图片失败:HTTP {} {}({})", resp.status().as_u16(), resp.status().canonical_reason().unwrap_or(""), image_url ); } // 流式写文件(避免大图 OOM),边下边累计,超 MAX_IMAGE_BYTES 立即中止 + 删半成品。 // 父目录不存在自动创建。原子写:tmp → rename(同目录不跨卷)。 let target = std::path::Path::new(output_path); if let Some(parent) = target.parent() { tokio::fs::create_dir_all(parent).await .map_err(|e| anyhow::anyhow!("创建目录失败: {}", e))?; } let tmp = format!("{}.tmp-img", output_path); let mut file = tokio::fs::File::create(&tmp).await .map_err(|e| anyhow::anyhow!("创建临时文件失败 ({}): {}", tmp, e))?; let mut stream = resp.bytes_stream(); let mut written: u64 = 0; let mut exceeded = false; use futures::StreamExt; while let Some(chunk) = stream.next().await { let chunk = chunk.map_err(|e| anyhow::anyhow!("读取图片流失败: {}", e))?; written += chunk.len() as u64; if written > MAX_IMAGE_BYTES { exceeded = true; break; } file.write_all(&chunk).await .map_err(|e| anyhow::anyhow!("写入图片失败: {}", e))?; } file.flush().await .map_err(|e| anyhow::anyhow!("flush 图片文件失败: {}", e))?; drop(file); if exceeded { let _ = tokio::fs::remove_file(&tmp).await; anyhow::bail!( "图片超过 {} 字节上限,已中止并清理临时文件", MAX_IMAGE_BYTES ); } if let Err(e) = tokio::fs::rename(&tmp, output_path).await { let _ = tokio::fs::remove_file(&tmp).await; return Err(anyhow::anyhow!("原子替换失败: {}", e)); } Ok(written) } /// 原子写字节到指定路径(tmp → rename,同目录不跨卷)。 /// b64 路径用此函数(已全部解码在内存,无需流式)。 async fn write_atomic(output_path: &str, data: &[u8]) -> anyhow::Result<()> { let target = std::path::Path::new(output_path); if let Some(parent) = target.parent() { tokio::fs::create_dir_all(parent).await .map_err(|e| anyhow::anyhow!("创建目录失败: {}", e))?; } let tmp = format!("{}.tmp-img", output_path); tokio::fs::write(&tmp, data).await .map_err(|e| { let _ = std::fs::remove_file(&tmp); anyhow::anyhow!("写入图片失败: {}", e) })?; tokio::fs::rename(&tmp, output_path).await .map_err(|e| { let _ = std::fs::remove_file(&tmp); anyhow::anyhow!("原子替换失败: {}", e) })?; Ok(()) } /// 响应 body 截断为最多 500 chars 的字符串(用于错误信息),供定位用。 /// 走 audit truncate_chars 正典(截断时尾部补 …);错误文案无硬长度约束,允许多 1 字符。 fn body_snippet(v: &Value) -> String { let s = v.to_string(); truncate_chars_ellipsis(&s, 500) } // ============================================================ // 单元测试 // // 覆盖: // ① build_images_url:端点拼接规则(纯函数,零依赖) // ② select_image_provider:provider 筛选逻辑(纯函数,内存构造 record) // ③ estimate_decoded_bytes:b64 解码字节数估算(纯函数,OOM 防护核心) // ④ handler 参数边界(prompt 缺失/空、provider_id 不存在) // ⑤ handler SSRF 防护(endpoint 走 validate_url 前置拒,不发真实请求,CI 可跑) // 真实 API 调用走 #[ignore](CI 无凭证/无网时跳过,本地手跑)。 // ============================================================ #[cfg(test)] mod tests { use super::*; use std::path::PathBuf; // ── build_images_url:端点拼接规则 ── #[test] fn build_images_url_v1_suffix() { // /v1 结尾 → 直接补 /images/generations assert_eq!( build_images_url("https://api.openai.com/v1"), "https://api.openai.com/v1/images/generations" ); // 不应产生 /v1/v1/images/... assert!(!build_images_url("https://api.openai.com/v1").contains("/v1/v1/")); } #[test] fn build_images_url_v4_suffix() { // GLM v4 风格 base assert_eq!( build_images_url("https://open.bigmodel.cn/api/paas/v4"), "https://open.bigmodel.cn/api/paas/v4/images/generations" ); } #[test] fn build_images_url_bare_base() { // 用户只填根 → 补 /v1/images/generations assert_eq!( build_images_url("https://api.sensenova.cn"), "https://api.sensenova.cn/v1/images/generations" ); } #[test] fn build_images_url_strips_trailing_slash() { assert_eq!( build_images_url("https://api.x.com/v1/"), "https://api.x.com/v1/images/generations" ); } // ── select_image_provider:筛选 + 匹配逻辑 ── /// 构造测试用 AiProviderRecord(只填筛选关键字段)。 fn mk_provider(id: &str, enabled: bool, ptype: &str, base_url: &str) -> AiProviderRecord { AiProviderRecord { id: id.into(), name: id.into(), provider_type: ptype.into(), api_key: String::new(), base_url: base_url.into(), default_model: String::new(), models: None, model_configs: Vec::new(), is_default: false, config: None, created_at: "0".into(), updated_at: "0".into(), enabled, weight: 50, } } #[test] fn select_provider_no_candidates_errors() { // 全 disabled 或非 openai_compat → 报错 let providers = vec![ mk_provider("p1", false, "openai_compat", "https://x"), mk_provider("p2", true, "anthropic", "https://y"), ]; let err = select_image_provider(&providers, "any-model", None).unwrap_err(); assert!(format!("{}", err).contains("无可用")); } #[test] fn select_provider_provider_id_match() { let providers = vec![ mk_provider("p1", true, "openai_compat", "https://api.sensenova.cn"), mk_provider("p2", true, "openai_compat", "https://api.openai.com"), ]; let got = select_image_provider(&providers, "any", Some("p2")).unwrap(); assert_eq!(got.id, "p2"); } #[test] fn select_provider_provider_id_not_found_errors() { let providers = vec![mk_provider("p1", true, "openai_compat", "https://x")]; let err = select_image_provider(&providers, "any", Some("nonexistent")).unwrap_err(); assert!(format!("{}", err).contains("未找到指定 provider")); } #[test] fn select_provider_host_match_sensenova() { // model 含 sensenova → 匹配 base_url 含 sensenova 的 provider let providers = vec![ mk_provider("openai", true, "openai_compat", "https://api.openai.com"), mk_provider("sense", true, "openai_compat", "https://api.sensenova.cn"), ]; let got = select_image_provider(&providers, "sensenova-u1-fast", None).unwrap(); assert_eq!(got.id, "sense"); } #[test] fn select_provider_host_match_dalle() { // model 含 dall-e → 匹配 base_url 含 openai 的 provider let providers = vec![ mk_provider("glm", true, "openai_compat", "https://open.bigmodel.cn"), mk_provider("oai", true, "openai_compat", "https://api.openai.com"), ]; let got = select_image_provider(&providers, "dall-e-3", None).unwrap(); assert_eq!(got.id, "oai"); } #[test] fn select_provider_host_unmatched_falls_back_first() { // model 名无厂商提示 → 兜底取筛选结果首个 let providers = vec![ mk_provider("first", true, "openai_compat", "https://api.unknown.com"), mk_provider("second", true, "openai_compat", "https://api.other.com"), ]; let got = select_image_provider(&providers, "my-custom-model", None).unwrap(); assert_eq!(got.id, "first"); } // ── estimate_decoded_bytes:b64 解码字节数估算(OOM 防护核心纯函数) ── #[test] fn estimate_decoded_bytes_normal_small() { // 1MB base64 → 估算 ~0.75MB(base64 每 4 字符编码 3 字节) let one_mb = 1024 * 1024; assert_eq!(estimate_decoded_bytes(one_mb), (one_mb as u64) * 3 / 4); assert_eq!(estimate_decoded_bytes(one_mb), 786_432); // 1MB * 3 / 4 } #[test] fn estimate_decoded_bytes_over_limit() { // 200MB 解码后场景:b64 长约 267MB → 估算 ~200MB > 50MB 上限,OOM 防护应挡 let two_hundred_mb_b64 = (200 * 1024 * 1024) as usize; // 假设这是 b64 长度 let estimated = estimate_decoded_bytes(two_hundred_mb_b64); assert!( estimated > MAX_IMAGE_BYTES, "200MB b64 估算 {} 应超 {} 上限,OOM 防护靠此拒", estimated, MAX_IMAGE_BYTES ); } #[test] fn estimate_decoded_bytes_boundary_at_max() { // 边界:估算正好等于 MAX_IMAGE_BYTES → 比较用 `>`(超过才拒),正好等于算合法通过 // 反推 b64 长度:要使 len*3/4 == MAX_IMAGE_BYTES,len = MAX_IMAGE_BYTES*4/3 let b64_len = (MAX_IMAGE_BYTES as u128 * 4 / 3) as usize; let estimated = estimate_decoded_bytes(b64_len); // 估算值应 ≤ MAX_IMAGE_BYTES(因整数除法向下取整,可能略小于上限) assert!( estimated <= MAX_IMAGE_BYTES, "估算 {} 应 ≤ 上限 {}(正好或略小于,`>` 比较下不拒)", estimated, MAX_IMAGE_BYTES ); } #[test] fn estimate_decoded_bytes_just_over_max() { // 边界:估算值 = MAX_IMAGE_BYTES + 1 → 应被拒 // 反推 b64 长度:要使 len*3/4 == MAX+1,len = (MAX+1)*4/3 let target = MAX_IMAGE_BYTES + 1; // 因整数除法,估算可能略小于 target,但应保证 > MAX_IMAGE_BYTES // 用更精确的 b64_len:target*4/3 + 4(向上补一个 base64 块),保证估算严格 > MAX let b64_len_safe = ((target as u128 * 4 + 2) / 3) as usize; let estimated_safe = estimate_decoded_bytes(b64_len_safe); assert!( estimated_safe > MAX_IMAGE_BYTES, "b64_len={} 估算 {} 应严格 > 上限 {}(超限应拒)", b64_len_safe, estimated_safe, MAX_IMAGE_BYTES ); } #[test] fn estimate_decoded_bytes_saturating_no_overflow() { // 极大值(u64::MAX 的 usize 等价)→ saturating_mul 防 panic,返回 u64 内合理值 // 注:usize 在 64 位平台 = u64,saturating_mul(u64::MAX, 3) 会 saturate 到 u64::MAX let huge = usize::MAX; let estimated = estimate_decoded_bytes(huge); // saturating_mul(MAX, 3) = MAX,再 / 4 = MAX/4,不 panic 即通过 assert_eq!(estimated, u64::MAX / 4); } #[test] fn estimate_decoded_bytes_zero() { // 空输入 → 0(边界,虽 b64_opt 已 filter 空串不进此路径,纯函数仍应优雅) assert_eq!(estimate_decoded_bytes(0), 0); } // ── handler 参数边界(走 db 需 in-memory db) ── #[tokio::test] async fn handler_missing_prompt_errors() { // 缺 prompt → 参数解析阶段报错(不发请求) let db = Arc::new(Database::open_in_memory().await.unwrap()); let allowed = Arc::new(RwLock::new(AllowedDirs::default())); let args = json!({ "model": "sensenova-u1-fast" }); let err = execute_generate_image(args, &db, &allowed).await.unwrap_err(); assert!(format!("{}", err).contains("prompt")); } #[tokio::test] async fn handler_empty_prompt_errors() { let db = Arc::new(Database::open_in_memory().await.unwrap()); let allowed = Arc::new(RwLock::new(AllowedDirs::default())); let args = json!({ "prompt": " " }); let err = execute_generate_image(args, &db, &allowed).await.unwrap_err(); assert!(format!("{}", err).contains("prompt 不能为空")); } #[tokio::test] async fn handler_no_provider_errors() { // 无 openai_compat provider → 报错(不发请求) let db = Arc::new(Database::open_in_memory().await.unwrap()); let allowed = Arc::new(RwLock::new(AllowedDirs::default())); let args = json!({ "prompt": "一只猫" }); let err = execute_generate_image(args, &db, &allowed).await.unwrap_err(); assert!(format!("{}", err).contains("无可用") || format!("{}", err).contains("provider")); } // ── handler SSRF 防护(endpoint 走 validate_url 前置拒,不发真实请求) ── // // base_url 来自 DB 用户配置(设置页可填任意 URL),provider 攻陷 / 配置错误 / 恶意 base_url // 即可打内网。endpoint POST 须走 SSRF 防护(validate_url + resolve_and_check_host),与图片 // URL 下载同源(2026-08-02 走查修复)。本组测试构造 enabled + openai_compat + api_key 非空 // 的 provider 写入 in-memory db,handler 走到 validate_url 阶段拒,不发真实请求(CI 可跑)。 // 风格对齐 http.rs 的 test_handler_rejects_localhost_url / test_handler_rejects_private_ip_url。 /// 构造带 api_key 的 openai_compat provider(复用 mk_provider 但填 api_key,跳过 keyring)。 fn mk_provider_with_key(id: &str, base_url: &str) -> AiProviderRecord { AiProviderRecord { api_key: "sk-test-dummy-not-real".into(), ..mk_provider(id, true, "openai_compat", base_url) } } #[tokio::test] async fn handler_rejects_localhost_base_url() { // base_url = http://localhost:8080 → validate_url L117-119 拒 localhost(SSRF 防护) // 在 build_client 前拒,不发真实请求 let db = Arc::new(Database::open_in_memory().await.unwrap()); AiProviderRepo::new(&db) .insert(mk_provider_with_key("p1", "http://localhost:8080")) .await .unwrap(); let allowed = Arc::new(RwLock::new(AllowedDirs::default())); let args = json!({ "prompt": "一只猫", "provider_id": "p1" }); let err = execute_generate_image(args, &db, &allowed).await.unwrap_err(); let msg = format!("{}", err); assert!( msg.contains("localhost") || msg.contains("SSRF"), "期望 SSRF 拒绝 localhost,实际: {}", msg ); } #[tokio::test] async fn handler_rejects_private_ip_base_url() { // base_url 含 169.254.169.254(云元数据服务,SSRF 头号目标)→ validate_url 字面量 IP 私网拒 let db = Arc::new(Database::open_in_memory().await.unwrap()); AiProviderRepo::new(&db) .insert(mk_provider_with_key("p1", "http://169.254.169.254/latest")) .await .unwrap(); let allowed = Arc::new(RwLock::new(AllowedDirs::default())); let args = json!({ "prompt": "一只猫", "provider_id": "p1" }); let err = execute_generate_image(args, &db, &allowed).await.unwrap_err(); let msg = format!("{}", err); assert!( msg.contains("私网") || msg.contains("SSRF") || msg.contains("169.254"), "期望 SSRF 拒绝元数据 IP,实际: {}", msg ); } #[tokio::test] async fn handler_rejects_loopback_base_url() { // base_url = http://127.0.0.1 → validate_url 字面量 IP 环回拒 let db = Arc::new(Database::open_in_memory().await.unwrap()); AiProviderRepo::new(&db) .insert(mk_provider_with_key("p1", "http://127.0.0.1:9000")) .await .unwrap(); let allowed = Arc::new(RwLock::new(AllowedDirs::default())); let args = json!({ "prompt": "一只猫", "provider_id": "p1" }); let err = execute_generate_image(args, &db, &allowed).await.unwrap_err(); let msg = format!("{}", err); assert!( msg.contains("私网") || msg.contains("SSRF") || msg.contains("127.0.0.1"), "期望 SSRF 拒绝环回 IP,实际: {}", msg ); } #[tokio::test] async fn handler_rejects_non_http_scheme_base_url() { // base_url = file:///etc/passwd → validate_url L104-109 拒非 http 协议 let db = Arc::new(Database::open_in_memory().await.unwrap()); AiProviderRepo::new(&db) .insert(mk_provider_with_key("p1", "file:///etc/passwd")) .await .unwrap(); let allowed = Arc::new(RwLock::new(AllowedDirs::default())); let args = json!({ "prompt": "一只猫", "provider_id": "p1" }); let err = execute_generate_image(args, &db, &allowed).await.unwrap_err(); let msg = format!("{}", err); assert!( msg.contains("协议"), "期望拒绝非 http 协议,实际: {}", msg ); } // ── 真实 SenseNova ast 图像生成验证(#[ignore],手动跑) ── // // 端到端验证:用真实 db 副本(SenseNova provider 配置)+ OS keyring 解析 api_key + // 真实调 https://token.sensenova.cn/v1/images/generations 生成一张测试图并下载落地。 // 覆盖 generate_image 全链路:provider 选择(host 匹配)/ build_images_url 端点拼接 / // 请求构造 / 响应解析 / SSRF 防护下载落地。 // // 运行方式: // set DEVFLOW_LIVE_DB=<真实 devflow.db 路径(含 wal)> // cargo test -p devflow --lib -- --ignored live_sensenova_u1_fast --nocapture // // 需真实 API 凭证(付费),默认 #[ignore] 不随 CI 跑。落盘到 /tmp_generated/ 便于清理。 /// 构造带持久授权目录的 AllowedDirs,指向项目内 tmp_generated/ 输出目录。 fn mk_allowed_with_generated_dir() -> (AllowedDirs, PathBuf) { let out_dir = std::env::current_dir() .expect("读 cwd") .join("tmp_generated"); let mut persistent = std::collections::HashSet::new(); persistent.insert(out_dir.clone()); (AllowedDirs { persistent, session: Default::default(), once: Default::default() }, out_dir) } #[tokio::test] #[ignore] // 真实 API 调用(付费),需 DEVFLOW_LIVE_DB 环境变量,手动跑 async fn live_sensenova_u1_fast() { let db_path = std::env::var("DEVFLOW_LIVE_DB") .map(PathBuf::from) .expect("请设置 DEVFLOW_LIVE_DB=<真实 devflow.db 路径> 以跑真实 API 验证"); assert!( db_path.exists(), "DEVFLOW_LIVE_DB 指向的 db 不存在: {}", db_path.display() ); // 打开真实 db(会跑迁移,用真实路径;读 SenseNova provider + keyring 解析 key) let db = Arc::new(Database::open(&db_path).await.expect("打开真实 db")); let (allowed, out_dir) = mk_allowed_with_generated_dir(); // 走 execute_generate_image 全链路:model 默认 sensenova-u1-fast → host 匹配 SenseNova provider let args = json!({ "prompt": "一张极简科技风信息图,主题「DevFlow AI 工作流」,顶部大标题,三个步骤图标(规划→执行→核查),浅蓝灰配色", "size": "1536x2752", "n": 1 }); let result = execute_generate_image(args, &db, &Arc::new(RwLock::new(allowed))) .await .expect("execute_generate_image 应成功生成并下载图片"); let path = result["path"].as_str().expect("result 应有 path"); let bytes = result["bytes_written"].as_u64().expect("result 应有 bytes_written"); let model = result["model"].as_str().unwrap_or(""); let provider_id = result["provider_id"].as_str().unwrap_or(""); println!("✅ 生成成功:"); println!(" model : {}", model); println!(" provider : {}", provider_id); println!(" 本地路径 : {}", path); println!(" 字节数 : {}", bytes); // 断言:文件真实落盘 + 非空 + 是图片(PNG 魔数) assert!( std::path::Path::new(path).exists(), "图片应真实落盘: {}", path ); assert!(bytes > 1000, "图片应 > 1KB,实际 {} bytes", bytes); let head = std::fs::read(path).expect("读图片头").into_iter().take(8).collect::>(); assert!( head.starts_with(&[0x89, b'P', b'N', b'G']) || head.starts_with(&[0xFF, 0xD8, 0xFF]), "应为 PNG/JPEG 魔数,实际头字节: {:?}", head ); // 清理输出目录(测后不留临时图) let _ = std::fs::remove_dir_all(&out_dir); println!(" (输出目录已清理:{})", out_dir.display()); } }