876 lines
38 KiB
Rust
876 lines
38 KiB
Rust
//! 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::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<Database>,
|
|
allowed_dirs: &Arc<RwLock<AllowedDirs>>,
|
|
) -> anyhow::Result<Value> {
|
|
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 = err_text.chars().take(500).collect();
|
|
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<AiProviderRecord> {
|
|
// 筛 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<u64> {
|
|
// 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 的字符串(用于错误信息),供定位用。
|
|
fn body_snippet(v: &Value) -> String {
|
|
let s = v.to_string();
|
|
s.chars().take(500).collect()
|
|
}
|
|
|
|
// ============================================================
|
|
// 单元测试
|
|
//
|
|
// 覆盖:
|
|
// ① 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 跑。落盘到 <cwd>/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::<Vec<_>>();
|
|
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());
|
|
}
|
|
}
|