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

yggdrasil/api/
upload.rs
