Implement compression quota refunds and admin manual subscription
This commit is contained in:
1553
src/api/admin.rs
Normal file
1553
src/api/admin.rs
Normal file
File diff suppressed because it is too large
Load Diff
596
src/api/auth.rs
Normal file
596
src/api/auth.rs
Normal file
@@ -0,0 +1,596 @@
|
||||
use crate::auth;
|
||||
use crate::api::envelope::Envelope;
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::services::mail;
|
||||
use crate::state::AppState;
|
||||
|
||||
use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier};
|
||||
use axum::{
|
||||
extract::State,
|
||||
http::HeaderMap,
|
||||
routing::post,
|
||||
Json, Router,
|
||||
};
|
||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use rand::RngCore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use uuid::Uuid;
|
||||
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/register", post(register))
|
||||
.route("/login", post(login))
|
||||
.route("/send-verification", post(send_verification))
|
||||
.route("/verify-email", post(verify_email))
|
||||
.route("/forgot-password", post(forgot_password))
|
||||
.route("/reset-password", post(reset_password))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct RegisterRequest {
|
||||
email: String,
|
||||
password: String,
|
||||
username: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct RegisterResponse {
|
||||
user: UserView,
|
||||
token: String,
|
||||
message: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct LoginRequest {
|
||||
email: String,
|
||||
password: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct LoginResponse {
|
||||
token: String,
|
||||
expires_at: DateTime<Utc>,
|
||||
user: UserView,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct UserView {
|
||||
id: Uuid,
|
||||
email: String,
|
||||
username: String,
|
||||
role: String,
|
||||
email_verified: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, sqlx::FromRow)]
|
||||
struct UserRow {
|
||||
id: Uuid,
|
||||
email: String,
|
||||
username: String,
|
||||
password_hash: String,
|
||||
role: String,
|
||||
is_active: bool,
|
||||
email_verified_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
async fn register(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<RegisterRequest>,
|
||||
) -> Result<Json<Envelope<RegisterResponse>>, AppError> {
|
||||
validate_email(&req.email)?;
|
||||
validate_username(&req.username)?;
|
||||
validate_password(&req.password)?;
|
||||
|
||||
let password_hash = hash_password(&req.password)?;
|
||||
|
||||
let user = sqlx::query_as::<_, UserRow>(
|
||||
r#"
|
||||
INSERT INTO users (email, username, password_hash)
|
||||
VALUES ($1, $2, $3)
|
||||
RETURNING
|
||||
id,
|
||||
email,
|
||||
username,
|
||||
password_hash,
|
||||
role::text AS role,
|
||||
is_active,
|
||||
email_verified_at
|
||||
"#,
|
||||
)
|
||||
.bind(req.email.to_lowercase())
|
||||
.bind(&req.username)
|
||||
.bind(password_hash)
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(map_unique_violation)?;
|
||||
|
||||
let (token, _expires_at) =
|
||||
auth::issue_jwt(&state.config.jwt_secret, state.config.jwt_expiry_hours, user.id, &user.role)?;
|
||||
|
||||
let verification_token = generate_token();
|
||||
let token_hash = sha256_hex(&verification_token);
|
||||
let expires_at_db = Utc::now() + Duration::hours(24);
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO email_verifications (user_id, token_hash, expires_at)
|
||||
VALUES ($1, $2, $3)
|
||||
"#,
|
||||
)
|
||||
.bind(user.id)
|
||||
.bind(token_hash)
|
||||
.bind(expires_at_db)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "创建邮箱验证记录失败").with_source(err))?;
|
||||
|
||||
let verification_url = format!(
|
||||
"{}/verify-email?token={}",
|
||||
state.config.public_base_url, verification_token
|
||||
);
|
||||
|
||||
mail::send_verification_email(&state, &user.email, &user.username, &verification_url)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::MailSendFailed, "验证邮件发送失败").with_source(err))?;
|
||||
|
||||
let body = RegisterResponse {
|
||||
user: UserView {
|
||||
id: user.id,
|
||||
email: user.email,
|
||||
username: user.username,
|
||||
role: user.role,
|
||||
email_verified: user.email_verified_at.is_some(),
|
||||
},
|
||||
token,
|
||||
message: "注册成功，验证邮件已发送至您的邮箱".to_string(),
|
||||
};
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: body,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn login(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<LoginRequest>,
|
||||
) -> Result<Json<Envelope<LoginResponse>>, AppError> {
|
||||
let identity = req.email.trim();
|
||||
if identity.is_empty() {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "邮箱或用户名不能为空"));
|
||||
}
|
||||
|
||||
let user = if identity.contains('@') {
|
||||
validate_email(identity)?;
|
||||
sqlx::query_as::<_, UserRow>(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
email,
|
||||
username,
|
||||
password_hash,
|
||||
role::text AS role,
|
||||
is_active,
|
||||
email_verified_at
|
||||
FROM users
|
||||
WHERE email = $1
|
||||
"#,
|
||||
)
|
||||
.bind(identity.to_lowercase())
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
} else {
|
||||
validate_username(identity)?;
|
||||
sqlx::query_as::<_, UserRow>(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
email,
|
||||
username,
|
||||
password_hash,
|
||||
role::text AS role,
|
||||
is_active,
|
||||
email_verified_at
|
||||
FROM users
|
||||
WHERE username = $1
|
||||
"#,
|
||||
)
|
||||
.bind(identity)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
}
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?
|
||||
.ok_or_else(|| AppError::new(ErrorCode::Unauthorized, "账号或密码错误"))?;
|
||||
|
||||
if !user.is_active {
|
||||
return Err(AppError::new(ErrorCode::Forbidden, "账号已被禁用"));
|
||||
}
|
||||
|
||||
verify_password(&req.password, &user.password_hash)?;
|
||||
|
||||
let (token, expires_at) =
|
||||
auth::issue_jwt(&state.config.jwt_secret, state.config.jwt_expiry_hours, user.id, &user.role)?;
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: LoginResponse {
|
||||
token,
|
||||
expires_at,
|
||||
user: UserView {
|
||||
id: user.id,
|
||||
email: user.email,
|
||||
username: user.username,
|
||||
role: user.role,
|
||||
email_verified: user.email_verified_at.is_some(),
|
||||
},
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct MessageResponse {
|
||||
message: String,
|
||||
}
|
||||
|
||||
async fn send_verification(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<Json<Envelope<MessageResponse>>, AppError> {
|
||||
let claims = auth::require_jwt(&state.config.jwt_secret, &headers)?;
|
||||
|
||||
// Rate limit: 1 per minute per user
|
||||
let key = format!("rate:send_verification:{}:{}", claims.sub, Utc::now().format("%Y%m%d%H%M"));
|
||||
let mut redis = state.redis.clone();
|
||||
let count: i64 = redis::cmd("INCR")
|
||||
.arg(&key)
|
||||
.query_async(&mut redis)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "Redis 限流失败").with_source(err))?;
|
||||
if count == 1 {
|
||||
let _: () = redis::cmd("EXPIRE")
|
||||
.arg(&key)
|
||||
.arg(60)
|
||||
.query_async(&mut redis)
|
||||
.await
|
||||
.unwrap_or(());
|
||||
}
|
||||
if count > 1 {
|
||||
return Err(AppError::new(ErrorCode::RateLimited, "发送过于频繁，请稍后再试"));
|
||||
}
|
||||
|
||||
let user = sqlx::query_as::<_, UserRow>(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
email,
|
||||
username,
|
||||
password_hash,
|
||||
role::text AS role,
|
||||
is_active,
|
||||
email_verified_at
|
||||
FROM users
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(claims.sub)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?
|
||||
.ok_or_else(|| AppError::new(ErrorCode::Unauthorized, "用户不存在或未登录"))?;
|
||||
|
||||
if user.email_verified_at.is_some() {
|
||||
return Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: MessageResponse {
|
||||
message: "邮箱已验证，无需重复验证".to_string(),
|
||||
},
|
||||
}));
|
||||
}
|
||||
|
||||
let verification_token = generate_token();
|
||||
let token_hash = sha256_hex(&verification_token);
|
||||
let expires_at_db = Utc::now() + Duration::hours(24);
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO email_verifications (user_id, token_hash, expires_at)
|
||||
VALUES ($1, $2, $3)
|
||||
"#,
|
||||
)
|
||||
.bind(user.id)
|
||||
.bind(token_hash)
|
||||
.bind(expires_at_db)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "创建邮箱验证记录失败").with_source(err))?;
|
||||
|
||||
let verification_url = format!(
|
||||
"{}/verify-email?token={}",
|
||||
state.config.public_base_url, verification_token
|
||||
);
|
||||
|
||||
mail::send_verification_email(&state, &user.email, &user.username, &verification_url)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::MailSendFailed, "验证邮件发送失败").with_source(err))?;
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: MessageResponse {
|
||||
message: "验证邮件已发送，请查收".to_string(),
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct VerifyEmailRequest {
|
||||
token: String,
|
||||
}
|
||||
|
||||
async fn verify_email(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<VerifyEmailRequest>,
|
||||
) -> Result<Json<Envelope<MessageResponse>>, AppError> {
|
||||
if req.token.trim().is_empty() {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "token 不能为空"));
|
||||
}
|
||||
|
||||
let token_hash = sha256_hex(&req.token);
|
||||
let now = Utc::now();
|
||||
|
||||
let updated = sqlx::query(
|
||||
r#"
|
||||
WITH v AS (
|
||||
SELECT user_id
|
||||
FROM email_verifications
|
||||
WHERE token_hash = $1
|
||||
AND verified_at IS NULL
|
||||
AND expires_at > $2
|
||||
LIMIT 1
|
||||
),
|
||||
u AS (
|
||||
UPDATE users
|
||||
SET email_verified_at = $2
|
||||
WHERE id = (SELECT user_id FROM v)
|
||||
AND email_verified_at IS NULL
|
||||
RETURNING id
|
||||
)
|
||||
UPDATE email_verifications
|
||||
SET verified_at = $2
|
||||
WHERE token_hash = $1
|
||||
AND verified_at IS NULL
|
||||
AND expires_at > $2
|
||||
AND user_id IN (SELECT id FROM u)
|
||||
"#,
|
||||
)
|
||||
.bind(token_hash)
|
||||
.bind(now)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "邮箱验证失败").with_source(err))?;
|
||||
|
||||
if updated.rows_affected() == 0 {
|
||||
return Err(AppError::new(ErrorCode::InvalidToken, "Token 无效或已过期"));
|
||||
}
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: MessageResponse {
|
||||
message: "邮箱验证成功".to_string(),
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ForgotPasswordRequest {
|
||||
email: String,
|
||||
}
|
||||
|
||||
async fn forgot_password(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<ForgotPasswordRequest>,
|
||||
) -> Result<Json<Envelope<MessageResponse>>, AppError> {
|
||||
validate_email(&req.email)?;
|
||||
|
||||
let user = sqlx::query_as::<_, UserRow>(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
email,
|
||||
username,
|
||||
password_hash,
|
||||
role::text AS role,
|
||||
is_active,
|
||||
email_verified_at
|
||||
FROM users
|
||||
WHERE email = $1
|
||||
"#,
|
||||
)
|
||||
.bind(req.email.to_lowercase())
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?;
|
||||
|
||||
if let Some(user) = user {
|
||||
let reset_token = generate_token();
|
||||
let token_hash = sha256_hex(&reset_token);
|
||||
let expires_at_db = Utc::now() + Duration::hours(1);
|
||||
|
||||
let _ = sqlx::query(
|
||||
r#"
|
||||
INSERT INTO password_resets (user_id, token_hash, expires_at)
|
||||
VALUES ($1, $2, $3)
|
||||
"#,
|
||||
)
|
||||
.bind(user.id)
|
||||
.bind(token_hash)
|
||||
.bind(expires_at_db)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
let reset_url = format!(
|
||||
"{}/reset-password?token={}",
|
||||
state.config.public_base_url, reset_token
|
||||
);
|
||||
|
||||
let _ = mail::send_password_reset_email(&state, &user.email, &user.username, &reset_url).await;
|
||||
}
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: MessageResponse {
|
||||
message: "如果该邮箱已注册，您将收到重置邮件".to_string(),
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ResetPasswordRequest {
|
||||
token: String,
|
||||
new_password: String,
|
||||
}
|
||||
|
||||
async fn reset_password(
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<ResetPasswordRequest>,
|
||||
) -> Result<Json<Envelope<MessageResponse>>, AppError> {
|
||||
if req.token.trim().is_empty() {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "token 不能为空"));
|
||||
}
|
||||
validate_password(&req.new_password)?;
|
||||
|
||||
let token_hash = sha256_hex(&req.token);
|
||||
let now = Utc::now();
|
||||
|
||||
let user_id: Option<Uuid> = sqlx::query_scalar(
|
||||
r#"
|
||||
SELECT user_id
|
||||
FROM password_resets
|
||||
WHERE token_hash = $1
|
||||
AND used_at IS NULL
|
||||
AND expires_at > $2
|
||||
"#,
|
||||
)
|
||||
.bind(&token_hash)
|
||||
.bind(now)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "重置密码失败").with_source(err))?;
|
||||
|
||||
let Some(user_id) = user_id else {
|
||||
return Err(AppError::new(ErrorCode::InvalidToken, "Token 无效或已过期"));
|
||||
};
|
||||
|
||||
let password_hash = hash_password(&req.new_password)?;
|
||||
|
||||
let mut tx = state
|
||||
.db
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?;
|
||||
|
||||
sqlx::query("UPDATE users SET password_hash = $1 WHERE id = $2")
|
||||
.bind(password_hash)
|
||||
.bind(user_id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新密码失败").with_source(err))?;
|
||||
|
||||
sqlx::query("UPDATE password_resets SET used_at = $2 WHERE token_hash = $1 AND used_at IS NULL")
|
||||
.bind(token_hash)
|
||||
.bind(now)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新重置记录失败").with_source(err))?;
|
||||
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: MessageResponse {
|
||||
message: "密码重置成功".to_string(),
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
fn validate_email(email: &str) -> Result<(), AppError> {
|
||||
if email.trim().is_empty() || !email.contains('@') {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "邮箱格式不正确"));
|
||||
}
|
||||
if email.len() > 255 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "邮箱过长"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_username(username: &str) -> Result<(), AppError> {
|
||||
if username.trim().is_empty() {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "用户名不能为空"));
|
||||
}
|
||||
if username.len() > 50 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "用户名过长"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_password(password: &str) -> Result<(), AppError> {
|
||||
if password.len() < 8 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "密码至少 8 位"));
|
||||
}
|
||||
if password.len() > 128 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "密码过长"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn hash_password(password: &str) -> Result<String, AppError> {
|
||||
let salt = argon2::password_hash::SaltString::generate(&mut rand::rngs::OsRng);
|
||||
Argon2::default()
|
||||
.hash_password(password.as_bytes(), &salt)
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "密码哈希失败").with_source(err))?
|
||||
.to_string()
|
||||
.pipe(Ok)
|
||||
}
|
||||
|
||||
fn verify_password(password: &str, password_hash: &str) -> Result<(), AppError> {
|
||||
let parsed = PasswordHash::new(password_hash)
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "密码哈希格式错误").with_source(err))?;
|
||||
Argon2::default()
|
||||
.verify_password(password.as_bytes(), &parsed)
|
||||
.map_err(|_| AppError::new(ErrorCode::Unauthorized, "账号或密码错误"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn generate_token() -> String {
|
||||
let mut bytes = [0u8; 32];
|
||||
rand::rngs::OsRng.fill_bytes(&mut bytes);
|
||||
URL_SAFE_NO_PAD.encode(bytes)
|
||||
}
|
||||
|
||||
fn sha256_hex(token: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(token.as_bytes());
|
||||
hex::encode(hasher.finalize())
|
||||
}
|
||||
|
||||
fn map_unique_violation(err: sqlx::Error) -> AppError {
|
||||
if let sqlx::Error::Database(db_err) = &err {
|
||||
if let Some(code) = db_err.code() {
|
||||
if code == "23505" {
|
||||
return AppError::new(ErrorCode::InvalidRequest, "邮箱或用户名已存在");
|
||||
}
|
||||
}
|
||||
}
|
||||
AppError::new(ErrorCode::Internal, "数据库操作失败").with_source(err)
|
||||
}
|
||||
|
||||
trait Pipe: Sized {
|
||||
fn pipe<T>(self, f: impl FnOnce(Self) -> T) -> T {
|
||||
f(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Pipe for T {}
|
||||
738
src/api/billing.rs
Normal file
738
src/api/billing.rs
Normal file
@@ -0,0 +1,738 @@
|
||||
use crate::api::context;
|
||||
use crate::api::envelope::Envelope;
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::services::billing;
|
||||
use crate::services::idempotency;
|
||||
use crate::services::settings;
|
||||
use crate::state::AppState;
|
||||
|
||||
use axum::extract::{ConnectInfo, State};
|
||||
use axum::http::HeaderMap;
|
||||
use axum::routing::{get, post};
|
||||
use axum::{Json, Router};
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::FromRow;
|
||||
use std::net::SocketAddr;
|
||||
use uuid::Uuid;
|
||||
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/billing/plans", get(list_plans))
|
||||
.route("/billing/subscription", get(get_subscription))
|
||||
.route("/billing/usage", get(get_usage))
|
||||
.route("/billing/invoices", get(list_invoices))
|
||||
.route("/billing/checkout", post(create_checkout))
|
||||
.route("/billing/portal", post(create_portal))
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow, Serialize)]
|
||||
struct PlanView {
|
||||
id: Uuid,
|
||||
code: String,
|
||||
name: String,
|
||||
currency: String,
|
||||
amount_cents: i32,
|
||||
interval: String,
|
||||
included_units_per_period: i32,
|
||||
max_file_size_mb: i32,
|
||||
max_files_per_batch: i32,
|
||||
retention_days: i32,
|
||||
features: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct PlansResponse {
|
||||
plans: Vec<PlanView>,
|
||||
}
|
||||
|
||||
async fn list_plans(State(state): State<AppState>) -> Result<Json<Envelope<PlansResponse>>, AppError> {
|
||||
let plans = sqlx::query_as::<_, PlanView>(
|
||||
r#"
|
||||
SELECT
|
||||
id, code, name, currency, amount_cents, interval,
|
||||
included_units_per_period, max_file_size_mb, max_files_per_batch, retention_days,
|
||||
features
|
||||
FROM plans
|
||||
WHERE is_active = true
|
||||
ORDER BY amount_cents ASC
|
||||
"#,
|
||||
)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询套餐失败").with_source(err))?;
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: PlansResponse { plans },
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct SubscriptionPlanView {
|
||||
id: Uuid,
|
||||
code: String,
|
||||
name: String,
|
||||
currency: String,
|
||||
amount_cents: i32,
|
||||
interval: String,
|
||||
included_units_per_period: i32,
|
||||
max_file_size_mb: i32,
|
||||
max_files_per_batch: i32,
|
||||
retention_days: i32,
|
||||
features: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct SubscriptionView {
|
||||
status: String,
|
||||
current_period_start: DateTime<Utc>,
|
||||
current_period_end: DateTime<Utc>,
|
||||
cancel_at_period_end: bool,
|
||||
plan: SubscriptionPlanView,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct SubscriptionResponse {
|
||||
subscription: SubscriptionView,
|
||||
}
|
||||
|
||||
async fn get_subscription(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<Json<Envelope<SubscriptionResponse>>, AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (_jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let user_id = match principal {
|
||||
context::Principal::User { user_id, .. } => user_id,
|
||||
_ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
|
||||
};
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct SubRow {
|
||||
status: String,
|
||||
current_period_start: DateTime<Utc>,
|
||||
current_period_end: DateTime<Utc>,
|
||||
cancel_at_period_end: bool,
|
||||
plan_id: Uuid,
|
||||
plan_code: String,
|
||||
plan_name: String,
|
||||
currency: String,
|
||||
amount_cents: i32,
|
||||
interval: String,
|
||||
included_units_per_period: i32,
|
||||
max_file_size_mb: i32,
|
||||
max_files_per_batch: i32,
|
||||
retention_days: i32,
|
||||
features: serde_json::Value,
|
||||
}
|
||||
|
||||
let sub = sqlx::query_as::<_, SubRow>(
|
||||
r#"
|
||||
SELECT
|
||||
s.status::text AS status,
|
||||
s.current_period_start,
|
||||
s.current_period_end,
|
||||
s.cancel_at_period_end,
|
||||
p.id AS plan_id,
|
||||
p.code AS plan_code,
|
||||
p.name AS plan_name,
|
||||
p.currency,
|
||||
p.amount_cents,
|
||||
p.interval,
|
||||
p.included_units_per_period,
|
||||
p.max_file_size_mb,
|
||||
p.max_files_per_batch,
|
||||
p.retention_days,
|
||||
p.features
|
||||
FROM subscriptions s
|
||||
JOIN plans p ON p.id = s.plan_id
|
||||
WHERE s.user_id = $1
|
||||
AND s.status IN ('active', 'trialing', 'past_due', 'canceled', 'incomplete')
|
||||
ORDER BY s.current_period_end DESC
|
||||
LIMIT 1
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询订阅失败").with_source(err))?;
|
||||
|
||||
let (status, period_start, period_end, cancel_at_period_end, plan) = if let Some(sub) = sub {
|
||||
(
|
||||
sub.status,
|
||||
sub.current_period_start,
|
||||
sub.current_period_end,
|
||||
sub.cancel_at_period_end,
|
||||
SubscriptionPlanView {
|
||||
id: sub.plan_id,
|
||||
code: sub.plan_code,
|
||||
name: sub.plan_name,
|
||||
currency: sub.currency,
|
||||
amount_cents: sub.amount_cents,
|
||||
interval: sub.interval,
|
||||
included_units_per_period: sub.included_units_per_period,
|
||||
max_file_size_mb: sub.max_file_size_mb,
|
||||
max_files_per_batch: sub.max_files_per_batch,
|
||||
retention_days: sub.retention_days,
|
||||
features: sub.features,
|
||||
},
|
||||
)
|
||||
} else {
|
||||
let plan: PlanView = sqlx::query_as::<_, PlanView>(
|
||||
r#"
|
||||
SELECT
|
||||
id, code, name, currency, amount_cents, interval,
|
||||
included_units_per_period, max_file_size_mb, max_files_per_batch, retention_days,
|
||||
features
|
||||
FROM plans
|
||||
WHERE code = 'free'
|
||||
LIMIT 1
|
||||
"#,
|
||||
)
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "未找到 Free 套餐").with_source(err))?;
|
||||
|
||||
let (start, end) = billing::current_month_period_utc8(Utc::now());
|
||||
(
|
||||
"free".to_string(),
|
||||
start,
|
||||
end,
|
||||
false,
|
||||
SubscriptionPlanView {
|
||||
id: plan.id,
|
||||
code: plan.code,
|
||||
name: plan.name,
|
||||
currency: plan.currency,
|
||||
amount_cents: plan.amount_cents,
|
||||
interval: plan.interval,
|
||||
included_units_per_period: plan.included_units_per_period,
|
||||
max_file_size_mb: plan.max_file_size_mb,
|
||||
max_files_per_batch: plan.max_files_per_batch,
|
||||
retention_days: plan.retention_days,
|
||||
features: plan.features,
|
||||
},
|
||||
)
|
||||
};
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: SubscriptionResponse {
|
||||
subscription: SubscriptionView {
|
||||
status,
|
||||
current_period_start: period_start,
|
||||
current_period_end: period_end,
|
||||
cancel_at_period_end,
|
||||
plan,
|
||||
},
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct UsageResponse {
|
||||
period_start: DateTime<Utc>,
|
||||
period_end: DateTime<Utc>,
|
||||
used_units: i32,
|
||||
included_units: i32,
|
||||
bonus_units: i32,
|
||||
total_units: i32,
|
||||
remaining_units: i32,
|
||||
}
|
||||
|
||||
async fn get_usage(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<Json<Envelope<UsageResponse>>, AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (_jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let user_id = match principal {
|
||||
context::Principal::User { user_id, .. } => user_id,
|
||||
_ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
|
||||
};
|
||||
|
||||
let billing = billing::get_user_billing(&state, user_id).await?;
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct UsageRow {
|
||||
used_units: i32,
|
||||
bonus_units: i32,
|
||||
}
|
||||
|
||||
let usage = sqlx::query_as::<_, UsageRow>(
|
||||
r#"
|
||||
SELECT used_units, bonus_units
|
||||
FROM usage_periods
|
||||
WHERE user_id = $1 AND period_start = $2 AND period_end = $3
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(billing.period_start)
|
||||
.bind(billing.period_end)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用量失败").with_source(err))?
|
||||
.unwrap_or(UsageRow {
|
||||
used_units: 0,
|
||||
bonus_units: 0,
|
||||
});
|
||||
|
||||
let included = billing.plan.included_units_per_period;
|
||||
let total = included + usage.bonus_units;
|
||||
let remaining = (total - usage.used_units).max(0);
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: UsageResponse {
|
||||
period_start: billing.period_start,
|
||||
period_end: billing.period_end,
|
||||
used_units: usage.used_units,
|
||||
included_units: included,
|
||||
bonus_units: usage.bonus_units,
|
||||
total_units: total,
|
||||
remaining_units: remaining,
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct PagingQuery {
|
||||
page: Option<u32>,
|
||||
limit: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow, Serialize)]
|
||||
struct InvoiceView {
|
||||
invoice_number: String,
|
||||
status: String,
|
||||
currency: String,
|
||||
total_amount_cents: i32,
|
||||
period_start: Option<DateTime<Utc>>,
|
||||
period_end: Option<DateTime<Utc>>,
|
||||
hosted_invoice_url: Option<String>,
|
||||
pdf_url: Option<String>,
|
||||
paid_at: Option<DateTime<Utc>>,
|
||||
created_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct InvoicesResponse {
|
||||
invoices: Vec<InvoiceView>,
|
||||
page: u32,
|
||||
limit: u32,
|
||||
}
|
||||
|
||||
async fn list_invoices(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
axum::extract::Query(query): axum::extract::Query<PagingQuery>,
|
||||
) -> Result<Json<Envelope<InvoicesResponse>>, AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (_jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let user_id = match principal {
|
||||
context::Principal::User { user_id, .. } => user_id,
|
||||
_ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
|
||||
};
|
||||
|
||||
let limit = query.limit.unwrap_or(20).clamp(1, 100);
|
||||
let page = query.page.unwrap_or(1).max(1);
|
||||
let offset = (page - 1) * limit;
|
||||
|
||||
let invoices = sqlx::query_as::<_, InvoiceView>(
|
||||
r#"
|
||||
SELECT
|
||||
invoice_number,
|
||||
status::text AS status,
|
||||
currency,
|
||||
total_amount_cents,
|
||||
period_start,
|
||||
period_end,
|
||||
hosted_invoice_url,
|
||||
pdf_url,
|
||||
paid_at,
|
||||
created_at
|
||||
FROM invoices
|
||||
WHERE user_id = $1
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $2 OFFSET $3
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(limit as i64)
|
||||
.bind(offset as i64)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询发票失败").with_source(err))?;
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: InvoicesResponse {
|
||||
invoices,
|
||||
page,
|
||||
limit,
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct CheckoutRequest {
|
||||
plan_id: Uuid,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct CheckoutResponse {
|
||||
checkout_url: String,
|
||||
}
|
||||
|
||||
async fn create_checkout(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
Json(req): Json<CheckoutRequest>,
|
||||
) -> Result<Json<Envelope<CheckoutResponse>>, AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (_jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let user_id = match principal {
|
||||
context::Principal::User { user_id, .. } => user_id,
|
||||
_ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
|
||||
};
|
||||
|
||||
let idempotency_key = headers
|
||||
.get("idempotency-key")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|v| !v.is_empty())
|
||||
.map(str::to_string);
|
||||
|
||||
let request_hash = idempotency_key.as_ref().map(|_| {
|
||||
let plan = req.plan_id.to_string();
|
||||
idempotency::sha256_hex(&[b"billing_checkout", plan.as_bytes()])
|
||||
});
|
||||
|
||||
let mut idem_acquired = false;
|
||||
if let (Some(idem), Some(request_hash)) = (idempotency_key.as_deref(), request_hash.as_deref()) {
|
||||
match idempotency::begin(
|
||||
&state,
|
||||
idempotency::Scope::User(user_id),
|
||||
idem,
|
||||
request_hash,
|
||||
state.config.idempotency_ttl_hours as i64,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
idempotency::BeginResult::Replay { response_body, .. } => {
|
||||
let resp: CheckoutResponse = serde_json::from_value(response_body).map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "幂等结果解析失败").with_source(err)
|
||||
})?;
|
||||
return Ok(Json(Envelope { success: true, data: resp }));
|
||||
}
|
||||
idempotency::BeginResult::InProgress => {
|
||||
if let Some((_status, body)) = idempotency::wait_for_replay(
|
||||
&state,
|
||||
idempotency::Scope::User(user_id),
|
||||
idem,
|
||||
request_hash,
|
||||
10_000,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
let resp: CheckoutResponse = serde_json::from_value(body).map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "幂等结果解析失败").with_source(err)
|
||||
})?;
|
||||
return Ok(Json(Envelope { success: true, data: resp }));
|
||||
}
|
||||
return Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
"请求正在处理中，请稍后重试",
|
||||
));
|
||||
}
|
||||
idempotency::BeginResult::Acquired { .. } => {
|
||||
idem_acquired = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let session_result: Result<String, AppError> = (async {
|
||||
let stripe_secret = settings::get_stripe_secret(&state)
|
||||
.await
|
||||
.map_err(|err| err.with_source("stripe secret not configured"))?;
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct PlanStripeRow {
|
||||
stripe_price_id: Option<String>,
|
||||
amount_cents: i32,
|
||||
is_active: bool,
|
||||
}
|
||||
|
||||
let plan = sqlx::query_as::<_, PlanStripeRow>(
|
||||
r#"
|
||||
SELECT stripe_price_id, amount_cents, is_active
|
||||
FROM plans
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(req.plan_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询套餐失败").with_source(err))?
|
||||
.ok_or_else(|| AppError::new(ErrorCode::NotFound, "套餐不存在"))?;
|
||||
|
||||
if !plan.is_active {
|
||||
return Err(AppError::new(ErrorCode::Forbidden, "套餐不可用"));
|
||||
}
|
||||
let Some(price_id) = plan.stripe_price_id.filter(|v| !v.trim().is_empty()) else {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "该套餐不可订阅"));
|
||||
};
|
||||
if plan.amount_cents <= 0 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "该套餐不可订阅"));
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct UserStripeRow {
|
||||
email: String,
|
||||
billing_customer_id: Option<String>,
|
||||
}
|
||||
|
||||
let user = sqlx::query_as::<_, UserStripeRow>(
|
||||
"SELECT email, billing_customer_id FROM users WHERE id = $1",
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?;
|
||||
|
||||
let customer_id = if let Some(cus) = user.billing_customer_id {
|
||||
cus
|
||||
} else {
|
||||
let cus = stripe_create_customer(&stripe_secret, &user.email, user_id).await?;
|
||||
let _ = sqlx::query("UPDATE users SET billing_customer_id = $2 WHERE id = $1")
|
||||
.bind(user_id)
|
||||
.bind(&cus)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
cus
|
||||
};
|
||||
|
||||
let success_url =
|
||||
format!("{}/dashboard/billing?checkout=success", state.config.public_base_url);
|
||||
let cancel_url = format!("{}/pricing?checkout=cancel", state.config.public_base_url);
|
||||
|
||||
stripe_create_checkout_session(
|
||||
&stripe_secret,
|
||||
&customer_id,
|
||||
&price_id,
|
||||
&success_url,
|
||||
&cancel_url,
|
||||
user_id,
|
||||
)
|
||||
.await
|
||||
})
|
||||
.await;
|
||||
|
||||
match session_result {
|
||||
Ok(session) => {
|
||||
if let (Some(idem), Some(request_hash)) =
|
||||
(idempotency_key.as_deref(), request_hash.as_deref())
|
||||
{
|
||||
if idem_acquired {
|
||||
let _ = idempotency::complete(
|
||||
&state,
|
||||
idempotency::Scope::User(user_id),
|
||||
idem,
|
||||
request_hash,
|
||||
200,
|
||||
serde_json::to_value(&CheckoutResponse {
|
||||
checkout_url: session.clone(),
|
||||
})
|
||||
.unwrap_or(serde_json::Value::Null),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: CheckoutResponse {
|
||||
checkout_url: session,
|
||||
},
|
||||
}))
|
||||
}
|
||||
Err(err) => {
|
||||
if let (Some(idem), Some(request_hash)) =
|
||||
(idempotency_key.as_deref(), request_hash.as_deref())
|
||||
{
|
||||
if idem_acquired {
|
||||
let _ = idempotency::abort(
|
||||
&state,
|
||||
idempotency::Scope::User(user_id),
|
||||
idem,
|
||||
request_hash,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct PortalResponse {
|
||||
url: String,
|
||||
}
|
||||
|
||||
async fn create_portal(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<Json<Envelope<PortalResponse>>, AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (_jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let user_id = match principal {
|
||||
context::Principal::User { user_id, .. } => user_id,
|
||||
_ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
|
||||
};
|
||||
|
||||
let stripe_secret = settings::get_stripe_secret(&state)
|
||||
.await
|
||||
.map_err(|err| err.with_source("stripe secret not configured"))?;
|
||||
|
||||
let customer_id: Option<String> =
|
||||
sqlx::query_scalar("SELECT billing_customer_id FROM users WHERE id = $1")
|
||||
.bind(user_id)
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?;
|
||||
|
||||
let Some(customer_id) = customer_id.filter(|v| !v.trim().is_empty()) else {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "未找到 Stripe Customer"));
|
||||
};
|
||||
|
||||
let return_url = format!("{}/dashboard/billing", state.config.public_base_url);
|
||||
let url = stripe_create_portal_session(&stripe_secret, &customer_id, &return_url).await?;
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: PortalResponse { url },
|
||||
}))
|
||||
}
|
||||
|
||||
async fn stripe_create_customer(secret: &str, email: &str, user_id: Uuid) -> Result<String, AppError> {
|
||||
let resp: serde_json::Value = stripe_post_form(
|
||||
secret,
|
||||
"/v1/customers",
|
||||
vec![
|
||||
("email".to_string(), email.to_string()),
|
||||
("metadata[user_id]".to_string(), user_id.to_string()),
|
||||
],
|
||||
)
|
||||
.await?;
|
||||
|
||||
let id = resp
|
||||
.get("id")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| AppError::new(ErrorCode::Internal, "Stripe customer 创建失败"))?;
|
||||
|
||||
Ok(id.to_string())
|
||||
}
|
||||
|
||||
async fn stripe_create_checkout_session(
|
||||
secret: &str,
|
||||
customer_id: &str,
|
||||
price_id: &str,
|
||||
success_url: &str,
|
||||
cancel_url: &str,
|
||||
user_id: Uuid,
|
||||
) -> Result<String, AppError> {
|
||||
let resp: serde_json::Value = stripe_post_form(
|
||||
secret,
|
||||
"/v1/checkout/sessions",
|
||||
vec![
|
||||
("mode".to_string(), "subscription".to_string()),
|
||||
("customer".to_string(), customer_id.to_string()),
|
||||
("line_items[0][price]".to_string(), price_id.to_string()),
|
||||
("line_items[0][quantity]".to_string(), "1".to_string()),
|
||||
("success_url".to_string(), success_url.to_string()),
|
||||
("cancel_url".to_string(), cancel_url.to_string()),
|
||||
("allow_promotion_codes".to_string(), "true".to_string()),
|
||||
("client_reference_id".to_string(), user_id.to_string()),
|
||||
("metadata[user_id]".to_string(), user_id.to_string()),
|
||||
],
|
||||
)
|
||||
.await?;
|
||||
|
||||
let url = resp
|
||||
.get("url")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| AppError::new(ErrorCode::Internal, "Stripe checkout 创建失败"))?;
|
||||
|
||||
Ok(url.to_string())
|
||||
}
|
||||
|
||||
async fn stripe_create_portal_session(
|
||||
secret: &str,
|
||||
customer_id: &str,
|
||||
return_url: &str,
|
||||
) -> Result<String, AppError> {
|
||||
let resp: serde_json::Value = stripe_post_form(
|
||||
secret,
|
||||
"/v1/billing_portal/sessions",
|
||||
vec![
|
||||
("customer".to_string(), customer_id.to_string()),
|
||||
("return_url".to_string(), return_url.to_string()),
|
||||
],
|
||||
)
|
||||
.await?;
|
||||
|
||||
let url = resp
|
||||
.get("url")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| AppError::new(ErrorCode::Internal, "Stripe portal 创建失败"))?;
|
||||
|
||||
Ok(url.to_string())
|
||||
}
|
||||
|
||||
async fn stripe_post_form(
|
||||
secret: &str,
|
||||
path: &str,
|
||||
form: Vec<(String, String)>,
|
||||
) -> Result<serde_json::Value, AppError> {
|
||||
let url = format!("https://api.stripe.com{path}");
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let resp = client
|
||||
.post(url)
|
||||
.bearer_auth(secret)
|
||||
.form(&form)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "Stripe 请求失败").with_source(err))?;
|
||||
|
||||
let status = resp.status();
|
||||
let body = resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "Stripe 响应读取失败").with_source(err))?;
|
||||
|
||||
if !status.is_success() {
|
||||
tracing::error!(status = %status, body = %body, "Stripe API error");
|
||||
return Err(AppError::new(ErrorCode::Internal, "Stripe API 调用失败"));
|
||||
}
|
||||
|
||||
serde_json::from_str(&body)
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "Stripe 响应解析失败").with_source(err))
|
||||
}
|
||||
1218
src/api/compress.rs
Normal file
1218
src/api/compress.rs
Normal file
File diff suppressed because it is too large
Load Diff
229
src/api/context.rs
Normal file
229
src/api/context.rs
Normal file
@@ -0,0 +1,229 @@
|
||||
use crate::auth;
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::state::AppState;
|
||||
|
||||
use axum::http::HeaderMap;
|
||||
use axum_extra::extract::cookie::{Cookie, CookieJar, SameSite};
|
||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||
use chrono::{DateTime, Utc};
|
||||
use hmac::{Hmac, Mac};
|
||||
use rand::RngCore;
|
||||
use serde::Serialize;
|
||||
use sha2::Sha256;
|
||||
use sqlx::FromRow;
|
||||
use std::net::IpAddr;
|
||||
use time::Duration as TimeDuration;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum Principal {
|
||||
Anonymous { session_id: String },
|
||||
User { user_id: Uuid, role: String, email_verified: bool },
|
||||
ApiKey {
|
||||
user_id: Uuid,
|
||||
api_key_id: Uuid,
|
||||
role: String,
|
||||
email_verified: bool,
|
||||
},
|
||||
}
|
||||
|
||||
pub fn client_ip(headers: &HeaderMap, connect_ip: IpAddr) -> IpAddr {
|
||||
if let Some(ip) = parse_forwarded_for(headers) {
|
||||
return ip;
|
||||
}
|
||||
if let Some(ip) = headers
|
||||
.get("x-real-ip")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.parse::<IpAddr>().ok())
|
||||
{
|
||||
return ip;
|
||||
}
|
||||
connect_ip
|
||||
}
|
||||
|
||||
fn parse_forwarded_for(headers: &HeaderMap) -> Option<IpAddr> {
|
||||
let value = headers.get("x-forwarded-for")?.to_str().ok()?;
|
||||
value
|
||||
.split(',')
|
||||
.next()
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
.and_then(|s| s.parse::<IpAddr>().ok())
|
||||
}
|
||||
|
||||
pub async fn authenticate(
|
||||
state: &AppState,
|
||||
jar: CookieJar,
|
||||
headers: &HeaderMap,
|
||||
ip: IpAddr,
|
||||
) -> Result<(CookieJar, Principal), AppError> {
|
||||
if let Some(principal) = try_jwt(state, headers).await? {
|
||||
return Ok((jar, principal));
|
||||
}
|
||||
if let Some(principal) = try_api_key(state, headers, ip).await? {
|
||||
return Ok((jar, principal));
|
||||
}
|
||||
|
||||
if !state.config.allow_anonymous_upload {
|
||||
return Err(AppError::new(ErrorCode::Unauthorized, "未登录"));
|
||||
}
|
||||
|
||||
let (jar, session_id) = ensure_session_cookie(jar);
|
||||
Ok((jar, Principal::Anonymous { session_id }))
|
||||
}
|
||||
|
||||
async fn try_jwt(state: &AppState, headers: &HeaderMap) -> Result<Option<Principal>, AppError> {
|
||||
let auth_header = headers
|
||||
.get(axum::http::header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("");
|
||||
|
||||
if !auth_header.starts_with("Bearer ") {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let claims = auth::require_jwt(&state.config.jwt_secret, headers)?;
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct UserAuthRow {
|
||||
id: Uuid,
|
||||
role: String,
|
||||
is_active: bool,
|
||||
email_verified_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
let user = sqlx::query_as::<_, UserAuthRow>(
|
||||
r#"
|
||||
SELECT id, role::text AS role, is_active, email_verified_at
|
||||
FROM users
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(claims.sub)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?
|
||||
.ok_or_else(|| AppError::new(ErrorCode::Unauthorized, "用户不存在或未登录"))?;
|
||||
|
||||
if !user.is_active {
|
||||
return Err(AppError::new(ErrorCode::Forbidden, "账号已被禁用"));
|
||||
}
|
||||
|
||||
Ok(Some(Principal::User {
|
||||
user_id: user.id,
|
||||
role: user.role,
|
||||
email_verified: user.email_verified_at.is_some(),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn try_api_key(
|
||||
state: &AppState,
|
||||
headers: &HeaderMap,
|
||||
ip: IpAddr,
|
||||
) -> Result<Option<Principal>, AppError> {
|
||||
let key = headers
|
||||
.get("x-api-key")
|
||||
.or_else(|| headers.get("X-API-Key"))
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|v| !v.is_empty());
|
||||
|
||||
let Some(full_key) = key else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let key_prefix = full_key
|
||||
.get(0..16)
|
||||
.ok_or_else(|| AppError::new(ErrorCode::Unauthorized, "API Key 格式错误"))?;
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct ApiKeyAuthRow {
|
||||
id: Uuid,
|
||||
user_id: Uuid,
|
||||
key_hash: String,
|
||||
is_active: bool,
|
||||
user_role: String,
|
||||
user_is_active: bool,
|
||||
email_verified_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
let row = sqlx::query_as::<_, ApiKeyAuthRow>(
|
||||
r#"
|
||||
SELECT
|
||||
k.id,
|
||||
k.user_id,
|
||||
k.key_hash,
|
||||
k.is_active,
|
||||
u.role::text AS user_role,
|
||||
u.is_active AS user_is_active,
|
||||
u.email_verified_at
|
||||
FROM api_keys k
|
||||
JOIN users u ON u.id = k.user_id
|
||||
WHERE k.key_prefix = $1
|
||||
"#,
|
||||
)
|
||||
.bind(key_prefix)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询 API Key 失败").with_source(err))?
|
||||
.ok_or_else(|| AppError::new(ErrorCode::Unauthorized, "API Key 无效"))?;
|
||||
|
||||
if !row.user_is_active || !row.is_active {
|
||||
return Err(AppError::new(ErrorCode::Forbidden, "API Key 已禁用"));
|
||||
}
|
||||
|
||||
let expected = api_key_hash(full_key, &state.config.api_key_pepper)?;
|
||||
if expected != row.key_hash {
|
||||
return Err(AppError::new(ErrorCode::Unauthorized, "API Key 无效"));
|
||||
}
|
||||
|
||||
let _ = sqlx::query("UPDATE api_keys SET last_used_at = NOW(), last_used_ip = $2 WHERE id = $1")
|
||||
.bind(row.id)
|
||||
.bind(ip.to_string())
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
Ok(Some(Principal::ApiKey {
|
||||
user_id: row.user_id,
|
||||
api_key_id: row.id,
|
||||
role: row.user_role,
|
||||
email_verified: row.email_verified_at.is_some(),
|
||||
}))
|
||||
}
|
||||
|
||||
pub fn ensure_session_cookie(jar: CookieJar) -> (CookieJar, String) {
|
||||
if let Some(cookie) = jar.get("if_session") {
|
||||
let session_id = cookie.value().trim().to_string();
|
||||
if !session_id.is_empty() {
|
||||
return (jar, session_id);
|
||||
}
|
||||
}
|
||||
|
||||
let session_id = generate_session_id();
|
||||
let cookie = Cookie::build(("if_session", session_id.clone()))
|
||||
.path("/")
|
||||
.http_only(true)
|
||||
.same_site(SameSite::Lax)
|
||||
.max_age(TimeDuration::days(7))
|
||||
.build();
|
||||
|
||||
(jar.add(cookie), session_id)
|
||||
}
|
||||
|
||||
fn generate_session_id() -> String {
|
||||
let mut bytes = [0u8; 32];
|
||||
rand::rngs::OsRng.fill_bytes(&mut bytes);
|
||||
URL_SAFE_NO_PAD.encode(bytes)
|
||||
}
|
||||
|
||||
pub fn api_key_hash(full_key: &str, pepper: &str) -> Result<String, AppError> {
|
||||
type HmacSha256 = Hmac<Sha256>;
|
||||
|
||||
let mut mac = HmacSha256::new_from_slice(pepper.as_bytes())
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "API Key pepper 错误").with_source(err))?;
|
||||
mac.update(full_key.as_bytes());
|
||||
let result = mac.finalize().into_bytes();
|
||||
Ok(hex::encode(result))
|
||||
}
|
||||
|
||||
343
src/api/downloads.rs
Normal file
343
src/api/downloads.rs
Normal file
@@ -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(())
|
||||
}
|
||||
8
src/api/envelope.rs
Normal file
8
src/api/envelope.rs
Normal file
@@ -0,0 +1,8 @@
|
||||
use serde::Serialize;
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct Envelope<T> {
|
||||
pub success: bool,
|
||||
pub data: T,
|
||||
}
|
||||
|
||||
39
src/api/health.rs
Normal file
39
src/api/health.rs
Normal file
@@ -0,0 +1,39 @@
|
||||
use crate::state::AppState;
|
||||
|
||||
use axum::{extract::State, http::StatusCode, response::IntoResponse, Json};
|
||||
use serde::Serialize;
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct HealthResponse {
|
||||
status: &'static str,
|
||||
database: &'static str,
|
||||
redis: &'static str,
|
||||
}
|
||||
|
||||
pub async fn health(State(state): State<AppState>) -> impl IntoResponse {
|
||||
let database_ok = sqlx::query("SELECT 1")
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.is_ok();
|
||||
|
||||
let mut redis_conn = state.redis.clone();
|
||||
let redis_ok = redis::cmd("PING")
|
||||
.query_async::<_, String>(&mut redis_conn)
|
||||
.await
|
||||
.is_ok();
|
||||
|
||||
let status = if database_ok && redis_ok {
|
||||
StatusCode::OK
|
||||
} else {
|
||||
StatusCode::SERVICE_UNAVAILABLE
|
||||
};
|
||||
|
||||
let body = HealthResponse {
|
||||
status: if status == StatusCode::OK { "healthy" } else { "unhealthy" },
|
||||
database: if database_ok { "connected" } else { "unavailable" },
|
||||
redis: if redis_ok { "connected" } else { "unavailable" },
|
||||
};
|
||||
|
||||
(status, Json(body))
|
||||
}
|
||||
|
||||
66
src/api/mod.rs
Normal file
66
src/api/mod.rs
Normal file
@@ -0,0 +1,66 @@
|
||||
mod auth;
|
||||
mod context;
|
||||
mod envelope;
|
||||
mod compress;
|
||||
mod downloads;
|
||||
mod billing;
|
||||
mod webhooks;
|
||||
mod user;
|
||||
mod tasks;
|
||||
mod admin;
|
||||
mod health;
|
||||
mod response;
|
||||
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::state::AppState;
|
||||
|
||||
use axum::extract::DefaultBodyLimit;
|
||||
use axum::Router;
|
||||
use std::net::SocketAddr;
|
||||
use tower_http::services::{ServeDir, ServeFile};
|
||||
use tower_http::trace::TraceLayer;
|
||||
|
||||
pub async fn run(state: AppState) -> Result<(), AppError> {
|
||||
let addr = format!("{}:{}", state.config.host, state.config.port);
|
||||
|
||||
if let Err(err) = crate::services::bootstrap::ensure_schema(&state).await {
|
||||
tracing::error!(error = %err, "数据库结构初始化失败");
|
||||
}
|
||||
if let Err(err) = crate::services::bootstrap::ensure_admin_user(&state).await {
|
||||
tracing::error!(error = %err, "管理员账号初始化失败");
|
||||
}
|
||||
|
||||
let static_service = ServeDir::new("static").not_found_service(ServeFile::new("static/index.html"));
|
||||
|
||||
let v1 = v1_router().layer(DefaultBodyLimit::max(100 * 1024 * 1024));
|
||||
|
||||
let app = Router::new()
|
||||
.route("/health", axum::routing::get(health::health))
|
||||
.nest("/downloads", downloads::router())
|
||||
.nest("/api/v1", v1)
|
||||
.fallback_service(static_service)
|
||||
.layer(TraceLayer::new_for_http())
|
||||
.with_state(state);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(&addr)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "监听端口失败").with_source(err))?;
|
||||
|
||||
tracing::info!(addr = %addr, "API server listening");
|
||||
|
||||
axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>())
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "HTTP 服务异常退出").with_source(err))
|
||||
}
|
||||
|
||||
fn v1_router() -> Router<AppState> {
|
||||
Router::new()
|
||||
.nest("/auth", auth::router())
|
||||
.merge(compress::router())
|
||||
.merge(tasks::router())
|
||||
.merge(billing::router())
|
||||
.merge(webhooks::router())
|
||||
.merge(user::router())
|
||||
.merge(admin::router())
|
||||
.fallback(response::not_found)
|
||||
}
|
||||
6
src/api/response.rs
Normal file
6
src/api/response.rs
Normal file
@@ -0,0 +1,6 @@
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
|
||||
pub async fn not_found() -> AppError {
|
||||
AppError::new(ErrorCode::NotFound, "接口不存在")
|
||||
}
|
||||
|
||||
987
src/api/tasks.rs
Normal file
987
src/api/tasks.rs
Normal file
@@ -0,0 +1,987 @@
|
||||
use crate::api::context;
|
||||
use crate::api::envelope::Envelope;
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::services::billing;
|
||||
use crate::services::billing::{BillingContext, Plan};
|
||||
use crate::services::compress;
|
||||
use crate::services::compress::{CompressionLevel, ImageFmt};
|
||||
use crate::services::idempotency;
|
||||
use crate::state::AppState;
|
||||
|
||||
use axum::extract::{ConnectInfo, Multipart, Path, State};
|
||||
use axum::http::HeaderMap;
|
||||
use axum::routing::{delete, get, post};
|
||||
use axum::{Json, Router};
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use sqlx::FromRow;
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use uuid::Uuid;
|
||||
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/compress/batch", post(create_batch_task))
|
||||
.route("/compress/tasks/{task_id}", get(get_task))
|
||||
.route("/compress/tasks/{task_id}/cancel", post(cancel_task))
|
||||
.route("/compress/tasks/{task_id}", delete(delete_task))
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct BatchCreateResponse {
|
||||
task_id: Uuid,
|
||||
total_files: i32,
|
||||
status: String,
|
||||
status_url: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct BatchFileInput {
|
||||
file_id: Uuid,
|
||||
original_name: String,
|
||||
original_format: ImageFmt,
|
||||
output_format: ImageFmt,
|
||||
original_size: u64,
|
||||
storage_path: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct BatchOptions {
|
||||
level: CompressionLevel,
|
||||
compression_rate: Option<u8>,
|
||||
output_format: Option<ImageFmt>,
|
||||
max_width: Option<u32>,
|
||||
max_height: Option<u32>,
|
||||
preserve_metadata: bool,
|
||||
}
|
||||
|
||||
async fn create_batch_task(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
mut multipart: Multipart,
|
||||
) -> Result<(axum_extra::extract::cookie::CookieJar, Json<Envelope<BatchCreateResponse>>), AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
if state.config.storage_type.to_ascii_lowercase() != "local" {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::StorageUnavailable,
|
||||
"当前仅支持本地存储（STORAGE_TYPE=local）",
|
||||
));
|
||||
}
|
||||
|
||||
let idempotency_key = headers
|
||||
.get("idempotency-key")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|v| !v.is_empty())
|
||||
.map(str::to_string);
|
||||
let idempotency_scope = match &principal {
|
||||
context::Principal::User { user_id, .. } => Some(idempotency::Scope::User(*user_id)),
|
||||
context::Principal::ApiKey { api_key_id, .. } => Some(idempotency::Scope::ApiKey(*api_key_id)),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
let task_id = Uuid::new_v4();
|
||||
let (files, opts, request_hash) = parse_batch_request(&state, task_id, &mut multipart).await?;
|
||||
|
||||
if files.is_empty() {
|
||||
cleanup_file_paths(&files).await;
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "缺少 files[]"));
|
||||
}
|
||||
|
||||
let mut idem_acquired = false;
|
||||
if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) {
|
||||
match idempotency::begin(
|
||||
&state,
|
||||
scope,
|
||||
idem_key,
|
||||
&request_hash,
|
||||
state.config.idempotency_ttl_hours as i64,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
idempotency::BeginResult::Replay { response_body, .. } => {
|
||||
cleanup_file_paths(&files).await;
|
||||
let resp: BatchCreateResponse =
|
||||
serde_json::from_value(response_body).map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "幂等结果解析失败").with_source(err)
|
||||
})?;
|
||||
return Ok((jar, Json(Envelope { success: true, data: resp })));
|
||||
}
|
||||
idempotency::BeginResult::InProgress => {
|
||||
cleanup_file_paths(&files).await;
|
||||
if let Some((_status, body)) = idempotency::wait_for_replay(
|
||||
&state,
|
||||
scope,
|
||||
idem_key,
|
||||
&request_hash,
|
||||
10_000,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
let resp: BatchCreateResponse =
|
||||
serde_json::from_value(body).map_err(|err| {
|
||||
AppError::new(ErrorCode::Internal, "幂等结果解析失败").with_source(err)
|
||||
})?;
|
||||
return Ok((jar, Json(Envelope { success: true, data: resp })));
|
||||
}
|
||||
return Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
"请求正在处理中，请稍后重试",
|
||||
));
|
||||
}
|
||||
idempotency::BeginResult::Acquired { .. } => {
|
||||
idem_acquired = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let create_result: Result<BatchCreateResponse, AppError> = (async {
|
||||
let (retention, task_owner, source) = match &principal {
|
||||
context::Principal::Anonymous { session_id } => {
|
||||
enforce_batch_limits_anonymous(&state, &files)?;
|
||||
let remaining = anonymous_remaining_units(&state, session_id, ip).await?;
|
||||
if remaining < files.len() as i64 {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::QuotaExceeded,
|
||||
"匿名试用次数已用完（每日 10 次）",
|
||||
));
|
||||
}
|
||||
Ok((
|
||||
Duration::hours(state.config.anon_retention_hours as i64),
|
||||
TaskOwner::Anonymous {
|
||||
session_id: session_id.clone(),
|
||||
},
|
||||
"web",
|
||||
))
|
||||
}
|
||||
context::Principal::User {
|
||||
user_id,
|
||||
email_verified,
|
||||
..
|
||||
} => {
|
||||
if !email_verified {
|
||||
return Err(AppError::new(ErrorCode::EmailNotVerified, "请先验证邮箱"));
|
||||
}
|
||||
let billing = billing::get_user_billing(&state, *user_id).await?;
|
||||
enforce_batch_limits_plan(&billing.plan, &files)?;
|
||||
ensure_quota_available(&state, &billing, files.len() as i32).await?;
|
||||
Ok((
|
||||
Duration::days(billing.plan.retention_days as i64),
|
||||
TaskOwner::User { user_id: *user_id },
|
||||
"web",
|
||||
))
|
||||
}
|
||||
context::Principal::ApiKey {
|
||||
user_id,
|
||||
api_key_id,
|
||||
email_verified,
|
||||
..
|
||||
} => {
|
||||
if !email_verified {
|
||||
return Err(AppError::new(ErrorCode::EmailNotVerified, "请先验证邮箱"));
|
||||
}
|
||||
let billing = billing::get_user_billing(&state, *user_id).await?;
|
||||
if !billing.plan.feature_api_enabled {
|
||||
return Err(AppError::new(ErrorCode::Forbidden, "当前套餐未开通 API"));
|
||||
}
|
||||
enforce_batch_limits_plan(&billing.plan, &files)?;
|
||||
ensure_quota_available(&state, &billing, files.len() as i32).await?;
|
||||
Ok((
|
||||
Duration::days(billing.plan.retention_days as i64),
|
||||
TaskOwner::ApiKey {
|
||||
user_id: *user_id,
|
||||
api_key_id: *api_key_id,
|
||||
},
|
||||
"api",
|
||||
))
|
||||
}
|
||||
}?;
|
||||
|
||||
tokio::fs::create_dir_all(&state.config.storage_path)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
AppError::new(ErrorCode::StorageUnavailable, "创建存储目录失败").with_source(err)
|
||||
})?;
|
||||
|
||||
let expires_at = Utc::now() + retention;
|
||||
let (user_id, session_id, api_key_id) = match &task_owner {
|
||||
TaskOwner::Anonymous { session_id } => (None, Some(session_id.clone()), None),
|
||||
TaskOwner::User { user_id } => (Some(*user_id), None, None),
|
||||
TaskOwner::ApiKey { user_id, api_key_id } => (Some(*user_id), None, Some(*api_key_id)),
|
||||
};
|
||||
|
||||
let total_original_size: i64 = files.iter().map(|f| f.original_size as i64).sum();
|
||||
|
||||
let mut tx = state
|
||||
.db
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO tasks (
|
||||
id, user_id, session_id, api_key_id, client_ip, source, status,
|
||||
compression_level, output_format, max_width, max_height, preserve_metadata,
|
||||
compression_rate,
|
||||
total_files, completed_files, failed_files,
|
||||
total_original_size, total_compressed_size,
|
||||
expires_at
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5::inet, $6::task_source, 'pending',
|
||||
$7::compression_level, $8, $9, $10, $11, $12,
|
||||
$13, 0, 0,
|
||||
$14, 0,
|
||||
$15
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(user_id)
|
||||
.bind(session_id)
|
||||
.bind(api_key_id)
|
||||
.bind(ip.to_string())
|
||||
.bind(source)
|
||||
.bind(opts.level.as_str())
|
||||
.bind(opts.output_format.map(|f| f.as_str()))
|
||||
.bind(opts.max_width.map(|v| v as i32))
|
||||
.bind(opts.max_height.map(|v| v as i32))
|
||||
.bind(false)
|
||||
.bind(opts.compression_rate.map(|v| v as i16))
|
||||
.bind(files.len() as i32)
|
||||
.bind(total_original_size)
|
||||
.bind(expires_at)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "创建任务失败").with_source(err))?;
|
||||
|
||||
for file in &files {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO task_files (
|
||||
id, task_id,
|
||||
original_name, original_format, output_format,
|
||||
original_size,
|
||||
storage_path, status
|
||||
) VALUES (
|
||||
$1, $2,
|
||||
$3, $4, $5,
|
||||
$6,
|
||||
$7, 'pending'
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(file.file_id)
|
||||
.bind(task_id)
|
||||
.bind(&file.original_name)
|
||||
.bind(file.original_format.as_str())
|
||||
.bind(file.output_format.as_str())
|
||||
.bind(file.original_size as i64)
|
||||
.bind(&file.storage_path)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "创建文件记录失败").with_source(err))?;
|
||||
}
|
||||
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;
|
||||
|
||||
if let Err(err) = enqueue_task(&state, task_id).await {
|
||||
let _ = sqlx::query("UPDATE tasks SET status = 'failed', error_message = $2 WHERE id = $1")
|
||||
.bind(task_id)
|
||||
.bind("队列提交失败")
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
return Err(err);
|
||||
}
|
||||
|
||||
Ok(BatchCreateResponse {
|
||||
task_id,
|
||||
total_files: files.len() as i32,
|
||||
status: "pending".to_string(),
|
||||
status_url: format!("/api/v1/compress/tasks/{task_id}"),
|
||||
})
|
||||
})
|
||||
.await;
|
||||
|
||||
match create_result {
|
||||
Ok(resp) => {
|
||||
if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) {
|
||||
if idem_acquired {
|
||||
let _ = idempotency::complete(
|
||||
&state,
|
||||
scope,
|
||||
idem_key,
|
||||
&request_hash,
|
||||
200,
|
||||
serde_json::to_value(&resp).unwrap_or(serde_json::Value::Null),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
Ok((jar, Json(Envelope { success: true, data: resp })))
|
||||
}
|
||||
Err(err) => {
|
||||
if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) {
|
||||
if idem_acquired {
|
||||
let _ = idempotency::abort(&state, scope, idem_key, &request_hash).await;
|
||||
}
|
||||
}
|
||||
cleanup_file_paths(&files).await;
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum TaskOwner {
|
||||
Anonymous { session_id: String },
|
||||
User { user_id: Uuid },
|
||||
ApiKey { user_id: Uuid, api_key_id: Uuid },
|
||||
}
|
||||
|
||||
async fn enqueue_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
|
||||
let mut conn = state.redis.clone();
|
||||
let now = Utc::now().to_rfc3339();
|
||||
redis::cmd("XADD")
|
||||
.arg("stream:compress_jobs")
|
||||
.arg("*")
|
||||
.arg("task_id")
|
||||
.arg(task_id.to_string())
|
||||
.arg("created_at")
|
||||
.arg(now)
|
||||
.query_async::<_, redis::Value>(&mut conn)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "写入队列失败").with_source(err))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn parse_batch_request(
|
||||
state: &AppState,
|
||||
task_id: Uuid,
|
||||
multipart: &mut Multipart,
|
||||
) -> Result<(Vec<BatchFileInput>, BatchOptions, String), AppError> {
|
||||
let mut files: Vec<BatchFileInput> = Vec::new();
|
||||
let mut file_digests: Vec<String> = Vec::new();
|
||||
let mut opts = BatchOptions {
|
||||
level: CompressionLevel::Medium,
|
||||
compression_rate: None,
|
||||
output_format: None,
|
||||
max_width: None,
|
||||
max_height: None,
|
||||
preserve_metadata: false,
|
||||
};
|
||||
|
||||
let base_dir = format!("{}/orig/{task_id}", state.config.storage_path);
|
||||
tokio::fs::create_dir_all(&base_dir)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::StorageUnavailable, "创建存储目录失败").with_source(err))?;
|
||||
|
||||
loop {
|
||||
let next = multipart.next_field().await.map_err(|err| {
|
||||
AppError::new(ErrorCode::InvalidRequest, "读取上传内容失败").with_source(err)
|
||||
});
|
||||
|
||||
let field = match next {
|
||||
Ok(v) => v,
|
||||
Err(err) => {
|
||||
cleanup_file_paths(&files).await;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
|
||||
let Some(field) = field else { break };
|
||||
|
||||
let name = field.name().unwrap_or("").to_string();
|
||||
if name == "files" || name == "files[]" {
|
||||
let file_id = Uuid::new_v4();
|
||||
let original_name = field.file_name().unwrap_or("upload").to_string();
|
||||
let bytes = match field.bytes().await {
|
||||
Ok(v) => v,
|
||||
Err(err) => {
|
||||
cleanup_file_paths(&files).await;
|
||||
return Err(
|
||||
AppError::new(ErrorCode::InvalidRequest, "读取文件失败").with_source(err)
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
let original_size = bytes.len() as u64;
|
||||
let file_digest = {
|
||||
let mut h = Sha256::new();
|
||||
h.update(&bytes);
|
||||
h.update(original_name.as_bytes());
|
||||
hex::encode(h.finalize())
|
||||
};
|
||||
let original_format = match compress::detect_format(&bytes) {
|
||||
Ok(v) => v,
|
||||
Err(err) => {
|
||||
cleanup_file_paths(&files).await;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
|
||||
let output_format = opts.output_format.unwrap_or(original_format);
|
||||
let path = format!("{base_dir}/{file_id}.{}", original_format.extension());
|
||||
|
||||
let mut f = match tokio::fs::File::create(&path).await {
|
||||
Ok(v) => v,
|
||||
Err(err) => {
|
||||
cleanup_file_paths(&files).await;
|
||||
return Err(
|
||||
AppError::new(ErrorCode::StorageUnavailable, "写入文件失败").with_source(err),
|
||||
);
|
||||
}
|
||||
};
|
||||
if let Err(err) = f.write_all(&bytes).await {
|
||||
let _ = tokio::fs::remove_file(&path).await;
|
||||
cleanup_file_paths(&files).await;
|
||||
return Err(
|
||||
AppError::new(ErrorCode::StorageUnavailable, "写入文件失败").with_source(err),
|
||||
);
|
||||
}
|
||||
|
||||
files.push(BatchFileInput {
|
||||
file_id,
|
||||
original_name,
|
||||
original_format,
|
||||
output_format,
|
||||
original_size,
|
||||
storage_path: path,
|
||||
});
|
||||
file_digests.push(file_digest);
|
||||
continue;
|
||||
}
|
||||
|
||||
let text = match field.text().await {
|
||||
Ok(v) => v,
|
||||
Err(err) => {
|
||||
cleanup_file_paths(&files).await;
|
||||
return Err(
|
||||
AppError::new(ErrorCode::InvalidRequest, "读取字段失败").with_source(err),
|
||||
);
|
||||
}
|
||||
};
|
||||
match name.as_str() {
|
||||
"level" => {
|
||||
opts.level = match compress::parse_level(&text) {
|
||||
Ok(v) => v,
|
||||
Err(err) => {
|
||||
cleanup_file_paths(&files).await;
|
||||
return Err(err);
|
||||
}
|
||||
}
|
||||
}
|
||||
"output_format" => {
|
||||
let v = text.trim();
|
||||
if !v.is_empty() {
|
||||
opts.output_format = Some(match compress::parse_output_format(v) {
|
||||
Ok(v) => v,
|
||||
Err(err) => {
|
||||
cleanup_file_paths(&files).await;
|
||||
return Err(err);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
"compression_rate" | "quality" => {
|
||||
let v = text.trim();
|
||||
if !v.is_empty() {
|
||||
opts.compression_rate = Some(match compress::parse_compression_rate(v) {
|
||||
Ok(v) => v,
|
||||
Err(err) => {
|
||||
cleanup_file_paths(&files).await;
|
||||
return Err(err);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
"max_width" => {
|
||||
let v = text.trim();
|
||||
if !v.is_empty() {
|
||||
opts.max_width = Some(match v.parse::<u32>() {
|
||||
Ok(n) => n,
|
||||
Err(_) => {
|
||||
cleanup_file_paths(&files).await;
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "max_width 格式错误"));
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
"max_height" => {
|
||||
let v = text.trim();
|
||||
if !v.is_empty() {
|
||||
opts.max_height = Some(match v.parse::<u32>() {
|
||||
Ok(n) => n,
|
||||
Err(_) => {
|
||||
cleanup_file_paths(&files).await;
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "max_height 格式错误"));
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
"preserve_metadata" => {
|
||||
opts.preserve_metadata = matches!(
|
||||
text.trim().to_ascii_lowercase().as_str(),
|
||||
"1" | "true" | "yes" | "y" | "on"
|
||||
);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(rate) = opts.compression_rate {
|
||||
opts.level = compress::rate_to_level(rate);
|
||||
}
|
||||
|
||||
if opts.output_format.is_some() {
|
||||
cleanup_file_paths(&files).await;
|
||||
return Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
"当前仅支持保持原图片格式",
|
||||
));
|
||||
}
|
||||
|
||||
let mw = opts.max_width.map(|v| v.to_string()).unwrap_or_default();
|
||||
let mh = opts.max_height.map(|v| v.to_string()).unwrap_or_default();
|
||||
let out_fmt = opts.output_format.map(|f| f.as_str()).unwrap_or("");
|
||||
let rate_key = opts
|
||||
.compression_rate
|
||||
.map(|v| v.to_string())
|
||||
.unwrap_or_default();
|
||||
let preserve = if opts.preserve_metadata { "1" } else { "0" };
|
||||
|
||||
let mut h = Sha256::new();
|
||||
h.update(b"compress_batch_v1");
|
||||
h.update(opts.level.as_str().as_bytes());
|
||||
h.update(out_fmt.as_bytes());
|
||||
h.update(rate_key.as_bytes());
|
||||
h.update(mw.as_bytes());
|
||||
h.update(mh.as_bytes());
|
||||
h.update(preserve.as_bytes());
|
||||
for d in &file_digests {
|
||||
h.update(d.as_bytes());
|
||||
}
|
||||
let request_hash = hex::encode(h.finalize());
|
||||
|
||||
Ok((files, opts, request_hash))
|
||||
}
|
||||
|
||||
fn enforce_batch_limits_anonymous(state: &AppState, files: &[BatchFileInput]) -> Result<(), AppError> {
|
||||
let max_files = state.config.anon_max_files_per_batch as usize;
|
||||
if files.len() > max_files {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
format!("匿名试用单次最多 {} 个文件", max_files),
|
||||
));
|
||||
}
|
||||
|
||||
let max_bytes = state.config.anon_max_file_size_mb * 1024 * 1024;
|
||||
for f in files {
|
||||
if f.original_size > max_bytes {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::FileTooLarge,
|
||||
format!("匿名试用单文件最大 {} MB", state.config.anon_max_file_size_mb),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn enforce_batch_limits_plan(plan: &Plan, files: &[BatchFileInput]) -> Result<(), AppError> {
|
||||
let max_files = plan.max_files_per_batch as usize;
|
||||
if files.len() > max_files {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
format!("当前套餐单次最多 {} 个文件", plan.max_files_per_batch),
|
||||
));
|
||||
}
|
||||
|
||||
let max_bytes = (plan.max_file_size_mb as u64) * 1024 * 1024;
|
||||
for f in files {
|
||||
if f.original_size > max_bytes {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::FileTooLarge,
|
||||
format!("当前套餐单文件最大 {} MB", plan.max_file_size_mb),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn ensure_quota_available(
|
||||
state: &AppState,
|
||||
ctx: &BillingContext,
|
||||
needed_units: i32,
|
||||
) -> Result<(), AppError> {
|
||||
if needed_units <= 0 {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct UsageRow {
|
||||
used_units: i32,
|
||||
bonus_units: i32,
|
||||
}
|
||||
|
||||
let usage = sqlx::query_as::<_, UsageRow>(
|
||||
r#"
|
||||
SELECT used_units, bonus_units
|
||||
FROM usage_periods
|
||||
WHERE user_id = $1 AND period_start = $2 AND period_end = $3
|
||||
"#,
|
||||
)
|
||||
.bind(ctx.user_id)
|
||||
.bind(ctx.period_start)
|
||||
.bind(ctx.period_end)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用量失败").with_source(err))?
|
||||
.unwrap_or(UsageRow {
|
||||
used_units: 0,
|
||||
bonus_units: 0,
|
||||
});
|
||||
|
||||
let total_units = ctx.plan.included_units_per_period + usage.bonus_units;
|
||||
let remaining = total_units - usage.used_units;
|
||||
if remaining < needed_units {
|
||||
return Err(AppError::new(ErrorCode::QuotaExceeded, "当期配额已用完"));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn anonymous_remaining_units(
|
||||
state: &AppState,
|
||||
session_id: &str,
|
||||
ip: IpAddr,
|
||||
) -> Result<i64, AppError> {
|
||||
let date = utc8_date();
|
||||
let session_key = format!("anon_quota:{session_id}:{date}");
|
||||
let ip_key = format!("anon_quota_ip:{ip}:{date}");
|
||||
|
||||
let mut conn = state.redis.clone();
|
||||
let v1: Option<i64> = redis::cmd("GET")
|
||||
.arg(session_key)
|
||||
.query_async(&mut conn)
|
||||
.await
|
||||
.unwrap_or(None);
|
||||
let v2: Option<i64> = redis::cmd("GET")
|
||||
.arg(ip_key)
|
||||
.query_async(&mut conn)
|
||||
.await
|
||||
.unwrap_or(None);
|
||||
|
||||
let limit = state.config.anon_daily_units as i64;
|
||||
Ok(std::cmp::min(limit - v1.unwrap_or(0), limit - v2.unwrap_or(0)))
|
||||
}
|
||||
|
||||
fn utc8_date() -> String {
|
||||
let now = Utc::now() + Duration::hours(8);
|
||||
now.format("%Y-%m-%d").to_string()
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct TaskRow {
|
||||
id: Uuid,
|
||||
status: String,
|
||||
total_files: i32,
|
||||
completed_files: i32,
|
||||
failed_files: i32,
|
||||
created_at: DateTime<Utc>,
|
||||
completed_at: Option<DateTime<Utc>>,
|
||||
expires_at: DateTime<Utc>,
|
||||
user_id: Option<Uuid>,
|
||||
session_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct TaskFileRow {
|
||||
id: Uuid,
|
||||
original_name: String,
|
||||
original_size: i64,
|
||||
compressed_size: Option<i64>,
|
||||
saved_percent: Option<f64>,
|
||||
status: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct TaskFileView {
|
||||
file_id: Uuid,
|
||||
original_name: String,
|
||||
original_size: i64,
|
||||
compressed_size: Option<i64>,
|
||||
saved_percent: Option<f64>,
|
||||
status: String,
|
||||
download_url: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct TaskView {
|
||||
task_id: Uuid,
|
||||
status: String,
|
||||
progress: i32,
|
||||
total_files: i32,
|
||||
completed_files: i32,
|
||||
failed_files: i32,
|
||||
files: Vec<TaskFileView>,
|
||||
download_all_url: String,
|
||||
created_at: DateTime<Utc>,
|
||||
completed_at: Option<DateTime<Utc>>,
|
||||
expires_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
async fn get_task(
|
||||
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, Json<Envelope<TaskView>>), AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let task = sqlx::query_as::<_, TaskRow>(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
status::text AS status,
|
||||
total_files,
|
||||
completed_files,
|
||||
failed_files,
|
||||
created_at,
|
||||
completed_at,
|
||||
expires_at,
|
||||
user_id,
|
||||
session_id
|
||||
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, "任务已过期或不存在"));
|
||||
}
|
||||
|
||||
authorize_task(&principal, task.user_id, task.session_id.as_deref().unwrap_or(""))?;
|
||||
|
||||
let files = sqlx::query_as::<_, TaskFileRow>(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
original_name,
|
||||
original_size,
|
||||
compressed_size,
|
||||
saved_percent::float8 AS saved_percent,
|
||||
status::text AS status
|
||||
FROM task_files
|
||||
WHERE task_id = $1
|
||||
ORDER BY created_at ASC
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询任务文件失败").with_source(err))?;
|
||||
|
||||
let file_views = files
|
||||
.into_iter()
|
||||
.map(|f| TaskFileView {
|
||||
file_id: f.id,
|
||||
original_name: f.original_name,
|
||||
original_size: f.original_size,
|
||||
compressed_size: f.compressed_size,
|
||||
saved_percent: f.saved_percent,
|
||||
status: f.status.clone(),
|
||||
download_url: if f.status == "completed" {
|
||||
Some(format!("/downloads/{}", f.id))
|
||||
} else {
|
||||
None
|
||||
},
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let processed = task.completed_files + task.failed_files;
|
||||
let progress = if task.total_files <= 0 {
|
||||
0
|
||||
} else {
|
||||
((processed as f64) * 100.0 / (task.total_files as f64))
|
||||
.round()
|
||||
.clamp(0.0, 100.0) as i32
|
||||
};
|
||||
|
||||
Ok((
|
||||
jar,
|
||||
Json(Envelope {
|
||||
success: true,
|
||||
data: TaskView {
|
||||
task_id,
|
||||
status: task.status,
|
||||
progress,
|
||||
total_files: task.total_files,
|
||||
completed_files: task.completed_files,
|
||||
failed_files: task.failed_files,
|
||||
files: file_views,
|
||||
download_all_url: format!("/downloads/tasks/{task_id}"),
|
||||
created_at: task.created_at,
|
||||
completed_at: task.completed_at,
|
||||
expires_at: task.expires_at,
|
||||
},
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
async fn cancel_task(
|
||||
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, Json<Envelope<serde_json::Value>>), AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let task = sqlx::query_as::<_, TaskRow>(
|
||||
"SELECT id, status::text AS status, total_files, completed_files, failed_files, created_at, completed_at, expires_at, user_id, session_id 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, "任务不存在"))?;
|
||||
|
||||
authorize_task(&principal, task.user_id, task.session_id.as_deref().unwrap_or(""))?;
|
||||
|
||||
if matches!(task.status.as_str(), "completed" | "failed" | "cancelled") {
|
||||
return Ok((
|
||||
jar,
|
||||
Json(Envelope {
|
||||
success: true,
|
||||
data: serde_json::json!({ "message": "任务已结束" }),
|
||||
}),
|
||||
));
|
||||
}
|
||||
|
||||
let updated = sqlx::query(
|
||||
"UPDATE tasks SET status = 'cancelled', completed_at = NOW() WHERE id = $1 AND status IN ('pending', 'processing')",
|
||||
)
|
||||
.bind(task_id)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "取消任务失败").with_source(err))?;
|
||||
|
||||
if updated.rows_affected() == 0 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "任务状态不可取消"));
|
||||
}
|
||||
|
||||
Ok((
|
||||
jar,
|
||||
Json(Envelope {
|
||||
success: true,
|
||||
data: serde_json::json!({ "message": "已取消" }),
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
async fn delete_task(
|
||||
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, Json<Envelope<serde_json::Value>>), AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let task = sqlx::query_as::<_, TaskRow>(
|
||||
"SELECT id, status::text AS status, total_files, completed_files, failed_files, created_at, completed_at, expires_at, user_id, session_id 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, "任务不存在"))?;
|
||||
|
||||
authorize_task(&principal, task.user_id, task.session_id.as_deref().unwrap_or(""))?;
|
||||
|
||||
if task.status == "processing" {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
"任务处理中，请先取消后再删除",
|
||||
));
|
||||
}
|
||||
|
||||
let paths: Vec<Option<String>> =
|
||||
sqlx::query_scalar("SELECT storage_path FROM task_files WHERE task_id = $1")
|
||||
.bind(task_id)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询文件失败").with_source(err))?;
|
||||
|
||||
for p in paths.into_iter().flatten() {
|
||||
let _ = tokio::fs::remove_file(p).await;
|
||||
}
|
||||
|
||||
if state.config.storage_type.to_ascii_lowercase() == "local" {
|
||||
let zip_path = format!("{}/zips/{task_id}.zip", state.config.storage_path);
|
||||
let _ = tokio::fs::remove_file(zip_path).await;
|
||||
let orig_dir = format!("{}/orig/{task_id}", state.config.storage_path);
|
||||
let _ = tokio::fs::remove_dir_all(orig_dir).await;
|
||||
}
|
||||
|
||||
let deleted = sqlx::query("DELETE FROM tasks WHERE id = $1")
|
||||
.bind(task_id)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "删除任务失败").with_source(err))?;
|
||||
|
||||
if deleted.rows_affected() == 0 {
|
||||
return Err(AppError::new(ErrorCode::NotFound, "任务不存在"));
|
||||
}
|
||||
|
||||
Ok((
|
||||
jar,
|
||||
Json(Envelope {
|
||||
success: true,
|
||||
data: serde_json::json!({ "message": "已删除" }),
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
fn authorize_task(
|
||||
principal: &context::Principal,
|
||||
user_id: Option<Uuid>,
|
||||
session_id: &str,
|
||||
) -> Result<(), AppError> {
|
||||
if let Some(owner) = user_id {
|
||||
match principal {
|
||||
context::Principal::User { user_id: me, .. } if *me == owner => Ok(()),
|
||||
context::Principal::ApiKey { user_id: me, .. } if *me == owner => Ok(()),
|
||||
_ => Err(AppError::new(ErrorCode::Forbidden, "无权限访问该任务")),
|
||||
}
|
||||
} else {
|
||||
match principal {
|
||||
context::Principal::Anonymous { session_id: sid } if sid == session_id => Ok(()),
|
||||
_ => Err(AppError::new(ErrorCode::Forbidden, "无权限访问该任务")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn cleanup_file_paths(files: &[BatchFileInput]) {
|
||||
for f in files {
|
||||
let _ = tokio::fs::remove_file(&f.storage_path).await;
|
||||
}
|
||||
}
|
||||
912
src/api/user.rs
Normal file
912
src/api/user.rs
Normal file
@@ -0,0 +1,912 @@
|
||||
use crate::api::context;
|
||||
use crate::api::envelope::Envelope;
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::services::billing;
|
||||
use crate::services::mail;
|
||||
use crate::state::AppState;
|
||||
|
||||
use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier};
|
||||
use axum::extract::{ConnectInfo, Path, Query, State};
|
||||
use axum::http::HeaderMap;
|
||||
use axum::routing::{delete, get, post, put};
|
||||
use axum::{Json, Router};
|
||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use rand::RngCore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use sqlx::FromRow;
|
||||
use std::net::SocketAddr;
|
||||
use uuid::Uuid;
|
||||
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/user/profile", get(get_profile))
|
||||
.route("/user/profile", put(update_profile))
|
||||
.route("/user/password", put(update_password))
|
||||
.route("/user/history", get(list_history))
|
||||
.route("/user/api-keys", get(list_api_keys))
|
||||
.route("/user/api-keys", post(create_api_key))
|
||||
.route("/user/api-keys/{key_id}/rotate", post(rotate_api_key))
|
||||
.route("/user/api-keys/{key_id}", delete(disable_api_key))
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow, Serialize)]
|
||||
struct ApiKeyView {
|
||||
id: Uuid,
|
||||
name: String,
|
||||
key_prefix: String,
|
||||
permissions: serde_json::Value,
|
||||
rate_limit: i32,
|
||||
is_active: bool,
|
||||
last_used_at: Option<chrono::DateTime<chrono::Utc>>,
|
||||
last_used_ip: Option<String>,
|
||||
created_at: chrono::DateTime<chrono::Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct ApiKeyListResponse {
|
||||
api_keys: Vec<ApiKeyView>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct UserView {
|
||||
id: Uuid,
|
||||
email: String,
|
||||
username: String,
|
||||
role: String,
|
||||
email_verified: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct MessageResponse {
|
||||
message: String,
|
||||
}
|
||||
|
||||
async fn get_profile(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<Json<Envelope<UserView>>, AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (_jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let user_id = match principal {
|
||||
context::Principal::User { user_id, .. } => user_id,
|
||||
_ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
|
||||
};
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct UserRow {
|
||||
id: Uuid,
|
||||
email: String,
|
||||
username: String,
|
||||
role: String,
|
||||
email_verified_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
let user = sqlx::query_as::<_, UserRow>(
|
||||
r#"
|
||||
SELECT id, email, username, role::text AS role, email_verified_at
|
||||
FROM users
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?;
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: UserView {
|
||||
id: user.id,
|
||||
email: user.email,
|
||||
username: user.username,
|
||||
role: user.role,
|
||||
email_verified: user.email_verified_at.is_some(),
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct UpdateProfileRequest {
|
||||
email: Option<String>,
|
||||
username: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct UpdateProfileResponse {
|
||||
user: UserView,
|
||||
message: String,
|
||||
}
|
||||
|
||||
async fn update_profile(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
Json(req): Json<UpdateProfileRequest>,
|
||||
) -> Result<Json<Envelope<UpdateProfileResponse>>, AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (_jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let user_id = match principal {
|
||||
context::Principal::User { user_id, .. } => user_id,
|
||||
_ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
|
||||
};
|
||||
|
||||
if req.email.is_none() && req.username.is_none() {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "未提供可更新字段"));
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct UserRow {
|
||||
id: Uuid,
|
||||
email: String,
|
||||
username: String,
|
||||
role: String,
|
||||
email_verified_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
let user = sqlx::query_as::<_, UserRow>(
|
||||
r#"
|
||||
SELECT id, email, username, role::text AS role, email_verified_at
|
||||
FROM users
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?;
|
||||
|
||||
let mut next_email = user.email.clone();
|
||||
let mut next_username = user.username.clone();
|
||||
let mut email_changed = false;
|
||||
|
||||
if let Some(email) = req.email.as_ref() {
|
||||
let email = email.trim().to_lowercase();
|
||||
validate_email(&email)?;
|
||||
if email != user.email {
|
||||
next_email = email;
|
||||
email_changed = true;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(username) = req.username.as_ref() {
|
||||
let username = username.trim().to_string();
|
||||
validate_username(&username)?;
|
||||
if username != user.username {
|
||||
next_username = username;
|
||||
}
|
||||
}
|
||||
|
||||
if next_email == user.email && next_username == user.username {
|
||||
return Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: UpdateProfileResponse {
|
||||
user: UserView {
|
||||
id: user.id,
|
||||
email: user.email,
|
||||
username: user.username,
|
||||
role: user.role,
|
||||
email_verified: user.email_verified_at.is_some(),
|
||||
},
|
||||
message: "暂无更新".to_string(),
|
||||
},
|
||||
}));
|
||||
}
|
||||
|
||||
let mut tx = state
|
||||
.db
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?;
|
||||
|
||||
let email_verified_at = if email_changed { None } else { user.email_verified_at };
|
||||
|
||||
let updated = sqlx::query_as::<_, UserRow>(
|
||||
r#"
|
||||
UPDATE users
|
||||
SET email = $2,
|
||||
username = $3,
|
||||
email_verified_at = $4,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
RETURNING id, email, username, role::text AS role, email_verified_at
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(&next_email)
|
||||
.bind(&next_username)
|
||||
.bind(email_verified_at)
|
||||
.fetch_one(&mut *tx)
|
||||
.await
|
||||
.map_err(map_unique_violation)?;
|
||||
|
||||
let mut verification_link: Option<String> = None;
|
||||
if email_changed {
|
||||
let token = generate_token();
|
||||
let token_hash = sha256_hex(&token);
|
||||
let expires_at = Utc::now() + Duration::hours(24);
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO email_verifications (user_id, token_hash, expires_at)
|
||||
VALUES ($1, $2, $3)
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(token_hash)
|
||||
.bind(expires_at)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "创建邮箱验证记录失败").with_source(err))?;
|
||||
|
||||
verification_link = Some(format!(
|
||||
"{}/verify-email?token={}",
|
||||
state.config.public_base_url, token
|
||||
));
|
||||
}
|
||||
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;
|
||||
|
||||
if let Some(link) = verification_link.as_deref() {
|
||||
mail::send_verification_email(&state, &updated.email, &updated.username, link)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::MailSendFailed, "验证邮件发送失败").with_source(err))?;
|
||||
}
|
||||
|
||||
let message = if email_changed {
|
||||
"资料已更新，请验证新邮箱".to_string()
|
||||
} else {
|
||||
"资料已更新".to_string()
|
||||
};
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: UpdateProfileResponse {
|
||||
user: UserView {
|
||||
id: updated.id,
|
||||
email: updated.email,
|
||||
username: updated.username,
|
||||
role: updated.role,
|
||||
email_verified: updated.email_verified_at.is_some(),
|
||||
},
|
||||
message,
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct UpdatePasswordRequest {
|
||||
current_password: String,
|
||||
new_password: String,
|
||||
}
|
||||
|
||||
async fn update_password(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
Json(req): Json<UpdatePasswordRequest>,
|
||||
) -> Result<Json<Envelope<MessageResponse>>, AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (_jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let user_id = match principal {
|
||||
context::Principal::User { user_id, .. } => user_id,
|
||||
_ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
|
||||
};
|
||||
|
||||
validate_password(&req.new_password)?;
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct PasswordRow {
|
||||
password_hash: String,
|
||||
}
|
||||
|
||||
let row = sqlx::query_as::<_, PasswordRow>("SELECT password_hash FROM users WHERE id = $1")
|
||||
.bind(user_id)
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?;
|
||||
|
||||
verify_password(&req.current_password, &row.password_hash)?;
|
||||
|
||||
let new_hash = hash_password(&req.new_password)?;
|
||||
sqlx::query("UPDATE users SET password_hash = $2, updated_at = NOW() WHERE id = $1")
|
||||
.bind(user_id)
|
||||
.bind(new_hash)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新密码失败").with_source(err))?;
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: MessageResponse {
|
||||
message: "密码已更新，请重新登录以确保安全".to_string(),
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct HistoryQuery {
|
||||
page: Option<u32>,
|
||||
limit: Option<u32>,
|
||||
status: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct HistoryFileView {
|
||||
file_id: Uuid,
|
||||
original_name: String,
|
||||
original_size: i64,
|
||||
compressed_size: Option<i64>,
|
||||
saved_percent: Option<f64>,
|
||||
status: String,
|
||||
output_format: String,
|
||||
error_message: Option<String>,
|
||||
download_url: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct HistoryTaskView {
|
||||
task_id: Uuid,
|
||||
status: String,
|
||||
source: String,
|
||||
progress: i32,
|
||||
total_files: i32,
|
||||
completed_files: i32,
|
||||
failed_files: i32,
|
||||
created_at: DateTime<Utc>,
|
||||
completed_at: Option<DateTime<Utc>>,
|
||||
expires_at: DateTime<Utc>,
|
||||
download_all_url: Option<String>,
|
||||
files: Vec<HistoryFileView>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct HistoryResponse {
|
||||
tasks: Vec<HistoryTaskView>,
|
||||
page: u32,
|
||||
limit: u32,
|
||||
total: i64,
|
||||
}
|
||||
|
||||
async fn list_history(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<HistoryQuery>,
|
||||
) -> Result<Json<Envelope<HistoryResponse>>, AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (_jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let user_id = match principal {
|
||||
context::Principal::User { user_id, .. } => user_id,
|
||||
_ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
|
||||
};
|
||||
|
||||
let limit = query.limit.unwrap_or(20).clamp(1, 100);
|
||||
let page = query.page.unwrap_or(1).max(1);
|
||||
let offset = (page - 1) * limit;
|
||||
let status = query.status.map(|s| s.trim().to_string()).filter(|s| !s.is_empty());
|
||||
|
||||
let total: i64 = if let Some(status) = &status {
|
||||
sqlx::query_scalar(
|
||||
"SELECT COUNT(*) FROM tasks WHERE user_id = $1 AND status::text = $2",
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(status)
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询历史失败").with_source(err))?
|
||||
} else {
|
||||
sqlx::query_scalar("SELECT COUNT(*) FROM tasks WHERE user_id = $1")
|
||||
.bind(user_id)
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询历史失败").with_source(err))?
|
||||
};
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct TaskRow {
|
||||
id: Uuid,
|
||||
status: String,
|
||||
source: String,
|
||||
total_files: i32,
|
||||
completed_files: i32,
|
||||
failed_files: i32,
|
||||
created_at: DateTime<Utc>,
|
||||
completed_at: Option<DateTime<Utc>>,
|
||||
expires_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
let tasks: Vec<TaskRow> = if let Some(status) = &status {
|
||||
sqlx::query_as::<_, TaskRow>(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
status::text AS status,
|
||||
source::text AS source,
|
||||
total_files,
|
||||
completed_files,
|
||||
failed_files,
|
||||
created_at,
|
||||
completed_at,
|
||||
expires_at
|
||||
FROM tasks
|
||||
WHERE user_id = $1 AND status::text = $2
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $3 OFFSET $4
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(status)
|
||||
.bind(limit as i64)
|
||||
.bind(offset as i64)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询历史失败").with_source(err))?
|
||||
} else {
|
||||
sqlx::query_as::<_, TaskRow>(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
status::text AS status,
|
||||
source::text AS source,
|
||||
total_files,
|
||||
completed_files,
|
||||
failed_files,
|
||||
created_at,
|
||||
completed_at,
|
||||
expires_at
|
||||
FROM tasks
|
||||
WHERE user_id = $1
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $2 OFFSET $3
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(limit as i64)
|
||||
.bind(offset as i64)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询历史失败").with_source(err))?
|
||||
};
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct FileRow {
|
||||
id: Uuid,
|
||||
original_name: String,
|
||||
original_size: i64,
|
||||
compressed_size: Option<i64>,
|
||||
saved_percent: Option<f64>,
|
||||
status: String,
|
||||
output_format: String,
|
||||
error_message: Option<String>,
|
||||
storage_path: Option<String>,
|
||||
}
|
||||
|
||||
let now = Utc::now();
|
||||
let mut result_tasks = Vec::with_capacity(tasks.len());
|
||||
for task in tasks {
|
||||
let files: Vec<FileRow> = sqlx::query_as::<_, FileRow>(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
original_name,
|
||||
original_size,
|
||||
compressed_size,
|
||||
saved_percent::float8 AS saved_percent,
|
||||
status::text AS status,
|
||||
output_format,
|
||||
error_message,
|
||||
storage_path
|
||||
FROM task_files
|
||||
WHERE task_id = $1
|
||||
ORDER BY created_at ASC
|
||||
"#,
|
||||
)
|
||||
.bind(task.id)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询任务文件失败").with_source(err))?;
|
||||
|
||||
let file_views = files
|
||||
.into_iter()
|
||||
.map(|file| HistoryFileView {
|
||||
file_id: file.id,
|
||||
original_name: file.original_name,
|
||||
original_size: file.original_size,
|
||||
compressed_size: file.compressed_size,
|
||||
saved_percent: file.saved_percent,
|
||||
status: file.status.clone(),
|
||||
output_format: file.output_format,
|
||||
error_message: file.error_message,
|
||||
download_url: if file.status == "completed" && file.storage_path.is_some() && task.expires_at > now {
|
||||
Some(format!("/downloads/{}", file.id))
|
||||
} else {
|
||||
None
|
||||
},
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let progress = if task.total_files > 0 {
|
||||
((task.completed_files + task.failed_files) * 100 / task.total_files).clamp(0, 100)
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
let download_all_url = if task.status == "completed" && task.expires_at > now {
|
||||
Some(format!("/downloads/tasks/{}", task.id))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
result_tasks.push(HistoryTaskView {
|
||||
task_id: task.id,
|
||||
status: task.status,
|
||||
source: task.source,
|
||||
progress,
|
||||
total_files: task.total_files,
|
||||
completed_files: task.completed_files,
|
||||
failed_files: task.failed_files,
|
||||
created_at: task.created_at,
|
||||
completed_at: task.completed_at,
|
||||
expires_at: task.expires_at,
|
||||
download_all_url,
|
||||
files: file_views,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: HistoryResponse {
|
||||
tasks: result_tasks,
|
||||
page,
|
||||
limit,
|
||||
total,
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
async fn list_api_keys(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<Json<Envelope<ApiKeyListResponse>>, AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (_jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let (user_id, _email_verified) = match principal {
|
||||
context::Principal::User {
|
||||
user_id,
|
||||
email_verified,
|
||||
..
|
||||
} => (user_id, email_verified),
|
||||
_ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
|
||||
};
|
||||
|
||||
let rows = sqlx::query_as::<_, ApiKeyView>(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
name,
|
||||
key_prefix,
|
||||
permissions,
|
||||
rate_limit,
|
||||
is_active,
|
||||
last_used_at,
|
||||
last_used_ip::text AS last_used_ip,
|
||||
created_at
|
||||
FROM api_keys
|
||||
WHERE user_id = $1
|
||||
ORDER BY created_at DESC
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询 API Key 失败").with_source(err))?;
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: ApiKeyListResponse { api_keys: rows },
|
||||
}))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct CreateApiKeyRequest {
|
||||
name: String,
|
||||
permissions: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct CreateApiKeyResponse {
|
||||
id: Uuid,
|
||||
name: String,
|
||||
key_prefix: String,
|
||||
key: String,
|
||||
message: String,
|
||||
}
|
||||
|
||||
async fn create_api_key(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
Json(req): Json<CreateApiKeyRequest>,
|
||||
) -> Result<Json<Envelope<CreateApiKeyResponse>>, AppError> {
|
||||
if req.name.trim().is_empty() || req.name.len() > 100 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "name 不合法"));
|
||||
}
|
||||
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (_jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let (user_id, email_verified) = match principal {
|
||||
context::Principal::User {
|
||||
user_id,
|
||||
email_verified,
|
||||
..
|
||||
} => (user_id, email_verified),
|
||||
_ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
|
||||
};
|
||||
|
||||
if !email_verified {
|
||||
return Err(AppError::new(ErrorCode::EmailNotVerified, "请先验证邮箱"));
|
||||
}
|
||||
|
||||
let billing = billing::get_user_billing(&state, user_id).await?;
|
||||
if !billing.plan.feature_api_enabled {
|
||||
return Err(AppError::new(ErrorCode::Forbidden, "当前套餐未开通 API Key"));
|
||||
}
|
||||
|
||||
let permissions = normalize_permissions(req.permissions)?;
|
||||
|
||||
let (full_key, key_prefix) = generate_api_key();
|
||||
let key_hash = context::api_key_hash(&full_key, &state.config.api_key_pepper)?;
|
||||
|
||||
let row_id: Uuid = sqlx::query_scalar(
|
||||
r#"
|
||||
INSERT INTO api_keys (user_id, name, key_prefix, key_hash, permissions, rate_limit)
|
||||
VALUES ($1, $2, $3, $4, $5, 100)
|
||||
RETURNING id
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(req.name.trim())
|
||||
.bind(&key_prefix)
|
||||
.bind(key_hash)
|
||||
.bind(&permissions)
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "创建 API Key 失败").with_source(err))?;
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: CreateApiKeyResponse {
|
||||
id: row_id,
|
||||
name: req.name.trim().to_string(),
|
||||
key_prefix,
|
||||
key: full_key,
|
||||
message: "请保存此 Key，它只会显示一次".to_string(),
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
async fn disable_api_key(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
Path(key_id): Path<Uuid>,
|
||||
) -> Result<Json<Envelope<serde_json::Value>>, AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (_jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let user_id = match principal {
|
||||
context::Principal::User { user_id, .. } => user_id,
|
||||
_ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
|
||||
};
|
||||
|
||||
let result = sqlx::query("UPDATE api_keys SET is_active = false WHERE id = $1 AND user_id = $2")
|
||||
.bind(key_id)
|
||||
.bind(user_id)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新 API Key 失败").with_source(err))?;
|
||||
|
||||
if result.rows_affected() == 0 {
|
||||
return Err(AppError::new(ErrorCode::NotFound, "API Key 不存在"));
|
||||
}
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: serde_json::json!({ "message": "已禁用" }),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn rotate_api_key(
|
||||
State(state): State<AppState>,
|
||||
jar: axum_extra::extract::cookie::CookieJar,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
headers: HeaderMap,
|
||||
Path(key_id): Path<Uuid>,
|
||||
) -> Result<Json<Envelope<CreateApiKeyResponse>>, AppError> {
|
||||
let ip = context::client_ip(&headers, addr.ip());
|
||||
let (_jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
|
||||
|
||||
let (user_id, email_verified) = match principal {
|
||||
context::Principal::User {
|
||||
user_id,
|
||||
email_verified,
|
||||
..
|
||||
} => (user_id, email_verified),
|
||||
_ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
|
||||
};
|
||||
|
||||
if !email_verified {
|
||||
return Err(AppError::new(ErrorCode::EmailNotVerified, "请先验证邮箱"));
|
||||
}
|
||||
|
||||
let (full_key, key_prefix) = generate_api_key();
|
||||
let key_hash = context::api_key_hash(&full_key, &state.config.api_key_pepper)?;
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct RotateRow {
|
||||
id: Uuid,
|
||||
name: String,
|
||||
}
|
||||
|
||||
let row = sqlx::query_as::<_, RotateRow>(
|
||||
r#"
|
||||
UPDATE api_keys
|
||||
SET key_prefix = $1,
|
||||
key_hash = $2,
|
||||
is_active = true
|
||||
WHERE id = $3 AND user_id = $4
|
||||
RETURNING id, name
|
||||
"#,
|
||||
)
|
||||
.bind(&key_prefix)
|
||||
.bind(key_hash)
|
||||
.bind(key_id)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新 API Key 失败").with_source(err))?
|
||||
.ok_or_else(|| AppError::new(ErrorCode::NotFound, "API Key 不存在"))?;
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: CreateApiKeyResponse {
|
||||
id: row.id,
|
||||
name: row.name,
|
||||
key_prefix,
|
||||
key: full_key,
|
||||
message: "请保存此 Key，它只会显示一次".to_string(),
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
fn generate_api_key() -> (String, String) {
|
||||
let mut prefix_bytes = [0u8; 4];
|
||||
rand::rngs::OsRng.fill_bytes(&mut prefix_bytes);
|
||||
let prefix = hex::encode(prefix_bytes);
|
||||
let key_prefix = format!("if_live_{prefix}");
|
||||
|
||||
let mut secret_bytes = [0u8; 32];
|
||||
rand::rngs::OsRng.fill_bytes(&mut secret_bytes);
|
||||
let secret = URL_SAFE_NO_PAD.encode(secret_bytes);
|
||||
|
||||
let full = format!("{key_prefix}_{secret}");
|
||||
(full, key_prefix)
|
||||
}
|
||||
|
||||
fn normalize_permissions(input: Option<Vec<String>>) -> Result<serde_json::Value, AppError> {
|
||||
let allowed = ["compress", "batch_compress", "read_stats", "billing_read", "webhook_manage"];
|
||||
|
||||
let mut perms = Vec::<String>::new();
|
||||
if let Some(values) = input {
|
||||
for value in values {
|
||||
let v = value.trim().to_ascii_lowercase();
|
||||
if v.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if !allowed.contains(&v.as_str()) {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
format!("不支持的权限: {v}"),
|
||||
));
|
||||
}
|
||||
if !perms.contains(&v) {
|
||||
perms.push(v);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if perms.is_empty() {
|
||||
perms.push("compress".to_string());
|
||||
}
|
||||
|
||||
Ok(serde_json::json!(perms))
|
||||
}
|
||||
|
||||
fn validate_email(email: &str) -> Result<(), AppError> {
|
||||
if email.trim().is_empty() || !email.contains('@') {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "邮箱格式不正确"));
|
||||
}
|
||||
if email.len() > 255 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "邮箱过长"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_username(username: &str) -> Result<(), AppError> {
|
||||
if username.trim().is_empty() {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "用户名不能为空"));
|
||||
}
|
||||
if username.len() > 50 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "用户名过长"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_password(password: &str) -> Result<(), AppError> {
|
||||
if password.len() < 8 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "密码至少 8 位"));
|
||||
}
|
||||
if password.len() > 128 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "密码过长"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn hash_password(password: &str) -> Result<String, AppError> {
|
||||
let salt = argon2::password_hash::SaltString::generate(&mut rand::rngs::OsRng);
|
||||
let hashed = Argon2::default()
|
||||
.hash_password(password.as_bytes(), &salt)
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "密码哈希失败").with_source(err))?;
|
||||
Ok(hashed.to_string())
|
||||
}
|
||||
|
||||
fn verify_password(password: &str, password_hash: &str) -> Result<(), AppError> {
|
||||
let parsed = PasswordHash::new(password_hash)
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "密码哈希格式错误").with_source(err))?;
|
||||
Argon2::default()
|
||||
.verify_password(password.as_bytes(), &parsed)
|
||||
.map_err(|_| AppError::new(ErrorCode::Unauthorized, "密码错误"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn generate_token() -> String {
|
||||
let mut bytes = [0u8; 32];
|
||||
rand::rngs::OsRng.fill_bytes(&mut bytes);
|
||||
URL_SAFE_NO_PAD.encode(bytes)
|
||||
}
|
||||
|
||||
fn sha256_hex(token: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(token.as_bytes());
|
||||
hex::encode(hasher.finalize())
|
||||
}
|
||||
|
||||
fn map_unique_violation(err: sqlx::Error) -> AppError {
|
||||
if let sqlx::Error::Database(db_err) = &err {
|
||||
if let Some(code) = db_err.code() {
|
||||
if code == "23505" {
|
||||
return AppError::new(ErrorCode::InvalidRequest, "邮箱或用户名已存在");
|
||||
}
|
||||
}
|
||||
}
|
||||
AppError::new(ErrorCode::Internal, "数据库操作失败").with_source(err)
|
||||
}
|
||||
520
src/api/webhooks.rs
Normal file
520
src/api/webhooks.rs
Normal file
@@ -0,0 +1,520 @@
|
||||
use crate::api::envelope::Envelope;
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::services::settings;
|
||||
use crate::state::AppState;
|
||||
|
||||
use axum::body::Bytes;
|
||||
use axum::extract::State;
|
||||
use axum::http::HeaderMap;
|
||||
use axum::routing::post;
|
||||
use axum::{Json, Router};
|
||||
use chrono::{TimeZone, Utc};
|
||||
use hmac::{Hmac, Mac};
|
||||
use serde::Deserialize;
|
||||
use sha2::Sha256;
|
||||
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new().route("/webhooks/stripe", post(stripe_webhook))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct StripeEvent {
|
||||
id: String,
|
||||
#[serde(rename = "type")]
|
||||
type_: String,
|
||||
data: StripeEventData,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct StripeEventData {
|
||||
object: serde_json::Value,
|
||||
}
|
||||
|
||||
async fn stripe_webhook(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Result<Json<Envelope<serde_json::Value>>, AppError> {
|
||||
let secret = settings::get_stripe_webhook_secret(&state)
|
||||
.await
|
||||
.map_err(|err| err.with_source("stripe webhook secret not configured"))?;
|
||||
|
||||
let sig = headers
|
||||
.get("Stripe-Signature")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "缺少 Stripe-Signature"))?;
|
||||
|
||||
verify_stripe_signature(&body, sig, &secret)?;
|
||||
|
||||
let payload_str = std::str::from_utf8(&body)
|
||||
.map_err(|_| AppError::new(ErrorCode::InvalidRequest, "Webhook payload 非 UTF-8"))?;
|
||||
let event: StripeEvent = serde_json::from_str(payload_str)
|
||||
.map_err(|err| AppError::new(ErrorCode::InvalidRequest, "Webhook JSON 解析失败").with_source(err))?;
|
||||
|
||||
let inserted: Option<String> = sqlx::query_scalar(
|
||||
r#"
|
||||
INSERT INTO webhook_events (provider, provider_event_id, event_type, payload)
|
||||
VALUES ('stripe', $1, $2, $3)
|
||||
ON CONFLICT (provider, provider_event_id) DO NOTHING
|
||||
RETURNING provider_event_id
|
||||
"#,
|
||||
)
|
||||
.bind(&event.id)
|
||||
.bind(&event.type_)
|
||||
.bind(&event.data.object)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "Webhook 入库失败").with_source(err))?;
|
||||
|
||||
if inserted.is_none() {
|
||||
return Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: serde_json::json!({ "status": "duplicate" }),
|
||||
}));
|
||||
}
|
||||
|
||||
if let Err(err) = process_stripe_event(&state, &event).await {
|
||||
let _ = sqlx::query(
|
||||
"UPDATE webhook_events SET status = 'failed', error_message = $2, processed_at = NOW() WHERE provider = 'stripe' AND provider_event_id = $1",
|
||||
)
|
||||
.bind(&event.id)
|
||||
.bind(err.to_string())
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
return Err(err);
|
||||
}
|
||||
|
||||
let _ = sqlx::query(
|
||||
"UPDATE webhook_events SET status = 'processed', processed_at = NOW() WHERE provider = 'stripe' AND provider_event_id = $1",
|
||||
)
|
||||
.bind(&event.id)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
Ok(Json(Envelope {
|
||||
success: true,
|
||||
data: serde_json::json!({ "status": "ok" }),
|
||||
}))
|
||||
}
|
||||
|
||||
fn verify_stripe_signature(payload: &[u8], sig_header: &str, secret: &str) -> Result<(), AppError> {
|
||||
let mut timestamp: Option<i64> = None;
|
||||
let mut signatures = Vec::<String>::new();
|
||||
|
||||
for part in sig_header.split(',') {
|
||||
let part = part.trim();
|
||||
if let Some(v) = part.strip_prefix("t=") {
|
||||
timestamp = v.parse::<i64>().ok();
|
||||
} else if let Some(v) = part.strip_prefix("v1=") {
|
||||
signatures.push(v.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
let Some(ts) = timestamp else {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "Stripe-Signature 缺少 t"));
|
||||
};
|
||||
if signatures.is_empty() {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "Stripe-Signature 缺少 v1"));
|
||||
}
|
||||
|
||||
// 5 minutes tolerance
|
||||
let now = Utc::now().timestamp();
|
||||
if (now - ts).abs() > 300 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "Webhook 时间戳过期"));
|
||||
}
|
||||
|
||||
type HmacSha256 = Hmac<Sha256>;
|
||||
let mut mac = HmacSha256::new_from_slice(secret.as_bytes())
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "Webhook secret 错误").with_source(err))?;
|
||||
mac.update(ts.to_string().as_bytes());
|
||||
mac.update(b".");
|
||||
mac.update(payload);
|
||||
let expected = hex::encode(mac.finalize().into_bytes());
|
||||
|
||||
if signatures.iter().any(|sig| secure_eq(sig, &expected)) {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(AppError::new(ErrorCode::InvalidRequest, "Webhook 验签失败"))
|
||||
}
|
||||
}
|
||||
|
||||
fn secure_eq(a: &str, b: &str) -> bool {
|
||||
if a.len() != b.len() {
|
||||
return false;
|
||||
}
|
||||
let mut out = 0u8;
|
||||
for (x, y) in a.as_bytes().iter().zip(b.as_bytes().iter()) {
|
||||
out |= x ^ y;
|
||||
}
|
||||
out == 0
|
||||
}
|
||||
|
||||
async fn process_stripe_event(state: &AppState, event: &StripeEvent) -> Result<(), AppError> {
|
||||
match event.type_.as_str() {
|
||||
"checkout.session.completed" => {
|
||||
map_checkout_session_completed(state, &event.data.object).await
|
||||
}
|
||||
"customer.subscription.created" | "customer.subscription.updated" => {
|
||||
upsert_subscription(state, &event.data.object).await
|
||||
}
|
||||
"customer.subscription.deleted" => cancel_subscription(state, &event.data.object).await,
|
||||
"invoice.paid" | "invoice.payment_failed" => upsert_invoice(state, &event.data.object).await,
|
||||
_ => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn map_checkout_session_completed(
|
||||
state: &AppState,
|
||||
object: &serde_json::Value,
|
||||
) -> Result<(), AppError> {
|
||||
let customer_id = object
|
||||
.get("customer")
|
||||
.and_then(|v| v.as_str())
|
||||
.filter(|v| !v.trim().is_empty());
|
||||
|
||||
let user_id = object
|
||||
.get("client_reference_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.and_then(|v| v.parse::<uuid::Uuid>().ok())
|
||||
.or_else(|| {
|
||||
object
|
||||
.pointer("/metadata/user_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.and_then(|v| v.parse::<uuid::Uuid>().ok())
|
||||
});
|
||||
|
||||
let Some(customer_id) = customer_id else {
|
||||
tracing::warn!("checkout.session.completed missing customer");
|
||||
return Ok(());
|
||||
};
|
||||
let Some(user_id) = user_id else {
|
||||
tracing::warn!(customer = %customer_id, "checkout.session.completed missing user_id");
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let updated = sqlx::query(
|
||||
r#"
|
||||
UPDATE users
|
||||
SET billing_customer_id = $2,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1
|
||||
AND (billing_customer_id IS NULL OR billing_customer_id = '')
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(customer_id)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新 Stripe Customer 映射失败").with_source(err))?;
|
||||
|
||||
if updated.rows_affected() == 0 {
|
||||
let existing: Option<String> = sqlx::query_scalar::<_, Option<String>>(
|
||||
"SELECT billing_customer_id FROM users WHERE id = $1",
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?
|
||||
.flatten();
|
||||
|
||||
if let Some(existing) = existing.filter(|v| !v.trim().is_empty()) {
|
||||
if existing != customer_id {
|
||||
tracing::warn!(
|
||||
user_id = %user_id,
|
||||
existing_customer = %existing,
|
||||
new_customer = %customer_id,
|
||||
"user already mapped to different stripe customer"
|
||||
);
|
||||
}
|
||||
} else {
|
||||
tracing::warn!(user_id = %user_id, "user not found for checkout.session.completed");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn upsert_subscription(state: &AppState, object: &serde_json::Value) -> Result<(), AppError> {
|
||||
let provider_subscription_id = object
|
||||
.get("id")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "subscription.id 缺失"))?;
|
||||
let provider_customer_id = object
|
||||
.get("customer")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "subscription.customer 缺失"))?;
|
||||
|
||||
let status = object
|
||||
.get("status")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("incomplete");
|
||||
let mapped_status = map_subscription_status(status);
|
||||
|
||||
let cps = object
|
||||
.get("current_period_start")
|
||||
.and_then(|v| v.as_i64())
|
||||
.unwrap_or(0);
|
||||
let cpe = object
|
||||
.get("current_period_end")
|
||||
.and_then(|v| v.as_i64())
|
||||
.unwrap_or(0);
|
||||
let current_period_start = Utc.timestamp_opt(cps, 0).single().unwrap_or_else(Utc::now);
|
||||
let current_period_end = Utc.timestamp_opt(cpe, 0).single().unwrap_or_else(Utc::now);
|
||||
|
||||
let cancel_at_period_end = object
|
||||
.get("cancel_at_period_end")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
|
||||
let price_id = object
|
||||
.pointer("/items/data/0/price/id")
|
||||
.and_then(|v| v.as_str())
|
||||
.or_else(|| object.pointer("/items/data/0/plan/id").and_then(|v| v.as_str()))
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "subscription.price 缺失"))?;
|
||||
|
||||
let user_id: Option<uuid::Uuid> =
|
||||
sqlx::query_scalar("SELECT id FROM users WHERE billing_customer_id = $1 LIMIT 1")
|
||||
.bind(provider_customer_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?;
|
||||
|
||||
let Some(user_id) = user_id else {
|
||||
tracing::warn!(customer = %provider_customer_id, "stripe customer not mapped to user");
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let plan_id: Option<uuid::Uuid> =
|
||||
sqlx::query_scalar("SELECT id FROM plans WHERE stripe_price_id = $1 LIMIT 1")
|
||||
.bind(price_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询套餐失败").with_source(err))?;
|
||||
|
||||
let Some(plan_id) = plan_id else {
|
||||
tracing::warn!(price = %price_id, "stripe price not mapped to plan");
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let updated: Option<uuid::Uuid> = sqlx::query_scalar(
|
||||
r#"
|
||||
UPDATE subscriptions
|
||||
SET user_id = $1,
|
||||
plan_id = $2,
|
||||
status = $3::subscription_status,
|
||||
current_period_start = $4,
|
||||
current_period_end = $5,
|
||||
cancel_at_period_end = $6,
|
||||
provider = 'stripe',
|
||||
provider_customer_id = $7,
|
||||
updated_at = NOW()
|
||||
WHERE provider = 'stripe' AND provider_subscription_id = $8
|
||||
RETURNING id
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(plan_id)
|
||||
.bind(mapped_status)
|
||||
.bind(current_period_start)
|
||||
.bind(current_period_end)
|
||||
.bind(cancel_at_period_end)
|
||||
.bind(provider_customer_id)
|
||||
.bind(provider_subscription_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新订阅失败").with_source(err))?;
|
||||
|
||||
if updated.is_none() {
|
||||
let _ = sqlx::query(
|
||||
r#"
|
||||
INSERT INTO subscriptions (
|
||||
user_id, plan_id, status,
|
||||
current_period_start, current_period_end,
|
||||
cancel_at_period_end,
|
||||
provider, provider_customer_id, provider_subscription_id
|
||||
) VALUES (
|
||||
$1, $2, $3::subscription_status,
|
||||
$4, $5,
|
||||
$6,
|
||||
'stripe', $7, $8
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(plan_id)
|
||||
.bind(mapped_status)
|
||||
.bind(current_period_start)
|
||||
.bind(current_period_end)
|
||||
.bind(cancel_at_period_end)
|
||||
.bind(provider_customer_id)
|
||||
.bind(provider_subscription_id)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "创建订阅失败").with_source(err))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cancel_subscription(state: &AppState, object: &serde_json::Value) -> Result<(), AppError> {
|
||||
let provider_subscription_id = object
|
||||
.get("id")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "subscription.id 缺失"))?;
|
||||
|
||||
let _ = sqlx::query(
|
||||
r#"
|
||||
UPDATE subscriptions
|
||||
SET status = 'canceled',
|
||||
cancel_at_period_end = false,
|
||||
canceled_at = NOW(),
|
||||
updated_at = NOW()
|
||||
WHERE provider = 'stripe' AND provider_subscription_id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(provider_subscription_id)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn upsert_invoice(state: &AppState, object: &serde_json::Value) -> Result<(), AppError> {
|
||||
let provider_invoice_id = object
|
||||
.get("id")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "invoice.id 缺失"))?;
|
||||
let provider_customer_id = object
|
||||
.get("customer")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "invoice.customer 缺失"))?;
|
||||
|
||||
let user_id: Option<uuid::Uuid> =
|
||||
sqlx::query_scalar("SELECT id FROM users WHERE billing_customer_id = $1 LIMIT 1")
|
||||
.bind(provider_customer_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?;
|
||||
|
||||
let Some(user_id) = user_id else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let stripe_status = object.get("status").and_then(|v| v.as_str()).unwrap_or("open");
|
||||
let status = map_invoice_status(stripe_status);
|
||||
|
||||
let invoice_number = object
|
||||
.get("number")
|
||||
.and_then(|v| v.as_str())
|
||||
.filter(|v| !v.trim().is_empty())
|
||||
.map(|v| v.to_string())
|
||||
.unwrap_or_else(|| format!("stripe_{provider_invoice_id}"));
|
||||
|
||||
let currency = object
|
||||
.get("currency")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("cny")
|
||||
.to_uppercase();
|
||||
let total_amount_cents = object.get("total").and_then(|v| v.as_i64()).unwrap_or(0) as i32;
|
||||
|
||||
let hosted_invoice_url = object.get("hosted_invoice_url").and_then(|v| v.as_str()).map(|v| v.to_string());
|
||||
let pdf_url = object.get("invoice_pdf").and_then(|v| v.as_str()).map(|v| v.to_string());
|
||||
|
||||
let period_start = object.get("period_start").and_then(|v| v.as_i64()).and_then(|ts| Utc.timestamp_opt(ts, 0).single());
|
||||
let period_end = object.get("period_end").and_then(|v| v.as_i64()).and_then(|ts| Utc.timestamp_opt(ts, 0).single());
|
||||
|
||||
let paid_at = object
|
||||
.pointer("/status_transitions/paid_at")
|
||||
.and_then(|v| v.as_i64())
|
||||
.and_then(|ts| Utc.timestamp_opt(ts, 0).single());
|
||||
|
||||
let updated = sqlx::query(
|
||||
r#"
|
||||
UPDATE invoices
|
||||
SET status = $1::invoice_status,
|
||||
currency = $2,
|
||||
total_amount_cents = $3,
|
||||
hosted_invoice_url = $4,
|
||||
pdf_url = $5,
|
||||
period_start = $6,
|
||||
period_end = $7,
|
||||
paid_at = $8
|
||||
WHERE provider = 'stripe' AND provider_invoice_id = $9
|
||||
"#,
|
||||
)
|
||||
.bind(status)
|
||||
.bind(¤cy)
|
||||
.bind(total_amount_cents)
|
||||
.bind(hosted_invoice_url.as_deref())
|
||||
.bind(pdf_url.as_deref())
|
||||
.bind(period_start)
|
||||
.bind(period_end)
|
||||
.bind(paid_at)
|
||||
.bind(provider_invoice_id)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新发票失败").with_source(err))?;
|
||||
|
||||
if updated.rows_affected() == 0 {
|
||||
let invoice_number = truncate(invoice_number, 50);
|
||||
let _ = sqlx::query(
|
||||
r#"
|
||||
INSERT INTO invoices (
|
||||
user_id, invoice_number, status, currency, total_amount_cents,
|
||||
period_start, period_end,
|
||||
provider, provider_invoice_id, hosted_invoice_url, pdf_url,
|
||||
paid_at
|
||||
) VALUES (
|
||||
$1, $2, $3::invoice_status, $4, $5,
|
||||
$6, $7,
|
||||
'stripe', $8, $9, $10,
|
||||
$11
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(invoice_number)
|
||||
.bind(status)
|
||||
.bind(¤cy)
|
||||
.bind(total_amount_cents)
|
||||
.bind(period_start)
|
||||
.bind(period_end)
|
||||
.bind(provider_invoice_id)
|
||||
.bind(hosted_invoice_url.as_deref())
|
||||
.bind(pdf_url.as_deref())
|
||||
.bind(paid_at)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "创建发票失败").with_source(err))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn truncate(mut s: String, max: usize) -> String {
|
||||
if s.len() > max {
|
||||
s.truncate(max);
|
||||
}
|
||||
s
|
||||
}
|
||||
|
||||
fn map_subscription_status(status: &str) -> &'static str {
|
||||
match status {
|
||||
"trialing" => "trialing",
|
||||
"active" => "active",
|
||||
"past_due" => "past_due",
|
||||
"canceled" => "canceled",
|
||||
_ => "incomplete",
|
||||
}
|
||||
}
|
||||
|
||||
fn map_invoice_status(status: &str) -> &'static str {
|
||||
match status {
|
||||
"draft" => "draft",
|
||||
"paid" => "paid",
|
||||
"void" => "void",
|
||||
"uncollectible" => "uncollectible",
|
||||
_ => "open",
|
||||
}
|
||||
}
|
||||
60
src/auth.rs
Normal file
60
src/auth.rs
Normal file
@@ -0,0 +1,60 @@
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
|
||||
use axum::http::HeaderMap;
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize, Clone)]
|
||||
pub struct Claims {
|
||||
pub sub: Uuid,
|
||||
pub role: String,
|
||||
pub exp: usize,
|
||||
}
|
||||
|
||||
pub fn issue_jwt(
|
||||
jwt_secret: &str,
|
||||
jwt_expiry_hours: i64,
|
||||
user_id: Uuid,
|
||||
role: &str,
|
||||
) -> Result<(String, DateTime<Utc>), AppError> {
|
||||
let expires_at = Utc::now() + Duration::hours(jwt_expiry_hours);
|
||||
let claims = Claims {
|
||||
sub: user_id,
|
||||
role: role.to_string(),
|
||||
exp: expires_at.timestamp() as usize,
|
||||
};
|
||||
|
||||
let token = jsonwebtoken::encode(
|
||||
&Header::default(),
|
||||
&claims,
|
||||
&EncodingKey::from_secret(jwt_secret.as_bytes()),
|
||||
)
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "生成 Token 失败").with_source(err))?;
|
||||
|
||||
Ok((token, expires_at))
|
||||
}
|
||||
|
||||
pub fn require_jwt(jwt_secret: &str, headers: &HeaderMap) -> Result<Claims, AppError> {
|
||||
let auth = headers
|
||||
.get(axum::http::header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("");
|
||||
let token = auth.strip_prefix("Bearer ").ok_or_else(|| {
|
||||
AppError::new(ErrorCode::Unauthorized, "缺少 Authorization: Bearer <token>")
|
||||
})?;
|
||||
|
||||
decode_jwt(jwt_secret, token)
|
||||
}
|
||||
|
||||
pub fn decode_jwt(jwt_secret: &str, token: &str) -> Result<Claims, AppError> {
|
||||
jsonwebtoken::decode::<Claims>(
|
||||
token,
|
||||
&DecodingKey::from_secret(jwt_secret.as_bytes()),
|
||||
&Validation::default(),
|
||||
)
|
||||
.map(|data| data.claims)
|
||||
.map_err(|_| AppError::new(ErrorCode::Unauthorized, "Token 无效或已过期"))
|
||||
}
|
||||
|
||||
161
src/config.rs
Normal file
161
src/config.rs
Normal file
@@ -0,0 +1,161 @@
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Config {
|
||||
pub role: String,
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
pub public_base_url: String,
|
||||
|
||||
pub database_url: String,
|
||||
pub database_max_connections: u32,
|
||||
|
||||
pub redis_url: String,
|
||||
|
||||
pub jwt_secret: String,
|
||||
pub jwt_expiry_hours: i64,
|
||||
|
||||
pub api_key_pepper: String,
|
||||
|
||||
pub billing_provider: String,
|
||||
pub stripe_secret_key: Option<String>,
|
||||
pub stripe_webhook_secret: Option<String>,
|
||||
|
||||
pub storage_type: String,
|
||||
pub storage_path: String,
|
||||
pub signed_url_ttl_minutes: u64,
|
||||
|
||||
pub allow_anonymous_upload: bool,
|
||||
pub anon_max_file_size_mb: u64,
|
||||
pub anon_max_files_per_batch: u32,
|
||||
pub anon_daily_units: u32,
|
||||
pub anon_retention_hours: u64,
|
||||
|
||||
pub max_image_pixels: u64,
|
||||
pub idempotency_ttl_hours: u64,
|
||||
|
||||
pub mail_enabled: bool,
|
||||
pub mail_log_links_when_disabled: bool,
|
||||
pub mail_provider: String,
|
||||
pub mail_from: String,
|
||||
pub mail_password: String,
|
||||
pub mail_from_name: String,
|
||||
pub mail_smtp_host: Option<String>,
|
||||
pub mail_smtp_port: Option<u16>,
|
||||
pub mail_smtp_encryption: Option<String>,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
pub fn from_env() -> Result<Self, AppError> {
|
||||
let role = env_string("IMAGEFORGE_ROLE").unwrap_or_else(|| "api".to_string());
|
||||
let host = env_string("HOST").unwrap_or_else(|| "0.0.0.0".to_string());
|
||||
let port = env_u16("PORT").unwrap_or(8080);
|
||||
let public_base_url =
|
||||
env_string("PUBLIC_BASE_URL").unwrap_or_else(|| "http://localhost:8080".to_string());
|
||||
|
||||
let database_url = env_string("DATABASE_URL")
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "缺少环境变量 DATABASE_URL"))?;
|
||||
let database_max_connections = env_u32("DATABASE_MAX_CONNECTIONS").unwrap_or(10);
|
||||
|
||||
let redis_url = env_string("REDIS_URL")
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "缺少环境变量 REDIS_URL"))?;
|
||||
|
||||
let jwt_secret = env_string("JWT_SECRET")
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "缺少环境变量 JWT_SECRET"))?;
|
||||
let jwt_expiry_hours = env_i64("JWT_EXPIRY_HOURS").unwrap_or(168);
|
||||
|
||||
let api_key_pepper = env_string("API_KEY_PEPPER")
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "缺少环境变量 API_KEY_PEPPER"))?;
|
||||
|
||||
let billing_provider =
|
||||
env_string("BILLING_PROVIDER").unwrap_or_else(|| "stripe".to_string());
|
||||
let stripe_secret_key = env_string("STRIPE_SECRET_KEY");
|
||||
let stripe_webhook_secret = env_string("STRIPE_WEBHOOK_SECRET");
|
||||
|
||||
let storage_type = env_string("STORAGE_TYPE").unwrap_or_else(|| "local".to_string());
|
||||
let storage_path = env_string("STORAGE_PATH").unwrap_or_else(|| "./uploads".to_string());
|
||||
let signed_url_ttl_minutes = env_u64("SIGNED_URL_TTL_MINUTES").unwrap_or(60);
|
||||
|
||||
let allow_anonymous_upload = env_bool("ALLOW_ANONYMOUS_UPLOAD").unwrap_or(true);
|
||||
let anon_max_file_size_mb = env_u64("ANON_MAX_FILE_SIZE_MB").unwrap_or(5);
|
||||
let anon_max_files_per_batch = env_u32("ANON_MAX_FILES_PER_BATCH").unwrap_or(5);
|
||||
let anon_daily_units = env_u32("ANON_DAILY_UNITS").unwrap_or(10);
|
||||
let anon_retention_hours = env_u64("ANON_RETENTION_HOURS").unwrap_or(24);
|
||||
|
||||
let max_image_pixels = env_u64("MAX_IMAGE_PIXELS").unwrap_or(40_000_000);
|
||||
let idempotency_ttl_hours = env_u64("IDEMPOTENCY_TTL_HOURS").unwrap_or(24);
|
||||
|
||||
let mail_enabled = env_bool("MAIL_ENABLED").unwrap_or(false);
|
||||
let mail_log_links_when_disabled = env_bool("MAIL_LOG_LINKS_WHEN_DISABLED").unwrap_or(false);
|
||||
let mail_provider = env_string("MAIL_PROVIDER").unwrap_or_else(|| "qq".to_string());
|
||||
let mail_from = env_string("MAIL_FROM").unwrap_or_else(|| "noreply@example.com".to_string());
|
||||
let mail_password = env_string("MAIL_PASSWORD").unwrap_or_default();
|
||||
let mail_from_name = env_string("MAIL_FROM_NAME").unwrap_or_else(|| "ImageForge".to_string());
|
||||
let mail_smtp_host = env_string("MAIL_SMTP_HOST");
|
||||
let mail_smtp_port = env_u16("MAIL_SMTP_PORT");
|
||||
let mail_smtp_encryption = env_string("MAIL_SMTP_ENCRYPTION");
|
||||
|
||||
Ok(Self {
|
||||
role,
|
||||
host,
|
||||
port,
|
||||
public_base_url,
|
||||
database_url,
|
||||
database_max_connections,
|
||||
redis_url,
|
||||
jwt_secret,
|
||||
jwt_expiry_hours,
|
||||
api_key_pepper,
|
||||
billing_provider,
|
||||
stripe_secret_key,
|
||||
stripe_webhook_secret,
|
||||
storage_type,
|
||||
storage_path,
|
||||
signed_url_ttl_minutes,
|
||||
allow_anonymous_upload,
|
||||
anon_max_file_size_mb,
|
||||
anon_max_files_per_batch,
|
||||
anon_daily_units,
|
||||
anon_retention_hours,
|
||||
max_image_pixels,
|
||||
idempotency_ttl_hours,
|
||||
mail_enabled,
|
||||
mail_log_links_when_disabled,
|
||||
mail_provider,
|
||||
mail_from,
|
||||
mail_password,
|
||||
mail_from_name,
|
||||
mail_smtp_host,
|
||||
mail_smtp_port,
|
||||
mail_smtp_encryption,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn env_string(key: &str) -> Option<String> {
|
||||
std::env::var(key).ok().filter(|value| !value.trim().is_empty())
|
||||
}
|
||||
|
||||
fn env_u16(key: &str) -> Option<u16> {
|
||||
env_string(key).and_then(|v| v.parse::<u16>().ok())
|
||||
}
|
||||
|
||||
fn env_u32(key: &str) -> Option<u32> {
|
||||
env_string(key).and_then(|v| v.parse::<u32>().ok())
|
||||
}
|
||||
|
||||
fn env_i64(key: &str) -> Option<i64> {
|
||||
env_string(key).and_then(|v| v.parse::<i64>().ok())
|
||||
}
|
||||
|
||||
fn env_u64(key: &str) -> Option<u64> {
|
||||
env_string(key).and_then(|v| v.parse::<u64>().ok())
|
||||
}
|
||||
|
||||
fn env_bool(key: &str) -> Option<bool> {
|
||||
env_string(key).and_then(|v| match v.trim().to_ascii_lowercase().as_str() {
|
||||
"1" | "true" | "yes" | "y" | "on" => Some(true),
|
||||
"0" | "false" | "no" | "n" | "off" => Some(false),
|
||||
_ => None,
|
||||
})
|
||||
}
|
||||
136
src/error.rs
Normal file
136
src/error.rs
Normal file
@@ -0,0 +1,136 @@
|
||||
use axum::{http::StatusCode, response::IntoResponse, Json};
|
||||
use serde::Serialize;
|
||||
use std::fmt::{Display, Formatter};
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Serialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
|
||||
pub enum ErrorCode {
|
||||
InvalidRequest,
|
||||
InvalidImage,
|
||||
UnsupportedFormat,
|
||||
TooManyPixels,
|
||||
FileTooLarge,
|
||||
InvalidToken,
|
||||
Unauthorized,
|
||||
Forbidden,
|
||||
NotFound,
|
||||
IdempotencyConflict,
|
||||
RateLimited,
|
||||
QuotaExceeded,
|
||||
EmailNotVerified,
|
||||
CompressionFailed,
|
||||
StorageUnavailable,
|
||||
MailSendFailed,
|
||||
Internal,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct AppError {
|
||||
pub code: ErrorCode,
|
||||
pub message: String,
|
||||
pub source: Option<String>,
|
||||
}
|
||||
|
||||
impl AppError {
|
||||
pub fn new(code: ErrorCode, message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
code,
|
||||
message: message.into(),
|
||||
source: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_source(mut self, err: impl std::fmt::Display) -> Self {
|
||||
self.source = Some(err.to_string());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for AppError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}: {}", self.code.as_str(), self.message)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for AppError {}
|
||||
|
||||
impl ErrorCode {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
ErrorCode::InvalidRequest => "INVALID_REQUEST",
|
||||
ErrorCode::InvalidImage => "INVALID_IMAGE",
|
||||
ErrorCode::UnsupportedFormat => "UNSUPPORTED_FORMAT",
|
||||
ErrorCode::TooManyPixels => "TOO_MANY_PIXELS",
|
||||
ErrorCode::FileTooLarge => "FILE_TOO_LARGE",
|
||||
ErrorCode::InvalidToken => "INVALID_TOKEN",
|
||||
ErrorCode::Unauthorized => "UNAUTHORIZED",
|
||||
ErrorCode::Forbidden => "FORBIDDEN",
|
||||
ErrorCode::NotFound => "NOT_FOUND",
|
||||
ErrorCode::IdempotencyConflict => "IDEMPOTENCY_CONFLICT",
|
||||
ErrorCode::RateLimited => "RATE_LIMITED",
|
||||
ErrorCode::QuotaExceeded => "QUOTA_EXCEEDED",
|
||||
ErrorCode::EmailNotVerified => "EMAIL_NOT_VERIFIED",
|
||||
ErrorCode::CompressionFailed => "COMPRESSION_FAILED",
|
||||
ErrorCode::StorageUnavailable => "STORAGE_UNAVAILABLE",
|
||||
ErrorCode::MailSendFailed => "MAIL_SEND_FAILED",
|
||||
ErrorCode::Internal => "INTERNAL",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct ErrorEnvelope {
|
||||
success: bool,
|
||||
error: ErrorPayload,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct ErrorPayload {
|
||||
code: ErrorCode,
|
||||
message: String,
|
||||
request_id: String,
|
||||
}
|
||||
|
||||
impl IntoResponse for AppError {
|
||||
fn into_response(self) -> axum::response::Response {
|
||||
let request_id = format!("req_{}", Uuid::new_v4());
|
||||
|
||||
if let Some(source) = &self.source {
|
||||
tracing::error!(code = %self.code.as_str(), request_id = %request_id, message = %self.message, source = %source);
|
||||
} else {
|
||||
tracing::error!(code = %self.code.as_str(), request_id = %request_id, message = %self.message);
|
||||
}
|
||||
|
||||
let status = match self.code {
|
||||
ErrorCode::InvalidRequest => StatusCode::BAD_REQUEST,
|
||||
ErrorCode::InvalidImage => StatusCode::BAD_REQUEST,
|
||||
ErrorCode::UnsupportedFormat => StatusCode::BAD_REQUEST,
|
||||
ErrorCode::TooManyPixels => StatusCode::BAD_REQUEST,
|
||||
ErrorCode::FileTooLarge => StatusCode::PAYLOAD_TOO_LARGE,
|
||||
ErrorCode::InvalidToken => StatusCode::BAD_REQUEST,
|
||||
ErrorCode::Unauthorized => StatusCode::UNAUTHORIZED,
|
||||
ErrorCode::Forbidden => StatusCode::FORBIDDEN,
|
||||
ErrorCode::NotFound => StatusCode::NOT_FOUND,
|
||||
ErrorCode::IdempotencyConflict => StatusCode::CONFLICT,
|
||||
ErrorCode::RateLimited => StatusCode::TOO_MANY_REQUESTS,
|
||||
ErrorCode::QuotaExceeded => StatusCode::PAYMENT_REQUIRED,
|
||||
ErrorCode::EmailNotVerified => StatusCode::FORBIDDEN,
|
||||
ErrorCode::CompressionFailed => StatusCode::INTERNAL_SERVER_ERROR,
|
||||
ErrorCode::StorageUnavailable => StatusCode::SERVICE_UNAVAILABLE,
|
||||
ErrorCode::MailSendFailed => StatusCode::INTERNAL_SERVER_ERROR,
|
||||
ErrorCode::Internal => StatusCode::INTERNAL_SERVER_ERROR,
|
||||
};
|
||||
|
||||
let body = ErrorEnvelope {
|
||||
success: false,
|
||||
error: ErrorPayload {
|
||||
code: self.code,
|
||||
message: self.message,
|
||||
request_id,
|
||||
},
|
||||
};
|
||||
|
||||
(status, Json(body)).into_response()
|
||||
}
|
||||
}
|
||||
64
src/main.rs
Normal file
64
src/main.rs
Normal file
@@ -0,0 +1,64 @@
|
||||
mod api;
|
||||
mod auth;
|
||||
mod config;
|
||||
mod error;
|
||||
mod services;
|
||||
mod state;
|
||||
mod worker;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::services::mail::Mailer;
|
||||
use crate::state::AppState;
|
||||
|
||||
use sqlx::postgres::PgPoolOptions;
|
||||
use tracing::Level;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), AppError> {
|
||||
dotenvy::dotenv().ok();
|
||||
init_tracing();
|
||||
|
||||
let config = Config::from_env()?;
|
||||
let mailer = Mailer::new(&config)?;
|
||||
|
||||
let db = PgPoolOptions::new()
|
||||
.max_connections(config.database_max_connections)
|
||||
.connect(&config.database_url)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "数据库连接失败").with_source(err))?;
|
||||
|
||||
let redis = redis::Client::open(config.redis_url.clone())
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "Redis 配置错误").with_source(err))?
|
||||
.get_connection_manager()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "Redis 连接失败").with_source(err))?;
|
||||
|
||||
let state = AppState {
|
||||
config,
|
||||
db,
|
||||
redis,
|
||||
mailer: std::sync::Arc::new(mailer),
|
||||
};
|
||||
|
||||
match state.config.role.as_str() {
|
||||
"api" => api::run(state).await,
|
||||
"worker" => worker::run(state).await,
|
||||
other => Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
format!("未知 IMAGEFORGE_ROLE: {other}（仅支持 api/worker）"),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn init_tracing() {
|
||||
let env_filter =
|
||||
tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| {
|
||||
tracing_subscriber::EnvFilter::new("info,tower_http=info,imageforge=info")
|
||||
});
|
||||
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(env_filter)
|
||||
.with_max_level(Level::INFO)
|
||||
.init();
|
||||
}
|
||||
135
src/services/billing.rs
Normal file
135
src/services/billing.rs
Normal file
@@ -0,0 +1,135 @@
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::state::AppState;
|
||||
|
||||
use chrono::{DateTime, Datelike, TimeZone, Utc};
|
||||
use sqlx::FromRow;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Plan {
|
||||
pub id: Uuid,
|
||||
pub code: String,
|
||||
pub included_units_per_period: i32,
|
||||
pub max_file_size_mb: i32,
|
||||
pub max_files_per_batch: i32,
|
||||
pub retention_days: i32,
|
||||
pub feature_api_enabled: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BillingContext {
|
||||
pub user_id: Uuid,
|
||||
pub subscription_id: Option<Uuid>,
|
||||
pub plan: Plan,
|
||||
pub period_start: DateTime<Utc>,
|
||||
pub period_end: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct SubscriptionRow {
|
||||
id: Uuid,
|
||||
status: String,
|
||||
current_period_start: DateTime<Utc>,
|
||||
current_period_end: DateTime<Utc>,
|
||||
plan_id: Uuid,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct PlanRow {
|
||||
id: Uuid,
|
||||
code: String,
|
||||
included_units_per_period: i32,
|
||||
max_file_size_mb: i32,
|
||||
max_files_per_batch: i32,
|
||||
retention_days: i32,
|
||||
features: serde_json::Value,
|
||||
}
|
||||
|
||||
pub async fn get_user_billing(state: &AppState, user_id: Uuid) -> Result<BillingContext, AppError> {
|
||||
let subscription = sqlx::query_as::<_, SubscriptionRow>(
|
||||
r#"
|
||||
SELECT id, status::text AS status, current_period_start, current_period_end, plan_id
|
||||
FROM subscriptions
|
||||
WHERE user_id = $1
|
||||
AND status IN ('active', 'trialing', 'past_due')
|
||||
ORDER BY current_period_end DESC
|
||||
LIMIT 1
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询订阅失败").with_source(err))?;
|
||||
|
||||
let (subscription_id, period_start, period_end, plan_id) = if let Some(sub) = subscription {
|
||||
if sub.status == "past_due" {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::Forbidden,
|
||||
"订阅欠费，请先完成支付",
|
||||
));
|
||||
}
|
||||
(Some(sub.id), sub.current_period_start, sub.current_period_end, sub.plan_id)
|
||||
} else {
|
||||
let plan_id: Uuid = sqlx::query_scalar("SELECT id FROM plans WHERE code = 'free' LIMIT 1")
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "未找到 Free 套餐").with_source(err))?;
|
||||
let (start, end) = current_month_period_utc8(Utc::now());
|
||||
(None, start, end, plan_id)
|
||||
};
|
||||
|
||||
let plan_row = sqlx::query_as::<_, PlanRow>(
|
||||
r#"
|
||||
SELECT id, code, included_units_per_period, max_file_size_mb, max_files_per_batch, retention_days, features
|
||||
FROM plans
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(plan_id)
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询套餐失败").with_source(err))?;
|
||||
|
||||
let feature_api_enabled = plan_row
|
||||
.features
|
||||
.get("api")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
|
||||
Ok(BillingContext {
|
||||
user_id,
|
||||
subscription_id,
|
||||
plan: Plan {
|
||||
id: plan_row.id,
|
||||
code: plan_row.code,
|
||||
included_units_per_period: plan_row.included_units_per_period,
|
||||
max_file_size_mb: plan_row.max_file_size_mb,
|
||||
max_files_per_batch: plan_row.max_files_per_batch,
|
||||
retention_days: plan_row.retention_days,
|
||||
feature_api_enabled,
|
||||
},
|
||||
period_start,
|
||||
period_end,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn current_month_period_utc8(now_utc: DateTime<Utc>) -> (DateTime<Utc>, DateTime<Utc>) {
|
||||
let tz = chrono::FixedOffset::east_opt(8 * 3600).unwrap();
|
||||
let now = now_utc.with_timezone(&tz);
|
||||
let year = now.year();
|
||||
let month = now.month();
|
||||
|
||||
let start = tz.with_ymd_and_hms(year, month, 1, 0, 0, 0).single().unwrap();
|
||||
|
||||
let (next_year, next_month) = if month == 12 {
|
||||
(year + 1, 1)
|
||||
} else {
|
||||
(year, month + 1)
|
||||
};
|
||||
let end = tz
|
||||
.with_ymd_and_hms(next_year, next_month, 1, 0, 0, 0)
|
||||
.single()
|
||||
.unwrap();
|
||||
|
||||
(start.with_timezone(&Utc), end.with_timezone(&Utc))
|
||||
}
|
||||
204
src/services/bootstrap.rs
Normal file
204
src/services/bootstrap.rs
Normal file
@@ -0,0 +1,204 @@
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::state::AppState;
|
||||
|
||||
use argon2::{Argon2, PasswordHasher};
|
||||
use chrono::Utc;
|
||||
use sqlx::FromRow;
|
||||
use tracing::{info, warn};
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct AdminRow {
|
||||
id: Uuid,
|
||||
username: String,
|
||||
role: String,
|
||||
}
|
||||
|
||||
pub async fn ensure_admin_user(state: &AppState) -> Result<(), AppError> {
|
||||
let Some(admin_email) = env_string("ADMIN_EMAIL") else {
|
||||
return Ok(());
|
||||
};
|
||||
let Some(admin_password) = env_string("ADMIN_PASSWORD") else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let admin_email = admin_email.trim().to_lowercase();
|
||||
let admin_password = admin_password.trim().to_string();
|
||||
if admin_email.is_empty() || admin_password.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let admin_username = env_string("ADMIN_USERNAME").unwrap_or_else(|| {
|
||||
admin_email
|
||||
.split('@')
|
||||
.next()
|
||||
.unwrap_or("admin")
|
||||
.to_string()
|
||||
});
|
||||
let admin_username = admin_username.trim().to_string();
|
||||
|
||||
validate_email(&admin_email)?;
|
||||
validate_username(&admin_username)?;
|
||||
validate_password(&admin_password)?;
|
||||
|
||||
let existing = sqlx::query_as::<_, AdminRow>(
|
||||
r#"
|
||||
SELECT id, username, role::text AS role
|
||||
FROM users
|
||||
WHERE email = $1
|
||||
"#,
|
||||
)
|
||||
.bind(&admin_email)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询管理员账号失败").with_source(err))?;
|
||||
|
||||
let password_hash = hash_password(&admin_password)?;
|
||||
|
||||
if let Some(row) = existing {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE users
|
||||
SET password_hash = $1,
|
||||
role = 'admin',
|
||||
is_active = true,
|
||||
email_verified_at = COALESCE(email_verified_at, NOW()),
|
||||
updated_at = NOW()
|
||||
WHERE id = $2
|
||||
"#,
|
||||
)
|
||||
.bind(&password_hash)
|
||||
.bind(row.id)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新管理员账号失败").with_source(err))?;
|
||||
|
||||
if row.username != admin_username {
|
||||
let name_taken: bool = sqlx::query_scalar(
|
||||
r#"
|
||||
SELECT EXISTS(
|
||||
SELECT 1 FROM users WHERE username = $1 AND id <> $2
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(&admin_username)
|
||||
.bind(row.id)
|
||||
.fetch_one(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "校验管理员用户名失败").with_source(err))?;
|
||||
|
||||
if name_taken {
|
||||
warn!(
|
||||
admin_email = %admin_email,
|
||||
admin_username = %admin_username,
|
||||
"管理员用户名已被占用，保留原用户名"
|
||||
);
|
||||
} else {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE users
|
||||
SET username = $1, updated_at = NOW()
|
||||
WHERE id = $2
|
||||
"#,
|
||||
)
|
||||
.bind(&admin_username)
|
||||
.bind(row.id)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新管理员用户名失败").with_source(err))?;
|
||||
}
|
||||
}
|
||||
|
||||
if row.role != "admin" {
|
||||
info!(admin_email = %admin_email, "管理员权限已启用");
|
||||
}
|
||||
} else {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO users (email, username, password_hash, role, email_verified_at)
|
||||
VALUES ($1, $2, $3, 'admin', $4)
|
||||
"#,
|
||||
)
|
||||
.bind(&admin_email)
|
||||
.bind(&admin_username)
|
||||
.bind(&password_hash)
|
||||
.bind(Utc::now())
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "创建管理员账号失败").with_source(err))?;
|
||||
|
||||
info!(
|
||||
admin_email = %admin_email,
|
||||
admin_username = %admin_username,
|
||||
"管理员账号已创建"
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn ensure_schema(state: &AppState) -> Result<(), AppError> {
|
||||
sqlx::query(
|
||||
"ALTER TABLE tasks ADD COLUMN IF NOT EXISTS compression_rate SMALLINT",
|
||||
)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "初始化数据库结构失败").with_source(err))?;
|
||||
|
||||
sqlx::query(
|
||||
"ALTER TABLE usage_periods ADD COLUMN IF NOT EXISTS bonus_units INTEGER NOT NULL DEFAULT 0",
|
||||
)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "初始化数据库结构失败").with_source(err))?;
|
||||
|
||||
let _ = sqlx::query(
|
||||
"UPDATE usage_periods SET bonus_units = bonus_units + ABS(used_units), used_units = 0 WHERE used_units < 0",
|
||||
)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_email(email: &str) -> Result<(), AppError> {
|
||||
if email.trim().is_empty() || !email.contains('@') {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "管理员邮箱格式不正确"));
|
||||
}
|
||||
if email.len() > 255 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "管理员邮箱过长"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_username(username: &str) -> Result<(), AppError> {
|
||||
if username.trim().is_empty() {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "管理员用户名不能为空"));
|
||||
}
|
||||
if username.len() > 50 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "管理员用户名过长"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_password(password: &str) -> Result<(), AppError> {
|
||||
if password.len() < 8 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "管理员密码至少 8 位"));
|
||||
}
|
||||
if password.len() > 128 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "管理员密码过长"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn hash_password(password: &str) -> Result<String, AppError> {
|
||||
let salt = argon2::password_hash::SaltString::generate(&mut rand::rngs::OsRng);
|
||||
let hashed = Argon2::default()
|
||||
.hash_password(password.as_bytes(), &salt)
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "密码哈希失败").with_source(err))?;
|
||||
Ok(hashed.to_string())
|
||||
}
|
||||
|
||||
fn env_string(key: &str) -> Option<String> {
|
||||
std::env::var(key).ok().filter(|value| !value.trim().is_empty())
|
||||
}
|
||||
533
src/services/compress.rs
Normal file
533
src/services/compress.rs
Normal file
@@ -0,0 +1,533 @@
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::state::AppState;
|
||||
|
||||
use img_parts::{Bytes as ImgBytes, DynImage, ImageEXIF, ImageICC};
|
||||
use image::codecs::bmp::BmpEncoder;
|
||||
use image::codecs::gif::{GifDecoder, GifEncoder};
|
||||
use image::codecs::ico::IcoEncoder;
|
||||
use image::codecs::jpeg::JpegEncoder;
|
||||
use image::codecs::png::PngEncoder;
|
||||
use image::codecs::tiff::TiffEncoder;
|
||||
use image::{DynamicImage, ExtendedColorType, ImageEncoder};
|
||||
use image::{AnimationDecoder, GenericImageView};
|
||||
use oxipng::StripChunks;
|
||||
use rgb::FromSlice;
|
||||
use std::io::Cursor;
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum CompressionLevel {
|
||||
High,
|
||||
Medium,
|
||||
Low,
|
||||
}
|
||||
|
||||
impl CompressionLevel {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::High => "high",
|
||||
Self::Medium => "medium",
|
||||
Self::Low => "low",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ImageFmt {
|
||||
Png,
|
||||
Jpeg,
|
||||
Webp,
|
||||
Avif,
|
||||
Gif,
|
||||
Bmp,
|
||||
Tiff,
|
||||
Ico,
|
||||
}
|
||||
|
||||
impl ImageFmt {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Png => "png",
|
||||
Self::Jpeg => "jpeg",
|
||||
Self::Webp => "webp",
|
||||
Self::Avif => "avif",
|
||||
Self::Gif => "gif",
|
||||
Self::Bmp => "bmp",
|
||||
Self::Tiff => "tiff",
|
||||
Self::Ico => "ico",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn extension(self) -> &'static str {
|
||||
match self {
|
||||
Self::Png => "png",
|
||||
Self::Jpeg => "jpg",
|
||||
Self::Webp => "webp",
|
||||
Self::Avif => "avif",
|
||||
Self::Gif => "gif",
|
||||
Self::Bmp => "bmp",
|
||||
Self::Tiff => "tiff",
|
||||
Self::Ico => "ico",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn content_type(self) -> &'static str {
|
||||
match self {
|
||||
Self::Png => "image/png",
|
||||
Self::Jpeg => "image/jpeg",
|
||||
Self::Webp => "image/webp",
|
||||
Self::Avif => "image/avif",
|
||||
Self::Gif => "image/gif",
|
||||
Self::Bmp => "image/bmp",
|
||||
Self::Tiff => "image/tiff",
|
||||
Self::Ico => "image/x-icon",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse_level(value: &str) -> Result<CompressionLevel, AppError> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"" | "medium" => Ok(CompressionLevel::Medium),
|
||||
"high" => Ok(CompressionLevel::High),
|
||||
"low" => Ok(CompressionLevel::Low),
|
||||
_ => Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
"level 仅支持 high/medium/low",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse_compression_rate(value: &str) -> Result<u8, AppError> {
|
||||
let rate: u8 = value
|
||||
.trim()
|
||||
.parse()
|
||||
.map_err(|_| AppError::new(ErrorCode::InvalidRequest, "compression_rate 需为 1-100 的整数"))?;
|
||||
if !(1..=100).contains(&rate) {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
"compression_rate 需在 1-100 之间",
|
||||
));
|
||||
}
|
||||
Ok(rate)
|
||||
}
|
||||
|
||||
pub fn rate_to_level(rate: u8) -> CompressionLevel {
|
||||
match rate {
|
||||
1..=33 => CompressionLevel::Low,
|
||||
34..=66 => CompressionLevel::Medium,
|
||||
_ => CompressionLevel::High,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse_output_format(value: &str) -> Result<ImageFmt, AppError> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"png" => Ok(ImageFmt::Png),
|
||||
"jpeg" | "jpg" => Ok(ImageFmt::Jpeg),
|
||||
"webp" => Ok(ImageFmt::Webp),
|
||||
"avif" => Ok(ImageFmt::Avif),
|
||||
"gif" => Ok(ImageFmt::Gif),
|
||||
"bmp" => Ok(ImageFmt::Bmp),
|
||||
"tif" | "tiff" => Ok(ImageFmt::Tiff),
|
||||
"ico" => Ok(ImageFmt::Ico),
|
||||
_ => Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
"output_format 仅支持 png/jpeg/webp/avif/gif/bmp/tiff/ico",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn detect_format(bytes: &[u8]) -> Result<ImageFmt, AppError> {
|
||||
if bytes.starts_with(b"\x89PNG\r\n\x1a\n") {
|
||||
return Ok(ImageFmt::Png);
|
||||
}
|
||||
if bytes.len() >= 2 && bytes[0] == 0xFF && bytes[1] == 0xD8 {
|
||||
return Ok(ImageFmt::Jpeg);
|
||||
}
|
||||
if bytes.len() >= 12 && &bytes[0..4] == b"RIFF" && &bytes[8..12] == b"WEBP" {
|
||||
return Ok(ImageFmt::Webp);
|
||||
}
|
||||
if bytes.len() >= 12 && &bytes[4..8] == b"ftyp" {
|
||||
if bytes[8..12].eq_ignore_ascii_case(b"avif")
|
||||
|| bytes[8..12].eq_ignore_ascii_case(b"avis")
|
||||
{
|
||||
return Ok(ImageFmt::Avif);
|
||||
}
|
||||
}
|
||||
if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") {
|
||||
return Ok(ImageFmt::Gif);
|
||||
}
|
||||
if bytes.len() >= 2 && bytes[0] == 0x42 && bytes[1] == 0x4D {
|
||||
return Ok(ImageFmt::Bmp);
|
||||
}
|
||||
if bytes.len() >= 4 {
|
||||
if &bytes[0..4] == b"II*\x00" || &bytes[0..4] == b"MM\x00*" {
|
||||
return Ok(ImageFmt::Tiff);
|
||||
}
|
||||
if bytes[0] == 0x00
|
||||
&& bytes[1] == 0x00
|
||||
&& (bytes[2] == 0x01 || bytes[2] == 0x02)
|
||||
&& bytes[3] == 0x00
|
||||
{
|
||||
return Ok(ImageFmt::Ico);
|
||||
}
|
||||
}
|
||||
Err(AppError::new(
|
||||
ErrorCode::UnsupportedFormat,
|
||||
"不支持的图片格式",
|
||||
))
|
||||
}
|
||||
|
||||
pub async fn compress_image_bytes(
|
||||
state: &AppState,
|
||||
input: &[u8],
|
||||
format_in: ImageFmt,
|
||||
format_out: ImageFmt,
|
||||
level: CompressionLevel,
|
||||
compression_rate: Option<u8>,
|
||||
max_width: Option<u32>,
|
||||
max_height: Option<u32>,
|
||||
preserve_metadata: bool,
|
||||
) -> Result<Vec<u8>, AppError> {
|
||||
if format_in == ImageFmt::Gif {
|
||||
if is_animated_gif(input)? {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::UnsupportedFormat,
|
||||
"暂不支持动图 GIF",
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let rate = effective_rate(compression_rate, level);
|
||||
let (icc_profile, exif) = if preserve_metadata {
|
||||
extract_metadata(input)
|
||||
} else {
|
||||
(None, None)
|
||||
};
|
||||
|
||||
let mut resized = false;
|
||||
let mut output = if format_in == ImageFmt::Png
|
||||
&& format_out == ImageFmt::Png
|
||||
&& max_width.is_none()
|
||||
&& max_height.is_none()
|
||||
{
|
||||
let preset = png_preset_from_rate(rate);
|
||||
let mut opts = oxipng::Options::from_preset(preset);
|
||||
if !preserve_metadata {
|
||||
opts.strip = StripChunks::Safe;
|
||||
}
|
||||
oxipng::optimize_from_memory(input, &opts)
|
||||
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "PNG 压缩失败").with_source(err))?
|
||||
} else {
|
||||
let image = image::load_from_memory(input)
|
||||
.map_err(|err| AppError::new(ErrorCode::InvalidImage, "图片解码失败").with_source(err))?;
|
||||
|
||||
enforce_pixel_limit(state, &image)?;
|
||||
|
||||
let (image, did_resize) = resize_if_needed(image, max_width, max_height);
|
||||
resized = did_resize;
|
||||
|
||||
match format_out {
|
||||
ImageFmt::Png => encode_png(image, rate, preserve_metadata)?,
|
||||
ImageFmt::Jpeg => encode_jpeg(image, rate)?,
|
||||
ImageFmt::Webp => encode_webp(image, rate)?,
|
||||
ImageFmt::Avif => encode_avif(image, rate)?,
|
||||
ImageFmt::Gif => encode_gif(image, rate)?,
|
||||
ImageFmt::Bmp => encode_bmp(image)?,
|
||||
ImageFmt::Tiff => encode_tiff(image)?,
|
||||
ImageFmt::Ico => encode_ico(image)?,
|
||||
}
|
||||
};
|
||||
|
||||
if preserve_metadata {
|
||||
output = apply_metadata(output, icc_profile, exif)?;
|
||||
}
|
||||
|
||||
if !resized && output.len() >= input.len() {
|
||||
if preserve_metadata {
|
||||
return Ok(input.to_vec());
|
||||
}
|
||||
let stripped = strip_metadata(input).unwrap_or_else(|_| input.to_vec());
|
||||
return Ok(if stripped.len() <= input.len() {
|
||||
stripped
|
||||
} else {
|
||||
input.to_vec()
|
||||
});
|
||||
}
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn enforce_pixel_limit(state: &AppState, image: &DynamicImage) -> Result<(), AppError> {
|
||||
let (w, h) = image.dimensions();
|
||||
let pixels = (w as u64).saturating_mul(h as u64);
|
||||
if pixels > state.config.max_image_pixels {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::TooManyPixels,
|
||||
format!("图片像素过大（{}x{}）", w, h),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn resize_if_needed(
|
||||
image: DynamicImage,
|
||||
max_width: Option<u32>,
|
||||
max_height: Option<u32>,
|
||||
) -> (DynamicImage, bool) {
|
||||
if max_width.is_none() && max_height.is_none() {
|
||||
return (image, false);
|
||||
}
|
||||
|
||||
let (w, h) = image.dimensions();
|
||||
let (target_w, target_h) = fit_within(w, h, max_width, max_height);
|
||||
if target_w == w && target_h == h {
|
||||
return (image, false);
|
||||
}
|
||||
|
||||
(
|
||||
image.resize(target_w, target_h, image::imageops::FilterType::Lanczos3),
|
||||
true,
|
||||
)
|
||||
}
|
||||
|
||||
fn fit_within(w: u32, h: u32, max_width: Option<u32>, max_height: Option<u32>) -> (u32, u32) {
|
||||
let mut scale = 1.0_f64;
|
||||
if let Some(mw) = max_width.filter(|v| *v > 0) {
|
||||
scale = scale.min(mw as f64 / w as f64);
|
||||
}
|
||||
if let Some(mh) = max_height.filter(|v| *v > 0) {
|
||||
scale = scale.min(mh as f64 / h as f64);
|
||||
}
|
||||
if scale >= 1.0 {
|
||||
return (w, h);
|
||||
}
|
||||
let nw = (w as f64 * scale).round().max(1.0) as u32;
|
||||
let nh = (h as f64 * scale).round().max(1.0) as u32;
|
||||
(nw, nh)
|
||||
}
|
||||
|
||||
fn encode_png(
|
||||
image: DynamicImage,
|
||||
rate: u8,
|
||||
preserve_metadata: bool,
|
||||
) -> Result<Vec<u8>, AppError> {
|
||||
let rgba = image.to_rgba8();
|
||||
let (w, h) = rgba.dimensions();
|
||||
let mut out = Vec::new();
|
||||
|
||||
let encoder = PngEncoder::new(&mut out);
|
||||
encoder
|
||||
.write_image(rgba.as_raw(), w, h, ExtendedColorType::Rgba8)
|
||||
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "PNG 编码失败").with_source(err))?;
|
||||
|
||||
let preset = png_preset_from_rate(rate);
|
||||
let mut opts = oxipng::Options::from_preset(preset);
|
||||
if !preserve_metadata {
|
||||
opts.strip = StripChunks::Safe;
|
||||
}
|
||||
oxipng::optimize_from_memory(&out, &opts)
|
||||
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "PNG 优化失败").with_source(err))
|
||||
}
|
||||
|
||||
fn encode_jpeg(image: DynamicImage, rate: u8) -> Result<Vec<u8>, AppError> {
|
||||
let rgb = image.to_rgb8();
|
||||
let (w, h) = rgb.dimensions();
|
||||
let mut out = Vec::new();
|
||||
|
||||
let quality = jpeg_quality_from_rate(rate);
|
||||
|
||||
let mut encoder = JpegEncoder::new_with_quality(&mut out, quality);
|
||||
encoder
|
||||
.encode(rgb.as_raw(), w, h, ExtendedColorType::Rgb8)
|
||||
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "JPEG 编码失败").with_source(err))?;
|
||||
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn encode_webp(image: DynamicImage, rate: u8) -> Result<Vec<u8>, AppError> {
|
||||
let rgba = image.to_rgba8();
|
||||
let (w, h) = rgba.dimensions();
|
||||
let encoder = webp::Encoder::from_rgba(rgba.as_raw(), w, h);
|
||||
|
||||
let bytes = if rate <= 10 {
|
||||
encoder.encode_lossless()
|
||||
} else {
|
||||
encoder.encode(webp_quality_from_rate(rate))
|
||||
};
|
||||
|
||||
Ok(bytes.to_vec())
|
||||
}
|
||||
|
||||
fn encode_avif(image: DynamicImage, rate: u8) -> Result<Vec<u8>, AppError> {
|
||||
let rgba = image.to_rgba8();
|
||||
let (w, h) = rgba.dimensions();
|
||||
|
||||
let quality = avif_quality_from_rate(rate);
|
||||
|
||||
let raw = rgba.into_raw();
|
||||
let pixels = raw.as_rgba();
|
||||
let img = ravif::Img::new(pixels, w as usize, h as usize);
|
||||
|
||||
let encoder = ravif::Encoder::new().with_quality(quality);
|
||||
let encoded = encoder
|
||||
.encode_rgba(img)
|
||||
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "AVIF 编码失败").with_source(err))?;
|
||||
|
||||
Ok(encoded.avif_file)
|
||||
}
|
||||
|
||||
fn encode_gif(image: DynamicImage, rate: u8) -> Result<Vec<u8>, AppError> {
|
||||
let rgba = image.to_rgba8();
|
||||
let (w, h) = rgba.dimensions();
|
||||
let mut out = Vec::new();
|
||||
|
||||
let speed = gif_speed_from_rate(rate);
|
||||
{
|
||||
let mut encoder = GifEncoder::new_with_speed(&mut out, speed);
|
||||
encoder
|
||||
.encode(rgba.as_raw(), w, h, ExtendedColorType::Rgba8)
|
||||
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "GIF 编码失败").with_source(err))?;
|
||||
}
|
||||
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn encode_bmp(image: DynamicImage) -> Result<Vec<u8>, AppError> {
|
||||
let rgba = image.to_rgba8();
|
||||
let (w, h) = rgba.dimensions();
|
||||
let mut out = Vec::new();
|
||||
let encoder = BmpEncoder::new(&mut out);
|
||||
encoder
|
||||
.write_image(rgba.as_raw(), w, h, ExtendedColorType::Rgba8)
|
||||
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "BMP 编码失败").with_source(err))?;
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn encode_tiff(image: DynamicImage) -> Result<Vec<u8>, AppError> {
|
||||
let rgba = image.to_rgba8();
|
||||
let (w, h) = rgba.dimensions();
|
||||
let mut out = Cursor::new(Vec::new());
|
||||
let encoder = TiffEncoder::new(&mut out);
|
||||
encoder
|
||||
.write_image(rgba.as_raw(), w, h, ExtendedColorType::Rgba8)
|
||||
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "TIFF 编码失败").with_source(err))?;
|
||||
Ok(out.into_inner())
|
||||
}
|
||||
|
||||
fn encode_ico(image: DynamicImage) -> Result<Vec<u8>, AppError> {
|
||||
let rgba = image.to_rgba8();
|
||||
let (w, h) = rgba.dimensions();
|
||||
let mut out = Vec::new();
|
||||
let encoder = IcoEncoder::new(&mut out);
|
||||
encoder
|
||||
.write_image(rgba.as_raw(), w, h, ExtendedColorType::Rgba8)
|
||||
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "ICO 编码失败").with_source(err))?;
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn extract_metadata(input: &[u8]) -> (Option<ImgBytes>, Option<ImgBytes>) {
|
||||
let bytes = ImgBytes::copy_from_slice(input);
|
||||
match DynImage::from_bytes(bytes) {
|
||||
Ok(Some(img)) => (img.icc_profile(), img.exif()),
|
||||
_ => (None, None),
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_metadata(
|
||||
output: Vec<u8>,
|
||||
icc_profile: Option<ImgBytes>,
|
||||
exif: Option<ImgBytes>,
|
||||
) -> Result<Vec<u8>, AppError> {
|
||||
if icc_profile.is_none() && exif.is_none() {
|
||||
return Ok(output);
|
||||
}
|
||||
|
||||
let out_bytes = ImgBytes::from(output);
|
||||
let dyn_img = DynImage::from_bytes(out_bytes.clone())
|
||||
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "解析输出图片元数据失败").with_source(err))?;
|
||||
|
||||
let Some(mut img) = dyn_img else {
|
||||
return Ok(out_bytes.to_vec());
|
||||
};
|
||||
|
||||
img.set_icc_profile(icc_profile);
|
||||
img.set_exif(exif);
|
||||
|
||||
let mut buf = Vec::new();
|
||||
img.encoder()
|
||||
.write_to(&mut buf)
|
||||
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "写入图片元数据失败").with_source(err))?;
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
fn strip_metadata(input: &[u8]) -> Result<Vec<u8>, AppError> {
|
||||
let bytes = ImgBytes::copy_from_slice(input);
|
||||
let dyn_img = DynImage::from_bytes(bytes.clone())
|
||||
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "解析图片元数据失败").with_source(err))?;
|
||||
let Some(mut img) = dyn_img else {
|
||||
return Ok(bytes.to_vec());
|
||||
};
|
||||
|
||||
img.set_icc_profile(None);
|
||||
img.set_exif(None);
|
||||
|
||||
let mut buf = Vec::new();
|
||||
img.encoder()
|
||||
.write_to(&mut buf)
|
||||
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "写入图片元数据失败").with_source(err))?;
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
fn effective_rate(rate: Option<u8>, level: CompressionLevel) -> u8 {
|
||||
match rate {
|
||||
Some(value) => value.clamp(1, 100),
|
||||
None => match level {
|
||||
CompressionLevel::Low => 25,
|
||||
CompressionLevel::Medium => 55,
|
||||
CompressionLevel::High => 80,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn png_preset_from_rate(rate: u8) -> u8 {
|
||||
(((rate.saturating_sub(1)) as f32 / 99.0) * 6.0).round() as u8
|
||||
}
|
||||
|
||||
fn jpeg_quality_from_rate(rate: u8) -> u8 {
|
||||
quality_from_rate(rate, 35, 95)
|
||||
}
|
||||
|
||||
fn webp_quality_from_rate(rate: u8) -> f32 {
|
||||
quality_from_rate(rate, 40, 92) as f32
|
||||
}
|
||||
|
||||
fn avif_quality_from_rate(rate: u8) -> f32 {
|
||||
quality_from_rate(rate, 35, 90) as f32
|
||||
}
|
||||
|
||||
fn quality_from_rate(rate: u8, min_quality: u8, max_quality: u8) -> u8 {
|
||||
let min_q = min_quality as i32;
|
||||
let max_q = max_quality as i32;
|
||||
let rate = rate.clamp(1, 100) as i32;
|
||||
let span = max_q - min_q;
|
||||
let q = max_q - ((rate - 1) * span / 99);
|
||||
q.clamp(min_q, max_q) as u8
|
||||
}
|
||||
|
||||
fn gif_speed_from_rate(rate: u8) -> i32 {
|
||||
let rate = rate.clamp(1, 100) as i32;
|
||||
1 + ((rate - 1) * 29 / 99)
|
||||
}
|
||||
|
||||
fn is_animated_gif(input: &[u8]) -> Result<bool, AppError> {
|
||||
let decoder = GifDecoder::new(Cursor::new(input))
|
||||
.map_err(|err| AppError::new(ErrorCode::InvalidImage, "GIF 解码失败").with_source(err))?;
|
||||
let mut frames = decoder.into_frames().into_iter();
|
||||
if let Some(frame) = frames.next() {
|
||||
frame.map_err(|err| AppError::new(ErrorCode::InvalidImage, "GIF 解码失败").with_source(err))?;
|
||||
}
|
||||
if let Some(frame) = frames.next() {
|
||||
frame.map_err(|err| AppError::new(ErrorCode::InvalidImage, "GIF 解码失败").with_source(err))?;
|
||||
return Ok(true);
|
||||
}
|
||||
Ok(false)
|
||||
}
|
||||
341
src/services/idempotency.rs
Normal file
341
src/services/idempotency.rs
Normal file
@@ -0,0 +1,341 @@
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::state::AppState;
|
||||
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use serde_json::Value as JsonValue;
|
||||
use sha2::{Digest, Sha256};
|
||||
use sqlx::FromRow;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum Scope {
|
||||
User(Uuid),
|
||||
ApiKey(Uuid),
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum BeginResult {
|
||||
Acquired { expires_at: DateTime<Utc> },
|
||||
Replay { response_status: i32, response_body: JsonValue },
|
||||
InProgress,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct IdemRow {
|
||||
request_hash: String,
|
||||
response_status: i32,
|
||||
response_body: Option<JsonValue>,
|
||||
expires_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
pub fn sha256_hex(parts: &[&[u8]]) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
for p in parts {
|
||||
hasher.update(p);
|
||||
hasher.update([0u8]); // separator
|
||||
}
|
||||
hex::encode(hasher.finalize())
|
||||
}
|
||||
|
||||
pub async fn begin(
|
||||
state: &AppState,
|
||||
scope: Scope,
|
||||
idempotency_key: &str,
|
||||
request_hash: &str,
|
||||
ttl_hours: i64,
|
||||
) -> Result<BeginResult, AppError> {
|
||||
if idempotency_key.trim().is_empty() {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "Idempotency-Key 不能为空"));
|
||||
}
|
||||
if idempotency_key.len() > 128 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "Idempotency-Key 过长"));
|
||||
}
|
||||
if request_hash.len() != 64 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "request_hash 不合法"));
|
||||
}
|
||||
|
||||
let now = Utc::now();
|
||||
let expires_at = now + Duration::hours(ttl_hours.max(1));
|
||||
|
||||
cleanup_expired_for_key(state, scope, idempotency_key, now).await?;
|
||||
|
||||
let inserted = match scope {
|
||||
Scope::User(user_id) => {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO idempotency_keys (
|
||||
user_id, idempotency_key, request_hash,
|
||||
response_status, response_body,
|
||||
expires_at
|
||||
) VALUES (
|
||||
$1, $2, $3,
|
||||
0, NULL,
|
||||
$4
|
||||
)
|
||||
ON CONFLICT DO NOTHING
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(idempotency_key)
|
||||
.bind(request_hash)
|
||||
.bind(expires_at)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
}
|
||||
Scope::ApiKey(api_key_id) => {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO idempotency_keys (
|
||||
api_key_id, idempotency_key, request_hash,
|
||||
response_status, response_body,
|
||||
expires_at
|
||||
) VALUES (
|
||||
$1, $2, $3,
|
||||
0, NULL,
|
||||
$4
|
||||
)
|
||||
ON CONFLICT DO NOTHING
|
||||
"#,
|
||||
)
|
||||
.bind(api_key_id)
|
||||
.bind(idempotency_key)
|
||||
.bind(request_hash)
|
||||
.bind(expires_at)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
}
|
||||
}
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等记录失败").with_source(err))?;
|
||||
|
||||
if inserted.rows_affected() > 0 {
|
||||
return Ok(BeginResult::Acquired { expires_at });
|
||||
}
|
||||
|
||||
let row = get_row(state, scope, idempotency_key, now).await?;
|
||||
let Some(row) = row else {
|
||||
return Ok(BeginResult::Acquired { expires_at });
|
||||
};
|
||||
|
||||
if row.request_hash != request_hash {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::IdempotencyConflict,
|
||||
"同一个 Idempotency-Key 的请求参数不一致",
|
||||
));
|
||||
}
|
||||
|
||||
if row.response_status == 0 || row.response_body.is_none() {
|
||||
return Ok(BeginResult::InProgress);
|
||||
}
|
||||
|
||||
Ok(BeginResult::Replay {
|
||||
response_status: row.response_status,
|
||||
response_body: row.response_body.unwrap_or(JsonValue::Null),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn wait_for_replay(
|
||||
state: &AppState,
|
||||
scope: Scope,
|
||||
idempotency_key: &str,
|
||||
request_hash: &str,
|
||||
max_wait_ms: u64,
|
||||
) -> Result<Option<(i32, JsonValue)>, AppError> {
|
||||
let started = tokio::time::Instant::now();
|
||||
let now = Utc::now();
|
||||
|
||||
loop {
|
||||
let row = get_row(state, scope, idempotency_key, now).await?;
|
||||
let Some(row) = row else { return Ok(None) };
|
||||
|
||||
if row.request_hash != request_hash {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::IdempotencyConflict,
|
||||
"同一个 Idempotency-Key 的请求参数不一致",
|
||||
));
|
||||
}
|
||||
|
||||
if row.response_status != 0 {
|
||||
return Ok(Some((
|
||||
row.response_status,
|
||||
row.response_body.unwrap_or(JsonValue::Null),
|
||||
)));
|
||||
}
|
||||
|
||||
if started.elapsed().as_millis() as u64 >= max_wait_ms {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(200)).await;
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn complete(
|
||||
state: &AppState,
|
||||
scope: Scope,
|
||||
idempotency_key: &str,
|
||||
request_hash: &str,
|
||||
response_status: i32,
|
||||
response_body: JsonValue,
|
||||
) -> Result<(), AppError> {
|
||||
let updated = match scope {
|
||||
Scope::User(user_id) => {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE idempotency_keys
|
||||
SET response_status = $4,
|
||||
response_body = $5
|
||||
WHERE user_id = $1
|
||||
AND idempotency_key = $2
|
||||
AND request_hash = $3
|
||||
AND response_status = 0
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(idempotency_key)
|
||||
.bind(request_hash)
|
||||
.bind(response_status)
|
||||
.bind(response_body)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
}
|
||||
Scope::ApiKey(api_key_id) => {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE idempotency_keys
|
||||
SET response_status = $4,
|
||||
response_body = $5
|
||||
WHERE api_key_id = $1
|
||||
AND idempotency_key = $2
|
||||
AND request_hash = $3
|
||||
AND response_status = 0
|
||||
"#,
|
||||
)
|
||||
.bind(api_key_id)
|
||||
.bind(idempotency_key)
|
||||
.bind(request_hash)
|
||||
.bind(response_status)
|
||||
.bind(response_body)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
}
|
||||
}
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等结果失败").with_source(err))?;
|
||||
|
||||
if updated.rows_affected() == 0 {
|
||||
tracing::warn!("idempotency record not updated (maybe already completed?)");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn abort(
|
||||
state: &AppState,
|
||||
scope: Scope,
|
||||
idempotency_key: &str,
|
||||
request_hash: &str,
|
||||
) -> Result<(), AppError> {
|
||||
match scope {
|
||||
Scope::User(user_id) => {
|
||||
let _ = sqlx::query(
|
||||
"DELETE FROM idempotency_keys WHERE user_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND response_status = 0",
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(idempotency_key)
|
||||
.bind(request_hash)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
}
|
||||
Scope::ApiKey(api_key_id) => {
|
||||
let _ = sqlx::query(
|
||||
"DELETE FROM idempotency_keys WHERE api_key_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND response_status = 0",
|
||||
)
|
||||
.bind(api_key_id)
|
||||
.bind(idempotency_key)
|
||||
.bind(request_hash)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cleanup_expired_for_key(
|
||||
state: &AppState,
|
||||
scope: Scope,
|
||||
idempotency_key: &str,
|
||||
now: DateTime<Utc>,
|
||||
) -> Result<(), AppError> {
|
||||
match scope {
|
||||
Scope::User(user_id) => {
|
||||
let _ = sqlx::query(
|
||||
"DELETE FROM idempotency_keys WHERE user_id = $1 AND idempotency_key = $2 AND expires_at < $3",
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(idempotency_key)
|
||||
.bind(now)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
}
|
||||
Scope::ApiKey(api_key_id) => {
|
||||
let _ = sqlx::query(
|
||||
"DELETE FROM idempotency_keys WHERE api_key_id = $1 AND idempotency_key = $2 AND expires_at < $3",
|
||||
)
|
||||
.bind(api_key_id)
|
||||
.bind(idempotency_key)
|
||||
.bind(now)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_row(
|
||||
state: &AppState,
|
||||
scope: Scope,
|
||||
idempotency_key: &str,
|
||||
now: DateTime<Utc>,
|
||||
) -> Result<Option<IdemRow>, AppError> {
|
||||
let row = match scope {
|
||||
Scope::User(user_id) => {
|
||||
sqlx::query_as::<_, IdemRow>(
|
||||
r#"
|
||||
SELECT request_hash, response_status, response_body, expires_at
|
||||
FROM idempotency_keys
|
||||
WHERE user_id = $1
|
||||
AND idempotency_key = $2
|
||||
AND expires_at > $3
|
||||
ORDER BY created_at DESC
|
||||
LIMIT 1
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(idempotency_key)
|
||||
.bind(now)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
}
|
||||
Scope::ApiKey(api_key_id) => {
|
||||
sqlx::query_as::<_, IdemRow>(
|
||||
r#"
|
||||
SELECT request_hash, response_status, response_body, expires_at
|
||||
FROM idempotency_keys
|
||||
WHERE api_key_id = $1
|
||||
AND idempotency_key = $2
|
||||
AND expires_at > $3
|
||||
ORDER BY created_at DESC
|
||||
LIMIT 1
|
||||
"#,
|
||||
)
|
||||
.bind(api_key_id)
|
||||
.bind(idempotency_key)
|
||||
.bind(now)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
}
|
||||
}
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询幂等记录失败").with_source(err))?;
|
||||
|
||||
Ok(row)
|
||||
}
|
||||
|
||||
344
src/services/mail.rs
Normal file
344
src/services/mail.rs
Normal file
@@ -0,0 +1,344 @@
|
||||
use crate::config::Config;
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::services::settings;
|
||||
use crate::state::AppState;
|
||||
|
||||
use chrono::Datelike;
|
||||
use lettre::message::{header::ContentType, MultiPart, SinglePart};
|
||||
use lettre::transport::smtp::authentication::Credentials;
|
||||
use lettre::transport::smtp::client::{Tls, TlsParameters};
|
||||
use lettre::{AsyncSmtpTransport, AsyncTransport, Message, Tokio1Executor};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Mailer {
|
||||
enabled: bool,
|
||||
log_links_when_disabled: bool,
|
||||
from: String,
|
||||
from_name: String,
|
||||
transport: Option<AsyncSmtpTransport<Tokio1Executor>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MailSettings {
|
||||
pub enabled: bool,
|
||||
pub log_links_when_disabled: bool,
|
||||
pub provider: String,
|
||||
pub from: String,
|
||||
pub from_name: String,
|
||||
pub password: String,
|
||||
pub smtp_host: Option<String>,
|
||||
pub smtp_port: Option<u16>,
|
||||
pub smtp_encryption: Option<String>,
|
||||
}
|
||||
|
||||
impl MailSettings {
|
||||
pub fn from_env(config: &Config) -> Self {
|
||||
Self {
|
||||
enabled: config.mail_enabled,
|
||||
log_links_when_disabled: config.mail_log_links_when_disabled,
|
||||
provider: config.mail_provider.clone(),
|
||||
from: config.mail_from.clone(),
|
||||
from_name: config.mail_from_name.clone(),
|
||||
password: config.mail_password.clone(),
|
||||
smtp_host: config.mail_smtp_host.clone(),
|
||||
smtp_port: config.mail_smtp_port,
|
||||
smtp_encryption: config.mail_smtp_encryption.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Mailer {
|
||||
pub fn new(config: &Config) -> Result<Self, AppError> {
|
||||
Self::from_settings(MailSettings::from_env(config))
|
||||
}
|
||||
|
||||
pub fn from_settings(settings: MailSettings) -> Result<Self, AppError> {
|
||||
if !settings.enabled {
|
||||
return Ok(Self {
|
||||
enabled: false,
|
||||
log_links_when_disabled: settings.log_links_when_disabled,
|
||||
from: settings.from,
|
||||
from_name: settings.from_name,
|
||||
transport: None,
|
||||
});
|
||||
}
|
||||
|
||||
if settings.password.trim().is_empty() {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
"邮件服务已启用但未配置授权码/密码",
|
||||
));
|
||||
}
|
||||
|
||||
let smtp = SmtpConfig::from_settings(&settings)?;
|
||||
let creds = Credentials::new(settings.from.clone(), settings.password.clone());
|
||||
|
||||
let tls_params = if smtp.encryption == SmtpEncryption::None {
|
||||
None
|
||||
} else {
|
||||
Some(
|
||||
TlsParameters::new(smtp.host.clone())
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "SMTP TLS 参数错误").with_source(err))?,
|
||||
)
|
||||
};
|
||||
|
||||
let tls = match (smtp.encryption, tls_params) {
|
||||
(SmtpEncryption::Ssl, Some(params)) => Tls::Wrapper(params),
|
||||
(SmtpEncryption::StartTls, Some(params)) => Tls::Required(params),
|
||||
(SmtpEncryption::None, _) => Tls::None,
|
||||
_ => Tls::None,
|
||||
};
|
||||
|
||||
let transport = AsyncSmtpTransport::<Tokio1Executor>::builder_dangerous(&smtp.host)
|
||||
.port(smtp.port)
|
||||
.tls(tls)
|
||||
.credentials(creds)
|
||||
.build();
|
||||
|
||||
Ok(Self {
|
||||
enabled: true,
|
||||
log_links_when_disabled: false,
|
||||
from: settings.from,
|
||||
from_name: settings.from_name,
|
||||
transport: Some(transport),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn send_verification_email(
|
||||
&self,
|
||||
to: &str,
|
||||
username: &str,
|
||||
verification_url: &str,
|
||||
) -> Result<(), AppError> {
|
||||
if !self.enabled {
|
||||
if self.log_links_when_disabled {
|
||||
tracing::info!(
|
||||
to = %to,
|
||||
verification_url = %verification_url,
|
||||
"MAIL_ENABLED=false, verification email link"
|
||||
);
|
||||
} else {
|
||||
tracing::info!(to = %to, "MAIL_ENABLED=false, skip sending verification email");
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let year = chrono::Utc::now().year().to_string();
|
||||
|
||||
let html = render_template(
|
||||
include_str!("../../templates/email_verification.html"),
|
||||
&[
|
||||
("{{username}}", username),
|
||||
("{{verification_url}}", verification_url),
|
||||
("{{year}}", &year),
|
||||
],
|
||||
);
|
||||
let text = render_template(
|
||||
include_str!("../../templates/email_verification.txt"),
|
||||
&[
|
||||
("{{username}}", username),
|
||||
("{{verification_url}}", verification_url),
|
||||
("{{year}}", &year),
|
||||
],
|
||||
);
|
||||
|
||||
self.send_email(to, "验证您的 ImageForge 账号", &text, &html)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn send_password_reset_email(
|
||||
&self,
|
||||
to: &str,
|
||||
username: &str,
|
||||
reset_url: &str,
|
||||
) -> Result<(), AppError> {
|
||||
if !self.enabled {
|
||||
if self.log_links_when_disabled {
|
||||
tracing::info!(to = %to, reset_url = %reset_url, "MAIL_ENABLED=false, password reset link");
|
||||
} else {
|
||||
tracing::info!(to = %to, "MAIL_ENABLED=false, skip sending password reset email");
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let year = chrono::Utc::now().year().to_string();
|
||||
|
||||
let html = render_template(
|
||||
include_str!("../../templates/password_reset.html"),
|
||||
&[
|
||||
("{{username}}", username),
|
||||
("{{reset_url}}", reset_url),
|
||||
("{{year}}", &year),
|
||||
],
|
||||
);
|
||||
let text = render_template(
|
||||
include_str!("../../templates/password_reset.txt"),
|
||||
&[
|
||||
("{{username}}", username),
|
||||
("{{reset_url}}", reset_url),
|
||||
("{{year}}", &year),
|
||||
],
|
||||
);
|
||||
|
||||
self.send_email(to, "重置您的 ImageForge 密码", &text, &html)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_email(
|
||||
&self,
|
||||
to: &str,
|
||||
subject: &str,
|
||||
text_body: &str,
|
||||
html_body: &str,
|
||||
) -> Result<(), AppError> {
|
||||
if !self.enabled {
|
||||
tracing::info!(to = %to, subject = %subject, "MAIL_ENABLED=false, skip sending email");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let Some(transport) = &self.transport else {
|
||||
return Err(AppError::new(ErrorCode::Internal, "邮件服务未初始化"));
|
||||
};
|
||||
|
||||
let from = format!("{} <{}>", self.from_name, self.from);
|
||||
let email = Message::builder()
|
||||
.from(from.parse().map_err(|err| {
|
||||
AppError::new(ErrorCode::InvalidRequest, "MAIL_FROM/MAIL_FROM_NAME 格式错误")
|
||||
.with_source(err)
|
||||
})?)
|
||||
.to(to.parse().map_err(|err| {
|
||||
AppError::new(ErrorCode::InvalidRequest, "收件人邮箱格式错误").with_source(err)
|
||||
})?)
|
||||
.subject(subject)
|
||||
.multipart(
|
||||
MultiPart::alternative()
|
||||
.singlepart(SinglePart::builder()
|
||||
.header(ContentType::TEXT_PLAIN)
|
||||
.body(text_body.to_string()))
|
||||
.singlepart(SinglePart::builder()
|
||||
.header(ContentType::TEXT_HTML)
|
||||
.body(html_body.to_string())),
|
||||
)
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "构建邮件失败").with_source(err))?;
|
||||
|
||||
transport
|
||||
.send(email)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "邮件发送失败").with_source(err))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct SmtpConfig {
|
||||
host: String,
|
||||
port: u16,
|
||||
encryption: SmtpEncryption,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum SmtpEncryption {
|
||||
Ssl,
|
||||
StartTls,
|
||||
None,
|
||||
}
|
||||
|
||||
impl SmtpConfig {
|
||||
fn from_settings(settings: &MailSettings) -> Result<Self, AppError> {
|
||||
if settings.provider.eq_ignore_ascii_case("custom") {
|
||||
let host = settings.smtp_host.clone().ok_or_else(|| {
|
||||
AppError::new(ErrorCode::InvalidRequest, "自定义 SMTP 必须配置 host")
|
||||
})?;
|
||||
let port = settings
|
||||
.smtp_port
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "自定义 SMTP 必须配置端口"))?;
|
||||
let encryption = parse_encryption(settings.smtp_encryption.as_deref().unwrap_or("ssl"))?;
|
||||
return Ok(Self { host, port, encryption });
|
||||
}
|
||||
|
||||
let provider = settings.provider.to_ascii_lowercase();
|
||||
let (host, port, encryption) = match provider.as_str() {
|
||||
"qq" => ("smtp.qq.com", 465, SmtpEncryption::Ssl),
|
||||
"163" => ("smtp.163.com", 465, SmtpEncryption::Ssl),
|
||||
"aliyun_enterprise" => ("smtp.qiye.aliyun.com", 465, SmtpEncryption::Ssl),
|
||||
"tencent_enterprise" => ("smtp.exmail.qq.com", 465, SmtpEncryption::Ssl),
|
||||
"gmail" => ("smtp.gmail.com", 587, SmtpEncryption::StartTls),
|
||||
"outlook" => ("smtp.office365.com", 587, SmtpEncryption::StartTls),
|
||||
other => {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
format!("未知 MAIL_PROVIDER: {other}"),
|
||||
))
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
host: host.to_string(),
|
||||
port,
|
||||
encryption,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_encryption(value: &str) -> Result<SmtpEncryption, AppError> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"ssl" => Ok(SmtpEncryption::Ssl),
|
||||
"starttls" => Ok(SmtpEncryption::StartTls),
|
||||
"none" => Ok(SmtpEncryption::None),
|
||||
_ => Err(AppError::new(
|
||||
ErrorCode::InvalidRequest,
|
||||
"MAIL_SMTP_ENCRYPTION 仅支持 ssl/starttls/none",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn render_template(template: &str, vars: &[(&str, &str)]) -> String {
|
||||
let mut out = template.to_string();
|
||||
for (key, value) in vars {
|
||||
out = out.replace(key, value);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
pub async fn send_verification_email(
|
||||
state: &AppState,
|
||||
to: &str,
|
||||
username: &str,
|
||||
verification_url: &str,
|
||||
) -> Result<(), AppError> {
|
||||
let mailer = resolve_mailer(state).await?;
|
||||
mailer
|
||||
.send_verification_email(to, username, verification_url)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn send_password_reset_email(
|
||||
state: &AppState,
|
||||
to: &str,
|
||||
username: &str,
|
||||
reset_url: &str,
|
||||
) -> Result<(), AppError> {
|
||||
let mailer = resolve_mailer(state).await?;
|
||||
mailer.send_password_reset_email(to, username, reset_url).await
|
||||
}
|
||||
|
||||
pub async fn send_test_email(state: &AppState, to: &str) -> Result<(), AppError> {
|
||||
let mailer = resolve_mailer(state).await?;
|
||||
let year = chrono::Utc::now().year().to_string();
|
||||
let html = format!(
|
||||
"<h2>ImageForge 邮件测试</h2><p>这是一封测试邮件。</p><p>{}</p>",
|
||||
year
|
||||
);
|
||||
let text = format!("ImageForge 邮件测试\n\n这是一封测试邮件。\n{}\n", year);
|
||||
mailer
|
||||
.send_email(to, "ImageForge 邮件测试", &text, &html)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn resolve_mailer(state: &AppState) -> Result<Mailer, AppError> {
|
||||
if let Some(settings) = settings::load_mail_settings(state).await? {
|
||||
return Mailer::from_settings(settings);
|
||||
}
|
||||
Ok(state.mailer.as_ref().clone())
|
||||
}
|
||||
7
src/services/mod.rs
Normal file
7
src/services/mod.rs
Normal file
@@ -0,0 +1,7 @@
|
||||
pub mod mail;
|
||||
pub mod billing;
|
||||
pub mod quota;
|
||||
pub mod compress;
|
||||
pub mod idempotency;
|
||||
pub mod settings;
|
||||
pub mod bootstrap;
|
||||
73
src/services/quota.rs
Normal file
73
src/services/quota.rs
Normal file
@@ -0,0 +1,73 @@
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::state::AppState;
|
||||
|
||||
use chrono::{Duration, Utc};
|
||||
use std::net::IpAddr;
|
||||
|
||||
pub async fn consume_anonymous_units(
|
||||
state: &AppState,
|
||||
session_id: &str,
|
||||
ip: IpAddr,
|
||||
units: u32,
|
||||
) -> Result<(), AppError> {
|
||||
if units == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let date = utc8_date();
|
||||
let session_key = format!("anon_quota:{session_id}:{date}");
|
||||
let ip_key = format!("anon_quota_ip:{ip}:{date}");
|
||||
|
||||
let mut conn = state.redis.clone();
|
||||
|
||||
let limit = state.config.anon_daily_units as i64;
|
||||
let ttl_seconds = 48 * 60 * 60;
|
||||
let inc = units as i64;
|
||||
|
||||
let script = redis::Script::new(
|
||||
r#"
|
||||
local limit = tonumber(ARGV[1])
|
||||
local ttl = tonumber(ARGV[2])
|
||||
local inc = tonumber(ARGV[3])
|
||||
|
||||
local v1 = tonumber(redis.call('GET', KEYS[1]) or '0')
|
||||
local v2 = tonumber(redis.call('GET', KEYS[2]) or '0')
|
||||
|
||||
if v1 + inc > limit or v2 + inc > limit then
|
||||
return -1
|
||||
end
|
||||
|
||||
v1 = redis.call('INCRBY', KEYS[1], inc)
|
||||
v2 = redis.call('INCRBY', KEYS[2], inc)
|
||||
|
||||
if v1 == inc then redis.call('EXPIRE', KEYS[1], ttl) end
|
||||
if v2 == inc then redis.call('EXPIRE', KEYS[2], ttl) end
|
||||
|
||||
return v1
|
||||
"#,
|
||||
);
|
||||
|
||||
let new_value: i64 = script
|
||||
.key(session_key)
|
||||
.key(ip_key)
|
||||
.arg(limit)
|
||||
.arg(ttl_seconds)
|
||||
.arg(inc)
|
||||
.invoke_async(&mut conn)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "匿名配额检查失败").with_source(err))?;
|
||||
|
||||
if new_value < 0 {
|
||||
return Err(AppError::new(
|
||||
ErrorCode::QuotaExceeded,
|
||||
"匿名试用次数已用完（每日 10 次）",
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn utc8_date() -> String {
|
||||
let now = Utc::now() + Duration::hours(8);
|
||||
now.format("%Y-%m-%d").to_string()
|
||||
}
|
||||
227
src/services/settings.rs
Normal file
227
src/services/settings.rs
Normal file
@@ -0,0 +1,227 @@
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::services::mail::MailSettings;
|
||||
use crate::state::AppState;
|
||||
|
||||
use aes_gcm::aead::{Aead, KeyInit};
|
||||
use aes_gcm::{Aes256Gcm, Nonce};
|
||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||
use rand::RngCore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde::de::DeserializeOwned;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MailCustomSmtp {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
pub encryption: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MailConfigStored {
|
||||
pub enabled: bool,
|
||||
pub provider: String,
|
||||
pub from: String,
|
||||
pub from_name: String,
|
||||
pub password_encrypted: Option<String>,
|
||||
pub custom_smtp: Option<MailCustomSmtp>,
|
||||
pub log_links_when_disabled: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct StripeConfigStored {
|
||||
pub secret_key_encrypted: Option<String>,
|
||||
pub webhook_secret_encrypted: Option<String>,
|
||||
pub secret_key_prefix: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StripeSecrets {
|
||||
pub secret_key: String,
|
||||
pub webhook_secret: Option<String>,
|
||||
pub secret_key_prefix: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn load_system_config<T: DeserializeOwned>(
|
||||
state: &AppState,
|
||||
key: &str,
|
||||
) -> Result<Option<T>, AppError> {
|
||||
let value: Option<serde_json::Value> =
|
||||
sqlx::query_scalar("SELECT value FROM system_config WHERE key = $1")
|
||||
.bind(key)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询系统配置失败").with_source(err))?;
|
||||
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let parsed = serde_json::from_value::<T>(value)
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "解析系统配置失败").with_source(err))?;
|
||||
Ok(Some(parsed))
|
||||
}
|
||||
|
||||
pub async fn upsert_system_config(
|
||||
state: &AppState,
|
||||
key: &str,
|
||||
value: serde_json::Value,
|
||||
description: Option<&str>,
|
||||
updated_by: Option<uuid::Uuid>,
|
||||
) -> Result<(), AppError> {
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO system_config (key, value, description, updated_at, updated_by)
|
||||
VALUES ($1, $2, $3, NOW(), $4)
|
||||
ON CONFLICT (key) DO UPDATE
|
||||
SET value = EXCLUDED.value,
|
||||
description = COALESCE(EXCLUDED.description, system_config.description),
|
||||
updated_at = NOW(),
|
||||
updated_by = $4
|
||||
"#,
|
||||
)
|
||||
.bind(key)
|
||||
.bind(value)
|
||||
.bind(description)
|
||||
.bind(updated_by)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新系统配置失败").with_source(err))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn load_mail_settings(state: &AppState) -> Result<Option<MailSettings>, AppError> {
|
||||
let Some(cfg) = load_system_config::<MailConfigStored>(state, "mail").await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let password = match cfg.password_encrypted.as_deref() {
|
||||
Some(value) => decrypt_secret(state, value)?,
|
||||
None => String::new(),
|
||||
};
|
||||
|
||||
Ok(Some(MailSettings {
|
||||
enabled: cfg.enabled,
|
||||
log_links_when_disabled: cfg
|
||||
.log_links_when_disabled
|
||||
.unwrap_or(state.config.mail_log_links_when_disabled),
|
||||
provider: cfg.provider,
|
||||
from: cfg.from,
|
||||
from_name: cfg.from_name,
|
||||
password,
|
||||
smtp_host: cfg.custom_smtp.as_ref().map(|v| v.host.clone()),
|
||||
smtp_port: cfg.custom_smtp.as_ref().map(|v| v.port),
|
||||
smtp_encryption: cfg.custom_smtp.as_ref().map(|v| v.encryption.clone()),
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn load_stripe_secrets(state: &AppState) -> Result<Option<StripeSecrets>, AppError> {
|
||||
let Some(cfg) = load_system_config::<StripeConfigStored>(state, "stripe").await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let secret_key = match cfg.secret_key_encrypted.as_deref() {
|
||||
Some(value) => decrypt_secret(state, value)?,
|
||||
None => String::new(),
|
||||
};
|
||||
|
||||
if secret_key.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let webhook_secret = match cfg.webhook_secret_encrypted.as_deref() {
|
||||
Some(value) if !value.is_empty() => Some(decrypt_secret(state, value)?),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
Ok(Some(StripeSecrets {
|
||||
secret_key,
|
||||
webhook_secret,
|
||||
secret_key_prefix: cfg.secret_key_prefix,
|
||||
}))
|
||||
}
|
||||
|
||||
pub fn encrypt_secret(state: &AppState, plain: &str) -> Result<String, AppError> {
|
||||
let key = derive_key(&state.config.api_key_pepper);
|
||||
let cipher = Aes256Gcm::new_from_slice(&key)
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "加密密钥初始化失败").with_source(err))?;
|
||||
|
||||
let mut nonce_bytes = [0u8; 12];
|
||||
rand::rngs::OsRng.fill_bytes(&mut nonce_bytes);
|
||||
let nonce = Nonce::from_slice(&nonce_bytes);
|
||||
|
||||
let ciphertext = cipher
|
||||
.encrypt(nonce, plain.as_bytes())
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "加密失败").with_source(err))?;
|
||||
|
||||
let mut out = Vec::with_capacity(nonce_bytes.len() + ciphertext.len());
|
||||
out.extend_from_slice(&nonce_bytes);
|
||||
out.extend_from_slice(&ciphertext);
|
||||
|
||||
Ok(URL_SAFE_NO_PAD.encode(out))
|
||||
}
|
||||
|
||||
pub fn decrypt_secret(state: &AppState, encoded: &str) -> Result<String, AppError> {
|
||||
let key = derive_key(&state.config.api_key_pepper);
|
||||
let cipher = Aes256Gcm::new_from_slice(&key)
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "解密密钥初始化失败").with_source(err))?;
|
||||
|
||||
let decoded = URL_SAFE_NO_PAD
|
||||
.decode(encoded.as_bytes())
|
||||
.map_err(|err| AppError::new(ErrorCode::InvalidRequest, "密文格式错误").with_source(err))?;
|
||||
|
||||
if decoded.len() < 12 {
|
||||
return Err(AppError::new(ErrorCode::InvalidRequest, "密文长度错误"));
|
||||
}
|
||||
|
||||
let (nonce_bytes, cipher_bytes) = decoded.split_at(12);
|
||||
let nonce = Nonce::from_slice(nonce_bytes);
|
||||
let plain = cipher
|
||||
.decrypt(nonce, cipher_bytes)
|
||||
.map_err(|err| AppError::new(ErrorCode::InvalidRequest, "密文解密失败").with_source(err))?;
|
||||
|
||||
String::from_utf8(plain)
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "解密文本编码错误").with_source(err))
|
||||
}
|
||||
|
||||
fn derive_key(secret: &str) -> [u8; 32] {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(secret.as_bytes());
|
||||
let result = hasher.finalize();
|
||||
let mut out = [0u8; 32];
|
||||
out.copy_from_slice(&result);
|
||||
out
|
||||
}
|
||||
|
||||
pub async fn get_stripe_secret(state: &AppState) -> Result<String, AppError> {
|
||||
if let Some(cfg) = load_stripe_secrets(state).await? {
|
||||
if !cfg.secret_key.trim().is_empty() {
|
||||
return Ok(cfg.secret_key);
|
||||
}
|
||||
}
|
||||
|
||||
state
|
||||
.config
|
||||
.stripe_secret_key
|
||||
.clone()
|
||||
.filter(|v| !v.trim().is_empty())
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "未配置 Stripe Secret Key"))
|
||||
}
|
||||
|
||||
pub async fn get_stripe_webhook_secret(state: &AppState) -> Result<String, AppError> {
|
||||
if let Some(cfg) = load_stripe_secrets(state).await? {
|
||||
if let Some(secret) = cfg.webhook_secret {
|
||||
if !secret.trim().is_empty() {
|
||||
return Ok(secret);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
state
|
||||
.config
|
||||
.stripe_webhook_secret
|
||||
.clone()
|
||||
.filter(|v| !v.trim().is_empty())
|
||||
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "未配置 Stripe Webhook Secret"))
|
||||
}
|
||||
10
src/state.rs
Normal file
10
src/state.rs
Normal file
@@ -0,0 +1,10 @@
|
||||
use crate::config::Config;
|
||||
use crate::services::mail::Mailer;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub config: Config,
|
||||
pub db: sqlx::PgPool,
|
||||
pub redis: redis::aio::ConnectionManager,
|
||||
pub mailer: std::sync::Arc<Mailer>,
|
||||
}
|
||||
738
src/worker/mod.rs
Normal file
738
src/worker/mod.rs
Normal file
@@ -0,0 +1,738 @@
|
||||
use crate::error::{AppError, ErrorCode};
|
||||
use crate::services::billing;
|
||||
use crate::services::compress;
|
||||
use crate::services::quota;
|
||||
use crate::state::AppState;
|
||||
|
||||
use redis::streams::StreamReadOptions;
|
||||
use redis::AsyncCommands;
|
||||
use sqlx::FromRow;
|
||||
use std::net::IpAddr;
|
||||
use std::time::Instant;
|
||||
use uuid::Uuid;
|
||||
|
||||
const STREAM_KEY: &str = "stream:compress_jobs";
|
||||
const GROUP_NAME: &str = "compress_workers";
|
||||
|
||||
pub async fn run(state: AppState) -> Result<(), AppError> {
|
||||
tracing::info!("Worker started");
|
||||
|
||||
if let Err(err) = crate::services::bootstrap::ensure_schema(&state).await {
|
||||
tracing::error!(error = %err, "数据库结构初始化失败");
|
||||
}
|
||||
|
||||
let consumer = format!("worker_{}", Uuid::new_v4());
|
||||
ensure_group(&state, &consumer).await?;
|
||||
|
||||
let mut last_maintenance = Instant::now();
|
||||
|
||||
loop {
|
||||
if let Err(err) = poll_once(&state, &consumer).await {
|
||||
tracing::error!(error = ?err, "worker poll error");
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
|
||||
}
|
||||
|
||||
if last_maintenance.elapsed().as_secs() >= 300 {
|
||||
if let Err(err) = maintenance(&state).await {
|
||||
tracing::error!(error = ?err, "maintenance failed");
|
||||
}
|
||||
last_maintenance = Instant::now();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn ensure_group(state: &AppState, _consumer: &str) -> Result<(), AppError> {
|
||||
let mut conn = state.redis.clone();
|
||||
|
||||
let res: Result<redis::Value, redis::RedisError> = redis::cmd("XGROUP")
|
||||
.arg("CREATE")
|
||||
.arg(STREAM_KEY)
|
||||
.arg(GROUP_NAME)
|
||||
.arg("0")
|
||||
.arg("MKSTREAM")
|
||||
.query_async(&mut conn)
|
||||
.await;
|
||||
|
||||
match res {
|
||||
Ok(_) => Ok(()),
|
||||
Err(err) => {
|
||||
let msg = err.to_string();
|
||||
if msg.contains("BUSYGROUP") {
|
||||
return Ok(());
|
||||
}
|
||||
Err(AppError::new(ErrorCode::Internal, "初始化队列失败").with_source(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn poll_once(state: &AppState, consumer: &str) -> Result<(), AppError> {
|
||||
let mut conn = state.redis.clone();
|
||||
|
||||
let opts = StreamReadOptions::default()
|
||||
.group(GROUP_NAME, consumer)
|
||||
.count(1)
|
||||
.block(5000);
|
||||
|
||||
let reply: redis::streams::StreamReadReply = conn
|
||||
.xread_options(&[STREAM_KEY], &[">"], &opts)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "读取队列失败").with_source(err))?;
|
||||
|
||||
if reply.keys.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
for key in reply.keys {
|
||||
for msg in key.ids {
|
||||
let Some(task_id_str) = msg.get::<String>("task_id") else {
|
||||
ack_message(&mut conn, &msg.id).await?;
|
||||
continue;
|
||||
};
|
||||
|
||||
let task_id = match Uuid::parse_str(&task_id_str) {
|
||||
Ok(v) => v,
|
||||
Err(_) => {
|
||||
ack_message(&mut conn, &msg.id).await?;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(err) = process_task(state, task_id).await {
|
||||
tracing::error!(task_id = %task_id, error = %err, "task processing failed");
|
||||
}
|
||||
|
||||
ack_message(&mut conn, &msg.id).await?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn ack_message(conn: &mut redis::aio::ConnectionManager, msg_id: &str) -> Result<(), AppError> {
|
||||
let _: i64 = redis::cmd("XACK")
|
||||
.arg(STREAM_KEY)
|
||||
.arg(GROUP_NAME)
|
||||
.arg(msg_id)
|
||||
.query_async(conn)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "确认队列消息失败").with_source(err))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct TaskProcRow {
|
||||
id: Uuid,
|
||||
status: String,
|
||||
compression_level: String,
|
||||
compression_rate: Option<i16>,
|
||||
max_width: Option<i32>,
|
||||
max_height: Option<i32>,
|
||||
preserve_metadata: bool,
|
||||
total_files: i32,
|
||||
completed_files: i32,
|
||||
failed_files: i32,
|
||||
user_id: Option<Uuid>,
|
||||
session_id: Option<String>,
|
||||
api_key_id: Option<Uuid>,
|
||||
source: String,
|
||||
client_ip: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
struct TaskFileProcRow {
|
||||
id: Uuid,
|
||||
storage_path: Option<String>,
|
||||
original_name: String,
|
||||
original_format: String,
|
||||
output_format: String,
|
||||
original_size: i64,
|
||||
status: String,
|
||||
}
|
||||
|
||||
async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
|
||||
let mut task: TaskProcRow = sqlx::query_as(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
status::text AS status,
|
||||
compression_level::text AS compression_level,
|
||||
compression_rate,
|
||||
max_width,
|
||||
max_height,
|
||||
preserve_metadata,
|
||||
total_files,
|
||||
completed_files,
|
||||
failed_files,
|
||||
user_id,
|
||||
session_id,
|
||||
api_key_id,
|
||||
source::text AS source,
|
||||
client_ip::text AS client_ip
|
||||
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 matches!(task.status.as_str(), "completed" | "failed" | "cancelled") {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let updated = sqlx::query(
|
||||
r#"
|
||||
UPDATE tasks
|
||||
SET status = 'processing', started_at = NOW()
|
||||
WHERE id = $1 AND status = 'pending'
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新任务状态失败").with_source(err))?;
|
||||
|
||||
if updated.rows_affected() == 0 && task.status == "pending" {
|
||||
// Another worker might have taken it.
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Refresh task row after status change
|
||||
task.status = "processing".to_string();
|
||||
|
||||
let compression_rate = task
|
||||
.compression_rate
|
||||
.and_then(|v| u8::try_from(v).ok());
|
||||
let level = compression_rate
|
||||
.map(compress::rate_to_level)
|
||||
.unwrap_or(compress::parse_level(&task.compression_level)?);
|
||||
let max_width = task.max_width.and_then(|v| u32::try_from(v).ok());
|
||||
let max_height = task.max_height.and_then(|v| u32::try_from(v).ok());
|
||||
|
||||
let mut files: Vec<TaskFileProcRow> = sqlx::query_as(
|
||||
r#"
|
||||
SELECT
|
||||
id,
|
||||
storage_path,
|
||||
original_name,
|
||||
original_format,
|
||||
output_format,
|
||||
original_size,
|
||||
status::text AS status
|
||||
FROM task_files
|
||||
WHERE task_id = $1
|
||||
ORDER BY created_at ASC
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询任务文件失败").with_source(err))?;
|
||||
|
||||
let billing_ctx = if let Some(user_id) = task.user_id {
|
||||
Some(billing::get_user_billing(state, user_id).await?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let anon_ip: Option<IpAddr> = task
|
||||
.client_ip
|
||||
.as_deref()
|
||||
.and_then(|s| s.parse::<IpAddr>().ok());
|
||||
|
||||
for file in &mut files {
|
||||
// Stop early if cancelled.
|
||||
let status: Option<String> = sqlx::query_scalar("SELECT status::text FROM tasks WHERE id = $1")
|
||||
.bind(task_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.unwrap_or(None);
|
||||
if matches!(status.as_deref(), Some("cancelled")) {
|
||||
break;
|
||||
}
|
||||
|
||||
if file.status != "pending" {
|
||||
continue;
|
||||
}
|
||||
|
||||
let updated = sqlx::query("UPDATE task_files SET status = 'processing' WHERE id = $1 AND status = 'pending'")
|
||||
.bind(file.id)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.unwrap_or_else(|_| sqlx::postgres::PgQueryResult::default());
|
||||
if updated.rows_affected() == 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(input_path) = file.storage_path.clone() else {
|
||||
mark_file_failed(state, task_id, file.id, "原文件不存在").await?;
|
||||
continue;
|
||||
};
|
||||
|
||||
let input_bytes = match tokio::fs::read(&input_path).await {
|
||||
Ok(v) => v,
|
||||
Err(_) => {
|
||||
mark_file_failed(state, task_id, file.id, "读取原文件失败").await?;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let format_in = parse_image_fmt(&file.original_format)?;
|
||||
let format_out = parse_image_fmt(&file.output_format)?;
|
||||
|
||||
let compressed = match compress::compress_image_bytes(
|
||||
state,
|
||||
&input_bytes,
|
||||
format_in,
|
||||
format_out,
|
||||
level,
|
||||
compression_rate,
|
||||
max_width,
|
||||
max_height,
|
||||
task.preserve_metadata,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(v) => v,
|
||||
Err(err) => {
|
||||
mark_file_failed(state, task_id, file.id, &err.message).await?;
|
||||
let _ = tokio::fs::remove_file(&input_path).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let original_size = input_bytes.len() as u64;
|
||||
let compressed_size = compressed.len() as u64;
|
||||
let saved_percent = if original_size == 0 {
|
||||
0.0
|
||||
} else {
|
||||
(original_size.saturating_sub(compressed_size) as f64) * 100.0 / (original_size as f64)
|
||||
};
|
||||
let charge_units = compressed_size < original_size;
|
||||
|
||||
// Anonymous quota enforcement requires session_id + client_ip.
|
||||
if task.user_id.is_none() && charge_units {
|
||||
let Some(session_id) = task.session_id.as_deref() else {
|
||||
mark_file_failed(state, task_id, file.id, "匿名任务缺少 session_id").await?;
|
||||
let _ = tokio::fs::remove_file(&input_path).await;
|
||||
continue;
|
||||
};
|
||||
let Some(ip) = anon_ip else {
|
||||
mark_file_failed(state, task_id, file.id, "匿名任务缺少 client_ip").await?;
|
||||
let _ = tokio::fs::remove_file(&input_path).await;
|
||||
continue;
|
||||
};
|
||||
if let Err(err) = quota::consume_anonymous_units(state, session_id, ip, 1).await {
|
||||
mark_file_failed(state, task_id, file.id, &err.message).await?;
|
||||
let _ = tokio::fs::remove_file(&input_path).await;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
let output_path = format!(
|
||||
"{}/{}.{}",
|
||||
state.config.storage_path,
|
||||
file.id,
|
||||
format_out.extension()
|
||||
);
|
||||
if let Err(err) = tokio::fs::write(&output_path, &compressed).await {
|
||||
mark_file_failed(state, task_id, file.id, "写入压缩文件失败").await?;
|
||||
let _ = tokio::fs::remove_file(&input_path).await;
|
||||
return Err(AppError::new(ErrorCode::StorageUnavailable, "写入压缩文件失败").with_source(err));
|
||||
}
|
||||
|
||||
if let Err(err) = finalize_file(
|
||||
state,
|
||||
&billing_ctx,
|
||||
task.api_key_id,
|
||||
&task.source,
|
||||
task_id,
|
||||
file.id,
|
||||
&output_path,
|
||||
original_size as i64,
|
||||
compressed_size as i64,
|
||||
saved_percent,
|
||||
format_in,
|
||||
format_out,
|
||||
charge_units,
|
||||
)
|
||||
.await
|
||||
{
|
||||
// If quota exceeded for paid users, don't leave output behind.
|
||||
if err.code == ErrorCode::QuotaExceeded {
|
||||
let _ = tokio::fs::remove_file(&output_path).await;
|
||||
mark_file_failed(state, task_id, file.id, &err.message).await?;
|
||||
} else {
|
||||
mark_file_failed(state, task_id, file.id, &err.message).await?;
|
||||
}
|
||||
let _ = tokio::fs::remove_file(&input_path).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Success: remove original.
|
||||
let _ = tokio::fs::remove_file(&input_path).await;
|
||||
}
|
||||
|
||||
finalize_task_status(state, task_id).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_image_fmt(value: &str) -> Result<compress::ImageFmt, AppError> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"png" => Ok(compress::ImageFmt::Png),
|
||||
"jpeg" | "jpg" => Ok(compress::ImageFmt::Jpeg),
|
||||
"webp" => Ok(compress::ImageFmt::Webp),
|
||||
"avif" => Ok(compress::ImageFmt::Avif),
|
||||
"gif" => Ok(compress::ImageFmt::Gif),
|
||||
"bmp" => Ok(compress::ImageFmt::Bmp),
|
||||
"tif" | "tiff" => Ok(compress::ImageFmt::Tiff),
|
||||
"ico" => Ok(compress::ImageFmt::Ico),
|
||||
_ => Err(AppError::new(ErrorCode::InvalidRequest, "未知图片格式")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn finalize_file(
|
||||
state: &AppState,
|
||||
billing_ctx: &Option<billing::BillingContext>,
|
||||
api_key_id: Option<Uuid>,
|
||||
source: &str,
|
||||
task_id: Uuid,
|
||||
task_file_id: Uuid,
|
||||
output_path: &str,
|
||||
bytes_in: i64,
|
||||
bytes_out: i64,
|
||||
saved_percent: f64,
|
||||
format_in: compress::ImageFmt,
|
||||
format_out: compress::ImageFmt,
|
||||
charge_units: bool,
|
||||
) -> Result<(), AppError> {
|
||||
let mut tx = state
|
||||
.db
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?;
|
||||
|
||||
// Paid users: charge before marking file completed (atomic w/ status update).
|
||||
if charge_units {
|
||||
if let Some(billing) = billing_ctx {
|
||||
charge_one_unit(
|
||||
&mut tx,
|
||||
billing,
|
||||
api_key_id,
|
||||
source,
|
||||
task_id,
|
||||
task_file_id,
|
||||
format_in,
|
||||
format_out,
|
||||
bytes_in as u64,
|
||||
bytes_out as u64,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE task_files
|
||||
SET storage_path = $2,
|
||||
compressed_size = $3,
|
||||
saved_percent = $4,
|
||||
status = 'completed',
|
||||
completed_at = NOW()
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(task_file_id)
|
||||
.bind(output_path)
|
||||
.bind(bytes_out)
|
||||
.bind(saved_percent)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新文件失败").with_source(err))?;
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE tasks
|
||||
SET completed_files = completed_files + 1,
|
||||
total_compressed_size = total_compressed_size + $2
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(bytes_out)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新任务统计失败").with_source(err))?;
|
||||
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn mark_file_failed(state: &AppState, task_id: Uuid, task_file_id: Uuid, message: &str) -> Result<(), AppError> {
|
||||
let mut tx = state
|
||||
.db
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?;
|
||||
|
||||
let updated = sqlx::query(
|
||||
r#"
|
||||
UPDATE task_files
|
||||
SET status = 'failed',
|
||||
error_message = $2,
|
||||
storage_path = NULL,
|
||||
completed_at = NOW()
|
||||
WHERE id = $1
|
||||
AND status NOT IN ('completed', 'failed')
|
||||
"#,
|
||||
)
|
||||
.bind(task_file_id)
|
||||
.bind(message)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新文件失败").with_source(err))?;
|
||||
|
||||
if updated.rows_affected() > 0 {
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE tasks
|
||||
SET failed_files = failed_files + 1
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(task_id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新任务统计失败").with_source(err))?;
|
||||
}
|
||||
|
||||
tx.commit()
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn mark_task_failed(state: &AppState, task_id: Uuid, message: &str) -> Result<(), AppError> {
|
||||
sqlx::query("UPDATE tasks SET status = 'failed', error_message = $2, completed_at = NOW() WHERE id = $1")
|
||||
.bind(task_id)
|
||||
.bind(message)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新任务失败").with_source(err))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn finalize_task_status(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
|
||||
let row: Option<(i32, i32, i32, String)> = sqlx::query_as(
|
||||
"SELECT total_files, completed_files, failed_files, status::text AS status FROM tasks WHERE id = $1",
|
||||
)
|
||||
.bind(task_id)
|
||||
.fetch_optional(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "查询任务失败").with_source(err))?;
|
||||
|
||||
let Some((total, completed, failed, status)) = row else { return Ok(()); };
|
||||
if status == "cancelled" {
|
||||
let paths: Vec<Option<String>> = sqlx::query_scalar(
|
||||
"SELECT storage_path FROM task_files WHERE task_id = $1 AND status IN ('pending','processing')",
|
||||
)
|
||||
.bind(task_id)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
for p in paths.into_iter().flatten() {
|
||||
let _ = tokio::fs::remove_file(p).await;
|
||||
}
|
||||
|
||||
let _ = sqlx::query(
|
||||
"UPDATE task_files SET status = 'failed', error_message = '已取消', storage_path = NULL, completed_at = NOW() WHERE task_id = $1 AND status IN ('pending','processing')",
|
||||
)
|
||||
.bind(task_id)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
let _ = sqlx::query(
|
||||
"UPDATE tasks SET failed_files = GREATEST(total_files - completed_files, 0), completed_at = NOW() WHERE id = $1 AND completed_at IS NULL",
|
||||
)
|
||||
.bind(task_id)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if completed + failed >= total && total > 0 {
|
||||
let final_status = if completed == 0 && failed == total {
|
||||
"failed"
|
||||
} else {
|
||||
"completed"
|
||||
};
|
||||
sqlx::query(
|
||||
"UPDATE tasks SET status = $2::task_status, completed_at = NOW() WHERE id = $1",
|
||||
)
|
||||
.bind(task_id)
|
||||
.bind(final_status)
|
||||
.execute(&state.db)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "更新任务状态失败").with_source(err))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn charge_one_unit(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
|
||||
billing: &billing::BillingContext,
|
||||
api_key_id: Option<Uuid>,
|
||||
source: &str,
|
||||
task_id: Uuid,
|
||||
task_file_id: Uuid,
|
||||
format_in: compress::ImageFmt,
|
||||
format_out: compress::ImageFmt,
|
||||
bytes_in: u64,
|
||||
bytes_out: u64,
|
||||
) -> Result<(), AppError> {
|
||||
// Ensure usage period row exists.
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO usage_periods (user_id, subscription_id, period_start, period_end)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
ON CONFLICT (user_id, period_start, period_end) DO NOTHING
|
||||
"#,
|
||||
)
|
||||
.bind(billing.user_id)
|
||||
.bind(billing.subscription_id)
|
||||
.bind(billing.period_start)
|
||||
.bind(billing.period_end)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "初始化用量周期失败").with_source(err))?;
|
||||
|
||||
let updated: Option<i32> = sqlx::query_scalar(
|
||||
r#"
|
||||
UPDATE usage_periods
|
||||
SET used_units = used_units + 1,
|
||||
bytes_in = bytes_in + $1,
|
||||
bytes_out = bytes_out + $2,
|
||||
updated_at = NOW()
|
||||
WHERE user_id = $3
|
||||
AND period_start = $4
|
||||
AND period_end = $5
|
||||
AND used_units + 1 <= $6 + bonus_units
|
||||
RETURNING used_units
|
||||
"#,
|
||||
)
|
||||
.bind(bytes_in as i64)
|
||||
.bind(bytes_out as i64)
|
||||
.bind(billing.user_id)
|
||||
.bind(billing.period_start)
|
||||
.bind(billing.period_end)
|
||||
.bind(billing.plan.included_units_per_period)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "扣减配额失败").with_source(err))?;
|
||||
|
||||
if updated.is_none() {
|
||||
return Err(AppError::new(ErrorCode::QuotaExceeded, "当期配额已用完"));
|
||||
}
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO usage_events (
|
||||
user_id, api_key_id, source,
|
||||
task_id, task_file_id,
|
||||
units, bytes_in, bytes_out, format_in, format_out
|
||||
) VALUES (
|
||||
$1, $2, $3::task_source,
|
||||
$4, $5,
|
||||
1, $6, $7, $8, $9
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(billing.user_id)
|
||||
.bind(api_key_id)
|
||||
.bind(source)
|
||||
.bind(task_id)
|
||||
.bind(task_file_id)
|
||||
.bind(bytes_in as i64)
|
||||
.bind(bytes_out as i64)
|
||||
.bind(format_in.as_str())
|
||||
.bind(format_out.as_str())
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_err(|err| AppError::new(ErrorCode::Internal, "写入用量明细失败").with_source(err))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn maintenance(state: &AppState) -> Result<(), AppError> {
|
||||
cleanup_expired_tasks(state).await?;
|
||||
cleanup_expired_records(state).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> {
|
||||
let _ = sqlx::query("DELETE FROM idempotency_keys WHERE expires_at < NOW()")
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
let _ = sqlx::query("DELETE FROM email_verifications WHERE expires_at < NOW() AND verified_at IS NULL")
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
let _ = sqlx::query("DELETE FROM password_resets WHERE expires_at < NOW() - INTERVAL '7 days'")
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
let _ = sqlx::query("DELETE FROM webhook_events WHERE received_at < NOW() - INTERVAL '90 days'")
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cleanup_expired_tasks(state: &AppState) -> Result<(), AppError> {
|
||||
let task_ids: Vec<Uuid> = sqlx::query_scalar("SELECT id FROM tasks WHERE expires_at < NOW() LIMIT 200")
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
if task_ids.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if state.config.storage_type.to_ascii_lowercase() == "local" {
|
||||
for task_id in &task_ids {
|
||||
let paths: Vec<Option<String>> =
|
||||
sqlx::query_scalar("SELECT storage_path FROM task_files WHERE task_id = $1")
|
||||
.bind(task_id)
|
||||
.fetch_all(&state.db)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
for p in paths.into_iter().flatten() {
|
||||
let _ = tokio::fs::remove_file(p).await;
|
||||
}
|
||||
|
||||
let zip_path = format!("{}/zips/{task_id}.zip", state.config.storage_path);
|
||||
let _ = tokio::fs::remove_file(zip_path).await;
|
||||
|
||||
let orig_dir = format!("{}/orig/{task_id}", state.config.storage_path);
|
||||
let _ = tokio::fs::remove_dir_all(orig_dir).await;
|
||||
}
|
||||
}
|
||||
|
||||
for task_id in task_ids {
|
||||
let _ = sqlx::query("DELETE FROM tasks WHERE id = $1")
|
||||
.bind(task_id)
|
||||
.execute(&state.db)
|
||||
.await;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user