feat: add configurable S3 object storage

2026-07-25 13:23:11 +08:00
parent 61fa9cb820
commit d1f093685d
29 changed files with 3703 additions and 284 deletions
--- a/src/api/downloads.rs
+++ b/src/api/downloads.rs
@@ -1,11 +1,12 @@
 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, Response};
+use axum::response::{IntoResponse, Redirect, Response};
 use axum::routing::get;
 use axum::Router;
 use chrono::{DateTime, Utc};
@@ -25,7 +26,9 @@ pub fn router() -> Router<AppState> {

 #[derive(Debug, FromRow)]
 struct DownloadRow {
-    storage_path: Option<String>,
+    storage_backend: String,
+    storage_endpoint_id: Option<Uuid>,
+    storage_key: Option<String>,
    output_format: String,
    original_name: String,
    file_status: String,
@@ -47,7 +50,9 @@ async fn download_file(
    let row = sqlx::query_as::<_, DownloadRow>(
        r#"
        SELECT
-          f.storage_path,
+          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,
@@ -75,32 +80,51 @@ async fn download_file(

    authorize_download(&principal, &row)?;

-    let Some(path) = &row.storage_path else {
+    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
+}

-    let file = tokio::fs::File::open(path).await.map_err(|err| {
+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(&row.output_format).parse().unwrap(),
-    );
-    resp_headers.insert(
-        header::CONTENT_DISPOSITION,
-        content_disposition(&row.original_name)?,
-    );
+    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()))
 }

@@ -126,10 +150,36 @@ fn content_type(format: &str) -> &'static str {
        "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() {
@@ -177,11 +227,17 @@ struct TaskZipRow {
    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_path: Option<String>,
+    storage_backend: String,
+    storage_endpoint_id: Option<Uuid>,
+    storage_key: Option<String>,
    original_name: String,
    output_format: String,
 }
@@ -203,7 +259,11 @@ async fn download_task_zip(
          session_id,
          status::text AS status,
          completed_at,
-          expires_at
+          expires_at,
+          retention_hours,
+          zip_storage_backend,
+          zip_storage_endpoint_id,
+          zip_storage_key
        FROM tasks
        WHERE id = $1
        "#,
@@ -235,26 +295,29 @@ async fn download_task_zip(
        }
    }

-    if !state.config.storage_type.eq_ignore_ascii_case("local") {
-        return Err(AppError::new(
-            ErrorCode::StorageUnavailable,
-            "当前仅支持本地存储（STORAGE_TYPE=local）",
-        ));
-    }
-
-    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)
-    })?;
-    let zip_path = PathBuf::from(format!("{zip_dir}/{task_id}.zip"));
-
-    if tokio::fs::try_exists(&zip_path).await.unwrap_or(false) {
-        return stream_zip(jar, zip_path, task_id).await;
+    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 rows = sqlx::query_as::<_, TaskZipFileRow>(
        r#"
-        SELECT storage_path, original_name, output_format
+        SELECT storage_backend, storage_endpoint_id,
+               COALESCE(storage_key, storage_path) AS storage_key,
+               original_name, output_format
        FROM task_files
        WHERE task_id = $1 AND status = 'completed'
        ORDER BY created_at ASC
@@ -269,53 +332,92 @@ async fn download_task_zip(
        return Err(AppError::new(ErrorCode::NotFound, "没有可打包的文件"));
    }

-    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 name = build_zip_entry_name(&row.original_name, &row.output_format, &mut used_names);
-        entries.push((name, path));
-    }
-    if entries.is_empty() {
-        return Err(AppError::new(ErrorCode::NotFound, "没有可打包的文件"));
-    }
-
-    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))?;
-
-    stream_zip(jar, zip_path, task_id).await
-}
-
-async fn stream_zip(
-    jar: axum_extra::extract::cookie::CookieJar,
-    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 temp_dir = PathBuf::from(format!(
+        "{}/tmp/zips/{task_id}-{}",
+        state.config.storage_path,
+        Uuid::new_v4()
+    ));
+    tokio::fs::create_dir_all(&temp_dir).await.map_err(|err| {
+        AppError::new(ErrorCode::StorageUnavailable, "创建 ZIP 临时目录失败").with_source(err)
    })?;
+    let zip_path = temp_dir.join(format!("task_{task_id}.zip"));

-    let stream = ReaderStream::new(file);
-    let body = Body::from_stream(stream);
+    let build_result: Result<storage::StoredObject, AppError> = async {
+        let mut used_names: HashMap<String, usize> = HashMap::new();
+        let mut entries: Vec<(String, String)> = Vec::new();
+        for (index, row) in rows.into_iter().enumerate() {
+            let Some(key) = row.storage_key else {
+                continue;
+            };
+            let path = temp_dir.join(format!("entry-{index}"));
+            storage::download_to_file(
+                &state,
+                &storage::ObjectLocator {
+                    backend: row.storage_backend,
+                    endpoint_id: row.storage_endpoint_id,
+                    key,
+                },
+                &path,
+            )
+            .await?;
+            let name =
+                build_zip_entry_name(&row.original_name, &row.output_format, &mut used_names);
+            entries.push((name, path.to_string_lossy().to_string()));
+        }
+        if entries.is_empty() {
+            return Err(AppError::new(ErrorCode::NotFound, "没有可打包的文件"));
+        }

-    let mut resp_headers = HeaderMap::new();
-    resp_headers.insert(header::CONTENT_TYPE, "application/zip".parse().unwrap());
-    resp_headers.insert(
-        header::CONTENT_DISPOSITION,
-        format!("attachment; filename=\"task_{task_id}.zip\"")
-            .parse()
-            .unwrap(),
-    );
+        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))?;

-    Ok((jar, (resp_headers, body).into_response()))
+        let object_key = storage::archive_key(task.retention_hours as i64, task_id);
+        storage::store_file(&state, &object_key, &zip_path, "application/zip").await
+    }
+    .await;
+
+    let _ = tokio::fs::remove_dir_all(&temp_dir).await;
+    let stored = build_result?;
+
+    sqlx::query(
+        r#"
+        UPDATE tasks
+        SET zip_storage_backend = $2,
+            zip_storage_endpoint_id = $3,
+            zip_storage_key = $4,
+            zip_storage_etag = $5,
+            zip_size = $6
+        WHERE id = $1 AND zip_storage_key IS NULL
+        "#,
+    )
+    .bind(task_id)
+    .bind(&stored.backend)
+    .bind(stored.endpoint_id)
+    .bind(&stored.key)
+    .bind(&stored.etag)
+    .bind(stored.size as i64)
+    .execute(&state.db)
+    .await
+    .map_err(|err| AppError::new(ErrorCode::Internal, "记录 ZIP 对象失败").with_source(err))?;
+
+    respond_object(
+        &state,
+        jar,
+        &storage::ObjectLocator {
+            backend: stored.backend,
+            endpoint_id: stored.endpoint_id,
+            key: stored.key,
+        },
+        &format!("task_{task_id}.zip"),
+        "application/zip",
+    )
+    .await
 }

 fn build_zip_entry_name(
@@ -333,6 +435,10 @@ fn build_zip_entry_name(
        "png" => "png",
        "webp" => "webp",
        "avif" => "avif",
+        "gif" => "gif",
+        "bmp" => "bmp",
+        "tif" | "tiff" => "tiff",
+        "ico" => "ico",
        _ => "bin",
    };

@@ -411,4 +517,10 @@ mod tests {
        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");
+    }
 }