From f3c7a77a37e3513f9e9d11fec53db64f9c4c1449 Mon Sep 17 00:00:00 2001 From: 237899745 <237899745@users.noreply.git.workyai.cn> Date: Sat, 25 Jul 2026 18:42:51 +0800 Subject: [PATCH] feat: harden auth and stream uploads --- Cargo.lock | 1 - Cargo.toml | 1 - README.md | 55 +--- docker/.env.production.example | 7 + docker/docker-compose.prod.yml | 10 +- docs/api.md | 13 +- docs/deployment.md | 2 + docs/security.md | 17 +- frontend/src/pages/DocsPage.vue | 1 + migrations/008_auth_token_version.sql | 2 + src/api/auth.rs | 314 ++++++++++++--------- src/api/compress.rs | 259 +++++++++-------- src/api/context.rs | 81 +++++- src/api/downloads.rs | 2 + src/api/mod.rs | 1 + src/api/multipart.rs | 23 ++ src/api/tasks.rs | 381 ++++++++++++++++---------- src/api/user.rs | 122 +++------ src/auth.rs | 19 ++ src/error.rs | 35 ++- src/services/bootstrap.rs | 12 +- src/services/credentials.rs | 99 +++++++ src/services/filename.rs | 34 +++ src/services/mod.rs | 3 + src/services/rate_limit.rs | 62 +++++ src/services/storage.rs | 11 +- 26 files changed, 1006 insertions(+), 561 deletions(-) create mode 100644 migrations/008_auth_token_version.sql create mode 100644 src/api/multipart.rs create mode 100644 src/services/credentials.rs create mode 100644 src/services/filename.rs create mode 100644 src/services/rate_limit.rs diff --git a/Cargo.lock b/Cargo.lock index f9a78ff..eac0c62 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2190,7 +2190,6 @@ dependencies = [ "serde_json", "sha2", "sqlx", - "thiserror 1.0.69", "time", "tokio", "tokio-util", diff --git a/Cargo.toml b/Cargo.toml index 61b7e41..e751680 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,7 +15,6 @@ time = "0.3" serde = { version = "1", features = ["derive"] } serde_json = "1" -thiserror = "1" bytes = "1" tracing = "0.1" diff --git a/README.md b/README.md index 5c4b270..a141f65 100644 --- a/README.md +++ b/README.md @@ -30,53 +30,22 @@ ## 目录结构 -> 说明：以下为目标目录结构（规划），会随实现逐步补齐。 - ``` imageforge/ ├── Cargo.toml ├── src/ -│ ├── main.rs # 入口 -│ ├── config.rs # 配置管理 -│ ├── error.rs # 错误处理 -│ ├── lib.rs -│ │ -│ ├── api/ # API 路由 -│ │ ├── mod.rs -│ │ ├── auth.rs # 认证相关 -│ │ ├── compress.rs # 压缩相关 -│ │ ├── user.rs # 用户相关 -│ │ └── admin.rs # 管理员相关 -│ │ -│ ├── services/ # 业务逻辑 -│ │ ├── mod.rs -│ │ ├── compress.rs # 压缩服务 -│ │ ├── auth.rs # 认证服务 -│ │ ├── user.rs # 用户服务 -│ │ └── storage.rs # 存储服务 -│ │ -│ ├── models/ # 数据模型 -│ │ ├── mod.rs -│ │ ├── user.rs -│ │ ├── image.rs -│ │ └── api_key.rs -│ │ -│ ├── compress/ # 压缩核心 -│ │ ├── mod.rs -│ │ ├── png.rs # PNG 压缩（oxipng + pngquant） -│ │ ├── jpeg.rs # JPEG 压缩（mozjpeg） -│ │ ├── webp.rs # WebP 压缩 -│ │ └── avif.rs # AVIF 压缩 -│ │ -│ └── middleware/ # 中间件 -│ ├── mod.rs -│ ├── auth.rs # 认证中间件 -│ └── rate_limit.rs # 限流中间件 -│ -├── migrations/ # 数据库迁移 -├── static/ # 静态资源 -├── frontend/ # 前端项目 -└── docker/ # Docker 配置 +│ ├── main.rs # API / Worker 进程入口 +│ ├── auth.rs # JWT 签发与解析 +│ ├── config.rs # 环境变量配置 +│ ├── error.rs # 统一错误响应 +│ ├── api/ # HTTP 路由、鉴权、压缩、管理后台 +│ ├── services/ # 压缩、存储、计费、限速、邮件等服务 +│ └── worker/ # Redis Streams 消费与任务处理 +├── migrations/ # SQLx 数据库迁移 +├── frontend/src/ # Vue3 管理端与用户端 +├── docker/ # 镜像、Compose 与 Nginx 配置 +├── docs/ # 设计、API、部署与运维文档 +└── scripts/ # 部署和质量回归脚本 ``` ## 快速开始 diff --git a/docker/.env.production.example b/docker/.env.production.example index f9ac1c2..8709e74 100644 --- a/docker/.env.production.example +++ b/docker/.env.production.example @@ -26,6 +26,13 @@ DATABASE_MAX_CONNECTIONS=10 WORKER_CONCURRENCY=2 IMAGE_PROCESSING_CONCURRENCY=2 +# Resource ceilings tuned for an 8-core / 16 GB application host. +POSTGRES_MEMORY_LIMIT=2g +REDIS_MEMORY_LIMIT=1g +REDIS_MAXMEMORY=768mb +API_MEMORY_LIMIT=3g +WORKER_MEMORY_LIMIT=8g + ALLOW_ANONYMOUS_UPLOAD=true ANON_MAX_FILE_SIZE_MB=5 ANON_MAX_FILES_PER_BATCH=5 diff --git a/docker/docker-compose.prod.yml b/docker/docker-compose.prod.yml index 1eb4a95..2f236c4 100644 --- a/docker/docker-compose.prod.yml +++ b/docker/docker-compose.prod.yml @@ -25,6 +25,8 @@ x-imageforge-environment: &imageforge-environment services: postgres: image: postgres:16-alpine + mem_limit: ${POSTGRES_MEMORY_LIMIT:-2g} + pids_limit: 256 environment: POSTGRES_DB: imageforge POSTGRES_USER: imageforge @@ -40,7 +42,9 @@ services: redis: image: redis:7-alpine - command: ["redis-server", "--appendonly", "yes", "--maxmemory-policy", "noeviction"] + mem_limit: ${REDIS_MEMORY_LIMIT:-1g} + pids_limit: 128 + command: ["redis-server", "--appendonly", "yes", "--maxmemory", "${REDIS_MAXMEMORY:-768mb}", "--maxmemory-policy", "noeviction"] volumes: - redis_data:/data healthcheck: @@ -52,6 +56,8 @@ services: api: image: imageforge:${IMAGEFORGE_TAG:-local} + mem_limit: ${API_MEMORY_LIMIT:-3g} + pids_limit: 256 build: context: .. dockerfile: docker/Dockerfile @@ -102,6 +108,8 @@ services: worker: image: imageforge:${IMAGEFORGE_TAG:-local} + mem_limit: ${WORKER_MEMORY_LIMIT:-8g} + pids_limit: 512 init: true environment: <<: *imageforge-environment diff --git a/docs/api.md b/docs/api.md index 2a01738..0f872b0 100644 --- a/docs/api.md +++ b/docs/api.md @@ -38,7 +38,7 @@ X-API-Key: ### 2.3 匿名试用（仅网站场景） - 不提供 API Key； - 通过 Cookie 维持匿名会话（服务端签发），仅允许较小文件与较低频率。 -- 每日 10 次（以成功压缩文件数计）；超出返回 `QUOTA_EXCEEDED`（HTTP `402`）。 +- 每日 10 次；进入处理即预留次数，处理失败会自动归还，不可压缩或无体积收益的有效图片仍计次。超出返回 `QUOTA_EXCEEDED`（HTTP `402`）。 - 日界：自然日（UTC+8），次日 00:00 重置。 - **匿名试用硬限制：Cookie + IP 双限制**（两者任一超出都拒绝），降低刷会话绕过风险。 @@ -59,12 +59,9 @@ Idempotency-Key: ### 3.2 限流（Rate Limit） 超出限制返回： - HTTP `429` -- 头：`Retry-After: ` +- 错误码：`RATE_LIMITED` -建议头（可选）： -- `RateLimit-Limit` -- `RateLimit-Remaining` -- `RateLimit-Reset` +API Key 按其 `rate_limit` 字段执行每分钟限制；登录、注册、找回密码和 Token 验证按 IP/账号执行独立限制。 ### 3.3 配额（Quota / Billing） 配额不足（当期额度耗尽）返回： @@ -496,6 +493,8 @@ Content-Type: application/json { "name": "Production Server", "permissions": ["compress", "batch_compress"] } ``` +省略 `permissions` 时默认授予 `compress`；该权限覆盖同步压缩，并兼容批量任务、任务查询和结果下载。仅授予 `batch_compress` 时不能调用同步压缩接口。 + 响应： ```json { @@ -721,6 +720,8 @@ Content-Type: application/json } ``` +错误响应同时返回 `X-Request-Id`，其值与响应体 `request_id` 一致，可用于日志定位。 + ```http PUT /admin/storage/endpoints/{endpoint_id} POST /admin/storage/endpoints/{endpoint_id}/test diff --git a/docs/deployment.md b/docs/deployment.md index 2bf5f29..3c63ace 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -13,6 +13,8 @@ Debian 13、4 核 CPU、8GB 内存的实测起始值为 `WORKER_CONCURRENCY=2` 和 `IMAGE_PROCESSING_CONCURRENCY=2`。AVIF 是 CPU 密集型编码，不要直接把并发设置为 CPU 核数的数倍。 +生产 Compose 默认按 8 核 16GB 主机设置可覆盖的资源上限：API 3GB、Worker 8GB、PostgreSQL 2GB、Redis 1GB；对应变量为 `API_MEMORY_LIMIT`、`WORKER_MEMORY_LIMIT`、`POSTGRES_MEMORY_LIMIT` 和 `REDIS_MEMORY_LIMIT`。Redis 的 `REDIS_MAXMEMORY` 默认 768MB，达到上限后返回写入错误而不是继续挤占宿主机内存。 + ### 首次启动 ```bash diff --git a/docs/security.md b/docs/security.md index f3d807f..aaa1ec7 100644 --- a/docs/security.md +++ b/docs/security.md @@ -23,13 +23,15 @@ ## 2. 认证与会话 ### 2.1 用户登录 -- 密码哈希：`argon2id`（带独立 salt，参数可配置）。 -- 登录保护：基础限速 + 失败次数冷却；可选验证码（V1+）。 +- 密码哈希：`argon2id`（独立 salt），哈希与校验在阻塞线程池执行，不占用 Tokio I/O 线程。 +- 登录保护：登录按 IP 限制为 5 分钟 30 次、按账号限制为 5 分钟 10 次；注册按 IP 限制为每小时 10 次。 +- 找回密码按 IP/邮箱双维度限速，验证和重置链接也有独立尝试上限。 - 账号状态：`is_active=false` 直接拒绝登录与 API。 ### 2.2 JWT 使用建议 - 对外 API：支持 Bearer Token（适合 CLI/SDK）。 - 网站（Vue3）：优先使用 HttpOnly Cookie 承载会话（降低 XSS 泄露风险），如使用 localStorage 必须配合严格 CSP。 +- JWT 包含用户 `token_version`；修改或重置密码会递增版本，使此前签发的 JWT 立即失效。 --- @@ -48,7 +50,8 @@ - 避免 bcrypt/argon2 用在高频 key 校验导致性能瓶颈。 ### 3.3 权限与限制 -- 最小权限：permissions（compress/batch/read_stats/billing_read 等）。 +- API Key 请求会执行 `permissions` 校验；`compress` 权限兼容同步、批量、任务查询和结果下载。 +- 每个 API Key 使用数据库中的 `rate_limit` 执行 Redis 原子分钟限速。 - 支持禁用/轮换；可选 IP 白名单（Business/V1+）。 - 每次请求记录 `last_used_at/last_used_ip/user_agent`（审计）。 @@ -58,11 +61,13 @@ ### 4.1 输入校验 - 只依赖扩展名不安全：必须校验魔数/探测真实格式。 +- 单文件上传按套餐上限分片读取；批量上传边接收边落盘，不在 API 内存中累计整批文件。 +- multipart 文本字段与字段数量单独设限，文件名清除控制字符并按数据库长度安全截断。 - 设定上限： - `max_file_size_mb` - `max_pixels`（宽×高） - `max_dimension`（单边） - - 解码超时（Worker 层，避免卡死） + - 图片处理全局并发信号量 ### 4.2 资源隔离 - 压缩属 CPU 密集型：放到 Worker；API 只做编排与轻量校验。 @@ -79,7 +84,7 @@ - **幂等**：`Idempotency-Key` 防止重试导致重复扣费。 - **配额硬限制**：到达当期额度返回 `QUOTA_EXCEEDED`（HTTP 402）。 -- **匿名试用**：每日 10 次（成功文件数计），采用 **Cookie + IP** 双维度 Redis 计数做硬限制。 +- **匿名试用**：每日 10 次，采用 **Cookie + IP** 双维度 Redis 原子预留；处理失败归还，已进入正常处理的不可压缩图片同样计次。 - **异常检测**（告警即可，首期不必自动封禁）： - 短时间内用量突增 - 失败率异常升高（疑似 fuzzing/探测） @@ -114,6 +119,6 @@ ## 8. 数据安全与保留 -- 结果保留期：按套餐（Free 24h、Pro 7d、Business 30d 等），匿名更短。 +- 结果保留期：未登录/Free 1 天、Pro 7 天、Business 15 天，实际值由套餐的 `retention_days` 决定。 - 支持用户主动删除任务/文件（立即删除对象存储 + DB 标记/审计）。 - 审计日志留存与脱敏：保留必要字段（IP、UA、动作、对象 ID），避免写入明文密钥/Token。 diff --git a/frontend/src/pages/DocsPage.vue b/frontend/src/pages/DocsPage.vue index 222f6b2..73d446f 100644 --- a/frontend/src/pages/DocsPage.vue +++ b/frontend/src/pages/DocsPage.vue @@ -19,6 +19,7 @@
  • Base URL：https://ys.workyai.cn/api/v1
  • 认证方式：X-API-Key（推荐）或 Authorization: Bearer <token>
    • +
  • API Key 会执行权限与每分钟请求上限；超限返回 429 RATE_LIMITED
  • 支持格式：PNG / JPG / JPEG / WebP / AVIF / GIF / BMP / TIFF / ICO（仅静态图片，支持 output_format 转码）
  • 压缩率：compression_rate 1-100；JPEG/WebP/AVIF 以该比例为体积上限，无损格式为尽力优化
  • 目标体积：target_size_bytes 仅支持 JPEG/WebP/AVIF，且不能与 compression_rate 同时提交
    • diff --git a/migrations/008_auth_token_version.sql b/migrations/008_auth_token_version.sql new file mode 100644 index 0000000..b09e216 --- /dev/null +++ b/migrations/008_auth_token_version.sql @@ -0,0 +1,2 @@ +ALTER TABLE users + ADD COLUMN IF NOT EXISTS token_version INTEGER NOT NULL DEFAULT 0; diff --git a/src/api/auth.rs b/src/api/auth.rs index 5053502..1211210 100644 --- a/src/api/auth.rs +++ b/src/api/auth.rs @@ -1,17 +1,19 @@ +use crate::api::context; use crate::api::envelope::Envelope; use crate::auth; use crate::error::{AppError, ErrorCode}; -use crate::services::mail; -use crate::services::settings; +use crate::services::{credentials, mail, rate_limit, settings}; use crate::state::AppState; -use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier}; -use axum::{extract::State, http::HeaderMap, routing::post, Json, Router}; -use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; +use axum::{ + extract::{ConnectInfo, State}, + http::HeaderMap, + routing::post, + Json, Router, +}; use chrono::{DateTime, Duration, Utc}; -use rand::RngCore; use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; +use std::net::SocketAddr; use uuid::Uuid; pub fn router() -> Router { @@ -69,17 +71,31 @@ struct UserRow { role: String, is_active: bool, email_verified_at: Option>, + token_version: i32, } async fn register( State(state): State, + ConnectInfo(addr): ConnectInfo, + headers: HeaderMap, Json(req): Json, ) -> Result>, AppError> { - validate_email(&req.email)?; - validate_username(&req.username)?; - validate_password(&req.password)?; + let ip = context::client_ip(&headers, addr.ip()); + rate_limit::enforce( + &state, + "auth_register_ip", + &ip.to_string(), + 10, + 60 * 60, + "注册请求过于频繁，请稍后再试", + ) + .await?; - let password_hash = hash_password(&req.password)?; + credentials::validate_email(&req.email)?; + credentials::validate_username(&req.username)?; + credentials::validate_password(&req.password)?; + + let password_hash = credentials::hash_password(&req.password).await?; let verification_required = settings::email_verification_required(&state).await?; let verified_at = (!verification_required).then(Utc::now); @@ -94,7 +110,8 @@ async fn register( password_hash, role::text AS role, is_active, - email_verified_at + email_verified_at, + token_version "#, ) .bind(req.email.to_lowercase()) @@ -110,11 +127,12 @@ async fn register( state.config.jwt_expiry_hours, user.id, &user.role, + user.token_version, )?; if verification_required { - let verification_token = generate_token(); - let token_hash = sha256_hex(&verification_token); + let verification_token = credentials::generate_token(); + let token_hash = credentials::sha256_hex(&verification_token); let expires_at_db = Utc::now() + Duration::hours(24); sqlx::query( @@ -168,6 +186,8 @@ async fn register( async fn login( State(state): State, + ConnectInfo(addr): ConnectInfo, + headers: HeaderMap, Json(req): Json, ) -> Result>, AppError> { let identity = req.email.trim(); @@ -178,8 +198,28 @@ async fn login( )); } + let ip = context::client_ip(&headers, addr.ip()); + rate_limit::enforce( + &state, + "auth_login_ip", + &ip.to_string(), + 30, + 5 * 60, + "登录请求过于频繁，请稍后再试", + ) + .await?; + rate_limit::enforce( + &state, + "auth_login_identity", + &identity.to_lowercase(), + 10, + 5 * 60, + "该账号登录尝试过于频繁，请稍后再试", + ) + .await?; + let user = if identity.contains('@') { - validate_email(identity)?; + credentials::validate_email(identity)?; sqlx::query_as::<_, UserRow>( r#" SELECT @@ -189,7 +229,8 @@ async fn login( password_hash, role::text AS role, is_active, - email_verified_at + email_verified_at, + token_version FROM users WHERE email = $1 "#, @@ -198,7 +239,7 @@ async fn login( .fetch_optional(&state.db) .await } else { - validate_username(identity)?; + credentials::validate_username(identity)?; sqlx::query_as::<_, UserRow>( r#" SELECT @@ -208,7 +249,8 @@ async fn login( password_hash, role::text AS role, is_active, - email_verified_at + email_verified_at, + token_version FROM users WHERE username = $1 "#, @@ -224,7 +266,9 @@ async fn login( return Err(AppError::new(ErrorCode::Forbidden, "账号已被禁用")); } - verify_password(&req.password, &user.password_hash)?; + if !credentials::verify_password(&req.password, &user.password_hash).await? { + return Err(AppError::new(ErrorCode::Unauthorized, "账号或密码错误")); + } let verification_required = settings::email_verification_required(&state).await?; let (token, expires_at) = auth::issue_jwt( @@ -232,6 +276,7 @@ async fn login( state.config.jwt_expiry_hours, user.id, &user.role, + user.token_version, )?; Ok(Json(Envelope { @@ -270,32 +315,15 @@ async fn send_verification( })); } - // Rate limit: 1 per minute per user - let key = format!( - "rate:send_verification:{}:{}", - claims.sub, - Utc::now().format("%Y%m%d%H%M") - ); - let mut redis = state.redis.clone(); - let count: i64 = redis::cmd("INCR") - .arg(&key) - .query_async(&mut redis) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "Redis 限流失败").with_source(err))?; - if count == 1 { - let _: () = redis::cmd("EXPIRE") - .arg(&key) - .arg(60) - .query_async(&mut redis) - .await - .unwrap_or(()); - } - if count > 1 { - return Err(AppError::new( - ErrorCode::RateLimited, - "发送过于频繁，请稍后再试", - )); - } + rate_limit::enforce( + &state, + "auth_send_verification_user", + &claims.sub.to_string(), + 1, + 60, + "发送过于频繁，请稍后再试", + ) + .await?; let user = sqlx::query_as::<_, UserRow>( r#" @@ -306,7 +334,8 @@ async fn send_verification( password_hash, role::text AS role, is_active, - email_verified_at + email_verified_at, + token_version FROM users WHERE id = $1 "#, @@ -317,6 +346,13 @@ async fn send_verification( .map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))? .ok_or_else(|| AppError::new(ErrorCode::Unauthorized, "用户不存在或未登录"))?; + if claims.ver != user.token_version { + return Err(AppError::new( + ErrorCode::Unauthorized, + "登录状态已失效，请重新登录", + )); + } + if user.email_verified_at.is_some() { return Ok(Json(Envelope { success: true, @@ -326,8 +362,8 @@ async fn send_verification( })); } - let verification_token = generate_token(); - let token_hash = sha256_hex(&verification_token); + let verification_token = credentials::generate_token(); + let token_hash = credentials::sha256_hex(&verification_token); let expires_at_db = Utc::now() + Duration::hours(24); sqlx::query( @@ -369,13 +405,26 @@ struct VerifyEmailRequest { async fn verify_email( State(state): State, + ConnectInfo(addr): ConnectInfo, + headers: HeaderMap, Json(req): Json, ) -> Result>, AppError> { if req.token.trim().is_empty() { return Err(AppError::new(ErrorCode::InvalidRequest, "token 不能为空")); } - let token_hash = sha256_hex(&req.token); + let ip = context::client_ip(&headers, addr.ip()); + rate_limit::enforce( + &state, + "auth_verify_email_ip", + &ip.to_string(), + 20, + 15 * 60, + "验证请求过于频繁，请稍后再试", + ) + .await?; + + let token_hash = credentials::sha256_hex(&req.token); let now = Utc::now(); let updated = sqlx::query( @@ -428,9 +477,31 @@ struct ForgotPasswordRequest { async fn forgot_password( State(state): State, + ConnectInfo(addr): ConnectInfo, + headers: HeaderMap, Json(req): Json, ) -> Result>, AppError> { - validate_email(&req.email)?; + credentials::validate_email(&req.email)?; + + let ip = context::client_ip(&headers, addr.ip()); + rate_limit::enforce( + &state, + "auth_forgot_ip", + &ip.to_string(), + 5, + 15 * 60, + "找回密码请求过于频繁，请稍后再试", + ) + .await?; + rate_limit::enforce( + &state, + "auth_forgot_email", + &req.email.to_lowercase(), + 3, + 15 * 60, + "找回密码请求过于频繁，请稍后再试", + ) + .await?; let user = sqlx::query_as::<_, UserRow>( r#" @@ -441,7 +512,8 @@ async fn forgot_password( password_hash, role::text AS role, is_active, - email_verified_at + email_verified_at, + token_version FROM users WHERE email = $1 "#, @@ -452,8 +524,8 @@ async fn forgot_password( .map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?; if let Some(user) = user { - let reset_token = generate_token(); - let token_hash = sha256_hex(&reset_token); + let reset_token = credentials::generate_token(); + let token_hash = credentials::sha256_hex(&reset_token); let expires_at_db = Utc::now() + Duration::hours(1); let _ = sqlx::query( @@ -493,15 +565,64 @@ struct ResetPasswordRequest { async fn reset_password( State(state): State, + ConnectInfo(addr): ConnectInfo, + headers: HeaderMap, Json(req): Json, ) -> Result>, AppError> { if req.token.trim().is_empty() { return Err(AppError::new(ErrorCode::InvalidRequest, "token 不能为空")); } - validate_password(&req.new_password)?; + credentials::validate_password(&req.new_password)?; - let token_hash = sha256_hex(&req.token); + let ip = context::client_ip(&headers, addr.ip()); + rate_limit::enforce( + &state, + "auth_reset_ip", + &ip.to_string(), + 10, + 15 * 60, + "重置密码请求过于频繁，请稍后再试", + ) + .await?; + rate_limit::enforce( + &state, + "auth_reset_token", + &req.token, + 5, + 15 * 60, + "该重置链接尝试次数过多，请重新申请", + ) + .await?; + + let token_hash = credentials::sha256_hex(&req.token); let now = Utc::now(); + let token_exists: bool = sqlx::query_scalar( + r#" + SELECT EXISTS( + SELECT 1 + FROM password_resets + WHERE token_hash = $1 + AND used_at IS NULL + AND expires_at > $2 + ) + "#, + ) + .bind(&token_hash) + .bind(now) + .fetch_one(&state.db) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "重置密码失败").with_source(err))?; + if !token_exists { + return Err(AppError::new(ErrorCode::InvalidToken, "Token 无效或已过期")); + } + + let password_hash = credentials::hash_password(&req.new_password).await?; + + let mut tx = state + .db + .begin() + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?; let user_id: Option = sqlx::query_scalar( r#" @@ -510,11 +631,12 @@ async fn reset_password( WHERE token_hash = $1 AND used_at IS NULL AND expires_at > $2 + FOR UPDATE "#, ) .bind(&token_hash) .bind(now) - .fetch_optional(&state.db) + .fetch_optional(&mut *tx) .await .map_err(|err| AppError::new(ErrorCode::Internal, "重置密码失败").with_source(err))?; @@ -522,15 +644,9 @@ async fn reset_password( return Err(AppError::new(ErrorCode::InvalidToken, "Token 无效或已过期")); }; - let password_hash = hash_password(&req.new_password)?; - - let mut tx = state - .db - .begin() - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?; - - sqlx::query("UPDATE users SET password_hash = $1 WHERE id = $2") + sqlx::query( + "UPDATE users SET password_hash = $1, token_version = token_version + 1, updated_at = NOW() WHERE id = $2", + ) .bind(password_hash) .bind(user_id) .execute(&mut *tx) @@ -558,66 +674,6 @@ async fn reset_password( })) } -fn validate_email(email: &str) -> Result<(), AppError> { - if email.trim().is_empty() || !email.contains('@') { - return Err(AppError::new(ErrorCode::InvalidRequest, "邮箱格式不正确")); - } - if email.len() > 255 { - return Err(AppError::new(ErrorCode::InvalidRequest, "邮箱过长")); - } - Ok(()) -} - -fn validate_username(username: &str) -> Result<(), AppError> { - if username.trim().is_empty() { - return Err(AppError::new(ErrorCode::InvalidRequest, "用户名不能为空")); - } - if username.len() > 50 { - return Err(AppError::new(ErrorCode::InvalidRequest, "用户名过长")); - } - Ok(()) -} - -fn validate_password(password: &str) -> Result<(), AppError> { - if password.len() < 8 { - return Err(AppError::new(ErrorCode::InvalidRequest, "密码至少 8 位")); - } - if password.len() > 128 { - return Err(AppError::new(ErrorCode::InvalidRequest, "密码过长")); - } - Ok(()) -} - -fn hash_password(password: &str) -> Result { - let salt = argon2::password_hash::SaltString::generate(&mut rand::rngs::OsRng); - Argon2::default() - .hash_password(password.as_bytes(), &salt) - .map_err(|err| AppError::new(ErrorCode::Internal, "密码哈希失败").with_source(err))? - .to_string() - .pipe(Ok) -} - -fn verify_password(password: &str, password_hash: &str) -> Result<(), AppError> { - let parsed = PasswordHash::new(password_hash) - .map_err(|err| AppError::new(ErrorCode::Internal, "密码哈希格式错误").with_source(err))?; - Argon2::default() - .verify_password(password.as_bytes(), &parsed) - .map_err(|_| AppError::new(ErrorCode::Unauthorized, "账号或密码错误"))?; - Ok(()) -} - -fn generate_token() -> String { - let mut bytes = [0u8; 32]; - rand::rngs::OsRng.fill_bytes(&mut bytes); - URL_SAFE_NO_PAD.encode(bytes) -} - -fn sha256_hex(token: &str) -> String { - let mut hasher = Sha256::new(); - hasher.update(token.as_bytes()); - hex::encode(hasher.finalize()) -} - fn map_unique_violation(err: sqlx::Error) -> AppError { if let sqlx::Error::Database(db_err) = &err { if let Some(code) = db_err.code() { @@ -628,11 +684,3 @@ fn map_unique_violation(err: sqlx::Error) -> AppError { } AppError::new(ErrorCode::Internal, "数据库操作失败").with_source(err) } - -trait Pipe: Sized { - fn pipe(self, f: impl FnOnce(Self) -> T) -> T { - f(self) - } -} - -impl Pipe for T {} diff --git a/src/api/compress.rs b/src/api/compress.rs index 59644ae..280faef 100644 --- a/src/api/compress.rs +++ b/src/api/compress.rs @@ -1,10 +1,12 @@ use crate::api::context; use crate::api::envelope::Envelope; +use crate::api::multipart as multipart_utils; use crate::error::{AppError, ErrorCode}; use crate::services::billing; -use crate::services::billing::{BillingContext, Plan}; +use crate::services::billing::BillingContext; use crate::services::compress; use crate::services::compress::{CompressionLevel, ImageFmt}; +use crate::services::filename; use crate::services::idempotency; use crate::services::quota; use crate::services::storage; @@ -108,8 +110,16 @@ async fn compress_json( > { let ip = context::client_ip(&headers, addr.ip()); let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?; + context::require_api_permission(&principal, &["compress"])?; + let admission = prepare_single_admission(&state, &principal, ip, true).await?; - let mut req = parse_single_file_request(&mut multipart).await?; + let mut req = parse_single_file_request( + &mut multipart, + admission.max_file_size_bytes, + admission.max_file_size_mb, + admission.limit_label, + ) + .await?; let format_in = compress::detect_format(&req.file_bytes)?; let format_out = req.output_format.unwrap_or(format_in); @@ -161,52 +171,8 @@ async fn compress_json( None }; - let (retention, quota_ctx) = match &principal { - context::Principal::Anonymous { session_id } => { - enforce_file_limits_anonymous(&state, &req.file_bytes)?; - ( - Duration::hours(state.config.anon_retention_hours as i64), - QuotaContext::Anonymous { - session_id: session_id.clone(), - ip, - }, - ) - } - context::Principal::User { - user_id, - email_verified, - .. - } => { - if !email_verified { - return Err(AppError::new(ErrorCode::EmailNotVerified, "请先验证邮箱")); - } - let billing = billing::get_user_billing(&state, *user_id).await?; - enforce_file_limits_plan(&billing.plan, &req.file_bytes)?; - ( - Duration::days(billing.plan.retention_days as i64), - QuotaContext::User(billing), - ) - } - context::Principal::ApiKey { - user_id, - api_key_id, - email_verified, - .. - } => { - if !email_verified { - return Err(AppError::new(ErrorCode::EmailNotVerified, "请先验证邮箱")); - } - let billing = billing::get_user_billing(&state, *user_id).await?; - if !billing.plan.feature_api_enabled { - return Err(AppError::new(ErrorCode::Forbidden, "当前套餐未开通 API")); - } - enforce_file_limits_plan(&billing.plan, &req.file_bytes)?; - ( - Duration::days(billing.plan.retention_days as i64), - QuotaContext::ApiKey(billing, *api_key_id), - ) - } - }; + let retention = admission.retention; + let quota_ctx = admission.quota_ctx; let mut idem_acquired = false; if let (Some(scope), Some(idem_key), Some(request_hash)) = ( @@ -291,6 +257,7 @@ async fn compress_json( .await?; let compressed_size = compressed.len() as u64; + let compressed = bytes::Bytes::from(compressed); let saved_bytes = original_size.saturating_sub(compressed_size); let saved_percent = if original_size == 0 { 0.0 @@ -431,6 +398,7 @@ async fn compress_direct( > { let ip = context::client_ip(&headers, addr.ip()); let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?; + context::require_api_permission(&principal, &["compress"])?; if matches!(principal, context::Principal::Anonymous { .. }) { return Err(AppError::new( @@ -438,17 +406,15 @@ async fn compress_direct( "对外 API 不支持匿名调用", )); } + let admission = prepare_single_admission(&state, &principal, ip, false).await?; - let mut req = parse_single_file_request(&mut multipart).await?; - - let email_verified = match &principal { - context::Principal::User { email_verified, .. } => *email_verified, - context::Principal::ApiKey { email_verified, .. } => *email_verified, - context::Principal::Anonymous { .. } => false, - }; - if !email_verified { - return Err(AppError::new(ErrorCode::EmailNotVerified, "请先验证邮箱")); - } + let mut req = parse_single_file_request( + &mut multipart, + admission.max_file_size_bytes, + admission.max_file_size_mb, + admission.limit_label, + ) + .await?; let format_in = compress::detect_format(&req.file_bytes)?; let format_out = req.output_format.unwrap_or(format_in); @@ -500,32 +466,8 @@ async fn compress_direct( None }; - let (retention, quota_ctx) = match &principal { - context::Principal::User { user_id, .. } => { - let billing = billing::get_user_billing(&state, *user_id).await?; - enforce_file_limits_plan(&billing.plan, &req.file_bytes)?; - ( - Duration::days(billing.plan.retention_days as i64), - QuotaContext::User(billing), - ) - } - context::Principal::ApiKey { - user_id, - api_key_id, - .. - } => { - let billing = billing::get_user_billing(&state, *user_id).await?; - if !billing.plan.feature_api_enabled { - return Err(AppError::new(ErrorCode::Forbidden, "当前套餐未开通 API")); - } - enforce_file_limits_plan(&billing.plan, &req.file_bytes)?; - ( - Duration::days(billing.plan.retention_days as i64), - QuotaContext::ApiKey(billing, *api_key_id), - ) - } - context::Principal::Anonymous { .. } => unreachable!(), - }; + let retention = admission.retention; + let quota_ctx = admission.quota_ctx; let mut idem_acquired = false; if let (Some(scope), Some(idem_key), Some(request_hash)) = ( @@ -655,6 +597,7 @@ async fn compress_direct( .await?; let compressed_size = compressed.len() as u64; + let compressed = bytes::Bytes::from(compressed); let saved_bytes = original_size.saturating_sub(compressed_size); let saved_percent = if original_size == 0 { 0.0 @@ -883,7 +826,12 @@ async fn load_direct_replay_bytes( Ok((bytes, fmt)) } -async fn parse_single_file_request(multipart: &mut Multipart) -> Result { +async fn parse_single_file_request( + multipart: &mut Multipart, + max_file_size_bytes: u64, + max_file_size_mb: u64, + limit_label: &str, +) -> Result { let mut file_bytes: Option> = None; let mut file_name: Option = None; let mut level = CompressionLevel::Medium; @@ -893,10 +841,18 @@ async fn parse_single_file_request(multipart: &mut Multipart) -> Result = None; let mut max_height: Option = None; let mut preserve_metadata = false; + let mut field_count = 0usize; while let Some(field) = multipart.next_field().await.map_err(|err| { AppError::new(ErrorCode::InvalidRequest, "读取上传内容失败").with_source(err) })? { + field_count += 1; + if field_count > 16 { + return Err(AppError::new( + ErrorCode::InvalidRequest, + "multipart 字段数量过多", + )); + } let name = field.name().unwrap_or("").to_string(); if name == "file" { if file_bytes.is_some() { @@ -905,17 +861,28 @@ async fn parse_single_file_request(multipart: &mut Multipart) -> Result max_file_size_bytes { + return Err(AppError::new( + ErrorCode::FileTooLarge, + format!("{limit_label}单文件最大 {max_file_size_mb} MB"), + )); + } + bytes.extend_from_slice(&chunk); + } + file_bytes = Some(bytes); continue; } - let text = field.text().await.map_err(|err| { - AppError::new(ErrorCode::InvalidRequest, "读取字段失败").with_source(err) - })?; + let text = multipart_utils::read_text(field).await?; match name.as_str() { "level" => { @@ -999,31 +966,6 @@ async fn parse_single_file_request(multipart: &mut Multipart) -> Result Result<(), AppError> { - let max = state.config.anon_max_file_size_mb * 1024 * 1024; - if bytes.len() as u64 > max { - return Err(AppError::new( - ErrorCode::FileTooLarge, - format!( - "匿名试用单文件最大 {} MB", - state.config.anon_max_file_size_mb - ), - )); - } - Ok(()) -} - -fn enforce_file_limits_plan(plan: &Plan, bytes: &[u8]) -> Result<(), AppError> { - let max = (plan.max_file_size_mb as u64) * 1024 * 1024; - if bytes.len() as u64 > max { - return Err(AppError::new( - ErrorCode::FileTooLarge, - format!("当前套餐单文件最大 {} MB", plan.max_file_size_mb), - )); - } - Ok(()) -} - #[derive(Debug, Clone)] enum QuotaContext { Anonymous { session_id: String, ip: IpAddr }, @@ -1031,6 +973,89 @@ enum QuotaContext { ApiKey(BillingContext, Uuid), } +struct SingleAdmission { + retention: Duration, + quota_ctx: QuotaContext, + max_file_size_bytes: u64, + max_file_size_mb: u64, + limit_label: &'static str, +} + +async fn prepare_single_admission( + state: &AppState, + principal: &context::Principal, + ip: IpAddr, + allow_anonymous: bool, +) -> Result { + match principal { + context::Principal::Anonymous { session_id } if allow_anonymous => Ok(SingleAdmission { + retention: Duration::hours(state.config.anon_retention_hours as i64), + quota_ctx: QuotaContext::Anonymous { + session_id: session_id.clone(), + ip, + }, + max_file_size_bytes: state.config.anon_max_file_size_mb * 1024 * 1024, + max_file_size_mb: state.config.anon_max_file_size_mb, + limit_label: "匿名试用", + }), + context::Principal::Anonymous { .. } => Err(AppError::new( + ErrorCode::Unauthorized, + "对外 API 不支持匿名调用", + )), + context::Principal::User { + user_id, + email_verified, + .. + } => { + if !email_verified { + return Err(AppError::new(ErrorCode::EmailNotVerified, "请先验证邮箱")); + } + let billing = billing::get_user_billing(state, *user_id).await?; + single_plan_admission(billing, None) + } + context::Principal::ApiKey { + user_id, + api_key_id, + email_verified, + .. + } => { + if !email_verified { + return Err(AppError::new(ErrorCode::EmailNotVerified, "请先验证邮箱")); + } + let billing = billing::get_user_billing(state, *user_id).await?; + if !billing.plan.feature_api_enabled { + return Err(AppError::new(ErrorCode::Forbidden, "当前套餐未开通 API")); + } + single_plan_admission(billing, Some(*api_key_id)) + } + } +} + +fn single_plan_admission( + billing: BillingContext, + api_key_id: Option, +) -> Result { + let max_file_size_mb = billing.plan.max_file_size_mb; + let retention_days = billing.plan.retention_days; + if max_file_size_mb <= 0 || retention_days <= 0 { + return Err(AppError::new(ErrorCode::Internal, "套餐上传限制配置无效")); + } + + let max_file_size_mb = max_file_size_mb as u64; + let retention = Duration::days(retention_days as i64); + let quota_ctx = match api_key_id { + Some(api_key_id) => QuotaContext::ApiKey(billing, api_key_id), + None => QuotaContext::User(billing), + }; + Ok(SingleAdmission { + retention, + quota_ctx, + max_file_size_bytes: max_file_size_mb * 1024 * 1024, + max_file_size_mb, + limit_label: "当前套餐", + }) +} + async fn ensure_quota_available( state: &AppState, ctx: &BillingContext, diff --git a/src/api/context.rs b/src/api/context.rs index 7b042d4..8aacb84 100644 --- a/src/api/context.rs +++ b/src/api/context.rs @@ -31,9 +31,31 @@ pub enum Principal { api_key_id: Uuid, role: String, email_verified: bool, + permissions: Vec, }, } +pub fn require_api_permission( + principal: &Principal, + accepted_permissions: &[&str], +) -> Result<(), AppError> { + let Principal::ApiKey { permissions, .. } = principal else { + return Ok(()); + }; + + if accepted_permissions + .iter() + .any(|expected| permissions.iter().any(|actual| actual == expected)) + { + return Ok(()); + } + + Err(AppError::new( + ErrorCode::Forbidden, + "API Key 缺少调用此接口所需的权限", + )) +} + pub fn client_ip(headers: &HeaderMap, connect_ip: IpAddr) -> IpAddr { resolve_client_ip(headers, connect_ip, crate::config::trust_proxy_headers()) } @@ -111,11 +133,12 @@ async fn try_jwt(state: &AppState, headers: &HeaderMap) -> Result>, + token_version: i32, } let user = sqlx::query_as::<_, UserAuthRow>( r#" - SELECT id, role::text AS role, is_active, email_verified_at + SELECT id, role::text AS role, is_active, email_verified_at, token_version FROM users WHERE id = $1 "#, @@ -129,6 +152,12 @@ async fn try_jwt(state: &AppState, headers: &HeaderMap) -> Result>, + permissions: serde_json::Value, + rate_limit: i32, } let row = sqlx::query_as::<_, ApiKeyAuthRow>( @@ -180,7 +211,9 @@ async fn try_api_key( k.is_active, u.role::text AS user_role, u.is_active AS user_is_active, - u.email_verified_at + u.email_verified_at, + k.permissions, + k.rate_limit FROM api_keys k JOIN users u ON u.id = k.user_id WHERE k.key_prefix = $1 @@ -201,6 +234,29 @@ async fn try_api_key( return Err(AppError::new(ErrorCode::Unauthorized, "API Key 无效")); } + crate::services::rate_limit::enforce( + state, + "api_key", + &row.id.to_string(), + row.rate_limit.clamp(1, 100_000) as u32, + 60, + "API Key 请求频率已超过限制", + ) + .await?; + + let permissions = row + .permissions + .as_array() + .ok_or_else(|| AppError::new(ErrorCode::Internal, "API Key 权限配置无效"))? + .iter() + .map(|permission| { + permission + .as_str() + .map(str::to_owned) + .ok_or_else(|| AppError::new(ErrorCode::Internal, "API Key 权限配置无效")) + }) + .collect::, _>>()?; + let _ = sqlx::query("UPDATE api_keys SET last_used_at = NOW(), last_used_ip = $2 WHERE id = $1") .bind(row.id) @@ -216,6 +272,7 @@ async fn try_api_key( api_key_id: row.id, role: row.user_role, email_verified: row.email_verified_at.is_some() || !verification_required, + permissions, })) } @@ -281,4 +338,24 @@ mod tests { "203.0.113.9".parse::().unwrap() ); } + + #[test] + fn api_key_permissions_are_enforced_without_affecting_users() { + let key = Principal::ApiKey { + user_id: Uuid::new_v4(), + api_key_id: Uuid::new_v4(), + role: "user".to_string(), + email_verified: true, + permissions: vec!["compress".to_string()], + }; + assert!(require_api_permission(&key, &["compress"]).is_ok()); + assert!(require_api_permission(&key, &["billing_read"]).is_err()); + + let user = Principal::User { + user_id: Uuid::new_v4(), + role: "user".to_string(), + email_verified: true, + }; + assert!(require_api_permission(&user, &["billing_read"]).is_ok()); + } } diff --git a/src/api/downloads.rs b/src/api/downloads.rs index 7b68cbe..cbcaccb 100644 --- a/src/api/downloads.rs +++ b/src/api/downloads.rs @@ -46,6 +46,7 @@ async fn download_file( ) -> Result<(axum_extra::extract::cookie::CookieJar, Response), AppError> { let ip = context::client_ip(&headers, addr.ip()); let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?; + context::require_api_permission(&principal, &["compress", "batch_compress"])?; let row = sqlx::query_as::<_, DownloadRow>( r#" @@ -251,6 +252,7 @@ async fn download_task_zip( ) -> Result<(axum_extra::extract::cookie::CookieJar, Response), AppError> { let ip = context::client_ip(&headers, addr.ip()); let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?; + context::require_api_permission(&principal, &["compress", "batch_compress"])?; let task = sqlx::query_as::<_, TaskZipRow>( r#" diff --git a/src/api/mod.rs b/src/api/mod.rs index c1b9bdb..10b06b8 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -7,6 +7,7 @@ mod context; mod downloads; mod envelope; mod health; +mod multipart; mod redemption; mod response; mod tasks; diff --git a/src/api/multipart.rs b/src/api/multipart.rs new file mode 100644 index 0000000..f68be7e --- /dev/null +++ b/src/api/multipart.rs @@ -0,0 +1,23 @@ +use crate::error::{AppError, ErrorCode}; + +use axum::extract::multipart::Field; + +const MAX_TEXT_FIELD_BYTES: usize = 8 * 1024; + +pub async fn read_text(mut field: Field<'_>) -> Result { + let mut bytes = Vec::new(); + while let Some(chunk) = field + .chunk() + .await + .map_err(|err| AppError::new(ErrorCode::InvalidRequest, "读取字段失败").with_source(err))? + { + if bytes.len().saturating_add(chunk.len()) > MAX_TEXT_FIELD_BYTES { + return Err(AppError::new(ErrorCode::InvalidRequest, "上传参数内容过长")); + } + bytes.extend_from_slice(&chunk); + } + + String::from_utf8(bytes).map_err(|err| { + AppError::new(ErrorCode::InvalidRequest, "上传参数不是有效 UTF-8").with_source(err) + }) +} diff --git a/src/api/tasks.rs b/src/api/tasks.rs index b17b33c..afbfef1 100644 --- a/src/api/tasks.rs +++ b/src/api/tasks.rs @@ -1,10 +1,12 @@ use crate::api::context; use crate::api::envelope::Envelope; +use crate::api::multipart as multipart_utils; use crate::error::{AppError, ErrorCode}; use crate::services::billing; use crate::services::billing::{BillingContext, Plan}; use crate::services::compress; use crate::services::compress::{CompressionLevel, ImageFmt}; +use crate::services::filename; use crate::services::idempotency; use crate::services::quota; use crate::services::storage; @@ -58,6 +60,23 @@ struct BatchOptions { preserve_metadata: bool, } +#[derive(Debug)] +struct BatchUploadLimits { + max_files: usize, + max_file_size_bytes: u64, + max_file_size_mb: u64, + label: &'static str, +} + +#[derive(Debug)] +struct BatchAdmission { + retention: Duration, + task_owner: TaskOwner, + source: &'static str, + billing: Option, + limits: BatchUploadLimits, +} + async fn create_batch_task( State(state): State, jar: axum_extra::extract::cookie::CookieJar, @@ -73,6 +92,8 @@ async fn create_batch_task( > { let ip = context::client_ip(&headers, addr.ip()); let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?; + context::require_api_permission(&principal, &["compress", "batch_compress"])?; + let admission = prepare_batch_admission(&state, &principal).await?; let idempotency_key = headers .get("idempotency-key") @@ -89,24 +110,41 @@ async fn create_batch_task( }; let task_id = Uuid::new_v4(); - let (files, opts, request_hash) = parse_batch_request(&state, task_id, &mut multipart).await?; + let parsed = parse_batch_request(&state, task_id, &mut multipart, &admission.limits).await; + let (files, opts, request_hash) = match parsed { + Ok(parsed) => parsed, + Err(err) => { + let base_dir = format!("{}/orig/{task_id}", state.config.storage_path); + let _ = tokio::fs::remove_dir_all(base_dir).await; + return Err(err); + } + }; if files.is_empty() { cleanup_file_paths(&files).await; + let base_dir = format!("{}/orig/{task_id}", state.config.storage_path); + let _ = tokio::fs::remove_dir_all(base_dir).await; return Err(AppError::new(ErrorCode::InvalidRequest, "缺少 files[]")); } let mut idem_acquired = false; if let (Some(scope), Some(idem_key)) = (idempotency_scope, idempotency_key.as_deref()) { - match idempotency::begin( + let begin_result = idempotency::begin( &state, scope, idem_key, &request_hash, state.config.idempotency_ttl_hours as i64, ) - .await? - { + .await; + let begin_result = match begin_result { + Ok(result) => result, + Err(err) => { + cleanup_file_paths(&files).await; + return Err(err); + } + }; + match begin_result { idempotency::BeginResult::Replay { response_body, .. } => { cleanup_file_paths(&files).await; let resp: BatchCreateResponse = @@ -152,68 +190,26 @@ async fn create_batch_task( let mut anonymous_reserved_units = 0u32; let create_result: Result = (async { - let (retention, task_owner, source) = match &principal { - context::Principal::Anonymous { session_id } => { - enforce_batch_limits_anonymous(&state, &files)?; + match &admission.task_owner { + TaskOwner::Anonymous { session_id } => { let units = u32::try_from(files.len()).map_err(|_| { AppError::new(ErrorCode::InvalidRequest, "批量文件数量超出限制") })?; quota::consume_anonymous_units(&state, session_id, ip, units).await?; anonymous_reserved_units = units; - Ok(( - Duration::hours(state.config.anon_retention_hours as i64), - TaskOwner::Anonymous { - session_id: session_id.clone(), - }, - "web", - )) } - context::Principal::User { - user_id, - email_verified, - .. - } => { - if !email_verified { - return Err(AppError::new(ErrorCode::EmailNotVerified, "请先验证邮箱")); - } - let billing = billing::get_user_billing(&state, *user_id).await?; - enforce_batch_limits_plan(&billing.plan, &files)?; - ensure_quota_available(&state, &billing, files.len() as i32).await?; - Ok(( - Duration::days(billing.plan.retention_days as i64), - TaskOwner::User { user_id: *user_id }, - "web", - )) + TaskOwner::User { .. } | TaskOwner::ApiKey { .. } => { + let billing = admission + .billing + .as_ref() + .ok_or_else(|| AppError::new(ErrorCode::Internal, "批量任务计费上下文缺失"))?; + ensure_quota_available(&state, billing, files.len() as i32).await?; } - context::Principal::ApiKey { - user_id, - api_key_id, - email_verified, - .. - } => { - if !email_verified { - return Err(AppError::new(ErrorCode::EmailNotVerified, "请先验证邮箱")); - } - let billing = billing::get_user_billing(&state, *user_id).await?; - if !billing.plan.feature_api_enabled { - return Err(AppError::new(ErrorCode::Forbidden, "当前套餐未开通 API")); - } - enforce_batch_limits_plan(&billing.plan, &files)?; - ensure_quota_available(&state, &billing, files.len() as i32).await?; - Ok(( - Duration::days(billing.plan.retention_days as i64), - TaskOwner::ApiKey { - user_id: *user_id, - api_key_id: *api_key_id, - }, - "api", - )) - } - }?; + } - let expires_at = Utc::now() + retention; - let retention_hours = retention.num_hours(); - let (user_id, session_id, api_key_id) = match &task_owner { + let expires_at = Utc::now() + admission.retention; + let retention_hours = admission.retention.num_hours(); + let (user_id, session_id, api_key_id) = match &admission.task_owner { TaskOwner::Anonymous { session_id } => (None, Some(session_id.clone()), None), TaskOwner::User { user_id } => (Some(*user_id), None, None), TaskOwner::ApiKey { @@ -252,7 +248,7 @@ async fn create_batch_task( .bind(session_id) .bind(api_key_id) .bind(ip.to_string()) - .bind(source) + .bind(admission.source) .bind(opts.level.as_str()) .bind(opts.output_format.map(|f| f.as_str())) .bind(opts.max_width.map(|v| v as i32)) @@ -374,6 +370,85 @@ enum TaskOwner { ApiKey { user_id: Uuid, api_key_id: Uuid }, } +async fn prepare_batch_admission( + state: &AppState, + principal: &context::Principal, +) -> Result { + match principal { + context::Principal::Anonymous { session_id } => Ok(BatchAdmission { + retention: Duration::hours(state.config.anon_retention_hours as i64), + task_owner: TaskOwner::Anonymous { + session_id: session_id.clone(), + }, + source: "web", + billing: None, + limits: BatchUploadLimits { + max_files: state.config.anon_max_files_per_batch as usize, + max_file_size_bytes: state.config.anon_max_file_size_mb * 1024 * 1024, + max_file_size_mb: state.config.anon_max_file_size_mb, + label: "匿名试用", + }, + }), + context::Principal::User { + user_id, + email_verified, + .. + } => { + if !email_verified { + return Err(AppError::new(ErrorCode::EmailNotVerified, "请先验证邮箱")); + } + let billing = billing::get_user_billing(state, *user_id).await?; + let limits = plan_upload_limits(&billing.plan)?; + Ok(BatchAdmission { + retention: Duration::days(billing.plan.retention_days as i64), + task_owner: TaskOwner::User { user_id: *user_id }, + source: "web", + billing: Some(billing), + limits, + }) + } + context::Principal::ApiKey { + user_id, + api_key_id, + email_verified, + .. + } => { + if !email_verified { + return Err(AppError::new(ErrorCode::EmailNotVerified, "请先验证邮箱")); + } + let billing = billing::get_user_billing(state, *user_id).await?; + if !billing.plan.feature_api_enabled { + return Err(AppError::new(ErrorCode::Forbidden, "当前套餐未开通 API")); + } + let limits = plan_upload_limits(&billing.plan)?; + Ok(BatchAdmission { + retention: Duration::days(billing.plan.retention_days as i64), + task_owner: TaskOwner::ApiKey { + user_id: *user_id, + api_key_id: *api_key_id, + }, + source: "api", + billing: Some(billing), + limits, + }) + } + } +} + +fn plan_upload_limits(plan: &Plan) -> Result { + if plan.max_files_per_batch <= 0 || plan.max_file_size_mb <= 0 { + return Err(AppError::new(ErrorCode::Internal, "套餐上传限制配置无效")); + } + + let max_file_size_mb = plan.max_file_size_mb as u64; + Ok(BatchUploadLimits { + max_files: plan.max_files_per_batch as usize, + max_file_size_bytes: max_file_size_mb * 1024 * 1024, + max_file_size_mb, + label: "当前套餐", + }) +} + async fn enqueue_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> { let mut conn = state.redis.clone(); let now = Utc::now().to_rfc3339(); @@ -397,6 +472,7 @@ async fn parse_batch_request( state: &AppState, task_id: Uuid, multipart: &mut Multipart, + limits: &BatchUploadLimits, ) -> Result<(Vec, BatchOptions, String), AppError> { let mut files: Vec = Vec::new(); let mut file_digests: Vec = Vec::new(); @@ -408,6 +484,7 @@ async fn parse_batch_request( max_height: None, preserve_metadata: false, }; + let mut field_count = 0usize; let base_dir = format!("{}/orig/{task_id}", state.config.storage_path); tokio::fs::create_dir_all(&base_dir).await.map_err(|err| { @@ -428,40 +505,30 @@ async fn parse_batch_request( }; let Some(field) = field else { break }; + field_count += 1; + if field_count > limits.max_files.saturating_add(16) { + cleanup_file_paths(&files).await; + return Err(AppError::new( + ErrorCode::InvalidRequest, + "multipart 字段数量过多", + )); + } let name = field.name().unwrap_or("").to_string(); if name == "files" || name == "files[]" { + if files.len() >= limits.max_files { + cleanup_file_paths(&files).await; + return Err(AppError::new( + ErrorCode::InvalidRequest, + format!("{}单次最多 {} 个文件", limits.label, limits.max_files), + )); + } + + let mut field = field; let file_id = Uuid::new_v4(); - let original_name = field.file_name().unwrap_or("upload").to_string(); - let bytes = match field.bytes().await { - Ok(v) => v, - Err(err) => { - cleanup_file_paths(&files).await; - return Err( - AppError::new(ErrorCode::InvalidRequest, "读取文件失败").with_source(err) - ); - } - }; - - let original_size = bytes.len() as u64; - let file_digest = { - let mut h = Sha256::new(); - h.update(&bytes); - h.update(original_name.as_bytes()); - hex::encode(h.finalize()) - }; - let original_format = match compress::detect_format(&bytes) { - Ok(v) => v, - Err(err) => { - cleanup_file_paths(&files).await; - return Err(err); - } - }; - - let output_format = opts.output_format.unwrap_or(original_format); - let path = format!("{base_dir}/{file_id}.{}", original_format.extension()); - - let mut f = match tokio::fs::File::create(&path).await { + let original_name = filename::normalize_upload_name(field.file_name()); + let temp_path = format!("{base_dir}/{file_id}.upload"); + let mut f = match tokio::fs::File::create(&temp_path).await { Ok(v) => v, Err(err) => { cleanup_file_paths(&files).await; @@ -469,13 +536,83 @@ async fn parse_batch_request( .with_source(err)); } }; - if let Err(err) = f.write_all(&bytes).await { - let _ = tokio::fs::remove_file(&path).await; + + const FORMAT_PREFIX_LIMIT: usize = 64 * 1024; + let mut format_prefix = Vec::with_capacity(FORMAT_PREFIX_LIMIT); + let mut original_size = 0u64; + let mut digest = Sha256::new(); + + loop { + let chunk = match field.chunk().await { + Ok(chunk) => chunk, + Err(err) => { + drop(f); + let _ = tokio::fs::remove_file(&temp_path).await; + cleanup_file_paths(&files).await; + return Err(AppError::new(ErrorCode::InvalidRequest, "读取文件失败") + .with_source(err)); + } + }; + let Some(chunk) = chunk else { break }; + + original_size = original_size + .checked_add(chunk.len() as u64) + .ok_or_else(|| AppError::new(ErrorCode::FileTooLarge, "文件大小超出限制"))?; + if original_size > limits.max_file_size_bytes { + drop(f); + let _ = tokio::fs::remove_file(&temp_path).await; + cleanup_file_paths(&files).await; + return Err(AppError::new( + ErrorCode::FileTooLarge, + format!("{}单文件最大 {} MB", limits.label, limits.max_file_size_mb), + )); + } + + digest.update(&chunk); + if format_prefix.len() < FORMAT_PREFIX_LIMIT { + let remaining = FORMAT_PREFIX_LIMIT - format_prefix.len(); + format_prefix.extend_from_slice(&chunk[..chunk.len().min(remaining)]); + } + if let Err(err) = f.write_all(&chunk).await { + drop(f); + let _ = tokio::fs::remove_file(&temp_path).await; + cleanup_file_paths(&files).await; + return Err(AppError::new(ErrorCode::StorageUnavailable, "写入文件失败") + .with_source(err)); + } + } + + if let Err(err) = f.flush().await { + drop(f); + let _ = tokio::fs::remove_file(&temp_path).await; cleanup_file_paths(&files).await; return Err( AppError::new(ErrorCode::StorageUnavailable, "写入文件失败").with_source(err) ); } + drop(f); + + let original_format = match compress::detect_format(&format_prefix) { + Ok(v) => v, + Err(err) => { + let _ = tokio::fs::remove_file(&temp_path).await; + cleanup_file_paths(&files).await; + return Err(err); + } + }; + let output_format = opts.output_format.unwrap_or(original_format); + let path = format!("{base_dir}/{file_id}.{}", original_format.extension()); + if let Err(err) = tokio::fs::rename(&temp_path, &path).await { + let _ = tokio::fs::remove_file(&temp_path).await; + cleanup_file_paths(&files).await; + return Err( + AppError::new(ErrorCode::StorageUnavailable, "保存上传文件失败") + .with_source(err), + ); + } + + digest.update(original_name.as_bytes()); + let file_digest = hex::encode(digest.finalize()); files.push(BatchFileInput { file_id, @@ -489,7 +626,7 @@ async fn parse_batch_request( continue; } - let text = match field.text().await { + let text = match multipart_utils::read_text(field).await { Ok(v) => v, Err(err) => { cleanup_file_paths(&files).await; @@ -607,56 +744,6 @@ async fn parse_batch_request( Ok((files, opts, request_hash)) } -fn enforce_batch_limits_anonymous( - state: &AppState, - files: &[BatchFileInput], -) -> Result<(), AppError> { - let max_files = state.config.anon_max_files_per_batch as usize; - if files.len() > max_files { - return Err(AppError::new( - ErrorCode::InvalidRequest, - format!("匿名试用单次最多 {} 个文件", max_files), - )); - } - - let max_bytes = state.config.anon_max_file_size_mb * 1024 * 1024; - for f in files { - if f.original_size > max_bytes { - return Err(AppError::new( - ErrorCode::FileTooLarge, - format!( - "匿名试用单文件最大 {} MB", - state.config.anon_max_file_size_mb - ), - )); - } - } - - Ok(()) -} - -fn enforce_batch_limits_plan(plan: &Plan, files: &[BatchFileInput]) -> Result<(), AppError> { - let max_files = plan.max_files_per_batch as usize; - if files.len() > max_files { - return Err(AppError::new( - ErrorCode::InvalidRequest, - format!("当前套餐单次最多 {} 个文件", plan.max_files_per_batch), - )); - } - - let max_bytes = (plan.max_file_size_mb as u64) * 1024 * 1024; - for f in files { - if f.original_size > max_bytes { - return Err(AppError::new( - ErrorCode::FileTooLarge, - format!("当前套餐单文件最大 {} MB", plan.max_file_size_mb), - )); - } - } - - Ok(()) -} - async fn ensure_quota_available( state: &AppState, ctx: &BillingContext, @@ -729,6 +816,7 @@ async fn get_task( > { let ip = context::client_ip(&headers, addr.ip()); let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?; + context::require_api_permission(&principal, &["compress", "batch_compress"])?; let task = sqlx::query_as::<_, TaskRow>( r#" @@ -843,6 +931,7 @@ async fn cancel_task( > { let ip = context::client_ip(&headers, addr.ip()); let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?; + context::require_api_permission(&principal, &["compress", "batch_compress"])?; let task = sqlx::query_as::<_, TaskRow>( "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", @@ -905,6 +994,7 @@ async fn delete_task( > { let ip = context::client_ip(&headers, addr.ip()); let (jar, principal) = context::authenticate(&state, jar, &headers, ip).await?; + context::require_api_permission(&principal, &["compress", "batch_compress"])?; let task = sqlx::query_as::<_, TaskRow>( "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", @@ -1037,7 +1127,14 @@ fn authorize_task( } async fn cleanup_file_paths(files: &[BatchFileInput]) { + let parent = files + .first() + .and_then(|file| std::path::Path::new(&file.storage_path).parent()) + .map(std::path::Path::to_path_buf); for f in files { let _ = tokio::fs::remove_file(&f.storage_path).await; } + if let Some(parent) = parent { + let _ = tokio::fs::remove_dir(parent).await; + } } diff --git a/src/api/user.rs b/src/api/user.rs index 42964bf..50f2a03 100644 --- a/src/api/user.rs +++ b/src/api/user.rs @@ -2,11 +2,9 @@ use crate::api::context; use crate::api::envelope::Envelope; use crate::error::{AppError, ErrorCode}; use crate::services::billing; -use crate::services::mail; -use crate::services::settings; +use crate::services::{credentials, mail, settings}; use crate::state::AppState; -use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier}; use axum::extract::{ConnectInfo, Path, Query, State}; use axum::http::HeaderMap; use axum::routing::{delete, get, post, put}; @@ -15,8 +13,8 @@ use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use chrono::{DateTime, Duration, Utc}; use rand::RngCore; use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; use sqlx::FromRow; +use std::collections::HashMap; use std::net::SocketAddr; use uuid::Uuid; @@ -172,7 +170,7 @@ async fn update_profile( if let Some(email) = req.email.as_ref() { let email = email.trim().to_lowercase(); - validate_email(&email)?; + credentials::validate_email(&email)?; if email != user.email { next_email = email; email_changed = true; @@ -181,7 +179,7 @@ async fn update_profile( if let Some(username) = req.username.as_ref() { let username = username.trim().to_string(); - validate_username(&username)?; + credentials::validate_username(&username)?; if username != user.username { next_username = username; } @@ -238,8 +236,8 @@ async fn update_profile( let mut verification_link: Option = None; if email_changed && verification_required { - let token = generate_token(); - let token_hash = sha256_hex(&token); + let token = credentials::generate_token(); + let token_hash = credentials::sha256_hex(&token); let expires_at = Utc::now() + Duration::hours(24); sqlx::query( @@ -317,7 +315,7 @@ async fn update_password( _ => return Err(AppError::new(ErrorCode::Unauthorized, "未登录")), }; - validate_password(&req.new_password)?; + credentials::validate_password(&req.new_password)?; #[derive(Debug, FromRow)] struct PasswordRow { @@ -330,10 +328,14 @@ async fn update_password( .await .map_err(|err| AppError::new(ErrorCode::Internal, "查询用户失败").with_source(err))?; - verify_password(&req.current_password, &row.password_hash)?; + if !credentials::verify_password(&req.current_password, &row.password_hash).await? { + return Err(AppError::new(ErrorCode::Unauthorized, "密码错误")); + } - let new_hash = hash_password(&req.new_password)?; - sqlx::query("UPDATE users SET password_hash = $2, updated_at = NOW() WHERE id = $1") + let new_hash = credentials::hash_password(&req.new_password).await?; + sqlx::query( + "UPDATE users SET password_hash = $2, token_version = token_version + 1, updated_at = NOW() WHERE id = $1", + ) .bind(user_id) .bind(new_hash) .execute(&state.db) @@ -498,6 +500,7 @@ async fn list_history( #[derive(Debug, FromRow)] struct FileRow { + task_id: Uuid, id: Uuid, original_name: String, original_size: i64, @@ -510,11 +513,14 @@ async fn list_history( } let now = Utc::now(); - let mut result_tasks = Vec::with_capacity(tasks.len()); - for task in tasks { - let files: Vec = sqlx::query_as::<_, FileRow>( + let task_ids = tasks.iter().map(|task| task.id).collect::>(); + let files = if task_ids.is_empty() { + Vec::new() + } else { + sqlx::query_as::<_, FileRow>( r#" SELECT + task_id, id, original_name, original_size, @@ -525,16 +531,25 @@ async fn list_history( error_message, COALESCE(storage_key, storage_path) IS NOT NULL AS has_storage FROM task_files - WHERE task_id = $1 - ORDER BY created_at ASC + WHERE task_id = ANY($1) + ORDER BY task_id, created_at ASC "#, ) - .bind(task.id) + .bind(&task_ids) .fetch_all(&state.db) .await - .map_err(|err| AppError::new(ErrorCode::Internal, "查询任务文件失败").with_source(err))?; + .map_err(|err| AppError::new(ErrorCode::Internal, "查询任务文件失败").with_source(err))? + }; + let mut files_by_task = HashMap::>::new(); + for file in files { + files_by_task.entry(file.task_id).or_default().push(file); + } - let file_views = files + let mut result_tasks = Vec::with_capacity(tasks.len()); + for task in tasks { + let file_views = files_by_task + .remove(&task.id) + .unwrap_or_default() .into_iter() .map(|file| HistoryFileView { file_id: file.id, @@ -837,13 +852,7 @@ fn generate_api_key() -> (String, String) { } fn normalize_permissions(input: Option>) -> Result { - let allowed = [ - "compress", - "batch_compress", - "read_stats", - "billing_read", - "webhook_manage", - ]; + let allowed = ["compress", "batch_compress"]; let mut perms = Vec::::new(); if let Some(values) = input { @@ -871,65 +880,6 @@ fn normalize_permissions(input: Option>) -> Result Result<(), AppError> { - if email.trim().is_empty() || !email.contains('@') { - return Err(AppError::new(ErrorCode::InvalidRequest, "邮箱格式不正确")); - } - if email.len() > 255 { - return Err(AppError::new(ErrorCode::InvalidRequest, "邮箱过长")); - } - Ok(()) -} - -fn validate_username(username: &str) -> Result<(), AppError> { - if username.trim().is_empty() { - return Err(AppError::new(ErrorCode::InvalidRequest, "用户名不能为空")); - } - if username.len() > 50 { - return Err(AppError::new(ErrorCode::InvalidRequest, "用户名过长")); - } - Ok(()) -} - -fn validate_password(password: &str) -> Result<(), AppError> { - if password.len() < 8 { - return Err(AppError::new(ErrorCode::InvalidRequest, "密码至少 8 位")); - } - if password.len() > 128 { - return Err(AppError::new(ErrorCode::InvalidRequest, "密码过长")); - } - Ok(()) -} - -fn hash_password(password: &str) -> Result { - let salt = argon2::password_hash::SaltString::generate(&mut rand::rngs::OsRng); - let hashed = Argon2::default() - .hash_password(password.as_bytes(), &salt) - .map_err(|err| AppError::new(ErrorCode::Internal, "密码哈希失败").with_source(err))?; - Ok(hashed.to_string()) -} - -fn verify_password(password: &str, password_hash: &str) -> Result<(), AppError> { - let parsed = PasswordHash::new(password_hash) - .map_err(|err| AppError::new(ErrorCode::Internal, "密码哈希格式错误").with_source(err))?; - Argon2::default() - .verify_password(password.as_bytes(), &parsed) - .map_err(|_| AppError::new(ErrorCode::Unauthorized, "密码错误"))?; - Ok(()) -} - -fn generate_token() -> String { - let mut bytes = [0u8; 32]; - rand::rngs::OsRng.fill_bytes(&mut bytes); - URL_SAFE_NO_PAD.encode(bytes) -} - -fn sha256_hex(token: &str) -> String { - let mut hasher = Sha256::new(); - hasher.update(token.as_bytes()); - hex::encode(hasher.finalize()) -} - fn map_unique_violation(err: sqlx::Error) -> AppError { if let sqlx::Error::Database(db_err) = &err { if let Some(code) = db_err.code() { diff --git a/src/auth.rs b/src/auth.rs index 2d47ff7..5518526 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -11,6 +11,8 @@ pub struct Claims { pub sub: Uuid, pub role: String, pub exp: usize, + #[serde(default)] + pub ver: i32, } pub fn issue_jwt( @@ -18,12 +20,14 @@ pub fn issue_jwt( jwt_expiry_hours: i64, user_id: Uuid, role: &str, + token_version: i32, ) -> Result<(String, DateTime), AppError> { let expires_at = Utc::now() + Duration::hours(jwt_expiry_hours); let claims = Claims { sub: user_id, role: role.to_string(), exp: expires_at.timestamp() as usize, + ver: token_version, }; let token = jsonwebtoken::encode( @@ -60,3 +64,18 @@ pub fn decode_jwt(jwt_secret: &str, token: &str) -> Result { .map(|data| data.claims) .map_err(|_| AppError::new(ErrorCode::Unauthorized, "Token 无效或已过期")) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn jwt_round_trip_preserves_token_version() { + let user_id = Uuid::new_v4(); + let (token, _) = issue_jwt("test-secret", 1, user_id, "user", 7).unwrap(); + let claims = decode_jwt("test-secret", &token).unwrap(); + assert_eq!(claims.sub, user_id); + assert_eq!(claims.role, "user"); + assert_eq!(claims.ver, 7); + } +} diff --git a/src/error.rs b/src/error.rs index f2374eb..de7ebc4 100644 --- a/src/error.rs +++ b/src/error.rs @@ -1,4 +1,8 @@ -use axum::{http::StatusCode, response::IntoResponse, Json}; +use axum::{ + http::{HeaderValue, StatusCode}, + response::IntoResponse, + Json, +}; use serde::Serialize; use std::fmt::{Display, Formatter}; use uuid::Uuid; @@ -96,12 +100,6 @@ impl IntoResponse for AppError { fn into_response(self) -> axum::response::Response { let request_id = format!("req_{}", Uuid::new_v4()); - if let Some(source) = &self.source { - tracing::error!(code = %self.code.as_str(), request_id = %request_id, message = %self.message, source = %source); - } else { - tracing::error!(code = %self.code.as_str(), request_id = %request_id, message = %self.message); - } - let status = match self.code { ErrorCode::InvalidRequest => StatusCode::BAD_REQUEST, ErrorCode::InvalidImage => StatusCode::BAD_REQUEST, @@ -122,15 +120,34 @@ impl IntoResponse for AppError { ErrorCode::Internal => StatusCode::INTERNAL_SERVER_ERROR, }; + if status.is_server_error() { + if let Some(source) = &self.source { + tracing::error!(code = %self.code.as_str(), request_id = %request_id, message = %self.message, source = %source); + } else { + tracing::error!(code = %self.code.as_str(), request_id = %request_id, message = %self.message); + } + } else if matches!( + status, + StatusCode::UNAUTHORIZED | StatusCode::TOO_MANY_REQUESTS + ) { + tracing::warn!(code = %self.code.as_str(), request_id = %request_id, message = %self.message); + } else { + tracing::debug!(code = %self.code.as_str(), request_id = %request_id, message = %self.message); + } + let body = ErrorEnvelope { success: false, error: ErrorPayload { code: self.code, message: self.message, - request_id, + request_id: request_id.clone(), }, }; - (status, Json(body)).into_response() + let mut response = (status, Json(body)).into_response(); + if let Ok(value) = HeaderValue::from_str(&request_id) { + response.headers_mut().insert("x-request-id", value); + } + response } } diff --git a/src/services/bootstrap.rs b/src/services/bootstrap.rs index 0bed3a1..f3a33c8 100644 --- a/src/services/bootstrap.rs +++ b/src/services/bootstrap.rs @@ -1,7 +1,7 @@ use crate::error::{AppError, ErrorCode}; +use crate::services::credentials; use crate::state::AppState; -use argon2::{Argon2, PasswordHasher}; use chrono::Utc; use sqlx::FromRow; use tracing::{info, warn}; @@ -71,7 +71,7 @@ pub async fn ensure_admin_user(state: &AppState) -> Result<(), AppError> { } let existing = matching.pop(); - let password_hash = hash_password(&admin_password)?; + let password_hash = credentials::hash_password(&admin_password).await?; if let Some(row) = existing { sqlx::query( @@ -210,14 +210,6 @@ fn validate_password(password: &str) -> Result<(), AppError> { Ok(()) } -fn hash_password(password: &str) -> Result { - let salt = argon2::password_hash::SaltString::generate(&mut rand::rngs::OsRng); - let hashed = Argon2::default() - .hash_password(password.as_bytes(), &salt) - .map_err(|err| AppError::new(ErrorCode::Internal, "密码哈希失败").with_source(err))?; - Ok(hashed.to_string()) -} - fn env_string(key: &str) -> Option { std::env::var(key) .ok() diff --git a/src/services/credentials.rs b/src/services/credentials.rs new file mode 100644 index 0000000..5344bcd --- /dev/null +++ b/src/services/credentials.rs @@ -0,0 +1,99 @@ +use crate::error::{AppError, ErrorCode}; + +use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier}; +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; +use rand::RngCore; +use sha2::{Digest, Sha256}; + +pub fn validate_email(email: &str) -> Result<(), AppError> { + if email.trim().is_empty() || !email.contains('@') { + return Err(AppError::new(ErrorCode::InvalidRequest, "邮箱格式不正确")); + } + if email.len() > 255 { + return Err(AppError::new(ErrorCode::InvalidRequest, "邮箱过长")); + } + Ok(()) +} + +pub fn validate_username(username: &str) -> Result<(), AppError> { + if username.trim().is_empty() { + return Err(AppError::new(ErrorCode::InvalidRequest, "用户名不能为空")); + } + if username.len() > 50 { + return Err(AppError::new(ErrorCode::InvalidRequest, "用户名过长")); + } + Ok(()) +} + +pub fn validate_password(password: &str) -> Result<(), AppError> { + if password.len() < 8 { + return Err(AppError::new(ErrorCode::InvalidRequest, "密码至少 8 位")); + } + if password.len() > 128 { + return Err(AppError::new(ErrorCode::InvalidRequest, "密码过长")); + } + Ok(()) +} + +pub async fn hash_password(password: &str) -> Result { + let password = password.to_owned(); + tokio::task::spawn_blocking(move || { + let salt = argon2::password_hash::SaltString::generate(&mut rand::rngs::OsRng); + Argon2::default() + .hash_password(password.as_bytes(), &salt) + .map(|hash| hash.to_string()) + .map_err(|err| err.to_string()) + }) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "密码哈希任务失败").with_source(err))? + .map_err(|err| AppError::new(ErrorCode::Internal, "密码哈希失败").with_source(err)) +} + +pub async fn verify_password(password: &str, password_hash: &str) -> Result { + let password = password.to_owned(); + let password_hash = password_hash.to_owned(); + tokio::task::spawn_blocking(move || { + let parsed = PasswordHash::new(&password_hash).map_err(|err| err.to_string())?; + Ok::<_, String>( + Argon2::default() + .verify_password(password.as_bytes(), &parsed) + .is_ok(), + ) + }) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "密码校验任务失败").with_source(err))? + .map_err(|err| AppError::new(ErrorCode::Internal, "密码哈希格式错误").with_source(err)) +} + +pub fn generate_token() -> String { + let mut bytes = [0u8; 32]; + rand::rngs::OsRng.fill_bytes(&mut bytes); + URL_SAFE_NO_PAD.encode(bytes) +} + +pub fn sha256_hex(value: &str) -> String { + let mut hasher = Sha256::new(); + hasher.update(value.as_bytes()); + hex::encode(hasher.finalize()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn password_hash_round_trip_and_rejects_wrong_password() { + let hash = hash_password("correct-horse").await.unwrap(); + assert!(verify_password("correct-horse", &hash).await.unwrap()); + assert!(!verify_password("wrong-password", &hash).await.unwrap()); + } + + #[test] + fn generated_tokens_are_url_safe_and_unique() { + let first = generate_token(); + let second = generate_token(); + assert_ne!(first, second); + assert!(!first.contains('=')); + assert_eq!(sha256_hex(&first).len(), 64); + } +} diff --git a/src/services/filename.rs b/src/services/filename.rs new file mode 100644 index 0000000..8ed7d18 --- /dev/null +++ b/src/services/filename.rs @@ -0,0 +1,34 @@ +pub fn normalize_upload_name(name: Option<&str>) -> String { + let normalized = name + .unwrap_or("upload") + .trim() + .chars() + .filter(|ch| *ch != '\0') + .map(|ch| if matches!(ch, '\r' | '\n') { '_' } else { ch }) + .take(255) + .collect::(); + + if normalized.is_empty() { + "upload".to_string() + } else { + normalized + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn upload_names_are_utf8_safe_and_database_bounded() { + let name = format!("{}\r\n.png", "图".repeat(300)); + let normalized = normalize_upload_name(Some(&name)); + assert_eq!(normalized.chars().count(), 255); + assert!(!normalized.contains(['\r', '\n', '\0'])); + } + + #[test] + fn empty_upload_name_uses_fallback() { + assert_eq!(normalize_upload_name(Some(" \0 ")), "upload"); + } +} diff --git a/src/services/mod.rs b/src/services/mod.rs index 9a21c3d..11cd28e 100644 --- a/src/services/mod.rs +++ b/src/services/mod.rs @@ -1,8 +1,11 @@ pub mod billing; pub mod bootstrap; pub mod compress; +pub mod credentials; +pub mod filename; pub mod idempotency; pub mod mail; pub mod quota; +pub mod rate_limit; pub mod settings; pub mod storage; diff --git a/src/services/rate_limit.rs b/src/services/rate_limit.rs new file mode 100644 index 0000000..d56f9b3 --- /dev/null +++ b/src/services/rate_limit.rs @@ -0,0 +1,62 @@ +use crate::error::{AppError, ErrorCode}; +use crate::state::AppState; + +use sha2::{Digest, Sha256}; + +pub async fn enforce( + state: &AppState, + namespace: &str, + discriminator: &str, + limit: u32, + window_seconds: u32, + message: &str, +) -> Result<(), AppError> { + if limit == 0 || window_seconds == 0 { + return Err(AppError::new(ErrorCode::Internal, "限速配置无效")); + } + + let key = rate_limit_key(namespace, discriminator); + let mut conn = state.redis.clone(); + let script = redis::Script::new( + r#" + local count = redis.call('INCR', KEYS[1]) + if count == 1 then + redis.call('EXPIRE', KEYS[1], ARGV[1]) + end + return count + "#, + ); + + let count: i64 = script + .key(key) + .arg(window_seconds) + .invoke_async(&mut conn) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "请求限速检查失败").with_source(err))?; + + if count > i64::from(limit) { + return Err(AppError::new(ErrorCode::RateLimited, message)); + } + + Ok(()) +} + +fn rate_limit_key(namespace: &str, discriminator: &str) -> String { + let mut hasher = Sha256::new(); + hasher.update(discriminator.as_bytes()); + let digest = hex::encode(hasher.finalize()); + format!("rate:{namespace}:{}", &digest[..32]) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rate_limit_keys_do_not_expose_identifiers() { + let key = rate_limit_key("login", "person@example.com"); + assert!(key.starts_with("rate:login:")); + assert!(!key.contains("person")); + assert_eq!(key.len(), "rate:login:".len() + 32); + } +} diff --git a/src/services/storage.rs b/src/services/storage.rs index f6338ee..e4fd7b9 100644 --- a/src/services/storage.rs +++ b/src/services/storage.rs @@ -148,13 +148,16 @@ fn retention_prefix(hours: i64) -> String { } } -pub async fn store_bytes( +pub async fn store_bytes( state: &AppState, key: &str, - bytes: Vec, + bytes: B, content_type: &str, -) -> Result { - let bytes = Bytes::from(bytes); +) -> Result +where + B: Into, +{ + let bytes = bytes.into(); if let Some(endpoint) = active_endpoint(state).await? { match store_bytes_s3(state, &endpoint, key, bytes.clone(), content_type).await { Ok(stored) => return Ok(stored),