yggdrasil/src/api/upload.rs
xfy 5743c7a581
Some checks failed
CI / check (push) Has been cancelled
CI / build (push) Has been cancelled
feat(upload): 上传图片按内容 SHA-256 去重,重复上传复用已登记素材
assets 新增 content_hash 列(唯一索引,迁移 016;存量行保持 NULL
不参与去重)。上传在安全性校验后、转码前对原始字节算 SHA-256:
命中即刷新 created_at/updated_at(重启 7 天清理保护窗)并返回
已有 URL,跳过 GIF/WebP 验证、转码与落盘。并发同内容上传由唯一
索引 + INSERT ... ON CONFLICT DO NOTHING 兜底,落败者删落盘文件
复用胜出者路径。命中响应带 reused: true。仅精确去重,视觉相似
检测(pHash)有意不做。
2026-07-27 13:25:26 +08:00

531 lines
20 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

//! 图片上传的 Axum 处理器。
//!
//! 处理 multipart 上传,校验 MIME 类型、文件大小与 admin 权限,
//! JPEG/PNG 自动转 WebP若体积更小则保留原格式GIF/WebP 保持原样。
//! 文件按日期分目录存放于 `uploads/`。
//!
//! 内容去重CAS以原始上传字节的 SHA-256 为内容指纹(`assets.content_hash`
//! 唯一索引)。重复上传同一内容时复用已登记素材——同一行、同一文件,不重复
//! 落盘,响应带 `"reused": true`;并发同内容上传由唯一索引 + ON CONFLICT 兜底。
//! 仅精确去重:尺寸/压缩不同的视觉相似图不合并(那是感知哈希 pHash 的领域,
//! 有意不做)。
//! 本模块属于手动注册的 Axum 路由,仅在 `feature = "server"` 时可用。
#[cfg(feature = "server")]
use axum::{
extract::{ConnectInfo, Extension, Multipart},
http::{HeaderMap, StatusCode},
response::Json,
};
#[cfg(feature = "server")]
use serde_json::{json, Value};
#[cfg(feature = "server")]
use std::net::SocketAddr;
#[cfg(feature = "server")]
use crate::auth::session::parse_session_token;
#[cfg(feature = "server")]
const ALLOWED_MIME_TYPES: &[&str] = &["image/jpeg", "image/png", "image/gif", "image/webp"];
#[cfg(feature = "server")]
const MAX_FILE_SIZE: usize = 5 * 1024 * 1024; // 5MB
/// 构造统一的 JSON 错误响应:`{ "success": false, "error": msg }`。
#[cfg(feature = "server")]
fn upload_error<T: serde::Serialize>(status: StatusCode, msg: T) -> (StatusCode, Json<Value>) {
(status, Json(json!({ "success": false, "error": msg })))
}
#[cfg(feature = "server")]
fn mime_to_ext(mime: &str) -> &'static str {
match mime {
"image/jpeg" => "jpg",
"image/png" => "png",
"image/webp" => "webp",
"image/gif" => "gif",
_ => "bin",
}
}
#[cfg(feature = "server")]
/// 通过文件头 magic bytes 校验实际格式是否与声明 MIME 一致。
fn validate_image_magic_bytes(data: &[u8], mime_type: &str) -> bool {
if data.is_empty() {
return false;
}
match mime_type {
"image/jpeg" => data.starts_with(&[0xFF, 0xD8, 0xFF]),
"image/png" => data.starts_with(&[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]),
"image/gif" => data.starts_with(b"GIF87a") || data.starts_with(b"GIF89a"),
"image/webp" => {
// RIFF....WEBP
data.len() >= 12 && &data[0..4] == b"RIFF" && &data[8..12] == b"WEBP"
}
_ => false,
}
}
#[cfg(feature = "server")]
/// 解码验证 GIF/WebP 原始字节,确保不是伪造扩展名的恶意文件。
fn validate_raw_image(data: &[u8], mime_type: &str) -> bool {
match mime_type {
"image/webp" => crate::webp::decode(data).is_ok(),
"image/gif" => image::load_from_memory(data).is_ok(),
_ => true,
}
}
#[cfg(feature = "server")]
/// 处理图片上传的 Axum handler。
///
/// 流程:限流 → 解析 session → 校验 admin → 读取 multipart → 校验类型/大小 →
/// 转码(如适用)→ 按日期落盘 → 返回相对 URL。
///
/// `ConnectInfo` 以可选扩展注入:`dioxus::server::serve()` 接管了 listener
/// 无法调用 `into_make_service_with_connect_info::<SocketAddr>()`,所以这里
/// 与 `serve_image` 保持一致的优雅降级——扩展缺失时退回 `"unknown"` 限流桶。
/// 生产环境应在反向代理后部署并配置 `TRUSTED_PROXY_COUNT`,让限流拿到真实 IP。
pub async fn upload_image(
connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
headers: HeaderMap,
mut multipart: Multipart,
) -> Result<Json<Value>, (StatusCode, Json<Value>)> {
// 0. Rate limit check
let peer = connect_info.map(|Extension(ConnectInfo(addr))| addr);
let ip = crate::api::rate_limit::get_client_ip_with_peer(&headers, peer);
if let Err(msg) = crate::api::rate_limit::check_upload_limit(&ip) {
return Err(upload_error(StatusCode::TOO_MANY_REQUESTS, msg));
}
// 1. Extract session from cookie
let cookie_header = headers
.get("cookie")
.and_then(|h| h.to_str().ok())
.unwrap_or("");
let token = match parse_session_token(cookie_header) {
Some(t) => t,
None => {
return Err(upload_error(StatusCode::UNAUTHORIZED, "未登录"));
}
};
// 2. Verify admin
let user = match crate::api::auth::get_user_by_token(token).await {
Ok(Some(u)) => u,
_ => {
return Err(upload_error(StatusCode::UNAUTHORIZED, "会话已过期"));
}
};
if user.role != crate::models::user::UserRole::Admin {
return Err(upload_error(StatusCode::FORBIDDEN, "权限不足"));
}
// 3. Read multipart field
let field = match multipart.next_field().await {
Ok(Some(f)) => f,
Ok(None) => {
return Err(upload_error(StatusCode::BAD_REQUEST, "未找到文件"));
}
Err(e) => {
tracing::error!("Multipart error: {:?}", e);
return Err(upload_error(StatusCode::BAD_REQUEST, "文件读取失败"));
}
};
// 4. Validate mime type
let mime_type = field.content_type().unwrap_or("").to_string();
if !ALLOWED_MIME_TYPES.contains(&mime_type.as_str()) {
return Err(upload_error(StatusCode::BAD_REQUEST, "不支持的文件类型"));
}
// 原始文件名(客户端提供,仅作 assets 表展示字段);需在 bytes() 消耗 field 前取出。
let original_filename = field.file_name().map(|s| s.to_string());
// 5. Read file data
let data = match field.bytes().await {
Ok(d) => d,
Err(e) => {
tracing::error!("Read file error: {:?}", e);
return Err(upload_error(
StatusCode::INTERNAL_SERVER_ERROR,
"文件读取失败",
));
}
};
if data.len() > MAX_FILE_SIZE {
return Err(upload_error(
StatusCode::PAYLOAD_TOO_LARGE,
"文件超过大小限制",
));
}
// 校验文件头 magic bytes防止仅修改扩展名/Content-Type 上传非图片文件。
if !validate_image_magic_bytes(&data, mime_type.as_str()) {
return Err(upload_error(StatusCode::BAD_REQUEST, "文件类型与内容不符"));
}
// 仅读 header 统一校验尺寸/像素上限,并拿回 (w, h) 供 assets 表登记,避免二次解析。
// 超限直接拒绝,避免大图走 decode 后被静默降级(原 fallback 存原图)。
let (img_width, img_height) =
match crate::api::image::upload_dimensions(&data, mime_type.as_str()) {
Ok(dims) => dims,
Err(msg) => return Err(upload_error(StatusCode::BAD_REQUEST, msg)),
};
let is_gif = mime_type.as_str() == "image/gif";
let is_webp = mime_type.as_str() == "image/webp";
// 内容去重CAS对原始上传字节算 SHA-256命中已登记素材直接复用
// 跳过 GIF/WebP 解码验证、转码与落盘(省下整个流程最贵的 CPU
// 放在安全性校验magic bytes / 尺寸上限)之后、转码之前。
// 命中时刷新 created_at/updated_at重传代表使用意图重启 7 天清理保护窗
// PURGE_GRACE_DAYS 保护的是「刚上传还没被文章引用」的素材)。
let content_hash = {
use sha2::Digest;
hex::encode(sha2::Sha256::digest(&data))
};
{
let client = crate::db::pool::get_conn().await.map_err(|e| {
tracing::error!("Dedup check: get conn failed: {e}");
upload_error(StatusCode::INTERNAL_SERVER_ERROR, "文件保存失败")
})?;
let reused = client
.query_opt(
"UPDATE assets SET created_at = NOW(), updated_at = NOW() \
WHERE content_hash = $1 RETURNING path",
&[&content_hash],
)
.await
.map_err(|e| {
tracing::error!("Dedup check query failed: {e}");
upload_error(StatusCode::INTERNAL_SERVER_ERROR, "文件保存失败")
})?;
if let Some(row) = reused {
let path: String = row.get("path");
tracing::info!("Image upload deduped: reuse {} (hash {})", path, &content_hash[..12]);
return Ok(Json(json!({
"success": true,
"url": format!("/uploads/{}", path),
"reused": true
})));
}
}
// 对不经过重编码的格式做解码验证。GIF 走 image::load_from_memory 会完整解码,
// 移到阻塞线程池避免拖住 async 运行时。
// data 是引用计数的 Bytesclone 廉价(原子 +1避免 to_vec() 的全文件深拷贝。
if is_gif || is_webp {
let validate_data = data.clone();
let validate_mime = mime_type.clone();
let is_valid = tokio::task::spawn_blocking(move || {
validate_raw_image(&validate_data, validate_mime.as_str())
})
.await
.map_err(|_| upload_error(StatusCode::INTERNAL_SERVER_ERROR, "图片校验任务失败"))?;
if !is_valid {
return Err(upload_error(
StatusCode::BAD_REQUEST,
"图片文件损坏或格式不正确",
));
}
}
// GIF 与 WebP 保持原格式;其余格式尝试转 WebP。
// Bytes clone 廉价(引用计数 +1move 进阻塞闭包无需全文件深拷贝;
// 仅在「保留原格式」回退分支才 to_vec() 出落盘所需的 Vec<u8>。
let (final_data, final_ext): (Vec<u8>, String) = if is_gif {
(data.to_vec(), "gif".to_string())
} else if is_webp {
(data.to_vec(), "webp".to_string())
} else {
let original_data = data.clone();
let mime = mime_type.clone();
let config = crate::webp::WEBP_CONFIG.clone();
// 在阻塞线程中执行图片解码与 WebP 编码,避免阻塞异步运行时。
let result = tokio::task::spawn_blocking(move || -> (Vec<u8>, String, bool) {
let total_start = std::time::Instant::now();
let cursor = std::io::Cursor::new(&original_data);
let format = match mime.as_str() {
"image/jpeg" => image::ImageFormat::Jpeg,
"image/png" => image::ImageFormat::Png,
_ => image::ImageFormat::Jpeg,
};
let mut reader = image::ImageReader::with_format(cursor, format);
reader.limits(crate::api::image::image_reader_limits());
match reader.decode() {
Ok(img) => {
let decode_time = total_start.elapsed();
let enc_start = std::time::Instant::now();
let result = match crate::webp::encode(&img, config.quality, config.method) {
Ok(webp_data) => {
let enc_time = enc_start.elapsed();
let total_time = total_start.elapsed();
// WebP 更小才采用,否则回退原格式以节省带宽。
if webp_data.len() < original_data.len() {
tracing::info!(
"WebP conversion: decode={:?} encode={:?} total={:?} {}x{} {} bytes -> {} bytes",
decode_time, enc_time, total_time,
img.width(), img.height(),
original_data.len(), webp_data.len()
);
(webp_data, "webp".to_string(), true)
} else {
tracing::info!(
"WebP conversion larger, keeping original: decode={:?} encode={:?} total={:?} {}x{} original={} webp={}",
decode_time, enc_time, total_time,
img.width(), img.height(),
original_data.len(), webp_data.len()
);
(original_data.to_vec(), mime_to_ext(&mime).to_string(), false)
}
}
Err(e) => {
tracing::warn!("WebP encode failed ({}), keeping original format", e);
(original_data.to_vec(), mime_to_ext(&mime).to_string(), false)
}
};
result
}
Err(e) => {
// 到这里尺寸校验已通过(超限在 header 阶段被拒),decode 失败只能是真损坏。
tracing::warn!("Failed to decode image ({}), keeping original format", e);
(original_data.to_vec(), mime_to_ext(&mime).to_string(), false)
}
}
})
.await;
match result {
Ok((converted_data, ext, _was_converted)) => (converted_data, ext),
Err(_) => {
tracing::warn!("spawn_blocking task panicked, keeping original format");
(data.to_vec(), mime_to_ext(&mime_type).to_string())
}
}
};
// 按上传时间组织目录uploads/YYYY/MM/DD。
// chrono 的 DelayedFormat 实现 Display可直接进 format!,省掉 year/month/day 三个中间 String。
let now = chrono::Utc::now();
let date = now.format("%Y/%m/%d");
let uuid = uuid::Uuid::new_v4().to_string();
let dir_path = format!("uploads/{}", date);
let file_name = format!("{}.{}.{}", now.format("%H%M%S"), uuid, final_ext);
let file_path = format!("{}/{}", dir_path, file_name);
let url_path = format!("/uploads/{}/{}", date, file_name);
if let Err(e) = tokio::fs::create_dir_all(&dir_path).await {
tracing::error!("Create dir error: {:?}", e);
return Err(upload_error(
StatusCode::INTERNAL_SERVER_ERROR,
"文件保存失败",
));
}
if let Err(e) = tokio::fs::write(&file_path, &final_data).await {
tracing::error!("Write file error: {:?}", e);
return Err(upload_error(
StatusCode::INTERNAL_SERVER_ERROR,
"文件保存失败",
));
}
tracing::info!("Image uploaded: {} ({} bytes)", file_path, final_data.len());
// 登记 assets 注册表。失败时补偿删除已落盘文件,避免产生未登记的孤儿文件。
// ON CONFLICT (content_hash) DO NOTHING 兜底并发竞态:两个请求同时上传同一
// 新内容时会双双错过上面的去重检查,唯一索引保证只有一个 INSERT 成功;
// 落败者删自己的落盘文件、复用胜出者的路径(返回 Some(reused_path))。
let rel_path = format!("{}/{}", date, file_name);
let final_mime = match final_ext.as_str() {
"jpg" => "image/jpeg",
"png" => "image/png",
"gif" => "image/gif",
_ => "image/webp",
};
let registered: Result<Option<String>, String> = async {
let client = crate::db::pool::get_conn()
.await
.map_err(|e| e.to_string())?;
// id 用 Uuid 类型直连 uuid 列with-uuid-1 桥接),避免 String→uuid 序列化失败。
let asset_id = uuid::Uuid::new_v4();
let inserted = client
.execute(
"INSERT INTO assets (id, path, filename, mime, size_bytes, width, height, content_hash)\
VALUES ($1, $2, $3, $4, $5, $6, $7, $8) \
ON CONFLICT (content_hash) DO NOTHING",
&[
&asset_id,
&rel_path,
&original_filename.unwrap_or_else(|| file_name.clone()),
&final_mime,
&(final_data.len() as i64),
&(img_width as i32),
&(img_height as i32),
&content_hash,
],
)
.await
.map_err(|e| e.to_string())?;
if inserted == 0 {
// 竞态落败:胜出者的行必然已提交(唯一索引冲突即可见),取其路径复用。
let row = client
.query_one(
"SELECT path FROM assets WHERE content_hash = $1",
&[&content_hash],
)
.await
.map_err(|e| e.to_string())?;
return Ok(Some(row.get("path")));
}
Ok(None)
}
.await;
match registered {
Ok(Some(reused_path)) => {
let _ = tokio::fs::remove_file(&file_path).await;
tracing::info!("Image upload deduped (concurrent race): reuse {}", reused_path);
return Ok(Json(json!({
"success": true,
"url": format!("/uploads/{}", reused_path),
"reused": true
})));
}
Ok(None) => {}
Err(e) => {
tracing::error!("Register asset failed ({}), removing uploaded file", e);
let _ = tokio::fs::remove_file(&file_path).await;
return Err(upload_error(
StatusCode::INTERNAL_SERVER_ERROR,
"文件保存失败",
));
}
}
Ok(Json(json!({
"success": true,
"url": url_path
})))
}
#[cfg(all(test, feature = "server"))]
mod tests {
#[test]
fn filename_format_no_spaces() {
let now_str = "120000";
let uuid = "abc-123";
let ext = "jpg";
let file_name = format!("{}.{}.{}", now_str, uuid, ext);
assert!(
!file_name.contains(' '),
"filename should not contain spaces: got '{}'",
file_name
);
}
#[test]
fn should_use_webp_ext_for_non_gif() {
let ext = "jpg";
let mime = "image/jpeg";
let is_gif = mime == "image/gif";
let final_ext = if is_gif { ext } else { "webp" };
assert_eq!(final_ext, "webp");
}
#[test]
fn should_preserve_gif_ext() {
let ext = "gif";
let mime = "image/gif";
let is_gif = mime == "image/gif";
let final_ext = if is_gif { ext } else { "webp" };
assert_eq!(final_ext, "gif");
}
#[test]
fn convert_to_webp_produces_bytes() {
let img = image::DynamicImage::new_rgb8(10, 10);
let result = crate::webp::encode(&img, 85.0, 4).unwrap();
assert!(!result.is_empty());
}
#[test]
fn webp_roundtrip_from_rgba() {
let img = image::DynamicImage::new_rgba8(2, 2);
let webp_bytes = crate::webp::encode(&img, 85.0, 4).unwrap();
let loaded = crate::webp::decode(&webp_bytes);
assert!(loaded.is_ok());
}
#[test]
fn mime_to_ext_maps_jpeg() {
assert_eq!(super::mime_to_ext("image/jpeg"), "jpg");
}
#[test]
fn mime_to_ext_maps_png() {
assert_eq!(super::mime_to_ext("image/png"), "png");
}
#[test]
fn mime_to_ext_maps_gif() {
assert_eq!(super::mime_to_ext("image/gif"), "gif");
}
#[test]
fn mime_to_ext_maps_webp() {
assert_eq!(super::mime_to_ext("image/webp"), "webp");
}
#[test]
fn mime_to_ext_falls_back_for_unknown_mime() {
assert_eq!(super::mime_to_ext("image/avif"), "bin");
assert_eq!(super::mime_to_ext("application/octet-stream"), "bin");
}
#[test]
fn validate_jpeg_magic_bytes() {
assert!(super::validate_image_magic_bytes(
&[0xFF, 0xD8, 0xFF],
"image/jpeg"
));
assert!(!super::validate_image_magic_bytes(
&[0x89, 0x50],
"image/jpeg"
));
}
#[test]
fn validate_png_magic_bytes() {
assert!(super::validate_image_magic_bytes(
&[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A],
"image/png"
));
assert!(!super::validate_image_magic_bytes(
&[0xFF, 0xD8],
"image/png"
));
}
#[test]
fn validate_gif_magic_bytes() {
assert!(super::validate_image_magic_bytes(b"GIF89a", "image/gif"));
assert!(super::validate_image_magic_bytes(b"GIF87a", "image/gif"));
assert!(!super::validate_image_magic_bytes(b"GIF90a", "image/gif"));
}
#[test]
fn validate_webp_magic_bytes() {
let webp = b"RIFF\x00\x00\x00\x00WEBPVP8 ";
assert!(super::validate_image_magic_bytes(&webp[..12], "image/webp"));
assert!(!super::validate_image_magic_bytes(
&[0xFF, 0xD8],
"image/webp"
));
}
}