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

perf: improve compression reliability and deployment safety

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

121
src/api/admin.rs
View File

@@ -11,7 +11,7 @@ use axum::extract::{ConnectInfo, Path, Query, State};
use axum::http::HeaderMap;
use axum::routing::{get, post, put};
use axum::{Json, Router};
use chrono::{DateTime, Datelike, Duration, FixedOffset, Timelike, TimeZone, Utc};
use chrono::{DateTime, Datelike, Duration, FixedOffset, TimeZone, Timelike, Utc};
use serde::{Deserialize, Serialize};
use sqlx::FromRow;
use std::net::{IpAddr, SocketAddr};
@@ -108,7 +108,10 @@ async fn resolve_user_id(state: &AppState, identifier: &str) -> Result<Uuid, App
.map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?;
if ids.len() > 1 {
return Err(AppError::new(ErrorCode::InvalidRequest, "用户 ID 前缀不唯一"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"用户 ID 前缀不唯一",
));
}
}
@@ -260,7 +263,10 @@ async fn list_users(
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 search = query.search.map(|s| s.trim().to_string()).filter(|s| !s.is_empty());
let search = query
.search
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty());
let total: i64 = if let Some(search) = &search {
let keyword = format!("%{}%", search);
@@ -414,7 +420,10 @@ async fn list_tasks(
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 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 status::text = $1")
@@ -723,7 +732,12 @@ async fn grant_credits(
.map_err(|err| AppError::new(ErrorCode::Internal, "查询订阅失败").with_source(err))?;
let (subscription_id, period_start, period_end, plan_id) = if let Some(sub) = sub {
(Some(sub.id), sub.current_period_start, sub.current_period_end, Some(sub.plan_id))
(
Some(sub.id),
sub.current_period_start,
sub.current_period_end,
Some(sub.plan_id),
)
} else {
let (start, end) = billing::current_month_period_utc8(Utc::now());
(None, start, end, None)
@@ -853,7 +867,10 @@ async fn create_manual_subscription(
let months = req.months.unwrap_or(1);
if months <= 0 || months > 24 {
return Err(AppError::new(ErrorCode::InvalidRequest, "months 需在 1-24 之间"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"months 需在 1-24 之间",
));
}
let user_id = resolve_user_id(&state, &req.user_id).await?;
@@ -1012,7 +1029,10 @@ fn days_in_month(tz: FixedOffset, year: i32, month: u32) -> u32 {
} else {
(year, month + 1)
};
let first_next = tz.with_ymd_and_hms(next_year, next_month, 1, 0, 0, 0).single().unwrap();
let first_next = tz
.with_ymd_and_hms(next_year, next_month, 1, 0, 0, 0)
.single()
.unwrap();
let last = first_next - Duration::days(1);
last.day()
}
@@ -1155,29 +1175,30 @@ async fn get_stripe_config(
let (_jar, _admin_id) = require_admin(&state, jar, &headers, ip).await?;
let stored = settings::load_system_config::<StripeConfigStored>(&state, "stripe").await?;
let (secret_key_configured, webhook_secret_configured, secret_key_prefix) = if let Some(cfg) = stored {
(
cfg.secret_key_encrypted.as_ref().is_some(),
cfg.webhook_secret_encrypted.as_ref().is_some(),
cfg.secret_key_prefix,
)
} else {
let env_secret = state
.config
.stripe_secret_key
.as_ref()
.filter(|v| !v.trim().is_empty());
let env_webhook = state
.config
.stripe_webhook_secret
.as_ref()
.filter(|v| !v.trim().is_empty());
(
env_secret.is_some(),
env_webhook.is_some(),
env_secret.map(|value| mask_secret(value)),
)
};
let (secret_key_configured, webhook_secret_configured, secret_key_prefix) =
if let Some(cfg) = stored {
(
cfg.secret_key_encrypted.as_ref().is_some(),
cfg.webhook_secret_encrypted.as_ref().is_some(),
cfg.secret_key_prefix,
)
} else {
let env_secret = state
.config
.stripe_secret_key
.as_ref()
.filter(|v| !v.trim().is_empty());
let env_webhook = state
.config
.stripe_webhook_secret
.as_ref()
.filter(|v| !v.trim().is_empty());
(
env_secret.is_some(),
env_webhook.is_some(),
env_secret.map(|value| mask_secret(value)),
)
};
Ok(Json(Envelope {
success: true,
@@ -1227,7 +1248,8 @@ async fn update_stripe_config(
if webhook_secret.is_empty() {
stored.webhook_secret_encrypted = None;
} else {
stored.webhook_secret_encrypted = Some(settings::encrypt_secret(&state, &webhook_secret)?);
stored.webhook_secret_encrypted =
Some(settings::encrypt_secret(&state, &webhook_secret)?);
}
}
@@ -1343,10 +1365,16 @@ async fn update_mail_config(
if req.provider.eq_ignore_ascii_case("custom") {
let custom = req.custom_smtp.as_ref().ok_or_else(|| {
AppError::new(ErrorCode::InvalidRequest, "自定义 SMTP 需要填写 host/port/encryption")
AppError::new(
ErrorCode::InvalidRequest,
"自定义 SMTP 需要填写 host/port/encryption",
)
})?;
if custom.host.trim().is_empty() {
return Err(AppError::new(ErrorCode::InvalidRequest, "SMTP host 不能为空"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"SMTP host 不能为空",
));
}
}
@@ -1387,8 +1415,9 @@ async fn update_mail_config(
settings::upsert_system_config(
&state,
"mail",
serde_json::to_value(&stored)
.map_err(|err| AppError::new(ErrorCode::Internal, "序列化邮件配置失败").with_source(err))?,
serde_json::to_value(&stored).map_err(|err| {
AppError::new(ErrorCode::Internal, "序列化邮件配置失败").with_source(err)
})?,
Some("邮件服务配置"),
Some(admin_id),
)
@@ -1423,15 +1452,11 @@ async fn test_mail(
let ip = context::client_ip(&headers, addr.ip());
let (_jar, admin_id) = require_admin(&state, jar, &headers, ip).await?;
let to = if let Some(to) = req.to.as_ref().map(|v| v.trim().to_string()) {
if to.is_empty() {
None
} else {
Some(to)
}
} else {
None
};
let to = req
.to
.as_ref()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty());
let recipient = if let Some(to) = to {
to
@@ -1440,13 +1465,17 @@ async fn test_mail(
.bind(admin_id)
.fetch_one(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询管理员邮箱失败").with_source(err))?;
.map_err(|err| {
AppError::new(ErrorCode::Internal, "查询管理员邮箱失败").with_source(err)
})?;
email
};
mail::send_test_email(&state, &recipient)
.await
.map_err(|err| AppError::new(ErrorCode::MailSendFailed, "测试邮件发送失败").with_source(err))?;
.map_err(|err| {
AppError::new(ErrorCode::MailSendFailed, "测试邮件发送失败").with_source(err)
})?;
Ok(Json(Envelope {
success: true,

66
src/api/auth.rs
View File

@@ -1,16 +1,11 @@
use crate::auth;
use crate::api::envelope::Envelope;
use crate::auth;
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 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;
@@ -106,8 +101,12 @@ async fn register(
.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 (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);
@@ -133,7 +132,9 @@ async fn register(
mail::send_verification_email(&state, &user.email, &user.username, &verification_url)
.await
.map_err(|err| AppError::new(ErrorCode::MailSendFailed, "验证邮件发送失败").with_source(err))?;
.map_err(|err| {
AppError::new(ErrorCode::MailSendFailed, "验证邮件发送失败").with_source(err)
})?;
let body = RegisterResponse {
user: UserView {
@@ -159,7 +160,10 @@ async fn login(
) -> Result<Json<Envelope<LoginResponse>>, AppError> {
let identity = req.email.trim();
if identity.is_empty() {
return Err(AppError::new(ErrorCode::InvalidRequest, "邮箱或用户名不能为空"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"邮箱或用户名不能为空",
));
}
let user = if identity.contains('@') {
@@ -210,8 +214,12 @@ async fn login(
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)?;
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,
@@ -241,7 +249,11 @@ async fn send_verification(
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 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)
@@ -257,7 +269,10 @@ async fn send_verification(
.unwrap_or(());
}
if count > 1 {
return Err(AppError::new(ErrorCode::RateLimited, "发送过于频繁，请稍后再试"));
return Err(AppError::new(
ErrorCode::RateLimited,
"发送过于频繁，请稍后再试",
));
}
let user = sqlx::query_as::<_, UserRow>(
@@ -313,7 +328,9 @@ async fn send_verification(
mail::send_verification_email(&state, &user.email, &user.username, &verification_url)
.await
.map_err(|err| AppError::new(ErrorCode::MailSendFailed, "验证邮件发送失败").with_source(err))?;
.map_err(|err| {
AppError::new(ErrorCode::MailSendFailed, "验证邮件发送失败").with_source(err)
})?;
Ok(Json(Envelope {
success: true,
@@ -434,7 +451,8 @@ async fn forgot_password(
state.config.public_base_url, reset_token
);
let _ = mail::send_password_reset_email(&state, &user.email, &user.username, &reset_url).await;
let _ =
mail::send_password_reset_email(&state, &user.email, &user.username, &reset_url).await;
}
Ok(Json(Envelope {
@@ -497,12 +515,14 @@ async fn reset_password(
.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))?;
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

50
src/api/billing.rs
View File

@@ -46,7 +46,9 @@ struct PlansResponse {
plans: Vec<PlanView>,
}
async fn list_plans(State(state): State<AppState>) -> Result<Json<Envelope<PlansResponse>>, AppError> {
async fn list_plans(
State(state): State<AppState>,
) -> Result<Json<Envelope<PlansResponse>>, AppError> {
let plans = sqlx::query_as::<_, PlanView>(
r#"
SELECT
@@ -422,7 +424,8 @@ async fn create_checkout(
});
let mut idem_acquired = false;
if let (Some(idem), Some(request_hash)) = (idempotency_key.as_deref(), request_hash.as_deref()) {
if let (Some(idem), Some(request_hash)) = (idempotency_key.as_deref(), request_hash.as_deref())
{
match idempotency::begin(
&state,
idempotency::Scope::User(user_id),
@@ -433,10 +436,14 @@ async fn create_checkout(
.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 }));
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(
@@ -451,14 +458,17 @@ async fn create_checkout(
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 Ok(Json(Envelope {
success: true,
data: resp,
}));
}
return Err(AppError::new(
ErrorCode::InvalidRequest,
"请求正在处理中，请稍后重试",
));
}
idempotency::BeginResult::Acquired { .. } => {
idempotency::BeginResult::Acquired => {
idem_acquired = true;
}
}
@@ -525,8 +535,10 @@ async fn create_checkout(
cus
};
let success_url =
format!("{}/dashboard/billing?checkout=success", state.config.public_base_url);
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(
@@ -619,7 +631,10 @@ async fn create_portal(
.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"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"未找到 Stripe Customer",
));
};
let return_url = format!("{}/dashboard/billing", state.config.public_base_url);
@@ -631,7 +646,11 @@ async fn create_portal(
}))
}
async fn stripe_create_customer(secret: &str, email: &str, user_id: Uuid) -> Result<String, AppError> {
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",
@@ -723,10 +742,9 @@ async fn stripe_post_form(
.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))?;
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");

222
src/api/compress.rs
View File

@@ -54,7 +54,7 @@ struct CompressRequest {
file_bytes: Vec<u8>,
level: CompressionLevel,
compression_rate: Option<u8>,
target_size_bytes: Option<u64>, // 新增：直接指定目标大小（字节）
target_size_bytes: Option<u64>, // 新增：直接指定目标大小（字节）
output_format: Option<ImageFmt>,
max_width: Option<u32>,
max_height: Option<u32>,
@@ -99,11 +99,17 @@ async fn compress_json(
ConnectInfo(addr): ConnectInfo<SocketAddr>,
headers: HeaderMap,
mut multipart: Multipart,
) -> Result<(axum_extra::extract::cookie::CookieJar, Json<Envelope<CompressResponse>>), AppError> {
) -> Result<
(
axum_extra::extract::cookie::CookieJar,
Json<Envelope<CompressResponse>>,
),
AppError,
> {
let ip = context::client_ip(&headers, addr.ip());
let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
let req = parse_single_file_request(&mut multipart).await?;
let mut req = parse_single_file_request(&mut multipart).await?;
let format_in = compress::detect_format(&req.file_bytes)?;
let format_out = req.output_format.unwrap_or(format_in);
@@ -121,7 +127,9 @@ async fn compress_json(
.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)),
context::Principal::ApiKey { api_key_id, .. } => {
Some(idempotency::Scope::ApiKey(*api_key_id))
}
_ => None,
};
@@ -139,7 +147,12 @@ async fn compress_json(
.compression_rate
.map(|v| v.to_string())
.unwrap_or_default();
let target_size_key = req
.target_size_bytes
.map(|v| v.to_string())
.unwrap_or_default();
h.update(rate_key.as_bytes());
h.update(target_size_key.as_bytes());
h.update(mw.as_bytes());
h.update(mh.as_bytes());
h.update(preserve.as_bytes());
@@ -211,32 +224,40 @@ async fn compress_json(
.await?
{
idempotency::BeginResult::Replay { response_body, .. } => {
let resp: CompressResponse = 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 })));
let resp: CompressResponse =
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 => {
if let Some((_status, body)) = idempotency::wait_for_replay(
&state,
scope,
idem_key,
request_hash,
10_000,
)
.await?
if let Some((_status, body)) =
idempotency::wait_for_replay(&state, scope, idem_key, request_hash, 10_000)
.await?
{
let resp: CompressResponse = 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 Ok((
jar,
Json(Envelope {
success: true,
data: resp,
}),
));
}
return Err(AppError::new(
ErrorCode::InvalidRequest,
"请求正在处理中，请稍后重试",
));
}
idempotency::BeginResult::Acquired { .. } => {
idempotency::BeginResult::Acquired => {
idem_acquired = true;
}
}
@@ -249,21 +270,22 @@ async fn compress_json(
QuotaContext::Anonymous { .. } => {}
}
let original_size = req.file_bytes.len() as u64;
let input = std::mem::take(&mut req.file_bytes);
let compressed = compress::compress_image_bytes(
&state,
&req.file_bytes,
input,
format_in,
format_out,
effective_level,
req.compression_rate,
req.target_size_bytes, // 新增：目标大小
req.target_size_bytes, // 新增：目标大小
req.max_width,
req.max_height,
req.preserve_metadata,
)
.await?;
let original_size = req.file_bytes.len() as u64;
let compressed_size = compressed.len() as u64;
let saved_bytes = original_size.saturating_sub(compressed_size);
let saved_percent = if original_size == 0 {
@@ -280,7 +302,7 @@ async fn compress_json(
}
}
if state.config.storage_type.to_ascii_lowercase() != "local" {
if !state.config.storage_type.eq_ignore_ascii_case("local") {
return Err(AppError::new(
ErrorCode::StorageUnavailable,
"当前仅支持本地存储（STORAGE_TYPE=local）",
@@ -370,7 +392,13 @@ async fn compress_json(
.await;
}
}
Ok((jar, Json(Envelope { success: true, data: resp })))
Ok((
jar,
Json(Envelope {
success: true,
data: resp,
}),
))
}
Err(err) => {
if let (Some(scope), Some(idem_key), Some(request_hash)) = (
@@ -393,7 +421,13 @@ async fn compress_direct(
ConnectInfo(addr): ConnectInfo<SocketAddr>,
headers: HeaderMap,
mut multipart: Multipart,
) -> Result<(axum_extra::extract::cookie::CookieJar, axum::response::Response), AppError> {
) -> Result<
(
axum_extra::extract::cookie::CookieJar,
axum::response::Response,
),
AppError,
> {
let ip = context::client_ip(&headers, addr.ip());
let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?;
@@ -404,7 +438,7 @@ async fn compress_direct(
));
}
let req = parse_single_file_request(&mut multipart).await?;
let mut req = parse_single_file_request(&mut multipart).await?;
let email_verified = match &principal {
context::Principal::User { email_verified, .. } => *email_verified,
@@ -431,7 +465,9 @@ async fn compress_direct(
.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)),
context::Principal::ApiKey { api_key_id, .. } => {
Some(idempotency::Scope::ApiKey(*api_key_id))
}
context::Principal::Anonymous { .. } => None,
};
@@ -449,7 +485,12 @@ async fn compress_direct(
.compression_rate
.map(|v| v.to_string())
.unwrap_or_default();
let target_size_key = req
.target_size_bytes
.map(|v| v.to_string())
.unwrap_or_default();
h.update(rate_key.as_bytes());
h.update(target_size_key.as_bytes());
h.update(mw.as_bytes());
h.update(mh.as_bytes());
h.update(preserve.as_bytes());
@@ -505,7 +546,8 @@ async fn compress_direct(
serde_json::from_value(response_body).map_err(|err| {
AppError::new(ErrorCode::Internal, "幂等结果解析失败").with_source(err)
})?;
let (bytes, fmt) = load_direct_replay_bytes(&state, &principal, data.file_id).await?;
let (bytes, fmt) =
load_direct_replay_bytes(&state, &principal, data.file_id).await?;
let mut resp_headers = HeaderMap::new();
resp_headers.insert(
@@ -537,14 +579,9 @@ async fn compress_direct(
return Ok((jar, response));
}
idempotency::BeginResult::InProgress => {
if let Some((_status, body)) = idempotency::wait_for_replay(
&state,
scope,
idem_key,
request_hash,
10_000,
)
.await?
if let Some((_status, body)) =
idempotency::wait_for_replay(&state, scope, idem_key, request_hash, 10_000)
.await?
{
let data: DirectIdempotencyData =
serde_json::from_value(body).map_err(|err| {
@@ -587,7 +624,7 @@ async fn compress_direct(
"请求正在处理中，请稍后重试",
));
}
idempotency::BeginResult::Acquired { .. } => {
idempotency::BeginResult::Acquired => {
idem_acquired = true;
}
}
@@ -600,21 +637,22 @@ async fn compress_direct(
QuotaContext::Anonymous { .. } => {}
}
let original_size = req.file_bytes.len() as u64;
let input = std::mem::take(&mut req.file_bytes);
let compressed = compress::compress_image_bytes(
&state,
&req.file_bytes,
input,
format_in,
format_out,
effective_level,
req.compression_rate,
req.target_size_bytes, // 新增：目标大小
req.target_size_bytes, // 新增：目标大小
req.max_width,
req.max_height,
req.preserve_metadata,
)
.await?;
let original_size = req.file_bytes.len() as u64;
let compressed_size = compressed.len() as u64;
let saved_bytes = original_size.saturating_sub(compressed_size);
let saved_percent = if original_size == 0 {
@@ -625,7 +663,7 @@ async fn compress_direct(
let skip_charge = req.compression_rate == Some(100);
let charge_units = !skip_charge && compressed_size < original_size;
if state.config.storage_type.to_ascii_lowercase() != "local" {
if !state.config.storage_type.eq_ignore_ascii_case("local") {
return Err(AppError::new(
ErrorCode::StorageUnavailable,
"当前仅支持本地存储（STORAGE_TYPE=local）",
@@ -689,7 +727,10 @@ async fn compress_direct(
"ImageForge-Compressed-Size",
compressed_size.to_string().parse().unwrap(),
);
resp_headers.insert("ImageForge-Saved-Bytes", saved_bytes.to_string().parse().unwrap());
resp_headers.insert(
"ImageForge-Saved-Bytes",
saved_bytes.to_string().parse().unwrap(),
);
resp_headers.insert(
"ImageForge-Saved-Percent",
format!("{saved_percent:.2}").parse().unwrap(),
@@ -752,12 +793,12 @@ async fn compress_direct(
}
async fn write_file(path: &str, bytes: &[u8]) -> Result<(), AppError> {
let mut file = tokio::fs::File::create(path)
.await
.map_err(|err| AppError::new(ErrorCode::StorageUnavailable, "写入文件失败").with_source(err))?;
file.write_all(bytes)
.await
.map_err(|err| AppError::new(ErrorCode::StorageUnavailable, "写入文件失败").with_source(err))?;
let mut file = tokio::fs::File::create(path).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "写入文件失败").with_source(err)
})?;
file.write_all(bytes).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "写入文件失败").with_source(err)
})?;
Ok(())
}
@@ -811,7 +852,9 @@ async fn load_direct_replay_bytes(
.fetch_optional(&state.db)
.await
}
context::Principal::Anonymous { .. } => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")),
context::Principal::Anonymous { .. } => {
return Err(AppError::new(ErrorCode::Unauthorized, "未登录"))
}
}
.map_err(|err| AppError::new(ErrorCode::Internal, "查询文件失败").with_source(err))?
.ok_or_else(|| AppError::new(ErrorCode::NotFound, "文件不存在"))?;
@@ -826,9 +869,9 @@ async fn load_direct_replay_bytes(
return Err(AppError::new(ErrorCode::NotFound, "文件不存在"));
};
let bytes = tokio::fs::read(&path)
.await
.map_err(|err| AppError::new(ErrorCode::StorageUnavailable, "读取文件失败").with_source(err))?;
let bytes = tokio::fs::read(&path).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "读取文件失败").with_source(err)
})?;
let fmt = compress::parse_output_format(&row.output_format)?;
Ok((bytes, fmt))
@@ -840,36 +883,27 @@ async fn parse_single_file_request(multipart: &mut Multipart) -> Result<Compress
let mut level = CompressionLevel::Medium;
let mut output_format: Option<ImageFmt> = None;
let mut compression_rate: Option<u8> = None;
let mut target_size_bytes: Option<u64> = None; // 新增
let mut target_size_bytes: Option<u64> = None; // 新增
let mut max_width: Option<u32> = None;
let mut max_height: Option<u32> = None;
let mut preserve_metadata = false;
while let Some(field) = multipart
.next_field()
.await
.map_err(|err| AppError::new(ErrorCode::InvalidRequest, "读取上传内容失败").with_source(err))?
{
while let Some(field) = multipart.next_field().await.map_err(|err| {
AppError::new(ErrorCode::InvalidRequest, "读取上传内容失败").with_source(err)
})? {
let name = field.name().unwrap_or("").to_string();
if name == "file" {
file_name = Some(
field
.file_name()
.unwrap_or("upload")
.to_string(),
);
let bytes = field
.bytes()
.await
.map_err(|err| AppError::new(ErrorCode::InvalidRequest, "读取文件失败").with_source(err))?;
file_name = Some(field.file_name().unwrap_or("upload").to_string());
let bytes = field.bytes().await.map_err(|err| {
AppError::new(ErrorCode::InvalidRequest, "读取文件失败").with_source(err)
})?;
file_bytes = Some(bytes.to_vec());
continue;
}
let text = field
.text()
.await
.map_err(|err| AppError::new(ErrorCode::InvalidRequest, "读取字段失败").with_source(err))?;
let text = field.text().await.map_err(|err| {
AppError::new(ErrorCode::InvalidRequest, "读取字段失败").with_source(err)
})?;
match name.as_str() {
"level" => {
@@ -890,19 +924,17 @@ async fn parse_single_file_request(multipart: &mut Multipart) -> Result<Compress
"max_width" => {
let v = text.trim();
if !v.is_empty() {
max_width = Some(
v.parse::<u32>()
.map_err(|_| AppError::new(ErrorCode::InvalidRequest, "max_width 格式错误"))?,
);
max_width = Some(v.parse::<u32>().map_err(|_| {
AppError::new(ErrorCode::InvalidRequest, "max_width 格式错误")
})?);
}
}
"max_height" => {
let v = text.trim();
if !v.is_empty() {
max_height = Some(
v.parse::<u32>()
.map_err(|_| AppError::new(ErrorCode::InvalidRequest, "max_height 格式错误"))?,
);
max_height = Some(v.parse::<u32>().map_err(|_| {
AppError::new(ErrorCode::InvalidRequest, "max_height 格式错误")
})?);
}
}
"preserve_metadata" => {
@@ -914,10 +946,12 @@ async fn parse_single_file_request(multipart: &mut Multipart) -> Result<Compress
"target_size_bytes" | "target_size" => {
let v = text.trim();
if !v.is_empty() {
target_size_bytes = Some(
v.parse::<u64>()
.map_err(|_| AppError::new(ErrorCode::InvalidRequest, "target_size_bytes 格式错误，需为正整数（字节）"))?,
);
target_size_bytes = Some(v.parse::<u64>().map_err(|_| {
AppError::new(
ErrorCode::InvalidRequest,
"target_size_bytes 格式错误，需为正整数（字节）",
)
})?);
// 最小目标大小限制：1KB
if let Some(size) = target_size_bytes {
if size < 1024 {
@@ -933,7 +967,8 @@ async fn parse_single_file_request(multipart: &mut Multipart) -> Result<Compress
}
}
let file_bytes = file_bytes.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "缺少 file"))?;
let file_bytes =
file_bytes.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "缺少 file"))?;
let file_name = file_name.unwrap_or_else(|| "upload".to_string());
Ok(CompressRequest {
@@ -941,7 +976,7 @@ async fn parse_single_file_request(multipart: &mut Multipart) -> Result<Compress
file_bytes,
level,
compression_rate,
target_size_bytes, // 新增
target_size_bytes, // 新增
output_format,
max_width,
max_height,
@@ -954,7 +989,10 @@ fn enforce_file_limits_anonymous(state: &AppState, bytes: &[u8]) -> Result<(), A
if bytes.len() as u64 > max {
return Err(AppError::new(
ErrorCode::FileTooLarge,
format!("匿名试用单文件最大 {} MB", state.config.anon_max_file_size_mb),
format!(
"匿名试用单文件最大 {} MB",
state.config.anon_max_file_size_mb
),
));
}
Ok(())
@@ -1020,6 +1058,7 @@ async fn ensure_quota_available(
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn record_task_and_metering(
state: &AppState,
principal: &context::Principal,
@@ -1042,9 +1081,15 @@ async fn record_task_and_metering(
charge_units: bool,
) -> Result<(), AppError> {
let (user_id, session_id, api_key_id, source) = match principal {
context::Principal::Anonymous { session_id } => (None, Some(session_id.clone()), None, "web"),
context::Principal::Anonymous { session_id } => {
(None, Some(session_id.clone()), None, "web")
}
context::Principal::User { user_id, .. } => (Some(*user_id), None, None, "web"),
context::Principal::ApiKey { user_id, api_key_id, .. } => (Some(*user_id), None, Some(*api_key_id), "api"),
context::Principal::ApiKey {
user_id,
api_key_id,
..
} => (Some(*user_id), None, Some(*api_key_id), "api"),
};
let mut tx = state
@@ -1161,6 +1206,7 @@ async fn record_task_and_metering(
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn charge_one_unit(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
billing: &BillingContext,

71
src/api/context.rs
View File

@@ -18,8 +18,14 @@ 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 },
Anonymous {
session_id: String,
},
User {
user_id: Uuid,
role: String,
email_verified: bool,
},
ApiKey {
user_id: Uuid,
api_key_id: Uuid,
@@ -29,6 +35,14 @@ pub enum Principal {
}
pub fn client_ip(headers: &HeaderMap, connect_ip: IpAddr) -> IpAddr {
resolve_client_ip(headers, connect_ip, crate::config::trust_proxy_headers())
}
fn resolve_client_ip(headers: &HeaderMap, connect_ip: IpAddr, trust_proxy: bool) -> IpAddr {
if !trust_proxy {
return connect_ip;
}
if let Some(ip) = parse_forwarded_for(headers) {
return ip;
}
@@ -69,7 +83,13 @@ pub async fn authenticate(
return Err(AppError::new(ErrorCode::Unauthorized, "未登录"));
}
let (jar, session_id) = ensure_session_cookie(jar);
let cookie_secure = state
.config
.public_base_url
.trim()
.to_ascii_lowercase()
.starts_with("https://");
let (jar, session_id) = ensure_session_cookie(jar, cookie_secure);
Ok((jar, Principal::Anonymous { session_id }))
}
@@ -178,11 +198,12 @@ async fn try_api_key(
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;
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,
@@ -192,7 +213,7 @@ async fn try_api_key(
}))
}
pub fn ensure_session_cookie(jar: CookieJar) -> (CookieJar, String) {
pub fn ensure_session_cookie(jar: CookieJar, secure: bool) -> (CookieJar, String) {
if let Some(cookie) = jar.get("if_session") {
let session_id = cookie.value().trim().to_string();
if !session_id.is_empty() {
@@ -204,6 +225,7 @@ pub fn ensure_session_cookie(jar: CookieJar) -> (CookieJar, String) {
let cookie = Cookie::build(("if_session", session_id.clone()))
.path("/")
.http_only(true)
.secure(secure)
.same_site(SameSite::Lax)
.max_age(TimeDuration::days(7))
.build();
@@ -220,10 +242,37 @@ fn generate_session_id() -> String {
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))?;
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))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn untrusted_proxy_headers_cannot_override_peer_ip() {
let mut headers = HeaderMap::new();
headers.insert("x-forwarded-for", "203.0.113.9".parse().unwrap());
let peer = "192.0.2.10".parse().unwrap();
assert_eq!(resolve_client_ip(&headers, peer, false), peer);
}
#[test]
fn trusted_proxy_headers_use_forwarded_client_ip() {
let mut headers = HeaderMap::new();
headers.insert(
"x-forwarded-for",
"203.0.113.9, 192.0.2.20".parse().unwrap(),
);
let peer = "192.0.2.10".parse().unwrap();
assert_eq!(
resolve_client_ip(&headers, peer, true),
"203.0.113.9".parse::<IpAddr>().unwrap()
);
}
}

135
src/api/downloads.rs
View File

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

1
src/api/envelope.rs
View File

@@ -5,4 +5,3 @@ pub struct Envelope<T> {
pub success: bool,
pub data: T,
}

40
src/api/health.rs
View File

@@ -2,6 +2,9 @@ use crate::state::AppState;
use axum::{extract::State, http::StatusCode, response::IntoResponse, Json};
use serde::Serialize;
use std::time::Duration;
const DEPENDENCY_TIMEOUT: Duration = Duration::from_secs(2);
#[derive(Debug, Serialize)]
struct HealthResponse {
@@ -11,16 +14,24 @@ struct HealthResponse {
}
pub async fn health(State(state): State<AppState>) -> impl IntoResponse {
let database_ok = sqlx::query("SELECT 1")
.execute(&state.db)
let database_check = async {
tokio::time::timeout(
DEPENDENCY_TIMEOUT,
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)
.is_ok_and(|result| result.is_ok())
};
let redis_check = async {
let mut redis_conn = state.redis.clone();
tokio::time::timeout(
DEPENDENCY_TIMEOUT,
redis::cmd("PING").query_async::<_, String>(&mut redis_conn),
)
.await
.is_ok();
.is_ok_and(|result| result.is_ok())
};
let (database_ok, redis_ok) = tokio::join!(database_check, redis_check);
let status = if database_ok && redis_ok {
StatusCode::OK
@@ -29,11 +40,18 @@ pub async fn health(State(state): State<AppState>) -> impl IntoResponse {
};
let body = HealthResponse {
status: if status == StatusCode::OK { "healthy" } else { "unhealthy" },
database: if database_ok { "connected" } else { "unavailable" },
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))
}

38
src/api/mod.rs
View File

@@ -1,15 +1,15 @@
mod auth;
mod context;
mod envelope;
mod compress;
mod downloads;
mod billing;
mod webhooks;
mod user;
mod tasks;
mod admin;
mod auth;
mod billing;
mod compress;
mod context;
mod downloads;
mod envelope;
mod health;
mod response;
mod tasks;
mod user;
mod webhooks;
use crate::error::{AppError, ErrorCode};
use crate::state::AppState;
@@ -23,14 +23,11 @@ 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, "管理员账号初始化失败");
}
crate::services::bootstrap::ensure_schema(&state).await?;
crate::services::bootstrap::ensure_admin_user(&state).await?;
let static_service = ServeDir::new("static").not_found_service(ServeFile::new("static/index.html"));
let static_service =
ServeDir::new("static").not_found_service(ServeFile::new("static/index.html"));
let v1 = v1_router().layer(DefaultBodyLimit::max(100 * 1024 * 1024));
@@ -48,9 +45,12 @@ pub async fn run(state: AppState) -> Result<(), AppError> {
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))
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> {

1
src/api/response.rs
View File

@@ -3,4 +3,3 @@ use crate::error::{AppError, ErrorCode};
pub async fn not_found() -> AppError {
AppError::new(ErrorCode::NotFound, "接口不存在")
}

170
src/api/tasks.rs
View File

@@ -62,11 +62,17 @@ async fn create_batch_task(
ConnectInfo(addr): ConnectInfo<SocketAddr>,
headers: HeaderMap,
mut multipart: Multipart,
) -> Result<(axum_extra::extract::cookie::CookieJar, Json<Envelope<BatchCreateResponse>>), AppError> {
) -> 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" {
if !state.config.storage_type.eq_ignore_ascii_case("local") {
return Err(AppError::new(
ErrorCode::StorageUnavailable,
"当前仅支持本地存储（STORAGE_TYPE=local）",
@@ -81,7 +87,9 @@ async fn create_batch_task(
.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)),
context::Principal::ApiKey { api_key_id, .. } => {
Some(idempotency::Scope::ApiKey(*api_key_id))
}
_ => None,
};
@@ -110,31 +118,38 @@ async fn create_batch_task(
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 })));
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?
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 Ok((
jar,
Json(Envelope {
success: true,
data: resp,
}),
));
}
return Err(AppError::new(
ErrorCode::InvalidRequest,
"请求正在处理中，请稍后重试",
));
}
idempotency::BeginResult::Acquired { .. } => {
idempotency::BeginResult::Acquired => {
idem_acquired = true;
}
}
@@ -212,16 +227,18 @@ async fn create_batch_task(
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)),
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))?;
let mut tx =
state.db.begin().await.map_err(|err| {
AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err)
})?;
sqlx::query(
r#"
@@ -285,7 +302,9 @@ async fn create_batch_task(
.bind(&file.storage_path)
.execute(&mut *tx)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "创建文件记录失败").with_source(err))?;
.map_err(|err| {
AppError::new(ErrorCode::Internal, "创建文件记录失败").with_source(err)
})?;
}
tx.commit()
@@ -293,11 +312,12 @@ async fn create_batch_task(
.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;
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);
}
@@ -325,7 +345,13 @@ async fn create_batch_task(
.await;
}
}
Ok((jar, Json(Envelope { success: true, data: resp })))
Ok((
jar,
Json(Envelope {
success: true,
data: resp,
}),
))
}
Err(err) => {
if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) {
@@ -379,9 +405,9 @@ async fn parse_batch_request(
};
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))?;
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| {
@@ -434,16 +460,15 @@ async fn parse_batch_request(
Ok(v) => v,
Err(err) => {
cleanup_file_paths(&files).await;
return Err(
AppError::new(ErrorCode::StorageUnavailable, "写入文件失败").with_source(err),
);
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),
AppError::new(ErrorCode::StorageUnavailable, "写入文件失败").with_source(err)
);
}
@@ -464,7 +489,7 @@ async fn parse_batch_request(
Err(err) => {
cleanup_file_paths(&files).await;
return Err(
AppError::new(ErrorCode::InvalidRequest, "读取字段失败").with_source(err),
AppError::new(ErrorCode::InvalidRequest, "读取字段失败").with_source(err)
);
}
};
@@ -509,7 +534,10 @@ async fn parse_batch_request(
Ok(n) => n,
Err(_) => {
cleanup_file_paths(&files).await;
return Err(AppError::new(ErrorCode::InvalidRequest, "max_width 格式错误"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"max_width 格式错误",
));
}
});
}
@@ -521,7 +549,10 @@ async fn parse_batch_request(
Ok(n) => n,
Err(_) => {
cleanup_file_paths(&files).await;
return Err(AppError::new(ErrorCode::InvalidRequest, "max_height 格式错误"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"max_height 格式错误",
));
}
});
}
@@ -571,7 +602,10 @@ async fn parse_batch_request(
Ok((files, opts, request_hash))
}
fn enforce_batch_limits_anonymous(state: &AppState, files: &[BatchFileInput]) -> Result<(), AppError> {
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(
@@ -585,7 +619,10 @@ fn enforce_batch_limits_anonymous(state: &AppState, files: &[BatchFileInput]) ->
if f.original_size > max_bytes {
return Err(AppError::new(
ErrorCode::FileTooLarge,
format!("匿名试用单文件最大 {} MB", state.config.anon_max_file_size_mb),
format!(
"匿名试用单文件最大 {} MB",
state.config.anon_max_file_size_mb
),
));
}
}
@@ -679,7 +716,10 @@ async fn anonymous_remaining_units(
.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)))
Ok(std::cmp::min(
limit - v1.unwrap_or(0),
limit - v2.unwrap_or(0),
))
}
fn utc8_date() -> String {
@@ -689,7 +729,6 @@ fn utc8_date() -> String {
#[derive(Debug, FromRow)]
struct TaskRow {
id: Uuid,
status: String,
total_files: i32,
completed_files: i32,
@@ -743,14 +782,19 @@ async fn get_task(
ConnectInfo(addr): ConnectInfo<SocketAddr>,
headers: HeaderMap,
Path(task_id): Path<Uuid>,
) -> Result<(axum_extra::extract::cookie::CookieJar, Json<Envelope<TaskView>>), AppError> {
) -> 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,
@@ -774,7 +818,11 @@ async fn get_task(
return Err(AppError::new(ErrorCode::NotFound, "任务已过期或不存在"));
}
authorize_task(&principal, task.user_id, task.session_id.as_deref().unwrap_or(""))?;
authorize_task(
&principal,
task.user_id,
task.session_id.as_deref().unwrap_or(""),
)?;
let files = sqlx::query_as::<_, TaskFileRow>(
r#"
@@ -848,12 +896,18 @@ async fn cancel_task(
ConnectInfo(addr): ConnectInfo<SocketAddr>,
headers: HeaderMap,
Path(task_id): Path<Uuid>,
) -> Result<(axum_extra::extract::cookie::CookieJar, Json<Envelope<serde_json::Value>>), AppError> {
) -> 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",
"SELECT 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)
@@ -861,7 +915,11 @@ async fn cancel_task(
.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(""))?;
authorize_task(
&principal,
task.user_id,
task.session_id.as_deref().unwrap_or(""),
)?;
if matches!(task.status.as_str(), "completed" | "failed" | "cancelled") {
return Ok((
@@ -900,12 +958,18 @@ async fn delete_task(
ConnectInfo(addr): ConnectInfo<SocketAddr>,
headers: HeaderMap,
Path(task_id): Path<Uuid>,
) -> Result<(axum_extra::extract::cookie::CookieJar, Json<Envelope<serde_json::Value>>), AppError> {
) -> 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",
"SELECT 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)
@@ -913,7 +977,11 @@ async fn delete_task(
.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(""))?;
authorize_task(
&principal,
task.user_id,
task.session_id.as_deref().unwrap_or(""),
)?;
if task.status == "processing" {
return Err(AppError::new(
@@ -933,7 +1001,7 @@ async fn delete_task(
let _ = tokio::fs::remove_file(p).await;
}
if state.config.storage_type.to_ascii_lowercase() == "local" {
if state.config.storage_type.eq_ignore_ascii_case("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);

66
src/api/user.rs
View File

@@ -205,7 +205,11 @@ async fn update_profile(
.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 email_verified_at = if email_changed {
None
} else {
user.email_verified_at
};
let updated = sqlx::query_as::<_, UserRow>(
r#"
@@ -243,7 +247,9 @@ async fn update_profile(
.bind(expires_at)
.execute(&mut *tx)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "创建邮箱验证记录失败").with_source(err))?;
.map_err(|err| {
AppError::new(ErrorCode::Internal, "创建邮箱验证记录失败").with_source(err)
})?;
verification_link = Some(format!(
"{}/verify-email?token={}",
@@ -258,7 +264,9 @@ async fn update_profile(
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))?;
.map_err(|err| {
AppError::new(ErrorCode::MailSendFailed, "验证邮件发送失败").with_source(err)
})?;
}
let message = if email_changed {
@@ -396,17 +404,18 @@ async fn list_history(
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 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))?
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)
@@ -530,7 +539,10 @@ async fn list_history(
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 {
download_url: if file.status == "completed"
&& file.storage_path.is_some()
&& task.expires_at > now
{
Some(format!("/downloads/{}", file.id))
} else {
None
@@ -667,7 +679,10 @@ async fn create_api_key(
let billing = billing::get_user_billing(&state, user_id).await?;
if !billing.plan.feature_api_enabled {
return Err(AppError::new(ErrorCode::Forbidden, "当前套餐未开通 API Key"));
return Err(AppError::new(
ErrorCode::Forbidden,
"当前套餐未开通 API Key",
));
}
let permissions = normalize_permissions(req.permissions)?;
@@ -718,12 +733,15 @@ async fn disable_api_key(
_ => 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))?;
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 不存在"));
@@ -813,7 +831,13 @@ fn generate_api_key() -> (String, String) {
}
fn normalize_permissions(input: Option<Vec<String>>) -> Result<serde_json::Value, AppError> {
let allowed = ["compress", "batch_compress", "read_stats", "billing_read", "webhook_manage"];
let allowed = [
"compress",
"batch_compress",
"read_stats",
"billing_read",
"webhook_manage",
];
let mut perms = Vec::<String>::new();
if let Some(values) = input {

64
src/api/webhooks.rs
View File

@@ -48,8 +48,9 @@ async fn stripe_webhook(
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 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#"
@@ -112,21 +113,31 @@ fn verify_stripe_signature(payload: &[u8], sig_header: &str, secret: &str) -> Re
}
let Some(ts) = timestamp else {
return Err(AppError::new(ErrorCode::InvalidRequest, "Stripe-Signature 缺少 t"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"Stripe-Signature 缺少 t",
));
};
if signatures.is_empty() {
return Err(AppError::new(ErrorCode::InvalidRequest, "Stripe-Signature 缺少 v1"));
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 时间戳过期"));
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))?;
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);
@@ -159,7 +170,9 @@ async fn process_stripe_event(state: &AppState, event: &StripeEvent) -> Result<(
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,
"invoice.paid" | "invoice.payment_failed" => {
upsert_invoice(state, &event.data.object).await
}
_ => Ok(()),
}
}
@@ -206,7 +219,9 @@ async fn map_checkout_session_completed(
.bind(customer_id)
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "更新 Stripe Customer 映射失败").with_source(err))?;
.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>>(
@@ -270,7 +285,11 @@ async fn upsert_subscription(state: &AppState, object: &serde_json::Value) -> Re
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()))
.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> =
@@ -401,7 +420,10 @@ async fn upsert_invoice(state: &AppState, object: &serde_json::Value) -> Result<
return Ok(());
};
let stripe_status = object.get("status").and_then(|v| v.as_str()).unwrap_or("open");
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
@@ -418,11 +440,23 @@ async fn upsert_invoice(state: &AppState, object: &serde_json::Value) -> Result<
.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 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 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")

6
src/auth.rs
View File

@@ -42,7 +42,10 @@ pub fn require_jwt(jwt_secret: &str, headers: &HeaderMap) -> Result<Claims, AppE
.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>")
AppError::new(
ErrorCode::Unauthorized,
"缺少 Authorization: Bearer <token>",
)
})?;
decode_jwt(jwt_secret, token)
@@ -57,4 +60,3 @@ pub fn decode_jwt(jwt_secret: &str, token: &str) -> Result<Claims, AppError> {
.map(|data| data.claims)
.map_err(|_| AppError::new(ErrorCode::Unauthorized, "Token 无效或已过期"))
}

50
src/config.rs
View File

@@ -1,5 +1,7 @@
use crate::error::{AppError, ErrorCode};
static TRUST_PROXY_HEADERS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
#[derive(Debug, Clone)]
pub struct Config {
pub role: String,
@@ -13,19 +15,18 @@ pub struct Config {
pub redis_url: String,
pub worker_concurrency: u32,
pub image_processing_concurrency: u32,
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,
@@ -67,22 +68,41 @@ impl Config {
.map(|v| v.get() as u32)
.unwrap_or(4)
});
let image_processing_concurrency = env_u32("IMAGE_PROCESSING_CONCURRENCY")
.filter(|value| *value > 0)
.unwrap_or_else(|| {
std::thread::available_parallelism()
.map(|v| v.get() as u32)
.unwrap_or(4)
});
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 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());
if !billing_provider.eq_ignore_ascii_case("stripe") {
return Err(AppError::new(
ErrorCode::InvalidRequest,
"BILLING_PROVIDER 目前仅支持 stripe",
));
}
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());
if !storage_type.eq_ignore_ascii_case("local") {
return Err(AppError::new(
ErrorCode::InvalidRequest,
"STORAGE_TYPE 目前仅支持 local",
));
}
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);
@@ -94,11 +114,14 @@ impl Config {
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_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_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_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");
@@ -112,15 +135,14 @@ impl Config {
database_max_connections,
redis_url,
worker_concurrency,
image_processing_concurrency,
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,
@@ -142,7 +164,9 @@ impl Config {
}
fn env_string(key: &str) -> Option<String> {
std::env::var(key).ok().filter(|value| !value.trim().is_empty())
std::env::var(key)
.ok()
.filter(|value| !value.trim().is_empty())
}
fn env_u16(key: &str) -> Option<u16> {
@@ -168,3 +192,7 @@ fn env_bool(key: &str) -> Option<bool> {
_ => None,
})
}
pub fn trust_proxy_headers() -> bool {
*TRUST_PROXY_HEADERS.get_or_init(|| env_bool("TRUST_PROXY_HEADERS").unwrap_or(false))
}

18
src/main.rs
View File

@@ -12,7 +12,6 @@ use crate::services::mail::Mailer;
use crate::state::AppState;
use sqlx::postgres::PgPoolOptions;
use tracing::Level;
#[tokio::main]
async fn main() -> Result<(), AppError> {
@@ -34,11 +33,16 @@ async fn main() -> Result<(), AppError> {
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "Redis 连接失败").with_source(err))?;
let image_processing_semaphore = std::sync::Arc::new(tokio::sync::Semaphore::new(
config.image_processing_concurrency as usize,
));
let state = AppState {
config,
db,
redis,
mailer: std::sync::Arc::new(mailer),
image_processing_semaphore,
};
match state.config.role.as_str() {
@@ -52,13 +56,9 @@ async fn main() -> Result<(), AppError> {
}
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")
});
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();
tracing_subscriber::fmt().with_env_filter(env_filter).init();
}

24
src/services/billing.rs
View File

@@ -7,8 +7,6 @@ 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,
@@ -36,8 +34,6 @@ struct SubscriptionRow {
#[derive(Debug, FromRow)]
struct PlanRow {
id: Uuid,
code: String,
included_units_per_period: i32,
max_file_size_mb: i32,
max_files_per_batch: i32,
@@ -68,19 +64,26 @@ pub async fn get_user_billing(state: &AppState, user_id: Uuid) -> Result<Billing
"订阅欠费，请先完成支付",
));
}
(Some(sub.id), sub.current_period_start, sub.current_period_end, sub.plan_id)
(
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))?;
.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
SELECT included_units_per_period, max_file_size_mb, max_files_per_batch, retention_days, features
FROM plans
WHERE id = $1
"#,
@@ -100,8 +103,6 @@ pub async fn get_user_billing(state: &AppState, user_id: Uuid) -> Result<Billing
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,
@@ -119,7 +120,10 @@ pub fn current_month_period_utc8(now_utc: DateTime<Utc>) -> (DateTime<Utc>, Date
let year = now.year();
let month = now.month();
let start = tz.with_ymd_and_hms(year, month, 1, 0, 0, 0).single().unwrap();
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)

61
src/services/bootstrap.rs
View File

@@ -7,6 +7,8 @@ use sqlx::FromRow;
use tracing::{info, warn};
use uuid::Uuid;
static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("./migrations");
#[derive(Debug, FromRow)]
struct AdminRow {
id: Uuid,
@@ -28,13 +30,8 @@ pub async fn ensure_admin_user(state: &AppState) -> Result<(), AppError> {
return Ok(());
}
let admin_username = env_string("ADMIN_USERNAME").unwrap_or_else(|| {
admin_email
.split('@')
.next()
.unwrap_or("admin")
.to_string()
});
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)?;
@@ -85,7 +82,9 @@ pub async fn ensure_admin_user(state: &AppState) -> Result<(), AppError> {
.bind(row.id)
.fetch_one(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "校验管理员用户名失败").with_source(err))?;
.map_err(|err| {
AppError::new(ErrorCode::Internal, "校验管理员用户名失败").with_source(err)
})?;
if name_taken {
warn!(
@@ -105,7 +104,9 @@ pub async fn ensure_admin_user(state: &AppState) -> Result<(), AppError> {
.bind(row.id)
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "更新管理员用户名失败").with_source(err))?;
.map_err(|err| {
AppError::new(ErrorCode::Internal, "更新管理员用户名失败").with_source(err)
})?;
}
}
@@ -138,32 +139,20 @@ pub async fn ensure_admin_user(state: &AppState) -> Result<(), AppError> {
}
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;
MIGRATOR
.run(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "执行数据库迁移失败").with_source(err))?;
Ok(())
}
fn validate_email(email: &str) -> Result<(), AppError> {
if email.trim().is_empty() || !email.contains('@') {
return Err(AppError::new(ErrorCode::InvalidRequest, "管理员邮箱格式不正确"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"管理员邮箱格式不正确",
));
}
if email.len() > 255 {
return Err(AppError::new(ErrorCode::InvalidRequest, "管理员邮箱过长"));
@@ -173,7 +162,10 @@ fn validate_email(email: &str) -> Result<(), AppError> {
fn validate_username(username: &str) -> Result<(), AppError> {
if username.trim().is_empty() {
return Err(AppError::new(ErrorCode::InvalidRequest, "管理员用户名不能为空"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"管理员用户名不能为空",
));
}
if username.len() > 50 {
return Err(AppError::new(ErrorCode::InvalidRequest, "管理员用户名过长"));
@@ -183,7 +175,10 @@ fn validate_username(username: &str) -> Result<(), AppError> {
fn validate_password(password: &str) -> Result<(), AppError> {
if password.len() < 8 {
return Err(AppError::new(ErrorCode::InvalidRequest, "管理员密码至少 8 位"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"管理员密码至少 8 位",
));
}
if password.len() > 128 {
return Err(AppError::new(ErrorCode::InvalidRequest, "管理员密码过长"));
@@ -200,5 +195,7 @@ fn hash_password(password: &str) -> Result<String, AppError> {
}
fn env_string(key: &str) -> Option<String> {
std::env::var(key).ok().filter(|value| !value.trim().is_empty())
std::env::var(key)
.ok()
.filter(|value| !value.trim().is_empty())
}

341
src/services/compress.rs
View File

@@ -1,15 +1,15 @@
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 image::{DynamicImage, ExtendedColorType, ImageEncoder};
use img_parts::{Bytes as ImgBytes, DynImage, ImageEXIF, ImageICC};
use oxipng::StripChunks;
use rgb::FromSlice;
use std::io::Cursor;
@@ -105,15 +105,12 @@ pub fn parse_level(value: &str) -> Result<CompressionLevel, AppError> {
}
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 的整数（压缩后体积占比）",
)
})?;
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,
@@ -199,8 +196,7 @@ pub fn detect_format(bytes: &[u8]) -> Result<ImageFmt, AppError> {
if has_brand(
&brands,
&[
*b"heic", *b"heix", *b"hevc", *b"hevx", *b"heis", *b"heim", *b"mif1",
*b"msf1",
*b"heic", *b"heix", *b"hevc", *b"hevx", *b"heis", *b"heim", *b"mif1", *b"msf1",
],
) {
return Err(AppError::new(
@@ -233,27 +229,69 @@ pub fn detect_format(bytes: &[u8]) -> Result<ImageFmt, AppError> {
))
}
#[allow(clippy::too_many_arguments)]
pub async fn compress_image_bytes(
state: &AppState,
input: &[u8],
input: Vec<u8>,
format_in: ImageFmt,
format_out: ImageFmt,
level: CompressionLevel,
compression_rate: Option<u8>,
target_size_bytes: Option<u64>, // 新增：直接指定目标大小（字节）
target_size_bytes: Option<u64>, // 新增：直接指定目标大小（字节）
max_width: Option<u32>,
max_height: Option<u32>,
preserve_metadata: bool,
) -> Result<Vec<u8>, AppError> {
let max_image_pixels = state.config.max_image_pixels;
let permit = state
.image_processing_semaphore
.clone()
.acquire_owned()
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "图片处理并发控制器已关闭").with_source(err)
})?;
tokio::task::spawn_blocking(move || {
let _permit = permit;
compress_image_bytes_sync(
input,
format_in,
format_out,
level,
compression_rate,
target_size_bytes,
max_width,
max_height,
preserve_metadata,
max_image_pixels,
)
})
.await
.map_err(|err| {
AppError::new(ErrorCode::CompressionFailed, "图片处理任务异常退出").with_source(err)
})?
}
#[allow(clippy::too_many_arguments)]
fn compress_image_bytes_sync(
input: Vec<u8>,
format_in: ImageFmt,
format_out: ImageFmt,
level: CompressionLevel,
compression_rate: Option<u8>,
target_size_bytes: Option<u64>,
max_width: Option<u32>,
max_height: Option<u32>,
preserve_metadata: bool,
max_image_pixels: u64,
) -> Result<Vec<u8>, AppError> {
let original_size = input.len() as u64;
if format_in == ImageFmt::Gif {
if is_animated_gif(input)? {
return Err(AppError::new(
ErrorCode::UnsupportedFormat,
"暂不支持动图 GIF",
));
}
if format_in == ImageFmt::Gif && is_animated_gif(&input)? {
return Err(AppError::new(
ErrorCode::UnsupportedFormat,
"暂不支持动图 GIF",
));
}
let retention_rate = effective_rate(compression_rate, level);
@@ -269,14 +307,14 @@ pub async fn compress_image_bytes(
&& max_height.is_none()
{
if preserve_metadata {
return Ok(input.to_vec());
return Ok(input);
}
let stripped = strip_metadata(input).unwrap_or_else(|_| input.to_vec());
let stripped = strip_metadata(&input).unwrap_or_else(|_| input.clone());
return Ok(stripped);
}
let (icc_profile, exif) = if preserve_metadata {
extract_metadata(input)
extract_metadata(&input)
} else {
(None, None)
};
@@ -294,13 +332,23 @@ pub async fn compress_image_bytes(
if !preserve_metadata {
opts.strip = StripChunks::Safe;
}
oxipng::optimize_from_memory(input, &opts)
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "PNG 压缩失败").with_source(err))?
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))?;
#[cfg(not(target_os = "linux"))]
if format_in == ImageFmt::Avif {
return Err(AppError::new(
ErrorCode::UnsupportedFormat,
"当前平台构建不支持 AVIF 解码，请转换为 PNG/JPEG/WebP 后重试",
));
}
enforce_pixel_limit(state, &image)?;
let image = image::load_from_memory(&input).map_err(|err| {
AppError::new(ErrorCode::InvalidImage, "图片解码失败").with_source(err)
})?;
enforce_pixel_limit(max_image_pixels, &image)?;
let (image, did_resize) = resize_if_needed(image, max_width, max_height);
resized = did_resize;
@@ -332,23 +380,23 @@ pub async fn compress_image_bytes(
if !resized && output.len() >= input.len() {
if preserve_metadata {
return Ok(input.to_vec());
return Ok(input);
}
let stripped = strip_metadata(input).unwrap_or_else(|_| input.to_vec());
let stripped = strip_metadata(&input).unwrap_or_else(|_| input.clone());
return Ok(if stripped.len() <= input.len() {
stripped
} else {
input.to_vec()
input
});
}
Ok(output)
}
fn enforce_pixel_limit(state: &AppState, image: &DynamicImage) -> Result<(), AppError> {
fn enforce_pixel_limit(max_image_pixels: u64, 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 {
if pixels > max_image_pixels {
return Err(AppError::new(
ErrorCode::TooManyPixels,
format!("图片像素过大（{}x{}）", w, h),
@@ -394,11 +442,7 @@ fn fit_within(w: u32, h: u32, max_width: Option<u32>, max_height: Option<u32>) -
(nw, nh)
}
fn encode_png(
image: DynamicImage,
rate: u8,
preserve_metadata: bool,
) -> Result<Vec<u8>, AppError> {
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();
@@ -406,7 +450,9 @@ fn encode_png(
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))?;
.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);
@@ -433,7 +479,9 @@ fn encode_jpeg_raw(raw: &[u8], w: u32, h: u32, quality: u8) -> Result<Vec<u8>, A
let mut encoder = JpegEncoder::new_with_quality(&mut out, quality);
encoder
.encode(raw, w, h, ExtendedColorType::Rgb8)
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "JPEG 编码失败").with_source(err))?;
.map_err(|err| {
AppError::new(ErrorCode::CompressionFailed, "JPEG 编码失败").with_source(err)
})?;
Ok(out)
}
@@ -451,12 +499,6 @@ fn encode_webp(image: DynamicImage, rate: u8) -> Result<Vec<u8>, AppError> {
Ok(bytes.to_vec())
}
fn encode_webp_with_quality(image: DynamicImage, quality: u8) -> Result<Vec<u8>, AppError> {
let rgba = image.to_rgba8();
let (w, h) = rgba.dimensions();
encode_webp_raw(rgba.as_raw(), w, h, quality)
}
fn encode_webp_raw(raw: &[u8], w: u32, h: u32, quality: u8) -> Result<Vec<u8>, AppError> {
let encoder = webp::Encoder::from_rgba(raw, w, h);
Ok(encoder.encode(quality as f32).to_vec())
@@ -473,26 +515,20 @@ fn encode_avif(image: DynamicImage, rate: u8) -> Result<Vec<u8>, AppError> {
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))?;
let encoded = encoder.encode_rgba(img).map_err(|err| {
AppError::new(ErrorCode::CompressionFailed, "AVIF 编码失败").with_source(err)
})?;
Ok(encoded.avif_file)
}
fn encode_avif_with_quality(image: DynamicImage, quality: u8) -> Result<Vec<u8>, AppError> {
let rgba = image.to_rgba8();
let (w, h) = rgba.dimensions();
encode_avif_raw(rgba.as_raw(), w, h, quality)
}
fn encode_avif_raw(raw: &[u8], w: u32, h: u32, quality: u8) -> Result<Vec<u8>, AppError> {
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 as f32);
let encoded = encoder
.encode_rgba(img)
.map_err(|err| AppError::new(ErrorCode::CompressionFailed, "AVIF 编码失败").with_source(err))?;
let encoded = encoder.encode_rgba(img).map_err(|err| {
AppError::new(ErrorCode::CompressionFailed, "AVIF 编码失败").with_source(err)
})?;
Ok(encoded.avif_file)
}
@@ -579,7 +615,8 @@ where
Some((_bytes, best_w, best_h, best_size)) => {
let new_pixels = (new_w as u64).saturating_mul(new_h as u64);
let best_pixels = (*best_w as u64).saturating_mul(*best_h as u64);
new_pixels > best_pixels || (new_pixels == best_pixels && result_size > *best_size)
new_pixels > best_pixels
|| (new_pixels == best_pixels && result_size > *best_size)
}
};
@@ -587,9 +624,7 @@ where
best_under = Some((result, new_w, new_h, result_size));
}
if new_w == orig_w
&& new_h == orig_h
&& target_size.saturating_sub(result_size) <= 1024
if new_w == orig_w && new_h == orig_h && target_size.saturating_sub(result_size) <= 1024
{
break;
}
@@ -645,11 +680,7 @@ where
let mut consider = |bytes: Vec<u8>| {
let size = bytes.len() as u64;
let is_under = size <= target_size;
let diff = if size > target_size {
size - target_size
} else {
target_size - size
};
let diff = size.abs_diff(target_size);
let should_update = match (best_is_under, is_under) {
(false, true) => true,
@@ -689,87 +720,6 @@ where
best.ok_or_else(|| AppError::new(ErrorCode::CompressionFailed, "压缩失败"))
}
fn encode_target_quality<F>(
min_q: u8,
max_q: u8,
target_size: u64,
mut encode: F,
) -> Result<Vec<u8>, AppError>
where
F: FnMut(u8) -> Result<Vec<u8>, AppError>,
{
let mut best: Option<Vec<u8>> = None;
let mut best_diff = u64::MAX;
let mut best_is_under = false; // 记录最佳结果是否小于目标
let mut best_size = 0u64;
// 考虑一个候选结果
let consider = |bytes: Vec<u8>, best: &mut Option<Vec<u8>>, best_diff: &mut u64, best_is_under: &mut bool, best_size: &mut u64| {
let size = bytes.len() as u64;
let is_under = size <= target_size;
let diff = if size > target_size {
size - target_size
} else {
target_size - size
};
// 优先选择不超过目标大小的结果
let should_update = match (*best_is_under, is_under) {
(false, true) => true, // 当前小于目标，之前大于目标 -> 更新
(true, false) => false, // 当前大于目标，之前小于目标 -> 不更新
_ => diff < *best_diff, // 同类情况，选择更接近的
};
if should_update {
*best_diff = diff;
*best_is_under = is_under;
*best_size = size;
*best = Some(bytes);
}
};
// 先尝试两端
let bytes = encode(min_q)?;
consider(bytes, &mut best, &mut best_diff, &mut best_is_under, &mut best_size);
if min_q != max_q {
let bytes = encode(max_q)?;
consider(bytes, &mut best, &mut best_diff, &mut best_is_under, &mut best_size);
}
// 二分查找，增加迭代次数到 12 次以提高精度
let mut low = min_q;
let mut high = max_q;
for _ in 0..12 {
if low > high {
break;
}
let mid = (low + high) / 2;
let bytes = encode(mid)?;
let size = bytes.len() as u64;
consider(bytes, &mut best, &mut best_diff, &mut best_is_under, &mut best_size);
if size > target_size {
high = mid.saturating_sub(1);
} else {
low = mid.saturating_add(1);
}
}
// 精细调整：如果当前结果超出目标太多，尝试更低质量
if best_size > target_size {
let mut q = min_q;
while q <= min_q.saturating_add(5) && q <= max_q {
let bytes = encode(q)?;
consider(bytes, &mut best, &mut best_diff, &mut best_is_under, &mut best_size);
if best_size <= target_size {
break; // 已找到满足条件的结果
}
q = q.saturating_add(1);
}
}
best.ok_or_else(|| AppError::new(ErrorCode::CompressionFailed, "压缩失败"))
}
fn encode_gif(image: DynamicImage, rate: u8) -> Result<Vec<u8>, AppError> {
let rgba = image.to_rgba8();
let (w, h) = rgba.dimensions();
@@ -780,7 +730,9 @@ fn encode_gif(image: DynamicImage, rate: u8) -> Result<Vec<u8>, AppError> {
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))?;
.map_err(|err| {
AppError::new(ErrorCode::CompressionFailed, "GIF 编码失败").with_source(err)
})?;
}
Ok(out)
@@ -793,7 +745,9 @@ fn encode_bmp(image: DynamicImage) -> Result<Vec<u8>, AppError> {
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))?;
.map_err(|err| {
AppError::new(ErrorCode::CompressionFailed, "BMP 编码失败").with_source(err)
})?;
Ok(out)
}
@@ -804,7 +758,9 @@ fn encode_tiff(image: DynamicImage) -> Result<Vec<u8>, AppError> {
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))?;
.map_err(|err| {
AppError::new(ErrorCode::CompressionFailed, "TIFF 编码失败").with_source(err)
})?;
Ok(out.into_inner())
}
@@ -815,7 +771,9 @@ fn encode_ico(image: DynamicImage) -> Result<Vec<u8>, AppError> {
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))?;
.map_err(|err| {
AppError::new(ErrorCode::CompressionFailed, "ICO 编码失败").with_source(err)
})?;
Ok(out)
}
@@ -837,8 +795,9 @@ fn apply_metadata(
}
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 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());
@@ -848,16 +807,17 @@ fn apply_metadata(
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))?;
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 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());
};
@@ -866,9 +826,9 @@ fn strip_metadata(input: &[u8]) -> Result<Vec<u8>, AppError> {
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))?;
img.encoder().write_to(&mut buf).map_err(|err| {
AppError::new(ErrorCode::CompressionFailed, "写入图片元数据失败").with_source(err)
})?;
Ok(buf)
}
@@ -927,13 +887,58 @@ fn strength_from_rate(rate: u8) -> u8 {
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();
let mut frames = decoder.into_frames();
if let Some(frame) = frames.next() {
frame.map_err(|err| AppError::new(ErrorCode::InvalidImage, "GIF 解码失败").with_source(err))?;
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))?;
frame.map_err(|err| {
AppError::new(ErrorCode::InvalidImage, "GIF 解码失败").with_source(err)
})?;
return Ok(true);
}
Ok(false)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn detects_supported_formats_from_signatures() {
assert_eq!(detect_format(b"\x89PNG\r\n\x1a\n").unwrap(), ImageFmt::Png);
assert_eq!(detect_format(b"\xff\xd8").unwrap(), ImageFmt::Jpeg);
assert_eq!(
detect_format(b"RIFF\x00\x00\x00\x00WEBP").unwrap(),
ImageFmt::Webp
);
assert_eq!(detect_format(b"GIF89a").unwrap(), ImageFmt::Gif);
assert_eq!(detect_format(b"BM").unwrap(), ImageFmt::Bmp);
}
#[test]
fn detects_avif_and_rejects_heic() {
let avif = b"\x00\x00\x00\x18ftypavif\x00\x00\x00\x00avif";
assert_eq!(detect_format(avif).unwrap(), ImageFmt::Avif);
let heic = b"\x00\x00\x00\x18ftypheic\x00\x00\x00\x00mif1";
let error = detect_format(heic).unwrap_err();
assert_eq!(error.code, ErrorCode::UnsupportedFormat);
}
#[test]
fn fit_within_preserves_aspect_ratio_and_never_upscales() {
assert_eq!(fit_within(4000, 2000, Some(1000), None), (1000, 500));
assert_eq!(fit_within(4000, 2000, None, Some(250)), (500, 250));
assert_eq!(fit_within(400, 200, Some(800), Some(800)), (400, 200));
}
#[test]
fn compression_rate_maps_to_expected_target_size() {
assert_eq!(target_size_from_rate(10_000, 1), 100);
assert_eq!(target_size_from_rate(10_000, 55), 5_500);
assert_eq!(target_size_from_rate(10_000, 100), 10_000);
}
}

30
src/services/idempotency.rs
View File

@@ -15,8 +15,8 @@ pub enum Scope {
#[derive(Debug)]
pub enum BeginResult {
Acquired { expires_at: DateTime<Utc> },
Replay { response_status: i32, response_body: JsonValue },
Acquired,
Replay { response_body: JsonValue },
InProgress,
}
@@ -25,7 +25,6 @@ struct IdemRow {
request_hash: String,
response_status: i32,
response_body: Option<JsonValue>,
expires_at: DateTime<Utc>,
}
pub fn sha256_hex(parts: &[&[u8]]) -> String {
@@ -45,13 +44,22 @@ pub async fn begin(
ttl_hours: i64,
) -> Result<BeginResult, AppError> {
if idempotency_key.trim().is_empty() {
return Err(AppError::new(ErrorCode::InvalidRequest, "Idempotency-Key 不能为空"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"Idempotency-Key 不能为空",
));
}
if idempotency_key.len() > 128 {
return Err(AppError::new(ErrorCode::InvalidRequest, "Idempotency-Key 过长"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"Idempotency-Key 过长",
));
}
if request_hash.len() != 64 {
return Err(AppError::new(ErrorCode::InvalidRequest, "request_hash 不合法"));
return Err(AppError::new(
ErrorCode::InvalidRequest,
"request_hash 不合法",
));
}
let now = Utc::now();
@@ -108,12 +116,12 @@ pub async fn begin(
.map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等记录失败").with_source(err))?;
if inserted.rows_affected() > 0 {
return Ok(BeginResult::Acquired { expires_at });
return Ok(BeginResult::Acquired);
}
let row = get_row(state, scope, idempotency_key, now).await?;
let Some(row) = row else {
return Ok(BeginResult::Acquired { expires_at });
return Ok(BeginResult::Acquired);
};
if row.request_hash != request_hash {
@@ -128,7 +136,6 @@ pub async fn begin(
}
Ok(BeginResult::Replay {
response_status: row.response_status,
response_body: row.response_body.unwrap_or(JsonValue::Null),
})
}
@@ -300,7 +307,7 @@ async fn get_row(
Scope::User(user_id) => {
sqlx::query_as::<_, IdemRow>(
r#"
SELECT request_hash, response_status, response_body, expires_at
SELECT request_hash, response_status, response_body
FROM idempotency_keys
WHERE user_id = $1
AND idempotency_key = $2
@@ -318,7 +325,7 @@ async fn get_row(
Scope::ApiKey(api_key_id) => {
sqlx::query_as::<_, IdemRow>(
r#"
SELECT request_hash, response_status, response_body, expires_at
SELECT request_hash, response_status, response_body
FROM idempotency_keys
WHERE api_key_id = $1
AND idempotency_key = $2
@@ -338,4 +345,3 @@ async fn get_row(
Ok(row)
}

49
src/services/mail.rs
View File

@@ -76,10 +76,9 @@ impl Mailer {
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))?,
)
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) {
@@ -203,8 +202,11 @@ impl Mailer {
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)
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)
@@ -212,12 +214,16 @@ impl Mailer {
.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())),
.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))?;
@@ -250,11 +256,16 @@ impl SmtpConfig {
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 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();
@@ -320,7 +331,9 @@ pub async fn send_password_reset_email(
reset_url: &str,
) -> Result<(), AppError> {
let mailer = resolve_mailer(state).await?;
mailer.send_password_reset_email(to, username, reset_url).await
mailer
.send_password_reset_email(to, username, reset_url)
.await
}
pub async fn send_test_email(state: &AppState, to: &str) -> Result<(), AppError> {

6
src/services/mod.rs
View File

@@ -1,7 +1,7 @@
pub mod mail;
pub mod billing;
pub mod quota;
pub mod bootstrap;
pub mod compress;
pub mod idempotency;
pub mod mail;
pub mod quota;
pub mod settings;
pub mod bootstrap;

8
src/services/settings.rs
View File

@@ -6,8 +6,8 @@ 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 serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -39,7 +39,6 @@ pub struct StripeConfigStored {
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>(
@@ -51,7 +50,9 @@ pub async fn load_system_config<T: DeserializeOwned>(
.bind(key)
.fetch_optional(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询系统配置失败").with_source(err))?;
.map_err(|err| {
AppError::new(ErrorCode::Internal, "查询系统配置失败").with_source(err)
})?;
let Some(value) = value else {
return Ok(None);
@@ -138,7 +139,6 @@ pub async fn load_stripe_secrets(state: &AppState) -> Result<Option<StripeSecret
Ok(Some(StripeSecrets {
secret_key,
webhook_secret,
secret_key_prefix: cfg.secret_key_prefix,
}))
}

1
src/state.rs
View File

@@ -7,4 +7,5 @@ pub struct AppState {
pub db: sqlx::PgPool,
pub redis: redis::aio::ConnectionManager,
pub mailer: std::sync::Arc<Mailer>,
pub image_processing_semaphore: std::sync::Arc<tokio::sync::Semaphore>,
}

185
src/worker/mod.rs
View File

@@ -20,9 +20,7 @@ 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, "数据库结构初始化失败");
}
crate::services::bootstrap::ensure_schema(&state).await?;
let consumer = format!("worker_{}", Uuid::new_v4());
ensure_group(&state, &consumer).await?;
@@ -71,15 +69,26 @@ async fn ensure_group(state: &AppState, _consumer: &str) -> Result<(), AppError>
async fn poll_once(state: &AppState, consumer: &str) -> Result<(), AppError> {
let mut conn = state.redis.clone();
let opts = StreamReadOptions::default()
// Retry messages already delivered to this consumer before taking new work.
let pending_opts = StreamReadOptions::default()
.group(GROUP_NAME, consumer)
.count(1)
.block(5000);
let reply: redis::streams::StreamReadReply = conn
.xread_options(&[STREAM_KEY], &[">"], &opts)
.count(1);
let mut reply: redis::streams::StreamReadReply = conn
.xread_options(&[STREAM_KEY], &["0"], &pending_opts)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "读取队列失败").with_source(err))?;
.map_err(|err| AppError::new(ErrorCode::Internal, "读取待重试任务失败").with_source(err))?;
if !reply.keys.iter().any(|key| !key.ids.is_empty()) {
let opts = StreamReadOptions::default()
.group(GROUP_NAME, consumer)
.count(1)
.block(5000);
reply = conn
.xread_options(&[STREAM_KEY], &[">"], &opts)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "读取队列失败").with_source(err))?;
}
if reply.keys.is_empty() {
return Ok(());
@@ -100,9 +109,10 @@ async fn poll_once(state: &AppState, consumer: &str) -> Result<(), AppError> {
}
};
if let Err(err) = process_task(state, task_id).await {
tracing::error!(task_id = %task_id, error = %err, "task processing failed");
}
process_task(state, task_id).await.map_err(|err| {
tracing::error!(task_id = %task_id, error = %err, "task processing failed; message left pending for retry");
err
})?;
ack_message(&mut conn, &msg.id).await?;
}
@@ -111,7 +121,10 @@ async fn poll_once(state: &AppState, consumer: &str) -> Result<(), AppError> {
Ok(())
}
async fn ack_message(conn: &mut redis::aio::ConnectionManager, msg_id: &str) -> Result<(), AppError> {
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)
@@ -124,16 +137,12 @@ async fn ack_message(conn: &mut redis::aio::ConnectionManager, msg_id: &str) ->
#[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>,
@@ -145,10 +154,8 @@ struct TaskProcRow {
struct TaskFileProcRow {
id: Uuid,
storage_path: Option<String>,
original_name: String,
original_format: String,
output_format: String,
original_size: i64,
status: String,
}
@@ -166,16 +173,12 @@ 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,
@@ -195,6 +198,8 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
return Ok(());
}
let is_retry = task.status == "processing";
let updated = sqlx::query(
r#"
UPDATE tasks
@@ -215,9 +220,17 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
// 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());
if is_retry {
sqlx::query(
"UPDATE task_files SET status = 'pending' WHERE task_id = $1 AND status = 'processing'",
)
.bind(task_id)
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "恢复待重试文件失败").with_source(err))?;
}
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)?);
@@ -229,10 +242,8 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
SELECT
id,
storage_path,
original_name,
original_format,
output_format,
original_size,
status::text AS status
FROM task_files
WHERE task_id = $1
@@ -281,7 +292,7 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
join_set.spawn(async move {
let _permit = permit;
if let Err(err) = process_task_file(
let result = process_task_file(
state,
task_id,
file,
@@ -292,19 +303,35 @@ async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
ctx,
billing_ctx,
)
.await
{
.await;
if let Err(err) = &result {
tracing::error!(task_id = %task_id, file_id = %file_id, error = %err, "file processing failed");
}
result
});
}
let mut first_error = None;
while let Some(result) = join_set.join_next().await {
if let Err(err) = result {
tracing::error!(task_id = %task_id, error = %err, "file worker panicked");
match result {
Ok(Ok(())) => {}
Ok(Err(err)) => {
if first_error.is_none() {
first_error = Some(err);
}
}
Err(err) => {
return Err(
AppError::new(ErrorCode::Internal, "文件处理线程异常退出").with_source(err)
);
}
}
}
if let Some(err) = first_error {
return Err(err);
}
finalize_task_status(state, task_id).await?;
Ok(())
}
@@ -332,6 +359,7 @@ async fn is_task_cancelled(state: &AppState, task_id: Uuid) -> Result<bool, AppE
Ok(matches!(status.as_deref(), Some("cancelled")))
}
#[allow(clippy::too_many_arguments)]
async fn process_task_file(
state: AppState,
task_id: Uuid,
@@ -347,11 +375,13 @@ async fn process_task_file(
return Ok(());
}
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());
let updated = sqlx::query(
"UPDATE task_files SET status = 'processing' WHERE id = $1 AND status = 'pending'",
)
.bind(file.id)
.execute(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "更新文件处理状态失败").with_source(err))?;
if updated.rows_affected() == 0 {
return Ok(());
}
@@ -374,17 +404,32 @@ async fn process_task_file(
}
};
let format_in = parse_image_fmt(&file.original_format)?;
let format_out = parse_image_fmt(&file.output_format)?;
let format_in = match parse_image_fmt(&file.original_format) {
Ok(format) => format,
Err(err) => {
mark_file_failed(&state, task_id, file.id, &err.message).await?;
let _ = tokio::fs::remove_file(&input_path).await;
return Ok(());
}
};
let format_out = match parse_image_fmt(&file.output_format) {
Ok(format) => format,
Err(err) => {
mark_file_failed(&state, task_id, file.id, &err.message).await?;
let _ = tokio::fs::remove_file(&input_path).await;
return Ok(());
}
};
let original_size = input_bytes.len() as u64;
let compressed = match compress::compress_image_bytes(
&state,
&input_bytes,
input_bytes,
format_in,
format_out,
level,
compression_rate,
None, // target_size_bytes: worker 批量任务不支持精确大小
None, // target_size_bytes: worker 批量任务不支持精确大小
max_width,
max_height,
ctx.preserve_metadata,
@@ -405,7 +450,6 @@ async fn process_task_file(
return Ok(());
}
let original_size = input_bytes.len() as u64;
let compressed_size = compressed.len() as u64;
let saved_percent = if original_size == 0 {
0.0
@@ -442,7 +486,9 @@ async fn process_task_file(
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));
return Err(
AppError::new(ErrorCode::StorageUnavailable, "写入压缩文件失败").with_source(err),
);
}
if is_task_cancelled(&state, task_id).await? {
@@ -483,6 +529,7 @@ async fn process_task_file(
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn finalize_file(
state: &AppState,
billing_ctx: &Option<billing::BillingContext>,
@@ -563,7 +610,12 @@ async fn finalize_file(
Ok(())
}
async fn mark_file_failed(state: &AppState, task_id: Uuid, task_file_id: Uuid, message: &str) -> Result<(), AppError> {
async fn mark_file_failed(
state: &AppState,
task_id: Uuid,
task_file_id: Uuid,
message: &str,
) -> Result<(), AppError> {
let mut tx = state
.db
.begin()
@@ -608,16 +660,6 @@ async fn mark_file_failed(state: &AppState, task_id: Uuid, task_file_id: Uuid, m
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",
@@ -627,7 +669,9 @@ async fn finalize_task_status(state: &AppState, task_id: Uuid) -> Result<(), App
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询任务失败").with_source(err))?;
let Some((total, completed, failed, status)) = row else { return Ok(()); };
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')",
@@ -675,6 +719,7 @@ async fn finalize_task_status(state: &AppState, task_id: Uuid) -> Result<(), App
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn charge_one_unit(
tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
billing: &billing::BillingContext,
@@ -771,32 +816,36 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> {
.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 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;
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();
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" {
if state.config.storage_type.eq_ignore_ascii_case("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")