Implement compression quota refunds and admin manual subscription

src/api/downloads.rs

@@ -0,0 +1,343 @@
+use crate::api::context;
+use crate::error::{AppError, ErrorCode};
+use crate::state::AppState;
+
+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 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_path: 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?;
+
+    let row = sqlx::query_as::<_, DownloadRow>(
+        r#"
+        SELECT
+          f.storage_path,
+          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(path) = &row.storage_path else {
+        return Err(AppError::new(ErrorCode::NotFound, "文件不存在"));
+    };
+
+    let bytes = tokio::fs::read(path)
+        .await
+        .map_err(|err| AppError::new(ErrorCode::StorageUnavailable, "读取文件失败").with_source(err))?;
+
+    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,
+        format!("attachment; filename=\"{}\"", sanitize_filename(&row.original_name))
+            .parse()
+            .unwrap(),
+    );
+
+    Ok((jar, (resp_headers, bytes).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",
+        _ => "application/octet-stream",
+    }
+}
+
+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', '"', '\\'], "_");
+    if out.len() > 120 {
+        out.truncate(120);
+    }
+    out
+}
+
+#[derive(Debug, FromRow)]
+struct TaskZipRow {
+    user_id: Option<Uuid>,
+    session_id: Option<String>,
+    status: String,
+    completed_at: Option<DateTime<Utc>>,
+    expires_at: DateTime<Utc>,
+}
+
+#[derive(Debug, FromRow)]
+struct TaskZipFileRow {
+    id: Uuid,
+    storage_path: Option<String>,
+    original_name: String,
+    output_format: String,
+}
+
+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?;
+
+    let task = sqlx::query_as::<_, TaskZipRow>(
+        r#"
+        SELECT
+          user_id,
+          session_id,
+          status::text AS status,
+          completed_at,
+          expires_at
+        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 state.config.storage_type.to_ascii_lowercase() != "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;
+    }
+
+    let rows = sqlx::query_as::<_, TaskZipFileRow>(
+        r#"
+        SELECT id, storage_path, original_name, output_format
+        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, "没有可打包的文件"));
+    }
+
+    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 stream = ReaderStream::new(file);
+    let body = Body::from_stream(stream);
+
+    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(),
+    );
+
+    Ok((jar, (resp_headers, body).into_response()))
+}
+
+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",
+        _ => "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', '"', '\\', '/', ':'], "_");
+    if out.len() > 120 {
+        out.truncate(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(())
+}