yggdrasil/api/
upload.rs

1//! 图片上传：web 处理器 + 共享入库流水线。
2//!
3//! 三条入口共用 `process_image_upload`：
4//! - web `POST /api/upload`（cookie 鉴权 admin，multipart）—— 见 `upload_image`；
5//! - web `POST /api/comments/upload`（评论区，允许匿名，IP 双层限流）—— 见 `comment_upload_image`；
6//! - MCP `POST /api/mcp/upload`（bearer 鉴权，multipart）—— 见 `mcp_upload_image`；
7//! - MCP `upload_media` 工具（URL 抓取）—— 见 `src/mcp/tools/media.rs`。
8//!
9//! 流水线：magic bytes 检 MIME → 大小校验 → 尺寸/像素校验 → SHA-256 内容去重
10//! （命中即复用）→ GIF/WebP 解码校验 → `spawn_blocking` 转码（GIF/WebP 原样，
11//! JPEG/PNG 仅在更小时转 WebP）→ 按日期落盘 → assets 登记（含并发竞态补偿）。
12//! JPEG/PNG 自动转 WebP（若体积更小则保留原格式），GIF/WebP 保持原样。
13//! 文件按日期分目录存放于 `uploads/`。
14//!
15//! 内容去重（CAS）：以原始上传字节的 SHA-256 为内容指纹（`assets.content_hash`，
16//! 唯一索引）。重复上传同一内容时复用已登记素材——同一行、同一文件，不重复
17//! 落盘，响应带 `"reused": true`；并发同内容上传由唯一索引 + ON CONFLICT 兜底。
18//! 仅精确去重：尺寸/压缩不同的视觉相似图不合并（那是感知哈希 pHash 的领域，
19//! 有意不做）。
20//! 本模块属于手动注册的 Axum 路由，仅在 `feature = "server"` 时可用。
21
22#[cfg(feature = "server")]
23use axum::extract::{ConnectInfo, Extension, Multipart};
24#[cfg(feature = "server")]
25use axum::http::{HeaderMap, StatusCode};
26#[cfg(feature = "server")]
27use axum::response::Response;
28#[cfg(feature = "server")]
29use axum::{response::IntoResponse, Json};
30#[cfg(feature = "server")]
31use serde_json::{json, Value};
32#[cfg(feature = "server")]
33use std::net::SocketAddr;
34
35#[cfg(feature = "server")]
36use crate::auth::session::parse_session_token;
37
38#[cfg(feature = "server")]
39const ALLOWED_MIME_TYPES: &[&str] = &["image/jpeg", "image/png", "image/gif", "image/webp"];
40#[cfg(feature = "server")]
41use crate::utils::server::MAX_FILE_SIZE;
42
43// ===========================================================================
44// web 处理器（cookie 鉴权）
45// ===========================================================================
46
47/// 构造统一的 JSON 错误响应：`{ "success": false, "error": msg }`。
48/// pub(crate)：备份导入（api/database/backup.rs）复用同一契约。
49#[cfg(feature = "server")]
50pub(crate) fn upload_error<T: serde::Serialize>(
51    status: StatusCode,
52    msg: T,
53) -> (StatusCode, Json<Value>) {
54    (status, Json(json!({ "success": false, "error": msg })))
55}
56
57/// 处理图片上传的 Axum handler（web 端，cookie 鉴权）。
58///
59/// 流程：限流 → 解析 session → 校验 admin → 读取 multipart → 早拒非法声明类型 →
60/// 读取字节 → 交给共享流水线 `process_image_upload`。
61///
62/// `ConnectInfo` 以可选扩展注入：`dioxus::server::serve()` 接管了 listener，
63/// 无法调用 `into_make_service_with_connect_info::<SocketAddr>()`，所以这里
64/// 与 `serve_image` 保持一致的优雅降级——扩展缺失时退回 `"unknown"` 限流桶。
65/// 生产环境应在反向代理后部署并配置 `TRUSTED_PROXY_COUNT`，让限流拿到真实 IP。
66#[cfg(feature = "server")]
67pub async fn upload_image(
68    connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
69    headers: HeaderMap,
70    mut multipart: Multipart,
71) -> Result<Json<Value>, (StatusCode, Json<Value>)> {
72    // 0. Rate limit check
73    let peer = connect_info.map(|Extension(ConnectInfo(addr))| addr);
74    let ip = crate::api::rate_limit::get_client_ip_with_peer(&headers, peer).await;
75    if let Err(msg) = crate::api::rate_limit::check_upload_limit(&ip) {
76        return Err(upload_error(StatusCode::TOO_MANY_REQUESTS, msg));
77    }
78
79    // 1. Extract session from cookie
80    let cookie_header = headers
81        .get("cookie")
82        .and_then(|h| h.to_str().ok())
83        .unwrap_or("");
84
85    let token = match parse_session_token(cookie_header) {
86        Some(t) => t,
87        None => {
88            return Err(upload_error(StatusCode::UNAUTHORIZED, "未登录"));
89        }
90    };
91
92    // 2. Verify admin
93    let user = match crate::api::auth::get_user_by_token(token).await {
94        Ok(Some(u)) => u,
95        _ => {
96            return Err(upload_error(StatusCode::UNAUTHORIZED, "会话已过期"));
97        }
98    };
99
100    if user.role != crate::models::user::UserRole::Admin {
101        return Err(upload_error(StatusCode::FORBIDDEN, "权限不足"));
102    }
103
104    // 3. Read multipart field
105    let field = match multipart.next_field().await {
106        Ok(Some(f)) => f,
107        Ok(None) => {
108            return Err(upload_error(StatusCode::BAD_REQUEST, "未找到文件"));
109        }
110        Err(e) => {
111            tracing::error!("Multipart error: {:?}", e);
112            return Err(upload_error(StatusCode::BAD_REQUEST, "文件读取失败"));
113        }
114    };
115
116    // 4. 早拒非法声明类型（快速路径，避免读字节后再判）。
117    //    流水线仍以 magic bytes 为权威——声明 jpeg 但实为 png 会被识别为 png 接受。
118    let declared_mime = field.content_type().unwrap_or("").to_string();
119    if !ALLOWED_MIME_TYPES.contains(&declared_mime.as_str()) {
120        return Err(upload_error(StatusCode::BAD_REQUEST, "不支持的文件类型"));
121    }
122
123    // 原始文件名（客户端提供，仅作 assets 表展示字段）；需在 bytes() 消耗 field 前取出。
124    let original_filename = field.file_name().map(|s| s.to_string());
125
126    // 5. Read file data
127    let data = match field.bytes().await {
128        Ok(d) => d,
129        Err(e) => {
130            tracing::error!("Read file error: {:?}", e);
131            return Err(upload_error(
132                StatusCode::INTERNAL_SERVER_ERROR,
133                "文件读取失败",
134            ));
135        }
136    };
137
138    // 6. 共享入库流水线。
139    match process_image_upload(data, original_filename, None).await {
140        Ok(out) => Ok(Json(json!({
141            "success": true,
142            "url": out.url,
143            "reused": out.reused
144        }))),
145        Err(e) => {
146            let (status, msg) = e.status_and_msg();
147            Err(upload_error(status, msg))
148        }
149    }
150}
151
152// ===========================================================================
153// 评论图片处理器（允许匿名，IP 双层限流）
154// ===========================================================================
155
156/// 评论图片上传的 Axum handler（评论区用，**允许匿名**）。
157///
158/// 与 admin `upload_image` 的区别：
159/// - 不做会话/角色校验——匿名访客也能为评论传图；
160/// - 限流用独立的 comment_upload 双层桶（每秒突发 + 每日总额），与 admin 的
161///   upload 桶隔离，匿名滥用不会挤占 admin 额度；
162/// - 入库流水线与响应契约完全一致（magic bytes 检 MIME、5MiB、尺寸/像素校验、
163///   SHA-256 去重、转码、assets 登记）。
164///
165/// 未引用素材由 `tasks::orphan_asset_purge` 定期回收（评论未提交、评论被删等
166/// 场景产生的孤儿图）。
167#[cfg(feature = "server")]
168pub async fn comment_upload_image(
169    connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
170    headers: HeaderMap,
171    mut multipart: Multipart,
172) -> Result<Json<Value>, (StatusCode, Json<Value>)> {
173    // 0. 双层限流（先于一切 body 读取）。
174    let peer = connect_info.map(|Extension(ConnectInfo(addr))| addr);
175    let ip = crate::api::rate_limit::get_client_ip_with_peer(&headers, peer).await;
176    if let Err(msg) = crate::api::rate_limit::check_comment_upload_limit(&ip) {
177        return Err(upload_error(StatusCode::TOO_MANY_REQUESTS, msg));
178    }
179
180    // 1. Read multipart field
181    let field = match multipart.next_field().await {
182        Ok(Some(f)) => f,
183        Ok(None) => {
184            return Err(upload_error(StatusCode::BAD_REQUEST, "未找到文件"));
185        }
186        Err(e) => {
187            tracing::error!("Comment multipart error: {:?}", e);
188            return Err(upload_error(StatusCode::BAD_REQUEST, "文件读取失败"));
189        }
190    };
191
192    // 2. 早拒非法声明类型（快速路径；流水线仍以 magic bytes 为权威）。
193    let declared_mime = field.content_type().unwrap_or("").to_string();
194    if !ALLOWED_MIME_TYPES.contains(&declared_mime.as_str()) {
195        return Err(upload_error(StatusCode::BAD_REQUEST, "不支持的文件类型"));
196    }
197
198    let original_filename = field.file_name().map(|s| s.to_string());
199
200    // 3. Read file data
201    let data = match field.bytes().await {
202        Ok(d) => d,
203        Err(e) => {
204            tracing::error!("Comment read file error: {:?}", e);
205            return Err(upload_error(
206                StatusCode::INTERNAL_SERVER_ERROR,
207                "文件读取失败",
208            ));
209        }
210    };
211
212    // 4. 共享入库流水线。
213    match process_image_upload(data, original_filename, None).await {
214        Ok(out) => Ok(Json(json!({
215            "success": true,
216            "url": out.url,
217            "reused": out.reused
218        }))),
219        Err(e) => {
220            let (status, msg) = e.status_and_msg();
221            Err(upload_error(status, msg))
222        }
223    }
224}
225
226// ===========================================================================
227// MCP 处理器（bearer 鉴权，multipart 二进制，带外传输）
228// ===========================================================================
229
230/// MCP bearer 上传错误 → JSON 响应（与 web 端格式一致）。
231#[cfg(feature = "server")]
232fn mcp_upload_error<T: serde::Serialize>(status: StatusCode, msg: T) -> Response {
233    (status, Json(json!({ "success": false, "error": msg }))).into_response()
234}
235
236/// 处理图片上传的 Axum handler（MCP 端，bearer token 鉴权）。
237///
238/// 与 web `upload_image` 的区别：
239/// - 鉴权用 `Authorization: Bearer ygg_...`（不是 cookie），经
240///   [`crate::mcp::auth::resolve_bearer_principal`] 解析；
241/// - 不挂 CSRF 中间件——bearer 在请求头里，浏览器不会自动附带，无 CSRF 风险；
242/// - 限流按 token_id 计数（复用 MCP 的 token-keyed governor）。
243///
244/// 供 AI 客户端的 host/shell 直接 POST 二进制（Claude Code 的 Bash+curl 等），
245/// 二进制不经 JSON-RPC，绕开 rmcp 4MiB 请求体上限。返回可直接嵌入 Markdown 的
246/// `/uploads/...` URL。
247#[cfg(feature = "server")]
248pub async fn mcp_upload_image(headers: HeaderMap, mut multipart: Multipart) -> Response {
249    // 1. bearer → principal（含 scope 校验：media 需要 write）。
250    let principal = match crate::mcp::auth::resolve_bearer_principal(&headers).await {
251        Ok(p) => p,
252        Err(status) => return mcp_upload_error(status, "未授权或令牌无效"),
253    };
254    if !principal
255        .scope
256        .grants(crate::models::mcp_token::TokenScope::Write)
257    {
258        return mcp_upload_error(StatusCode::FORBIDDEN, "权限不足：需要 write 作用域");
259    }
260
261    // 2. token-keyed 限流（与 /mcp 中间件的 MCP_LIMITER 隔离：上传单独配额）。
262    if let Err(msg) = crate::mcp::auth::check_mcp_upload_limit(&principal.token_id) {
263        return mcp_upload_error(StatusCode::TOO_MANY_REQUESTS, msg);
264    }
265
266    // 3. 支持 file + 可选 alt，字段顺序不限；不接受重复字段。
267    let mut file = None;
268    let mut alt = None;
269    loop {
270        let field = match multipart.next_field().await {
271            Ok(Some(field)) => field,
272            Ok(None) => break,
273            Err(_) => return mcp_upload_error(StatusCode::BAD_REQUEST, "文件读取失败"),
274        };
275        if field.name() == Some("alt") {
276            if alt.is_some() {
277                return mcp_upload_error(StatusCode::BAD_REQUEST, "alt 字段重复");
278            }
279            alt = match field.text().await {
280                Ok(value) => Some(value),
281                Err(_) => return mcp_upload_error(StatusCode::BAD_REQUEST, "alt 读取失败"),
282            };
283        } else if field.name() == Some("file") || field.file_name().is_some() {
284            if file.is_some() {
285                return mcp_upload_error(StatusCode::BAD_REQUEST, "每次只允许上传一个文件");
286            }
287            if !ALLOWED_MIME_TYPES.contains(&field.content_type().unwrap_or("")) {
288                return mcp_upload_error(StatusCode::BAD_REQUEST, "不支持的文件类型");
289            }
290            let filename = field.file_name().map(str::to_string);
291            let data = match field.bytes().await {
292                Ok(data) => data,
293                Err(_) => return mcp_upload_error(StatusCode::BAD_REQUEST, "文件读取失败"),
294            };
295            file = Some((data, filename));
296        } else {
297            return mcp_upload_error(StatusCode::BAD_REQUEST, "未知的上传字段");
298        }
299    }
300    let Some((data, original_filename)) = file else {
301        return mcp_upload_error(StatusCode::BAD_REQUEST, "未找到文件");
302    };
303
304    // 4. 共享入库流水线。
305    match process_image_upload(data, original_filename, alt).await {
306        Ok(out) => Json(json!({
307            "success": true,
308            "url": out.url,
309            "asset_id": out.asset_id,
310            "alt": out.alt,
311            "reused": out.reused,
312            "width": out.width,
313            "height": out.height,
314            "mime": out.mime
315        }))
316        .into_response(),
317        Err(e) => {
318            let (status, msg) = e.status_and_msg();
319            mcp_upload_error(status, msg)
320        }
321    }
322}
323
324// ===========================================================================
325// 共享入库流水线
326// ===========================================================================
327
328/// 单条图片入库的结果。
329#[cfg(feature = "server")]
330#[derive(Debug, serde::Serialize)]
331pub(crate) struct UploadOutcome {
332    /// 可直接嵌入 Markdown 的相对 URL：`/uploads/YYYY/MM/DD/HHMMSS.uuid.ext`。
333    pub url: String,
334    /// 素材注册表 ID，重复上传返回同一 ID。
335    pub asset_id: String,
336    pub alt: Option<String>,
337    /// 是否命中已登记素材（内容去重或并发竞态复用）。
338    pub reused: bool,
339    pub width: u32,
340    pub height: u32,
341    /// 最终 MIME（转码后；JPEG→WebP 成功则为 image/webp）。
342    pub mime: String,
343}
344
345/// 流水线错误：映射到 HTTP 状态 + 脱敏消息（不泄露 SQL/路径细节）。
346#[cfg(feature = "server")]
347#[derive(Debug)]
348pub(crate) enum UploadError {
349    Empty,
350    BadType,   // magic bytes 无法识别为 JPEG/PNG/GIF/WebP
351    TooLarge,  // 超过 MAX_FILE_SIZE
352    Oversized, // 像素超过 MAX_IMAGE_PIXELS
353    Corrupt,   // GIF/WebP 解码失败
354    /// 内部错误：携带静态上下文标签供 Debug 诊断（status_and_msg 统一返回脱敏消息）。
355    #[allow(dead_code)]
356    Internal(&'static str),
357}
358
359#[cfg(feature = "server")]
360impl UploadError {
361    /// 包装底层错误：服务端日志记完整 `{e}`，客户端只见静态 `ctx`。
362    fn internal<E: std::fmt::Display>(e: E, ctx: &'static str) -> Self {
363        tracing::error!("upload {ctx}: {e}");
364        UploadError::Internal(ctx)
365    }
366
367    /// 映射到 (HTTP 状态, 脱敏消息)。
368    fn status_and_msg(&self) -> (StatusCode, &'static str) {
369        match self {
370            UploadError::Empty => (StatusCode::BAD_REQUEST, "空文件"),
371            UploadError::BadType => (StatusCode::BAD_REQUEST, "不支持的文件类型"),
372            UploadError::TooLarge => (StatusCode::PAYLOAD_TOO_LARGE, "文件超过大小限制"),
373            UploadError::Oversized => (StatusCode::BAD_REQUEST, "图片尺寸超过上限"),
374            UploadError::Corrupt => (StatusCode::BAD_REQUEST, "图片文件损坏或格式不正确"),
375            UploadError::Internal(_) => (StatusCode::INTERNAL_SERVER_ERROR, "文件保存失败"),
376        }
377    }
378}
379
380/// 单一图片入库流水线（web 上传 / MCP bearer 端点 / MCP URL 抓取共用）。
381///
382/// 输入：原始字节 + 可选展示文件名。**不信任客户端声明的 MIME**——以 magic
383/// bytes 为唯一真相。输出可直接嵌入 Markdown 的 `/uploads/...` URL。
384///
385/// 步骤：大小校验 → magic bytes 检 MIME → 尺寸/像素校验 → SHA-256 去重（命中
386/// 即复用，跳过最贵的转码）→ GIF/WebP 解码校验 → `spawn_blocking` 转码 →
387/// 按日期落盘 → assets 登记（含并发竞态补偿，落败者删自己的文件复用胜出者）。
388#[cfg(feature = "server")]
389pub(crate) async fn process_image_upload(
390    data: bytes::Bytes,
391    original_filename: Option<String>,
392    alt: Option<String>,
393) -> Result<UploadOutcome, UploadError> {
394    if data.is_empty() {
395        return Err(UploadError::Empty);
396    }
397    if data.len() > MAX_FILE_SIZE {
398        return Err(UploadError::TooLarge);
399    }
400
401    // 1. magic bytes 检 MIME（不信任声明类型/扩展名）。
402    let mime_type = detect_mime(&data).ok_or(UploadError::BadType)?;
403
404    // 2. 仅读 header 校验尺寸/像素上限，并拿回 (w,h) 供 assets 登记，避免二次解析。
405    //    超限直接拒绝，避免大图 decode 后被静默降级（原 fallback 存原图）。
406    let (img_width, img_height) =
407        crate::api::image::upload_dimensions(&data, mime_type).map_err(|msg| {
408            tracing::warn!("upload dimensions check failed: {msg}");
409            UploadError::Oversized
410        })?;
411
412    let alt = alt.map(|value| value.trim().to_string());
413    let is_gif = mime_type == "image/gif";
414    let is_webp = mime_type == "image/webp";
415
416    // 3. 内容去重（CAS）：对原始上传字节算 SHA-256，命中已登记素材直接复用，
417    //    跳过 GIF/WebP 解码验证、转码与落盘（省下整个流程最贵的 CPU）。
418    //    放在安全性校验之后、转码之前。命中时刷新 created_at/updated_at：
419    //    重传代表使用意图，重启 7 天清理保护窗（PURGE_GRACE_DAYS 保护的是
420    //    「刚上传还没被文章引用」的素材）。
421    let content_hash = {
422        use sha2::Digest;
423        hex::encode(sha2::Sha256::digest(&data))
424    };
425    {
426        let client = crate::db::pool::get_conn()
427            .await
428            .map_err(|e| UploadError::internal(e, "dedup conn"))?;
429        let reused = client
430            .query_opt(
431                "UPDATE assets SET created_at = NOW(), updated_at = NOW(), \
432                 alt = CASE WHEN $2::text IS NULL THEN alt ELSE NULLIF($2, '') END \
433                 WHERE content_hash = $1 RETURNING id, path, mime, width, height, alt",
434                &[&content_hash, &alt],
435            )
436            .await
437            .map_err(|e| UploadError::internal(e, "dedup check"))?;
438        if let Some(row) = reused {
439            let path: String = row.get("path");
440            tracing::info!(
441                "Image deduped: reuse {} (hash {})",
442                path,
443                &content_hash[..12]
444            );
445            return Ok(upload_outcome(&row, true));
446        }
447    }
448
449    // 4. GIF/WebP 解码校验（不经过重编码的格式必须验真，防伪造扩展名的恶意文件）。
450    //    GIF 走 image::load_from_memory 会完整解码，移到阻塞线程池避免拖住 async 运行时。
451    if is_gif || is_webp {
452        let validate_data = data.clone();
453        let validate_mime = mime_type.to_string();
454        let is_valid = tokio::task::spawn_blocking(move || {
455            validate_raw_image(&validate_data, validate_mime.as_str())
456        })
457        .await
458        .map_err(|e| UploadError::internal(e, "validate task"))?;
459        if !is_valid {
460            return Err(UploadError::Corrupt);
461        }
462    }
463
464    // 5. 转码：GIF/WebP 原样；JPEG/PNG 仅在 WebP 更小时转。
465    //    Bytes clone 廉价（引用计数 +1），move 进阻塞闭包无需全文件深拷贝。
466    let (final_data, final_ext) = transcode(data, mime_type, is_gif, is_webp).await;
467
468    // 6. 按上传时间组织目录：uploads/YYYY/MM/DD。
469    //    chrono 的 DelayedFormat 实现 Display，可直接进 format!，省掉中间 String。
470    let now = chrono::Utc::now();
471    let date = now.format("%Y/%m/%d");
472    let uuid_str = uuid::Uuid::new_v4().to_string();
473
474    let dir_path = format!("uploads/{}", date);
475    let file_name = format!("{}.{}.{}", now.format("%H%M%S"), uuid_str, final_ext);
476    let file_path = format!("{}/{}", dir_path, file_name);
477    let rel_path = format!("{}/{}", date, file_name);
478    let final_mime = mime_for_ext(&final_ext);
479
480    if let Err(e) = tokio::fs::create_dir_all(&dir_path).await {
481        return Err(UploadError::internal(e, "create dir"));
482    }
483    if let Err(e) = tokio::fs::write(&file_path, &final_data).await {
484        return Err(UploadError::internal(e, "write file"));
485    }
486
487    tracing::info!("Image uploaded: {} ({} bytes)", file_path, final_data.len());
488
489    // 7. 注册或复用同内容素材，原子返回最终元数据；并发落败者删除自己落盘的文件。
490    let registered: Result<tokio_postgres::Row, UploadError> = async {
491        let client = crate::db::pool::get_conn()
492            .await
493            .map_err(|e| UploadError::internal(e, "register conn"))?;
494        client.query_one(
495            "INSERT INTO assets (id, path, filename, mime, size_bytes, width, height, content_hash, alt)
496             VALUES ($1, $2, $3, $4, $5, $6, $7, $8, NULLIF($9, ''))
497             ON CONFLICT (content_hash) DO UPDATE SET
498                 created_at = NOW(), updated_at = NOW(),
499                 alt = CASE WHEN $9::text IS NULL THEN assets.alt ELSE NULLIF($9, '') END
500             RETURNING id, path, mime, width, height, alt",
501            &[&uuid::Uuid::new_v4(), &rel_path,
502              &original_filename.unwrap_or_else(|| file_name.clone()), &final_mime,
503              &(final_data.len() as i64), &(img_width as i32), &(img_height as i32),
504              &content_hash, &alt],
505        ).await.map_err(|e| UploadError::internal(e, "register asset"))
506    }.await;
507
508    match registered {
509        Ok(row) => {
510            let reused = row.get::<_, String>("path") != rel_path;
511            if reused {
512                let _ = tokio::fs::remove_file(&file_path).await;
513            }
514            Ok(upload_outcome(&row, reused))
515        }
516        Err(e) => {
517            // 登记失败：补偿删除已落盘文件。
518            let _ = tokio::fs::remove_file(&file_path).await;
519            Err(e)
520        }
521    }
522}
523
524#[cfg(feature = "server")]
525fn upload_outcome(row: &tokio_postgres::Row, reused: bool) -> UploadOutcome {
526    UploadOutcome {
527        url: format!("/uploads/{}", row.get::<_, String>("path")),
528        asset_id: row.get::<_, uuid::Uuid>("id").to_string(),
529        alt: row.get("alt"),
530        reused,
531        width: row.get::<_, i32>("width") as u32,
532        height: row.get::<_, i32>("height") as u32,
533        mime: row.get("mime"),
534    }
535}
536
537// ===========================================================================
538// 图片处理辅助
539// ===========================================================================
540
541/// 从 magic bytes 检测 MIME 类型（不信任客户端声明的扩展名/Content-Type）。
542#[cfg(feature = "server")]
543pub(crate) fn detect_mime(data: &[u8]) -> Option<&'static str> {
544    if data.starts_with(&[0xFF, 0xD8, 0xFF]) {
545        Some("image/jpeg")
546    } else if data.starts_with(&[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]) {
547        Some("image/png")
548    } else if data.starts_with(b"GIF87a") || data.starts_with(b"GIF89a") {
549        Some("image/gif")
550    } else if data.len() >= 12 && &data[0..4] == b"RIFF" && &data[8..12] == b"WEBP" {
551        Some("image/webp")
552    } else {
553        None
554    }
555}
556
557#[cfg(feature = "server")]
558fn mime_to_ext(mime: &str) -> &'static str {
559    match mime {
560        "image/jpeg" => "jpg",
561        "image/png" => "png",
562        "image/webp" => "webp",
563        "image/gif" => "gif",
564        _ => "bin",
565    }
566}
567
568#[cfg(feature = "server")]
569fn mime_for_ext(ext: &str) -> &'static str {
570    match ext {
571        "jpg" => "image/jpeg",
572        "png" => "image/png",
573        "gif" => "image/gif",
574        _ => "image/webp",
575    }
576}
577
578/// 解码验证 GIF/WebP 原始字节，确保不是伪造扩展名的恶意文件。
579#[cfg(feature = "server")]
580fn validate_raw_image(data: &[u8], mime_type: &str) -> bool {
581    match mime_type {
582        "image/webp" => crate::infra::webp::decode(data).is_ok(),
583        "image/gif" => image::load_from_memory(data).is_ok(),
584        _ => true,
585    }
586}
587
588/// 转码核心（同步）：GIF/WebP 保持原格式，JPEG/PNG 尝试转 WebP（更小才采用）。
589#[cfg(feature = "server")]
590fn transcode_image_blocking(
591    data: &[u8],
592    mime: &'static str,
593    is_gif: bool,
594    is_webp: bool,
595) -> (Vec<u8>, String) {
596    if is_gif {
597        return (data.to_vec(), "gif".to_string());
598    }
599    if is_webp {
600        return (data.to_vec(), "webp".to_string());
601    }
602
603    // JPEG/PNG → 尝试 WebP。
604    let format = match mime {
605        "image/jpeg" => image::ImageFormat::Jpeg,
606        "image/png" => image::ImageFormat::Png,
607        _ => image::ImageFormat::Jpeg,
608    };
609    let cursor = std::io::Cursor::new(data);
610    let mut reader = image::ImageReader::with_format(cursor, format);
611    reader.limits(crate::api::image::image_reader_limits());
612
613    match reader.decode() {
614        Ok(img) => {
615            let config = crate::infra::webp::WEBP_CONFIG.clone();
616            match crate::infra::webp::encode(&img, config.quality, config.method) {
617                Ok(webp_data) if webp_data.len() < data.len() => {
618                    tracing::info!(
619                        "WebP conversion: {}x{} {} -> {} bytes",
620                        img.width(),
621                        img.height(),
622                        data.len(),
623                        webp_data.len()
624                    );
625                    (webp_data, "webp".to_string())
626                }
627                Ok(_) => {
628                    // WebP 更大，保留原格式。
629                    (data.to_vec(), mime_to_ext(mime).to_string())
630                }
631                Err(e) => {
632                    tracing::warn!("WebP encode failed ({}), keeping original", e);
633                    (data.to_vec(), mime_to_ext(mime).to_string())
634                }
635            }
636        }
637        // 到这里尺寸校验已通过（超限在 header 阶段被拒），decode 失败只能是真损坏。
638        Err(e) => {
639            tracing::warn!("Failed to decode image ({}), keeping original format", e);
640            (data.to_vec(), mime_to_ext(mime).to_string())
641        }
642    }
643}
644
645/// 在阻塞线程中执行转码，避免阻塞 async 运行时。
646/// Bytes clone 廉价（引用计数 +1）；join 失败（panic）时回退原格式。
647#[cfg(feature = "server")]
648async fn transcode(
649    data: bytes::Bytes,
650    mime: &'static str,
651    is_gif: bool,
652    is_webp: bool,
653) -> (Vec<u8>, String) {
654    let for_task = data.clone();
655    match tokio::task::spawn_blocking(move || {
656        transcode_image_blocking(&for_task, mime, is_gif, is_webp)
657    })
658    .await
659    {
660        Ok(result) => result,
661        Err(e) => {
662            tracing::warn!("transcode task panicked ({}), keeping original", e);
663            (data.to_vec(), mime_to_ext(mime).to_string())
664        }
665    }
666}
667
668#[cfg(all(test, feature = "server"))]
669mod tests {
670
671    /// Uses the same disposable database as the private article workflow test.
672    #[test]
673    #[ignore = "requires disposable DATABASE_URL database ygg_mcp_test"]
674    fn upload_metadata_database_and_http() {
675        crate::db::TEST_DATABASE_RUNTIME.block_on(upload_metadata_database_and_http_impl());
676    }
677
678    async fn upload_metadata_database_and_http_impl() {
679        let _guard = crate::db::TEST_DATABASE_LOCK.lock().await;
680        use super::*;
681        let mut client = crate::db::pool::get_conn().await.unwrap();
682        let database: String = client
683            .query_one("SELECT current_database()", &[])
684            .await
685            .unwrap()
686            .get(0);
687        assert_eq!(database, "ygg_mcp_test", "requires isolated test database");
688        crate::db::migrate::run_on_conn(&mut client).await.unwrap();
689        let owner = client
690            .query_opt("SELECT id FROM users WHERE role = 'admin'", &[])
691            .await
692            .unwrap();
693        let user_id: i32 = match owner {
694        Some(row) => row.get(0),
695        None => client.query_one("INSERT INTO users (username, email, password_hash, role) VALUES ('upload-test', 'upload@test.invalid', 'unused', 'admin') RETURNING id", &[]).await.unwrap().get(0),
696    };
697        let nonce = uuid::Uuid::new_v4();
698        let color = nonce.as_bytes();
699        let image = image::RgbaImage::from_fn(64, 64, |x, _| {
700            image::Rgba([color[x as usize % 16], color[1], color[2], 255])
701        });
702        let mut png = std::io::Cursor::new(Vec::new());
703        image.write_to(&mut png, image::ImageFormat::Png).unwrap();
704        let data = bytes::Bytes::from(png.into_inner());
705        let (a, b) = tokio::join!(
706            process_image_upload(
707                data.clone(),
708                Some("test.png".into()),
709                Some("  首次图片  ".into())
710            ),
711            process_image_upload(
712                data.clone(),
713                Some("test.png".into()),
714                Some("  首次图片  ".into())
715            ),
716        );
717        let (a, b) = (a.unwrap(), b.unwrap());
718        assert_eq!(a.asset_id, b.asset_id);
719        assert_eq!(a.url, b.url);
720        assert_ne!(a.reused, b.reused);
721        assert_eq!(a.mime, b.mime);
722        assert_eq!(a.alt.as_deref(), Some("首次图片"));
723        let reused = process_image_upload(data.clone(), None, None)
724            .await
725            .unwrap();
726        assert_eq!(reused.asset_id, a.asset_id);
727        assert_eq!(reused.alt, a.alt);
728        assert_eq!(reused.mime, a.mime);
729        let id = uuid::Uuid::parse_str(&a.asset_id).unwrap();
730        let stored = client
731            .query_one("SELECT alt, mime FROM assets WHERE id = $1", &[&id])
732            .await
733            .unwrap();
734        assert_eq!(stored.get::<_, Option<String>>(0), a.alt);
735        assert_eq!(stored.get::<_, String>(1), a.mime);
736
737        // Exercise actual multipart parsing, authentication and serialized response.
738        let token = format!("ygg_{nonce}");
739        client.execute("INSERT INTO mcp_tokens (id, user_id, name, scope, token_enc, token_hash) VALUES ($1, $2, 'upload-test', 'write', 'unused', $3)",
740        &[&nonce, &user_id, &crate::mcp::auth::hash_token(&token)]).await.unwrap();
741        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
742        let addr = listener.local_addr().unwrap();
743        let router = axum::Router::new().route("/upload", axum::routing::post(mcp_upload_image));
744        let server = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() });
745        let http = reqwest::Client::new();
746        let body = |alt: &str, alt_first: bool, duplicate: bool| {
747            let alt_part = format!(
748                "--test-boundary\r\nContent-Disposition: form-data; name=\"alt\"\r\n\r\n{alt}\r\n"
749            )
750            .into_bytes();
751            let mut file_part = b"--test-boundary\r\nContent-Disposition: form-data; name=\"file\"; filename=\"test.png\"\r\nContent-Type: image/png\r\n\r\n".to_vec();
752            file_part.extend_from_slice(&data);
753            file_part.extend_from_slice(b"\r\n");
754            let mut out = Vec::new();
755            if alt_first {
756                out.extend_from_slice(&alt_part);
757            }
758            out.extend_from_slice(&file_part);
759            if duplicate {
760                out.extend_from_slice(&file_part);
761            }
762            if !alt_first {
763                out.extend_from_slice(&alt_part);
764            }
765            out.extend_from_slice(b"--test-boundary--\r\n");
766            out
767        };
768        for (alt, alt_first) in [("  HTTP 图片  ", true), ("   ", false)] {
769            let response = http
770                .post(format!("http://{addr}/upload"))
771                .bearer_auth(&token)
772                .header(
773                    "content-type",
774                    "multipart/form-data; boundary=test-boundary",
775                )
776                .body(body(alt, alt_first, false))
777                .send()
778                .await
779                .unwrap();
780            let status = response.status();
781            let text = response.text().await.unwrap();
782            assert_eq!(status, StatusCode::OK, "{text}");
783            let result: serde_json::Value = serde_json::from_str(&text).unwrap();
784            assert_eq!(result["asset_id"], a.asset_id);
785            assert_eq!(result["mime"], a.mime);
786            let expected = if alt.trim().is_empty() {
787                None
788            } else {
789                Some(alt.trim().to_string())
790            };
791            assert_eq!(result["alt"], serde_json::to_value(&expected).unwrap());
792            assert_eq!(
793                client
794                    .query_one("SELECT alt FROM assets WHERE id = $1", &[&id])
795                    .await
796                    .unwrap()
797                    .get::<_, Option<String>>(0),
798                expected
799            );
800        }
801        let response = http
802            .post(format!("http://{addr}/upload"))
803            .bearer_auth(&token)
804            .header(
805                "content-type",
806                "multipart/form-data; boundary=test-boundary",
807            )
808            .body(body("rejected", false, true))
809            .send()
810            .await
811            .unwrap();
812        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
813        assert!(client
814            .query_one("SELECT alt FROM assets WHERE id = $1", &[&id])
815            .await
816            .unwrap()
817            .get::<_, Option<String>>(0)
818            .is_none());
819        let response = http
820            .post(format!("http://{addr}/upload"))
821            .header(
822                "content-type",
823                "multipart/form-data; boundary=test-boundary",
824            )
825            .body(body("unauthorized", true, false))
826            .send()
827            .await
828            .unwrap();
829        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
830        server.abort();
831        client
832            .execute("DELETE FROM mcp_tokens WHERE id = $1", &[&nonce])
833            .await
834            .unwrap();
835        client
836            .execute("DELETE FROM assets WHERE id = $1", &[&id])
837            .await
838            .unwrap();
839        tokio::fs::remove_file(a.url.trim_start_matches('/'))
840            .await
841            .unwrap();
842    }
843
844    #[test]
845    fn filename_format_no_spaces() {
846        let now_str = "120000";
847        let uuid = "abc-123";
848        let ext = "jpg";
849        let file_name = format!("{}.{}.{}", now_str, uuid, ext);
850        assert!(
851            !file_name.contains(' '),
852            "filename should not contain spaces: got '{}'",
853            file_name
854        );
855    }
856
857    #[test]
858    fn should_use_webp_ext_for_non_gif() {
859        let ext = "jpg";
860        let mime = "image/jpeg";
861        let is_gif = mime == "image/gif";
862        let final_ext = if is_gif { ext } else { "webp" };
863        assert_eq!(final_ext, "webp");
864    }
865
866    #[test]
867    fn should_preserve_gif_ext() {
868        let ext = "gif";
869        let mime = "image/gif";
870        let is_gif = mime == "image/gif";
871        let final_ext = if is_gif { ext } else { "webp" };
872        assert_eq!(final_ext, "gif");
873    }
874
875    #[test]
876    fn convert_to_webp_produces_bytes() {
877        let img = image::DynamicImage::new_rgb8(10, 10);
878        let result = crate::infra::webp::encode(&img, 85.0, 4).unwrap();
879        assert!(!result.is_empty());
880    }
881
882    #[test]
883    fn webp_roundtrip_from_rgba() {
884        let img = image::DynamicImage::new_rgba8(2, 2);
885        let webp_bytes = crate::infra::webp::encode(&img, 85.0, 4).unwrap();
886        let loaded = crate::infra::webp::decode(&webp_bytes);
887        assert!(loaded.is_ok());
888    }
889
890    #[test]
891    fn mime_to_ext_maps_jpeg() {
892        assert_eq!(super::mime_to_ext("image/jpeg"), "jpg");
893    }
894
895    #[test]
896    fn mime_to_ext_maps_png() {
897        assert_eq!(super::mime_to_ext("image/png"), "png");
898    }
899
900    #[test]
901    fn mime_to_ext_maps_gif() {
902        assert_eq!(super::mime_to_ext("image/gif"), "gif");
903    }
904
905    #[test]
906    fn mime_to_ext_maps_webp() {
907        assert_eq!(super::mime_to_ext("image/webp"), "webp");
908    }
909
910    #[test]
911    fn mime_to_ext_falls_back_for_unknown_mime() {
912        assert_eq!(super::mime_to_ext("image/avif"), "bin");
913        assert_eq!(super::mime_to_ext("application/octet-stream"), "bin");
914    }
915
916    #[test]
917    fn mime_for_ext_roundtrip() {
918        assert_eq!(super::mime_for_ext("jpg"), "image/jpeg");
919        assert_eq!(super::mime_for_ext("png"), "image/png");
920        assert_eq!(super::mime_for_ext("gif"), "image/gif");
921        assert_eq!(super::mime_for_ext("webp"), "image/webp");
922    }
923
924    #[test]
925    fn detect_mime_jpeg() {
926        assert_eq!(
927            super::detect_mime(&[0xFF, 0xD8, 0xFF, 0xE0]),
928            Some("image/jpeg")
929        );
930        assert_eq!(super::detect_mime(&[0x89, 0x50]), None);
931    }
932
933    #[test]
934    fn detect_mime_png() {
935        assert_eq!(
936            super::detect_mime(&[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]),
937            Some("image/png")
938        );
939        assert_eq!(super::detect_mime(&[0xFF, 0xD8]), None);
940    }
941
942    #[test]
943    fn detect_mime_gif() {
944        assert_eq!(super::detect_mime(b"GIF89a"), Some("image/gif"));
945        assert_eq!(super::detect_mime(b"GIF87a"), Some("image/gif"));
946        assert_eq!(super::detect_mime(b"GIF90a"), None);
947    }
948
949    #[test]
950    fn detect_mime_webp() {
951        let webp = b"RIFF\x00\x00\x00\x00WEBPVP8 ";
952        assert_eq!(super::detect_mime(&webp[..12]), Some("image/webp"));
953        assert_eq!(super::detect_mime(&[0xFF, 0xD8]), None);
954    }
955
956    #[test]
957    fn detect_mime_unknown() {
958        assert_eq!(super::detect_mime(b"hello world"), None);
959        assert_eq!(super::detect_mime(&[]), None);
960    }
961}
yggdrasil/api/upload.rs

yggdrasil/api/
upload.rs
