1184 lines
40 KiB
Rust
1184 lines
40 KiB
Rust
use crate::api::context;
|
|
use crate::error::{AppError, ErrorCode};
|
|
use crate::services::storage;
|
|
use crate::state::AppState;
|
|
|
|
use axum::body::Body;
|
|
use axum::extract::{ConnectInfo, Path, State};
|
|
use axum::http::{header, HeaderMap};
|
|
use axum::response::{IntoResponse, Redirect, 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;
|
|
use std::path::PathBuf;
|
|
use tokio_util::io::ReaderStream;
|
|
use uuid::Uuid;
|
|
|
|
pub fn router() -> Router<AppState> {
|
|
Router::new()
|
|
.route("/tasks/{task_id}", get(download_task_zip))
|
|
.route("/{file_id}", get(download_file))
|
|
}
|
|
|
|
#[derive(Debug, FromRow)]
|
|
struct DownloadRow {
|
|
storage_backend: String,
|
|
storage_endpoint_id: Option<Uuid>,
|
|
storage_key: Option<String>,
|
|
output_format: String,
|
|
original_name: String,
|
|
file_status: String,
|
|
task_user_id: Option<Uuid>,
|
|
task_session_id: Option<String>,
|
|
expires_at: DateTime<Utc>,
|
|
}
|
|
|
|
async fn download_file(
|
|
State(state): State<AppState>,
|
|
jar: axum_extra::extract::cookie::CookieJar,
|
|
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
|
headers: HeaderMap,
|
|
Path(file_id): Path<Uuid>,
|
|
) -> Result<(axum_extra::extract::cookie::CookieJar, Response), AppError> {
|
|
let ip = context::client_ip(&headers, addr.ip());
|
|
let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
|
context::require_api_permission(&principal, &["compress", "batch_compress"])?;
|
|
|
|
let row = sqlx::query_as::<_, DownloadRow>(
|
|
r#"
|
|
SELECT
|
|
f.storage_backend,
|
|
f.storage_endpoint_id,
|
|
COALESCE(f.storage_key, f.storage_path) AS storage_key,
|
|
f.output_format,
|
|
f.original_name,
|
|
f.status::text AS file_status,
|
|
t.user_id AS task_user_id,
|
|
t.session_id AS task_session_id,
|
|
t.expires_at
|
|
FROM task_files f
|
|
JOIN tasks t ON t.id = f.task_id
|
|
WHERE f.id = $1
|
|
"#,
|
|
)
|
|
.bind(file_id)
|
|
.fetch_optional(&state.db)
|
|
.await
|
|
.map_err(|err| AppError::new(ErrorCode::Internal, "查询文件失败").with_source(err))?
|
|
.ok_or_else(|| AppError::new(ErrorCode::NotFound, "文件不存在"))?;
|
|
|
|
if row.expires_at <= Utc::now() {
|
|
return Err(AppError::new(ErrorCode::NotFound, "文件已过期或不存在"));
|
|
}
|
|
|
|
if row.file_status != "completed" {
|
|
return Err(AppError::new(ErrorCode::NotFound, "文件不存在"));
|
|
}
|
|
|
|
authorize_download(&principal, &row)?;
|
|
|
|
let Some(key) = row.storage_key else {
|
|
return Err(AppError::new(ErrorCode::NotFound, "文件不存在"));
|
|
};
|
|
let object = storage::ObjectLocator {
|
|
backend: row.storage_backend,
|
|
endpoint_id: row.storage_endpoint_id,
|
|
key,
|
|
};
|
|
respond_object(
|
|
&state,
|
|
jar,
|
|
&object,
|
|
&output_file_name(&row.original_name, &row.output_format),
|
|
content_type(&row.output_format),
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn respond_object(
|
|
state: &AppState,
|
|
jar: axum_extra::extract::cookie::CookieJar,
|
|
object: &storage::ObjectLocator,
|
|
file_name: &str,
|
|
content_type_value: &str,
|
|
) -> Result<(axum_extra::extract::cookie::CookieJar, Response), AppError> {
|
|
if let Some(url) =
|
|
storage::presign_download(state, object, file_name, content_type_value).await?
|
|
{
|
|
return Ok((jar, Redirect::temporary(&url).into_response()));
|
|
}
|
|
|
|
let file = tokio::fs::File::open(&object.key).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(header::CONTENT_TYPE, content_type_value.parse().unwrap());
|
|
resp_headers.insert(header::CONTENT_DISPOSITION, content_disposition(file_name)?);
|
|
if let Some(content_length) = content_length {
|
|
resp_headers.insert(
|
|
header::CONTENT_LENGTH,
|
|
content_length.to_string().parse().unwrap(),
|
|
);
|
|
}
|
|
Ok((jar, (resp_headers, body).into_response()))
|
|
}
|
|
|
|
fn authorize_download(principal: &context::Principal, row: &DownloadRow) -> Result<(), AppError> {
|
|
if let Some(user_id) = row.task_user_id {
|
|
match principal {
|
|
context::Principal::User { user_id: me, .. } if *me == user_id => Ok(()),
|
|
context::Principal::ApiKey { user_id: me, .. } if *me == user_id => Ok(()),
|
|
_ => Err(AppError::new(ErrorCode::Forbidden, "无权限下载该文件")),
|
|
}
|
|
} else {
|
|
let expected = row.task_session_id.as_deref().unwrap_or("");
|
|
match principal {
|
|
context::Principal::Anonymous { session_id } if session_id == expected => Ok(()),
|
|
_ => Err(AppError::new(ErrorCode::Forbidden, "无权限下载该文件")),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn content_type(format: &str) -> &'static str {
|
|
match format.trim().to_ascii_lowercase().as_str() {
|
|
"png" => "image/png",
|
|
"jpeg" | "jpg" => "image/jpeg",
|
|
"webp" => "image/webp",
|
|
"avif" => "image/avif",
|
|
"gif" => "image/gif",
|
|
"bmp" => "image/bmp",
|
|
"tif" | "tiff" => "image/tiff",
|
|
"ico" => "image/x-icon",
|
|
_ => "application/octet-stream",
|
|
}
|
|
}
|
|
|
|
fn output_file_name(original_name: &str, output_format: &str) -> String {
|
|
let sanitized = sanitize_filename(original_name);
|
|
let base = sanitized
|
|
.rsplit_once('.')
|
|
.map(|(value, _)| value)
|
|
.unwrap_or(&sanitized)
|
|
.trim_end_matches('.');
|
|
let base = if base.is_empty() { "download" } else { base };
|
|
let extension = match output_format.trim().to_ascii_lowercase().as_str() {
|
|
"jpeg" | "jpg" => "jpg",
|
|
"png" => "png",
|
|
"webp" => "webp",
|
|
"avif" => "avif",
|
|
"gif" => "gif",
|
|
"bmp" => "bmp",
|
|
"tif" | "tiff" => "tiff",
|
|
"ico" => "ico",
|
|
_ => "bin",
|
|
};
|
|
format!("{base}.{extension}")
|
|
}
|
|
|
|
fn sanitize_filename(name: &str) -> String {
|
|
let mut out = name.trim().to_string();
|
|
if out.is_empty() {
|
|
out = "download".to_string();
|
|
}
|
|
out = out.replace(['\r', '\n', '"', '\\'], "_");
|
|
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>,
|
|
session_id: Option<String>,
|
|
status: String,
|
|
completed_at: Option<DateTime<Utc>>,
|
|
expires_at: DateTime<Utc>,
|
|
retention_hours: i32,
|
|
zip_storage_backend: Option<String>,
|
|
zip_storage_endpoint_id: Option<Uuid>,
|
|
zip_storage_key: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, FromRow)]
|
|
struct TaskZipFileRow {
|
|
storage_backend: String,
|
|
storage_endpoint_id: Option<Uuid>,
|
|
storage_key: Option<String>,
|
|
original_name: String,
|
|
output_format: String,
|
|
compressed_size: Option<i64>,
|
|
}
|
|
|
|
#[derive(Debug, FromRow)]
|
|
struct ZipBuildStateRow {
|
|
zip_storage_backend: Option<String>,
|
|
zip_storage_endpoint_id: Option<Uuid>,
|
|
zip_storage_key: Option<String>,
|
|
expires_at: DateTime<Utc>,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
enum ZipBuildClaim {
|
|
Acquired { token: Uuid },
|
|
Cached(storage::ObjectLocator),
|
|
Busy,
|
|
}
|
|
|
|
const ZIP_BUILD_LEASE_SECONDS: i64 = 15 * 60;
|
|
const ZIP_BUILD_WAIT_SECONDS: u64 = 30;
|
|
|
|
async fn download_task_zip(
|
|
State(state): State<AppState>,
|
|
jar: axum_extra::extract::cookie::CookieJar,
|
|
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
|
headers: HeaderMap,
|
|
Path(task_id): Path<Uuid>,
|
|
) -> Result<(axum_extra::extract::cookie::CookieJar, Response), AppError> {
|
|
let ip = context::client_ip(&headers, addr.ip());
|
|
let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
|
context::require_api_permission(&principal, &["compress", "batch_compress"])?;
|
|
|
|
let task = sqlx::query_as::<_, TaskZipRow>(
|
|
r#"
|
|
SELECT
|
|
user_id,
|
|
session_id,
|
|
status::text AS status,
|
|
completed_at,
|
|
expires_at,
|
|
retention_hours,
|
|
zip_storage_backend,
|
|
zip_storage_endpoint_id,
|
|
zip_storage_key
|
|
FROM tasks
|
|
WHERE id = $1
|
|
"#,
|
|
)
|
|
.bind(task_id)
|
|
.fetch_optional(&state.db)
|
|
.await
|
|
.map_err(|err| AppError::new(ErrorCode::Internal, "查询任务失败").with_source(err))?
|
|
.ok_or_else(|| AppError::new(ErrorCode::NotFound, "任务不存在"))?;
|
|
|
|
if task.expires_at <= Utc::now() {
|
|
return Err(AppError::new(ErrorCode::NotFound, "任务已过期或不存在"));
|
|
}
|
|
if task.completed_at.is_none() || matches!(task.status.as_str(), "pending" | "processing") {
|
|
return Err(AppError::new(ErrorCode::InvalidRequest, "任务尚未完成"));
|
|
}
|
|
|
|
if let Some(user_id) = task.user_id {
|
|
match principal {
|
|
context::Principal::User { user_id: me, .. } if me == user_id => {}
|
|
context::Principal::ApiKey { user_id: me, .. } if me == user_id => {}
|
|
_ => return Err(AppError::new(ErrorCode::Forbidden, "无权限下载该任务")),
|
|
}
|
|
} else {
|
|
let expected = task.session_id.as_deref().unwrap_or("");
|
|
match principal {
|
|
context::Principal::Anonymous { session_id } if session_id == expected => {}
|
|
_ => return Err(AppError::new(ErrorCode::Forbidden, "无权限下载该任务")),
|
|
}
|
|
}
|
|
|
|
if let (Some(backend), Some(key)) = (
|
|
task.zip_storage_backend.clone(),
|
|
task.zip_storage_key.clone(),
|
|
) {
|
|
return respond_object(
|
|
&state,
|
|
jar,
|
|
&storage::ObjectLocator {
|
|
backend,
|
|
endpoint_id: task.zip_storage_endpoint_id,
|
|
key,
|
|
},
|
|
&format!("task_{task_id}.zip"),
|
|
"application/zip",
|
|
)
|
|
.await;
|
|
}
|
|
|
|
let object = resolve_task_zip(&state, task_id, task.retention_hours as i64).await?;
|
|
|
|
respond_object(
|
|
&state,
|
|
jar,
|
|
&object,
|
|
&format!("task_{task_id}.zip"),
|
|
"application/zip",
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn resolve_task_zip(
|
|
state: &AppState,
|
|
task_id: Uuid,
|
|
retention_hours: i64,
|
|
) -> Result<storage::ObjectLocator, AppError> {
|
|
if let Some(object) = load_published_zip(state, task_id).await? {
|
|
return Ok(object);
|
|
}
|
|
let rows = sqlx::query_as::<_, TaskZipFileRow>(
|
|
r#"
|
|
SELECT storage_backend, storage_endpoint_id,
|
|
COALESCE(storage_key, storage_path) AS storage_key,
|
|
original_name, output_format, compressed_size
|
|
FROM task_files
|
|
WHERE task_id = $1 AND status = 'completed'
|
|
ORDER BY created_at ASC
|
|
"#,
|
|
)
|
|
.bind(task_id)
|
|
.fetch_all(&state.db)
|
|
.await
|
|
.map_err(|err| AppError::new(ErrorCode::Internal, "查询任务文件失败").with_source(err))?;
|
|
|
|
if rows.is_empty() {
|
|
return Err(AppError::new(ErrorCode::NotFound, "没有可打包的文件"));
|
|
}
|
|
validate_zip_budget(state, &rows)?;
|
|
|
|
let deadline =
|
|
tokio::time::Instant::now() + std::time::Duration::from_secs(ZIP_BUILD_WAIT_SECONDS);
|
|
loop {
|
|
match claim_zip_build(state, task_id).await? {
|
|
ZipBuildClaim::Cached(object) => return Ok(object),
|
|
ZipBuildClaim::Acquired { token } => {
|
|
return build_claimed_zip(state, task_id, retention_hours, token, &rows).await;
|
|
}
|
|
ZipBuildClaim::Busy if tokio::time::Instant::now() < deadline => {
|
|
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
|
}
|
|
ZipBuildClaim::Busy => {
|
|
return Err(AppError::new(
|
|
ErrorCode::StorageUnavailable,
|
|
"ZIP 正在生成，请稍后重试",
|
|
));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn validate_zip_budget(state: &AppState, rows: &[TaskZipFileRow]) -> Result<u64, AppError> {
|
|
if rows.len() > state.config.zip_max_entries as usize {
|
|
return Err(AppError::new(
|
|
ErrorCode::FileTooLarge,
|
|
format!("ZIP 文件数量超过 {} 个上限", state.config.zip_max_entries),
|
|
));
|
|
}
|
|
rows.iter().try_fold(0_u64, |total, row| {
|
|
if row.storage_key.is_none() {
|
|
return Err(AppError::new(
|
|
ErrorCode::StorageUnavailable,
|
|
"ZIP 源文件存储信息不完整",
|
|
));
|
|
}
|
|
let size = row
|
|
.compressed_size
|
|
.and_then(|value| u64::try_from(value).ok())
|
|
.ok_or_else(|| AppError::new(ErrorCode::StorageUnavailable, "ZIP 源文件大小无效"))?;
|
|
let next = total
|
|
.checked_add(size)
|
|
.ok_or_else(|| AppError::new(ErrorCode::FileTooLarge, "ZIP 源文件总大小超出限制"))?;
|
|
if next > state.config.zip_max_uncompressed_bytes {
|
|
return Err(AppError::new(
|
|
ErrorCode::FileTooLarge,
|
|
format!(
|
|
"ZIP 源文件总大小超过 {} 字节上限",
|
|
state.config.zip_max_uncompressed_bytes
|
|
),
|
|
));
|
|
}
|
|
Ok(next)
|
|
})
|
|
}
|
|
|
|
async fn claim_zip_build(state: &AppState, task_id: Uuid) -> Result<ZipBuildClaim, AppError> {
|
|
let token = Uuid::new_v4();
|
|
let acquired: Option<i64> = sqlx::query_scalar(
|
|
r#"
|
|
UPDATE tasks
|
|
SET zip_build_token = $2,
|
|
zip_build_lease_until = NOW() + ($3 * INTERVAL '1 second'),
|
|
zip_build_attempt = zip_build_attempt + 1
|
|
WHERE id = $1
|
|
AND zip_storage_key IS NULL
|
|
AND completed_at IS NOT NULL
|
|
AND expires_at > NOW()
|
|
AND (
|
|
zip_build_token IS NULL
|
|
OR zip_build_lease_until <= NOW()
|
|
)
|
|
RETURNING zip_build_attempt
|
|
"#,
|
|
)
|
|
.bind(task_id)
|
|
.bind(token)
|
|
.bind(ZIP_BUILD_LEASE_SECONDS)
|
|
.fetch_optional(&state.db)
|
|
.await
|
|
.map_err(|err| AppError::new(ErrorCode::Internal, "获取 ZIP 构建租约失败").with_source(err))?;
|
|
if acquired.is_some() {
|
|
return Ok(ZipBuildClaim::Acquired { token });
|
|
}
|
|
|
|
let current = sqlx::query_as::<_, ZipBuildStateRow>(
|
|
r#"
|
|
SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key, expires_at
|
|
FROM tasks
|
|
WHERE id = $1
|
|
"#,
|
|
)
|
|
.bind(task_id)
|
|
.fetch_optional(&state.db)
|
|
.await
|
|
.map_err(|err| AppError::new(ErrorCode::Internal, "查询 ZIP 构建状态失败").with_source(err))?
|
|
.ok_or_else(|| AppError::new(ErrorCode::NotFound, "任务不存在"))?;
|
|
if current.expires_at <= Utc::now() {
|
|
return Err(AppError::new(ErrorCode::NotFound, "任务已过期或不存在"));
|
|
}
|
|
if let (Some(backend), Some(key)) = (current.zip_storage_backend, current.zip_storage_key) {
|
|
return Ok(ZipBuildClaim::Cached(storage::ObjectLocator {
|
|
backend,
|
|
endpoint_id: current.zip_storage_endpoint_id,
|
|
key,
|
|
}));
|
|
}
|
|
// A lease may have been released between the UPDATE and this read. The
|
|
// caller retries the atomic claim after a short wait in either case.
|
|
Ok(ZipBuildClaim::Busy)
|
|
}
|
|
|
|
async fn renew_zip_build(state: &AppState, task_id: Uuid, token: Uuid) -> Result<(), AppError> {
|
|
let updated = sqlx::query(
|
|
r#"
|
|
UPDATE tasks
|
|
SET zip_build_lease_until = NOW() + ($3 * INTERVAL '1 second')
|
|
WHERE id = $1
|
|
AND zip_build_token = $2
|
|
AND zip_storage_key IS NULL
|
|
AND expires_at > NOW()
|
|
"#,
|
|
)
|
|
.bind(task_id)
|
|
.bind(token)
|
|
.bind(ZIP_BUILD_LEASE_SECONDS)
|
|
.execute(&state.db)
|
|
.await
|
|
.map_err(|err| AppError::new(ErrorCode::Internal, "续租 ZIP 构建失败").with_source(err))?;
|
|
if updated.rows_affected() != 1 {
|
|
return Err(AppError::new(
|
|
ErrorCode::StorageUnavailable,
|
|
"ZIP 构建租约已失效，请重试",
|
|
));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn release_zip_build(state: &AppState, task_id: Uuid, token: Uuid) {
|
|
if let Err(err) = sqlx::query(
|
|
r#"
|
|
UPDATE tasks
|
|
SET zip_build_token = NULL,
|
|
zip_build_lease_until = NULL
|
|
WHERE id = $1
|
|
AND zip_build_token = $2
|
|
AND zip_storage_key IS NULL
|
|
"#,
|
|
)
|
|
.bind(task_id)
|
|
.bind(token)
|
|
.execute(&state.db)
|
|
.await
|
|
{
|
|
tracing::warn!(task_id = %task_id, zip_build_token = %token, error = %err, "failed to release ZIP build lease");
|
|
}
|
|
}
|
|
|
|
async fn build_claimed_zip(
|
|
state: &AppState,
|
|
task_id: Uuid,
|
|
retention_hours: i64,
|
|
token: Uuid,
|
|
rows: &[TaskZipFileRow],
|
|
) -> Result<storage::ObjectLocator, AppError> {
|
|
let permit = match tokio::time::timeout(
|
|
std::time::Duration::from_secs(ZIP_BUILD_WAIT_SECONDS),
|
|
state.zip_build_semaphore.clone().acquire_owned(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(Ok(permit)) => permit,
|
|
Ok(Err(err)) => {
|
|
release_zip_build(state, task_id, token).await;
|
|
return Err(AppError::new(ErrorCode::Internal, "ZIP 并发闸门已关闭").with_source(err));
|
|
}
|
|
Err(_) => {
|
|
release_zip_build(state, task_id, token).await;
|
|
return Err(AppError::new(
|
|
ErrorCode::StorageUnavailable,
|
|
"ZIP 生成繁忙，请稍后重试",
|
|
));
|
|
}
|
|
};
|
|
|
|
let temp_dir = PathBuf::from(format!(
|
|
"{}/tmp/zips/{task_id}-{token}",
|
|
state.config.storage_path
|
|
));
|
|
let zip_path = temp_dir.join(format!("task_{task_id}.zip"));
|
|
let build_result = build_zip_attempt(
|
|
state,
|
|
task_id,
|
|
token,
|
|
retention_hours,
|
|
rows,
|
|
&temp_dir,
|
|
&zip_path,
|
|
)
|
|
.await;
|
|
drop(permit);
|
|
let _ = tokio::fs::remove_dir_all(&temp_dir).await;
|
|
|
|
let stored = match build_result {
|
|
Ok(stored) => stored,
|
|
Err(err) => {
|
|
release_zip_build(state, task_id, token).await;
|
|
return Err(err);
|
|
}
|
|
};
|
|
publish_zip_attempt(state, task_id, token, stored).await
|
|
}
|
|
|
|
async fn build_zip_attempt(
|
|
state: &AppState,
|
|
task_id: Uuid,
|
|
token: Uuid,
|
|
retention_hours: i64,
|
|
rows: &[TaskZipFileRow],
|
|
temp_dir: &std::path::Path,
|
|
zip_path: &std::path::Path,
|
|
) -> Result<storage::StoredObject, AppError> {
|
|
tokio::fs::create_dir_all(temp_dir).await.map_err(|err| {
|
|
AppError::new(ErrorCode::StorageUnavailable, "创建 ZIP 临时目录失败").with_source(err)
|
|
})?;
|
|
let mut used_names: HashMap<String, usize> = HashMap::new();
|
|
let mut entries: Vec<(String, String)> = Vec::with_capacity(rows.len());
|
|
let mut actual_bytes = 0_u64;
|
|
for (index, row) in rows.iter().enumerate() {
|
|
renew_zip_build(state, task_id, token).await?;
|
|
let key = row.storage_key.as_ref().ok_or_else(|| {
|
|
AppError::new(ErrorCode::StorageUnavailable, "ZIP 源文件存储信息不完整")
|
|
})?;
|
|
let path = temp_dir.join(format!("entry-{index}"));
|
|
storage::download_to_file(
|
|
state,
|
|
&storage::ObjectLocator {
|
|
backend: row.storage_backend.clone(),
|
|
endpoint_id: row.storage_endpoint_id,
|
|
key: key.clone(),
|
|
},
|
|
&path,
|
|
)
|
|
.await?;
|
|
let size = tokio::fs::metadata(&path)
|
|
.await
|
|
.map_err(|err| {
|
|
AppError::new(ErrorCode::StorageUnavailable, "读取 ZIP 临时文件失败")
|
|
.with_source(err)
|
|
})?
|
|
.len();
|
|
actual_bytes = actual_bytes
|
|
.checked_add(size)
|
|
.ok_or_else(|| AppError::new(ErrorCode::FileTooLarge, "ZIP 实际文件总大小超出限制"))?;
|
|
if actual_bytes > state.config.zip_max_uncompressed_bytes {
|
|
return Err(AppError::new(
|
|
ErrorCode::FileTooLarge,
|
|
"ZIP 实际文件总大小超出限制",
|
|
));
|
|
}
|
|
let name = build_zip_entry_name(&row.original_name, &row.output_format, &mut used_names);
|
|
entries.push((name, path.to_string_lossy().to_string()));
|
|
}
|
|
|
|
renew_zip_build(state, task_id, token).await?;
|
|
let zip_path_cloned = zip_path.to_path_buf();
|
|
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))?;
|
|
|
|
renew_zip_build(state, task_id, token).await?;
|
|
let object_key = storage::archive_attempt_key(retention_hours, task_id, token);
|
|
storage::store_file(state, &object_key, zip_path, "application/zip").await
|
|
}
|
|
|
|
async fn publish_zip_attempt(
|
|
state: &AppState,
|
|
task_id: Uuid,
|
|
token: Uuid,
|
|
stored: storage::StoredObject,
|
|
) -> Result<storage::ObjectLocator, AppError> {
|
|
let published = sqlx::query(
|
|
r#"
|
|
UPDATE tasks
|
|
SET zip_storage_backend = $3,
|
|
zip_storage_endpoint_id = $4,
|
|
zip_storage_key = $5,
|
|
zip_storage_etag = $6,
|
|
zip_size = $7,
|
|
zip_build_token = NULL,
|
|
zip_build_lease_until = NULL
|
|
WHERE id = $1
|
|
AND zip_build_token = $2
|
|
AND zip_storage_key IS NULL
|
|
AND expires_at > NOW()
|
|
"#,
|
|
)
|
|
.bind(task_id)
|
|
.bind(token)
|
|
.bind(&stored.backend)
|
|
.bind(stored.endpoint_id)
|
|
.bind(&stored.key)
|
|
.bind(&stored.etag)
|
|
.bind(stored.size as i64)
|
|
.execute(&state.db)
|
|
.await;
|
|
|
|
match published {
|
|
Ok(result) if result.rows_affected() == 1 => Ok(storage::ObjectLocator {
|
|
backend: stored.backend,
|
|
endpoint_id: stored.endpoint_id,
|
|
key: stored.key,
|
|
}),
|
|
Ok(_) => {
|
|
delete_unpublished_zip(state, task_id, token, &stored).await;
|
|
let current = load_published_zip(state, task_id).await?;
|
|
current.ok_or_else(|| {
|
|
AppError::new(ErrorCode::StorageUnavailable, "ZIP 发布租约已失效，请重试")
|
|
})
|
|
}
|
|
Err(err) => {
|
|
delete_unpublished_zip(state, task_id, token, &stored).await;
|
|
Err(AppError::new(ErrorCode::Internal, "记录 ZIP 对象失败").with_source(err))
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn delete_unpublished_zip(
|
|
state: &AppState,
|
|
task_id: Uuid,
|
|
token: Uuid,
|
|
stored: &storage::StoredObject,
|
|
) {
|
|
let object = storage::ObjectLocator {
|
|
backend: stored.backend.clone(),
|
|
endpoint_id: stored.endpoint_id,
|
|
key: stored.key.clone(),
|
|
};
|
|
let mut last_error = None;
|
|
for attempt in 1..=3_u64 {
|
|
match storage::delete_object(state, &object).await {
|
|
Ok(()) => {
|
|
last_error = None;
|
|
break;
|
|
}
|
|
Err(err) => {
|
|
last_error = Some(err);
|
|
tokio::time::sleep(std::time::Duration::from_millis(100 * attempt)).await;
|
|
}
|
|
}
|
|
}
|
|
if let Some(err) = last_error {
|
|
tracing::error!(task_id = %task_id, zip_build_token = %token, object_key = %stored.key, error = %err, "failed to delete unpublished ZIP attempt after retries");
|
|
}
|
|
release_zip_build(state, task_id, token).await;
|
|
}
|
|
|
|
async fn load_published_zip(
|
|
state: &AppState,
|
|
task_id: Uuid,
|
|
) -> Result<Option<storage::ObjectLocator>, AppError> {
|
|
let row: Option<(Option<String>, Option<Uuid>, Option<String>)> = sqlx::query_as(
|
|
"SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key FROM tasks WHERE id = $1",
|
|
)
|
|
.bind(task_id)
|
|
.fetch_optional(&state.db)
|
|
.await
|
|
.map_err(|err| AppError::new(ErrorCode::Internal, "查询已发布 ZIP 失败").with_source(err))?;
|
|
Ok(row.and_then(|(backend, endpoint_id, key)| {
|
|
Some(storage::ObjectLocator {
|
|
backend: backend?,
|
|
endpoint_id,
|
|
key: key?,
|
|
})
|
|
}))
|
|
}
|
|
|
|
fn build_zip_entry_name(
|
|
original_name: &str,
|
|
output_format: &str,
|
|
used: &mut HashMap<String, usize>,
|
|
) -> String {
|
|
let mut base = sanitize_zip_name(original_name);
|
|
if let Some((head, _ext)) = base.rsplit_once('.') {
|
|
base = head.to_string();
|
|
}
|
|
|
|
let ext = match output_format.trim().to_ascii_lowercase().as_str() {
|
|
"jpeg" | "jpg" => "jpg",
|
|
"png" => "png",
|
|
"webp" => "webp",
|
|
"avif" => "avif",
|
|
"gif" => "gif",
|
|
"bmp" => "bmp",
|
|
"tif" | "tiff" => "tiff",
|
|
"ico" => "ico",
|
|
_ => "bin",
|
|
};
|
|
|
|
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 {
|
|
*counter = 1;
|
|
return candidate;
|
|
}
|
|
|
|
let name = format!("{base} ({counter}).{ext}");
|
|
*counter += 1;
|
|
name
|
|
}
|
|
|
|
fn sanitize_zip_name(name: &str) -> String {
|
|
let mut out = name.trim().to_string();
|
|
out = out.replace(['\r', '\n', '"', '\\', '/', ':'], "_");
|
|
truncate_utf8(&mut out, 120);
|
|
out
|
|
}
|
|
|
|
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);
|
|
|
|
for (name, path) in entries {
|
|
zip.start_file(name, options)
|
|
.map_err(|e| format!("zip start_file: {e}"))?;
|
|
let mut f = std::fs::File::open(path).map_err(|e| format!("open file: {e}"))?;
|
|
std::io::copy(&mut f, &mut zip).map_err(|e| format!("copy: {e}"))?;
|
|
}
|
|
|
|
zip.finish().map_err(|e| format!("finish: {e}"))?;
|
|
|
|
std::fs::rename(&tmp, zip_path).map_err(|e| format!("rename: {e}"))?;
|
|
tracing::info!(task_id = %task_id, path = %zip_path.to_string_lossy(), "ZIP generated");
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::config::Config;
|
|
use crate::services::mail::Mailer;
|
|
use crate::services::settings;
|
|
use sqlx::postgres::PgPoolOptions;
|
|
use std::sync::Arc;
|
|
use tokio::sync::{Barrier, Semaphore};
|
|
|
|
#[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()));
|
|
}
|
|
|
|
#[test]
|
|
fn output_file_name_matches_converted_format() {
|
|
assert_eq!(output_file_name("photo.png", "webp"), "photo.webp");
|
|
assert_eq!(output_file_name("没有扩展名", "jpeg"), "没有扩展名.jpg");
|
|
}
|
|
|
|
async fn build_zip_test_state(
|
|
pool: sqlx::PgPool,
|
|
database_url: String,
|
|
redis_url: String,
|
|
storage_path: String,
|
|
) -> AppState {
|
|
let mut config = Config::from_env().expect("load ZIP test config");
|
|
config.database_url = database_url;
|
|
config.redis_url = redis_url;
|
|
config.storage_path = storage_path;
|
|
config.zip_build_concurrency = 2;
|
|
config.zip_max_entries = 200;
|
|
config.zip_max_uncompressed_bytes = 2 * 1024 * 1024;
|
|
config.mail_enabled = false;
|
|
config.mail_log_links_when_disabled = false;
|
|
let redis = redis::Client::open(config.redis_url.clone())
|
|
.expect("create ZIP test Redis client")
|
|
.get_connection_manager()
|
|
.await
|
|
.expect("connect ZIP test Redis");
|
|
AppState {
|
|
mailer: Arc::new(Mailer::new(&config).expect("create disabled ZIP test mailer")),
|
|
image_processing_semaphore: Arc::new(Semaphore::new(2)),
|
|
zip_build_semaphore: Arc::new(Semaphore::new(2)),
|
|
runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(),
|
|
storage_cache: storage::StorageCache::new(),
|
|
config,
|
|
db: pool,
|
|
redis,
|
|
}
|
|
}
|
|
|
|
async fn insert_zip_task(
|
|
pool: &sqlx::PgPool,
|
|
task_id: Uuid,
|
|
marker: &str,
|
|
input_path: &std::path::Path,
|
|
recorded_size: i64,
|
|
) {
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO tasks (
|
|
id, session_id, status, total_files, completed_files,
|
|
total_original_size, total_compressed_size,
|
|
started_at, completed_at, expires_at, retention_hours
|
|
) VALUES (
|
|
$1, $2, 'completed', 1, 1,
|
|
$3, $3, NOW(), NOW(), NOW() + INTERVAL '1 day', 24
|
|
)
|
|
"#,
|
|
)
|
|
.bind(task_id)
|
|
.bind(format!("zip-session-{marker}"))
|
|
.bind(recorded_size)
|
|
.execute(pool)
|
|
.await
|
|
.expect("insert ZIP test task");
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO task_files (
|
|
id, task_id, original_name, original_format, output_format,
|
|
original_size, compressed_size, saved_percent,
|
|
storage_path, storage_backend, storage_key,
|
|
status, completed_at
|
|
) VALUES (
|
|
$1, $2, $3, 'png', 'png',
|
|
$4, $4, 0,
|
|
$5, 'local', $5,
|
|
'completed', NOW()
|
|
)
|
|
"#,
|
|
)
|
|
.bind(Uuid::new_v4())
|
|
.bind(task_id)
|
|
.bind(format!("{marker}.png"))
|
|
.bind(recorded_size)
|
|
.bind(input_path.to_string_lossy().to_string())
|
|
.execute(pool)
|
|
.await
|
|
.expect("insert ZIP test file");
|
|
}
|
|
|
|
async fn configure_test_s3(state: &AppState, marker: &str) -> Option<Uuid> {
|
|
let endpoint = std::env::var("IMAGEFORGE_TEST_S3_ENDPOINT").ok()?;
|
|
let bucket = std::env::var("IMAGEFORGE_TEST_S3_BUCKET").ok()?;
|
|
let access_key = std::env::var("IMAGEFORGE_TEST_S3_ACCESS_KEY").ok()?;
|
|
let secret_key = std::env::var("IMAGEFORGE_TEST_S3_SECRET_KEY").ok()?;
|
|
let endpoint_id = Uuid::new_v4();
|
|
let encrypted_access =
|
|
settings::encrypt_secret(state, &access_key).expect("encrypt ZIP test S3 access key");
|
|
let encrypted_secret =
|
|
settings::encrypt_secret(state, &secret_key).expect("encrypt ZIP test S3 secret key");
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO storage_endpoints (
|
|
id, name, internal_endpoint, public_endpoint,
|
|
bucket, region,
|
|
access_key_encrypted, secret_key_encrypted, access_key_hint,
|
|
force_path_style, is_active
|
|
) VALUES (
|
|
$1, $2, $3, $3,
|
|
$4, 'us-east-1',
|
|
$5, $6, 'test',
|
|
true, true
|
|
)
|
|
"#,
|
|
)
|
|
.bind(endpoint_id)
|
|
.bind(format!("zip-test-{marker}"))
|
|
.bind(endpoint)
|
|
.bind(bucket)
|
|
.bind(encrypted_access)
|
|
.bind(encrypted_secret)
|
|
.execute(&state.db)
|
|
.await
|
|
.expect("insert ZIP test S3 endpoint");
|
|
Some(endpoint_id)
|
|
}
|
|
|
|
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
#[ignore = "requires isolated IMAGEFORGE_TEST_DATABASE_URL and IMAGEFORGE_TEST_REDIS_URL; optional IMAGEFORGE_TEST_S3_* uses MinIO"]
|
|
async fn zip_build_is_single_flight_bounded_and_fenced() {
|
|
let database_url = std::env::var("IMAGEFORGE_TEST_DATABASE_URL")
|
|
.expect("IMAGEFORGE_TEST_DATABASE_URL must be set");
|
|
assert!(
|
|
database_url.to_ascii_lowercase().contains("test"),
|
|
"refusing to run destructive integration test outside a test database"
|
|
);
|
|
let redis_url = std::env::var("IMAGEFORGE_TEST_REDIS_URL")
|
|
.expect("IMAGEFORGE_TEST_REDIS_URL must be set");
|
|
let pool = PgPoolOptions::new()
|
|
.max_connections(32)
|
|
.connect(&database_url)
|
|
.await
|
|
.expect("connect ZIP test database");
|
|
sqlx::migrate!().run(&pool).await.expect("run migrations");
|
|
let marker = Uuid::new_v4().simple().to_string();
|
|
let storage_root = std::env::temp_dir().join(format!("imageforge-zip-test-{marker}"));
|
|
tokio::fs::create_dir_all(&storage_root)
|
|
.await
|
|
.expect("create ZIP test storage root");
|
|
let state = build_zip_test_state(
|
|
pool.clone(),
|
|
database_url,
|
|
redis_url,
|
|
storage_root.to_string_lossy().to_string(),
|
|
)
|
|
.await;
|
|
let endpoint_id = configure_test_s3(&state, &marker).await;
|
|
|
|
let input_path = storage_root.join("single-flight-input.png");
|
|
tokio::fs::write(&input_path, b"single-flight-payload")
|
|
.await
|
|
.expect("write ZIP input");
|
|
let task_id = Uuid::new_v4();
|
|
insert_zip_task(&pool, task_id, &marker, &input_path, 21).await;
|
|
let barrier = Arc::new(Barrier::new(20));
|
|
let mut joins = Vec::new();
|
|
for _ in 0..20 {
|
|
let state = state.clone();
|
|
let barrier = barrier.clone();
|
|
joins.push(tokio::spawn(async move {
|
|
barrier.wait().await;
|
|
resolve_task_zip(&state, task_id, 24).await
|
|
}));
|
|
}
|
|
let mut locators = Vec::new();
|
|
for join in joins {
|
|
locators.push(
|
|
join.await
|
|
.expect("join concurrent ZIP request")
|
|
.expect("resolve concurrent ZIP request"),
|
|
);
|
|
}
|
|
assert!(locators
|
|
.iter()
|
|
.all(|locator| locator.key == locators[0].key));
|
|
assert!(locators[0]
|
|
.key
|
|
.replace('\\', "/")
|
|
.contains(&format!("/attempts/{task_id}/")));
|
|
if endpoint_id.is_some() {
|
|
assert_eq!(locators[0].backend, "s3");
|
|
}
|
|
let attempts: i64 = sqlx::query_scalar("SELECT zip_build_attempt FROM tasks WHERE id = $1")
|
|
.bind(task_id)
|
|
.fetch_one(&pool)
|
|
.await
|
|
.expect("query ZIP build attempts");
|
|
assert_eq!(attempts, 1, "concurrent ZIP requests built more than once");
|
|
|
|
let over_budget_task = Uuid::new_v4();
|
|
let missing_path = storage_root.join("must-not-be-downloaded.png");
|
|
insert_zip_task(
|
|
&pool,
|
|
over_budget_task,
|
|
&format!("{marker}-over-budget"),
|
|
&missing_path,
|
|
state.config.zip_max_uncompressed_bytes as i64 + 1,
|
|
)
|
|
.await;
|
|
let over_budget = resolve_task_zip(&state, over_budget_task, 24)
|
|
.await
|
|
.expect_err("over-budget ZIP reached the download phase");
|
|
assert_eq!(over_budget.code, ErrorCode::FileTooLarge);
|
|
let over_budget_attempts: i64 =
|
|
sqlx::query_scalar("SELECT zip_build_attempt FROM tasks WHERE id = $1")
|
|
.bind(over_budget_task)
|
|
.fetch_one(&pool)
|
|
.await
|
|
.expect("query over-budget ZIP attempts");
|
|
assert_eq!(over_budget_attempts, 0);
|
|
|
|
let takeover_task = Uuid::new_v4();
|
|
insert_zip_task(
|
|
&pool,
|
|
takeover_task,
|
|
&format!("{marker}-takeover"),
|
|
&input_path,
|
|
21,
|
|
)
|
|
.await;
|
|
let first_token = match claim_zip_build(&state, takeover_task)
|
|
.await
|
|
.expect("claim simulated failing ZIP builder")
|
|
{
|
|
ZipBuildClaim::Acquired { token } => token,
|
|
other => panic!("unexpected initial ZIP claim: {other:?}"),
|
|
};
|
|
let waiter_state = state.clone();
|
|
let waiter =
|
|
tokio::spawn(async move { resolve_task_zip(&waiter_state, takeover_task, 24).await });
|
|
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
|
|
release_zip_build(&state, takeover_task, first_token).await;
|
|
let takeover_locator = waiter
|
|
.await
|
|
.expect("join ZIP takeover waiter")
|
|
.expect("waiter safely took over ZIP build");
|
|
let takeover_attempts: i64 =
|
|
sqlx::query_scalar("SELECT zip_build_attempt FROM tasks WHERE id = $1")
|
|
.bind(takeover_task)
|
|
.fetch_one(&pool)
|
|
.await
|
|
.expect("query takeover attempts");
|
|
assert_eq!(takeover_attempts, 2);
|
|
|
|
let deleted_task = Uuid::new_v4();
|
|
insert_zip_task(
|
|
&pool,
|
|
deleted_task,
|
|
&format!("{marker}-deleted"),
|
|
&input_path,
|
|
21,
|
|
)
|
|
.await;
|
|
let deleted_token = match claim_zip_build(&state, deleted_task)
|
|
.await
|
|
.expect("claim deleted-task ZIP builder")
|
|
{
|
|
ZipBuildClaim::Acquired { token } => token,
|
|
other => panic!("unexpected deleted-task ZIP claim: {other:?}"),
|
|
};
|
|
let unpublished_path = storage_root.join("unpublished.zip");
|
|
tokio::fs::write(&unpublished_path, b"unpublished-zip")
|
|
.await
|
|
.expect("write unpublished ZIP fixture");
|
|
let unpublished = storage::store_file(
|
|
&state,
|
|
&storage::archive_attempt_key(24, deleted_task, deleted_token),
|
|
&unpublished_path,
|
|
"application/zip",
|
|
)
|
|
.await
|
|
.expect("store unpublished ZIP attempt");
|
|
sqlx::query("DELETE FROM tasks WHERE id = $1")
|
|
.bind(deleted_task)
|
|
.execute(&pool)
|
|
.await
|
|
.expect("delete task before ZIP publish");
|
|
let publish_error =
|
|
publish_zip_attempt(&state, deleted_task, deleted_token, unpublished.clone())
|
|
.await
|
|
.expect_err("published ZIP after task deletion");
|
|
assert_eq!(publish_error.code, ErrorCode::StorageUnavailable);
|
|
let orphan_read = storage::read_bytes(
|
|
&state,
|
|
&storage::ObjectLocator {
|
|
backend: unpublished.backend,
|
|
endpoint_id: unpublished.endpoint_id,
|
|
key: unpublished.key,
|
|
},
|
|
)
|
|
.await;
|
|
assert!(orphan_read.is_err(), "unpublished ZIP object was orphaned");
|
|
|
|
for locator in [&locators[0], &takeover_locator] {
|
|
storage::delete_object(&state, locator)
|
|
.await
|
|
.expect("delete published ZIP test object");
|
|
}
|
|
sqlx::query("DELETE FROM tasks WHERE id = ANY($1)")
|
|
.bind(&[task_id, over_budget_task, takeover_task][..])
|
|
.execute(&pool)
|
|
.await
|
|
.expect("delete ZIP test tasks");
|
|
if let Some(endpoint_id) = endpoint_id {
|
|
sqlx::query("DELETE FROM storage_endpoints WHERE id = $1")
|
|
.bind(endpoint_id)
|
|
.execute(&pool)
|
|
.await
|
|
.expect("delete ZIP test S3 endpoint");
|
|
}
|
|
tokio::fs::remove_dir_all(&storage_root)
|
|
.await
|
|
.expect("remove ZIP test storage root");
|
|
}
|
|
}
|