237899745/ystp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf: improve compression reliability and deployment safety

Browse Source
This commit is contained in:
237899745
2026-07-25 10:29:49 +08:00
parent 06220ca921
commit 9d7668bdee
34 changed files with 1391 additions and 1042 deletions
Download Patch File Download Diff File Expand all files Collapse all files

135
src/api/downloads.rs
View File

@@ -2,13 +2,14 @@ use crate::api::context;
use crate::error::{AppError, ErrorCode};
use crate::state::AppState;
use axum::body::Body;
use axum::extract::{ConnectInfo, Path, State};
use axum::http::{header, HeaderMap};
use axum::body::Body;
use axum::response::{IntoResponse, Response};
use axum::routing::get;
use axum::Router;
use chrono::{DateTime, Utc};
use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
use sqlx::FromRow;
use std::collections::HashMap;
use std::net::SocketAddr;
@@ -78,9 +79,11 @@ async fn download_file(
return Err(AppError::new(ErrorCode::NotFound, "文件不存在"));
};
let bytes = tokio::fs::read(path)
.await
.map_err(|err| AppError::new(ErrorCode::StorageUnavailable, "读取文件失败").with_source(err))?;
let file = tokio::fs::File::open(path).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "读取文件失败").with_source(err)
})?;
let content_length = file.metadata().await.ok().map(|metadata| metadata.len());
let body = Body::from_stream(ReaderStream::new(file));
let mut resp_headers = HeaderMap::new();
resp_headers.insert(
@@ -89,12 +92,16 @@ async fn download_file(
);
resp_headers.insert(
header::CONTENT_DISPOSITION,
format!("attachment; filename=\"{}\"", sanitize_filename(&row.original_name))
.parse()
.unwrap(),
content_disposition(&row.original_name)?,
);
if let Some(content_length) = content_length {
resp_headers.insert(
header::CONTENT_LENGTH,
content_length.to_string().parse().unwrap(),
);
}
Ok((jar, (resp_headers, bytes).into_response()))
Ok((jar, (resp_headers, body).into_response()))
}
fn authorize_download(principal: &context::Principal, row: &DownloadRow) -> Result<(), AppError> {
@@ -129,12 +136,40 @@ fn sanitize_filename(name: &str) -> String {
out = "download".to_string();
}
out = out.replace(['\r', '\n', '"', '\\'], "_");
if out.len() > 120 {
out.truncate(120);
}
truncate_utf8(&mut out, 120);
out
}
fn truncate_utf8(value: &mut String, max_bytes: usize) {
if value.len() <= max_bytes {
return;
}
let mut end = max_bytes;
while !value.is_char_boundary(end) {
end -= 1;
}
value.truncate(end);
}
fn content_disposition(name: &str) -> Result<axum::http::HeaderValue, AppError> {
let sanitized = sanitize_filename(name);
let ascii_fallback: String = sanitized
.chars()
.map(|ch| {
if ch.is_ascii_alphanumeric() || matches!(ch, '.' | '-' | '_') {
ch
} else {
'_'
}
})
.collect();
let encoded = utf8_percent_encode(&sanitized, NON_ALPHANUMERIC);
let value = format!("attachment; filename=\"{ascii_fallback}\"; filename*=UTF-8''{encoded}");
axum::http::HeaderValue::from_str(&value)
.map_err(|err| AppError::new(ErrorCode::Internal, "生成下载文件名失败").with_source(err))
}
#[derive(Debug, FromRow)]
struct TaskZipRow {
user_id: Option<Uuid>,
@@ -146,7 +181,6 @@ struct TaskZipRow {
#[derive(Debug, FromRow)]
struct TaskZipFileRow {
id: Uuid,
storage_path: Option<String>,
original_name: String,
output_format: String,
@@ -201,7 +235,7 @@ async fn download_task_zip(
}
}
if state.config.storage_type.to_ascii_lowercase() != "local" {
if !state.config.storage_type.eq_ignore_ascii_case("local") {
return Err(AppError::new(
ErrorCode::StorageUnavailable,
"当前仅支持本地存储（STORAGE_TYPE=local）",
@@ -209,9 +243,9 @@ async fn download_task_zip(
}
let zip_dir = format!("{}/zips", state.config.storage_path);
tokio::fs::create_dir_all(&zip_dir)
.await
.map_err(|err| AppError::new(ErrorCode::StorageUnavailable, "创建存储目录失败").with_source(err))?;
tokio::fs::create_dir_all(&zip_dir).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "创建存储目录失败").with_source(err)
})?;
let zip_path = PathBuf::from(format!("{zip_dir}/{task_id}.zip"));
if tokio::fs::try_exists(&zip_path).await.unwrap_or(false) {
@@ -220,7 +254,7 @@ async fn download_task_zip(
let rows = sqlx::query_as::<_, TaskZipFileRow>(
r#"
SELECT id, storage_path, original_name, output_format
SELECT storage_path, original_name, output_format
FROM task_files
WHERE task_id = $1 AND status = 'completed'
ORDER BY created_at ASC
@@ -238,7 +272,9 @@ async fn download_task_zip(
let mut used_names: HashMap<String, usize> = HashMap::new();
let mut entries: Vec<(String, String)> = Vec::new();
for row in rows {
let Some(path) = row.storage_path else { continue };
let Some(path) = row.storage_path else {
continue;
};
let name = build_zip_entry_name(&row.original_name, &row.output_format, &mut used_names);
entries.push((name, path));
}
@@ -248,10 +284,12 @@ async fn download_task_zip(
let zip_path_cloned = zip_path.clone();
let task_id_str = task_id.to_string();
tokio::task::spawn_blocking(move || generate_zip_file(&zip_path_cloned, &task_id_str, &entries))
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "生成 ZIP 失败").with_source(err))?
.map_err(|err| AppError::new(ErrorCode::Internal, "生成 ZIP 失败").with_source(err))?;
tokio::task::spawn_blocking(move || {
generate_zip_file(&zip_path_cloned, &task_id_str, &entries)
})
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "生成 ZIP 失败").with_source(err))?
.map_err(|err| AppError::new(ErrorCode::Internal, "生成 ZIP 失败").with_source(err))?;
stream_zip(jar, zip_path, task_id).await
}
@@ -261,9 +299,9 @@ async fn stream_zip(
zip_path: PathBuf,
task_id: Uuid,
) -> Result<(axum_extra::extract::cookie::CookieJar, Response), AppError> {
let file = tokio::fs::File::open(&zip_path)
.await
.map_err(|err| AppError::new(ErrorCode::StorageUnavailable, "读取 ZIP 失败").with_source(err))?;
let file = tokio::fs::File::open(&zip_path).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "读取 ZIP 失败").with_source(err)
})?;
let stream = ReaderStream::new(file);
let body = Body::from_stream(stream);
@@ -298,7 +336,11 @@ fn build_zip_entry_name(
_ => "bin",
};
let base = if base.is_empty() { "file".to_string() } else { base };
let base = if base.is_empty() {
"file".to_string()
} else {
base
};
let candidate = format!("{base}.{ext}");
let counter = used.entry(candidate.clone()).or_insert(0);
if *counter == 0 {
@@ -314,19 +356,21 @@ fn build_zip_entry_name(
fn sanitize_zip_name(name: &str) -> String {
let mut out = name.trim().to_string();
out = out.replace(['\r', '\n', '"', '\\', '/', ':'], "_");
if out.len() > 120 {
out.truncate(120);
}
truncate_utf8(&mut out, 120);
out
}
fn generate_zip_file(zip_path: &PathBuf, task_id: &str, entries: &[(String, String)]) -> Result<(), String> {
fn generate_zip_file(
zip_path: &PathBuf,
task_id: &str,
entries: &[(String, String)],
) -> Result<(), String> {
let tmp = PathBuf::from(format!("{}.tmp", zip_path.to_string_lossy()));
let file = std::fs::File::create(&tmp).map_err(|e| format!("create zip: {e}"))?;
let mut zip = zip::ZipWriter::new(file);
let options = zip::write::FileOptions::<()>::default()
.compression_method(zip::CompressionMethod::Stored);
let options =
zip::write::FileOptions::<()>::default().compression_method(zip::CompressionMethod::Stored);
for (name, path) in entries {
zip.start_file(name, options)
@@ -341,3 +385,30 @@ fn generate_zip_file(zip_path: &PathBuf, task_id: &str, entries: &[(String, Stri
tracing::info!(task_id = %task_id, path = %zip_path.to_string_lossy(), "ZIP generated");
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn content_disposition_supports_unicode_names() {
let value = content_disposition("测试 图片.jpg").unwrap();
let value = value.to_str().unwrap();
assert!(value.contains("filename=\"_____.jpg\""));
assert!(value.contains("filename*=UTF-8''"));
assert!(value.contains("%E6%B5%8B%E8%AF%95"));
}
#[test]
fn sanitize_filename_blocks_header_injection() {
assert_eq!(sanitize_filename("a\r\n\"b\\c.png"), "a___b_c.png");
}
#[test]
fn sanitize_filename_truncates_at_utf8_boundary() {
let name = format!("{}中.png", "a".repeat(119));
let sanitized = sanitize_filename(&name);
assert_eq!(sanitized, "a".repeat(119));
assert!(sanitized.is_char_boundary(sanitized.len()));
}
}