- 增加文件大小限制和流式写入

This commit is contained in:
Pine
2026-05-13 02:47:03 +08:00
parent 1f3bd690f0
commit 68ddb6446b
5 changed files with 26 additions and 15 deletions
+23 -7
View File
@@ -3,7 +3,7 @@ use crate::models::*;
use crate::storage::{get_media_type, is_allowed_extension, STORAGE};
use axum::{
body::Body,
extract::{Form, Multipart, State},
extract::{DefaultBodyLimit, Form, Multipart, State},
http::{Response, StatusCode, Uri},
response::sse::{Event, KeepAlive, Sse},
response::Json,
@@ -17,6 +17,7 @@ use std::convert::Infallible;
use std::path::PathBuf;
use std::sync::Arc;
use std::fs;
use tokio::io::AsyncWriteExt;
use tokio_stream::wrappers::BroadcastStream;
use tokio_stream::StreamExt;
use tower_http::cors::CorsLayer;
@@ -181,6 +182,7 @@ pub fn create_router(state: AppState) -> Router {
.nest_service("/file", ServeDir::new(media_dir_str()))
// SPA 回退 → 嵌入式文件(跨平台,编译到二进制中)
.fallback_service(get(spa_fallback))
.layer(DefaultBodyLimit::max(2048 * 1024 * 1024)) // 2GB 上传限制
.layer(CorsLayer::permissive())
.with_state(state)
}
@@ -290,7 +292,7 @@ async fn list_media() -> Json<FilesResponse> {
async fn upload_handler(
mut multipart: Multipart,
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
while let Ok(Some(field)) = multipart.next_field().await {
while let Ok(Some(mut field)) = multipart.next_field().await {
let file_name = field.file_name().unwrap_or("file").to_string();
let ext = std::path::Path::new(&file_name)
.extension()
@@ -310,12 +312,26 @@ async fn upload_handler(
let new_name = format!("{}_{}{}", stem, ts, ext);
let save_path = media_dir().join(&new_name);
let data = field.bytes().await.map_err(|e| {
(StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
})?;
fs::write(&save_path, &data).map_err(|e| {
(StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
// 流式写入文件,避免大视频文件占用过多内存
let mut file = tokio::fs::File::create(&save_path).await.map_err(|e| {
(StatusCode::INTERNAL_SERVER_ERROR, format!("文件创建失败: {}", e))
})?;
loop {
match field.chunk().await {
Ok(Some(chunk)) => {
file.write_all(&chunk).await.map_err(|e| {
(StatusCode::INTERNAL_SERVER_ERROR, format!("写入文件失败: {}", e))
})?;
}
Ok(None) => break,
Err(e) => {
return Err((
StatusCode::INTERNAL_SERVER_ERROR,
format!("读取上传数据失败: {}", e),
))
}
}
}
}
Ok(Json(serde_json::json!({"ok": true})))
}