From e6048ea01021e4db885832cf7390aad4a0b32174 Mon Sep 17 00:00:00 2001 From: Milky0217 Date: Fri, 24 Apr 2026 16:13:49 +0800 Subject: [PATCH] =?UTF-8?q?feat(auth):=20=E5=AE=9E=E7=8E=B0=E7=99=BB?= =?UTF-8?q?=E5=BD=95=E9=99=90=E6=B5=81=E5=92=8C=E6=80=A7=E8=83=BD=E4=BC=98?= =?UTF-8?q?=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 性能优化: - 添加数据库索引优化查询性能 (006) - weather_data: user_id, date, is_favorite 索引 - users: openid 索引 - payment_orders: status 索引 - 新增 refresh_tokens 表支持双 Token 机制 (004) - 新增 web_login_codes 表支持网页端扫码登录 (005) 安全增强: - 实现基于 IP 的登录限流 (rate_limiter.rs) - 滑动窗口算法: 5次/分钟/IP - 自动清理过期记录 - 429 TooManyRequests 响应 新模块: - src/rate_limiter.rs: 限流模块 - src/alipay.rs: 支付宝签名模块 (RSA2) - src/error.rs: 统一错误类型 (含 TooManyRequests) - src/handlers/meta.rs: 元数据处理器 代码清理: - 修复 .gitignore 规则，正确跟踪 src/ 和 migrations/ --- .gitignore | 41 +++- Cargo.lock | 1 + migrations/004_add_refresh_tokens.sql | 27 +++ migrations/005_add_web_login_codes.sql | 27 +++ migrations/006_add_performance_indexes.sql | 50 +++++ src/alipay.rs | 241 +++++++++++++++++++++ src/error.rs | 180 +++++++++++++++ src/handlers/auth.rs | 24 +- src/handlers/meta.rs | 112 ++++++++++ src/main.rs | 1 + src/rate_limiter.rs | 157 ++++++++++++++ 11 files changed, 852 insertions(+), 9 deletions(-) create mode 100644 migrations/004_add_refresh_tokens.sql create mode 100644 migrations/005_add_web_login_codes.sql create mode 100644 migrations/006_add_performance_indexes.sql create mode 100644 src/alipay.rs create mode 100644 src/error.rs create mode 100644 src/handlers/meta.rs create mode 100644 src/rate_limiter.rs diff --git a/.gitignore b/.gitignore index 9482034..2bccce0 100644 --- a/.gitignore +++ b/.gitignore @@ -1,9 +1,34 @@ -# Ignore all files -* +# Build output +target/ -# But track these -!AGENTS.md -!*.service -!*.sh -!deploy.sh -!test_deployment.sh \ No newline at end of file +# IDE +.vscode/ + +# Logs +logs/ + +# Environment (contains secrets) +.env +.env.example + +# Clawhub +.clawhub/ + +# Sisyphus +.sisyphus/ + +# Docs +docs/ + +# Scripts +scripts/ + +# Parent directory reference (if accidentally included) +ASD-backend/ + +# Generated files +static/css/report.css +templates/ + +# Rust lock file (optional - uncomment if you want to track it) +# Cargo.lock diff --git a/Cargo.lock b/Cargo.lock index d01efc1..8186b98 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1976,6 +1976,7 @@ dependencies = [ "log", "openssl", "pkcs8", + "rand 0.8.5", "reqwest", "rsa", "serde", diff --git a/migrations/004_add_refresh_tokens.sql b/migrations/004_add_refresh_tokens.sql new file mode 100644 index 0000000..6e669e5 --- /dev/null +++ b/migrations/004_add_refresh_tokens.sql @@ -0,0 +1,27 @@ +-- ============================================ +-- 迁移: 004_add_refresh_tokens.sql +-- 目的: 添加 refresh_tokens 表支持双 Token 机制 +-- 日期: 2026-04-18 +-- 依赖: 无（users 表已在 001 创建） +-- 说明: 用于存储 refresh_token，支持 access_token 过期后自动续期 +-- ============================================ + +-- 迁移: 添加 refresh_tokens 表 +-- 用于支持 access_token 刷新机制 + +CREATE TABLE IF NOT EXISTS refresh_tokens ( + id SERIAL PRIMARY KEY, + user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE, + token VARCHAR(512) NOT NULL, + expires_at TIMESTAMPTZ NOT NULL, + created_at TIMESTAMPTZ DEFAULT NOW() +); + +-- 索引：加速 token 查询和用户清理 +CREATE INDEX IF NOT EXISTS idx_refresh_tokens_token ON refresh_tokens(token); +CREATE INDEX IF NOT EXISTS idx_refresh_tokens_user_id ON refresh_tokens(user_id); +CREATE INDEX IF NOT EXISTS idx_refresh_tokens_expires_at ON refresh_tokens(expires_at); + +COMMENT ON TABLE refresh_tokens IS 'Refresh token 存储表，支持 access_token 续期'; +COMMENT ON COLUMN refresh_tokens.token IS 'Refresh token 字符串'; +COMMENT ON COLUMN refresh_tokens.expires_at IS '过期时间'; diff --git a/migrations/005_add_web_login_codes.sql b/migrations/005_add_web_login_codes.sql new file mode 100644 index 0000000..90ac6ed --- /dev/null +++ b/migrations/005_add_web_login_codes.sql @@ -0,0 +1,27 @@ +-- ============================================ +-- 迁移: 005_add_web_login_codes.sql +-- 目的: 添加网页端微信扫码登录临时码表 +-- 日期: 2026-04-20 +-- 依赖: users 表（001 创建） +-- 说明: 用户扫码后生成临时登录码，前端轮询验证登录状态 +-- ============================================ + +-- 05_add_web_login_codes.sql +-- 网页端微信扫码登录：临时登录码表 + +CREATE TABLE IF NOT EXISTS web_login_codes ( + id SERIAL PRIMARY KEY, + code VARCHAR(32) UNIQUE NOT NULL, -- 登录码，如 ASD-ABC123 + user_id INTEGER, -- 关联用户（确认登录后写入） + openid VARCHAR(128), -- 用户 openid + token TEXT, -- 生成的 JWT（确认后写入） + expires_at TIMESTAMPTZ NOT NULL, -- 过期时间（10分钟内有效） + created_at TIMESTAMPTZ DEFAULT NOW() +); + +-- 登录码索引 +CREATE INDEX IF NOT EXISTS idx_web_login_codes_code ON web_login_codes(code); +CREATE INDEX IF NOT EXISTS idx_web_login_codes_expires ON web_login_codes(expires_at); + +-- 定期清理过期登录码（可由 cron 或服务启动时触发） +-- DELETE FROM web_login_codes WHERE expires_at < NOW(); diff --git a/migrations/006_add_performance_indexes.sql b/migrations/006_add_performance_indexes.sql new file mode 100644 index 0000000..a5139fb --- /dev/null +++ b/migrations/006_add_performance_indexes.sql @@ -0,0 +1,50 @@ +-- 迁移: 006_add_performance_indexes.sql +-- 目的: 为常用查询字段添加索引，提升查询性能 +-- 日期: 2026-04-24 + +-- ============================================ +-- weather_data 表索引 +-- ============================================ + +-- 索引：加速按用户查询天气数据列表 +-- 用途：GET /weather 接口按 user_id + date 排序查询 +CREATE INDEX IF NOT EXISTS idx_weather_data_user_id ON weather_data(user_id); + +-- 索引：加速按日期排序查询 +-- 用途：列表查询默认按 date DESC, hour DESC, min DESC 排序 +CREATE INDEX IF NOT EXISTS idx_weather_data_date ON weather_data(date DESC); + +-- 索引：加速用户收藏列表查询 +-- 用途：GET /api/favorites 查询 user_id + is_favorite = true +CREATE INDEX IF NOT EXISTS idx_weather_data_user_favorite ON weather_data(user_id, is_favorite); + +-- 索引：复合索引加速用户数据列表（无需额外排序） +-- 用途：用户历史记录查询 +CREATE INDEX IF NOT EXISTS idx_weather_data_user_date ON weather_data(user_id, date DESC); + +-- ============================================ +-- users 表索引 +-- ============================================ + +-- 索引：加速 openid 查询（登录时高频使用） +-- 注意：openid 可能已有 UNIQUE 约束自动创建索引，此处确保存在 +CREATE INDEX IF NOT EXISTS idx_users_openid ON users(openid); + +-- ============================================ +-- payment_orders 表索引（补充） +-- ============================================ + +-- 索引：加速按状态查询待处理订单 +-- 用途：管理员查询 pending 状态订单 +CREATE INDEX IF NOT EXISTS idx_payment_orders_status ON payment_orders(status); + +-- ============================================ +-- web_login_codes 表索引（来自 005） +-- ============================================ + +-- 确保 web_login_codes 表索引存在（05已创建，此处确保兼容性） +-- 登录码查询和过期清理 +CREATE INDEX IF NOT EXISTS idx_web_login_codes_code ON web_login_codes(code); +CREATE INDEX IF NOT EXISTS idx_web_login_codes_expires_at ON web_login_codes(expires_at); + +COMMENT ON TABLE weather_data IS '添加性能索引：user_id、date、is_favorite'; diff --git a/src/alipay.rs b/src/alipay.rs new file mode 100644 index 0000000..aff192f --- /dev/null +++ b/src/alipay.rs @@ -0,0 +1,241 @@ +// src/alipay.rs — 支付宝 RSA2 签名与请求封装 +use rsa::pkcs8::{DecodePrivateKey, EncodePrivateKey, LineEnding}; +use rsa::{Pkcs1v15Sign, RsaPrivateKey}; +use sha2::Sha256; +use std::collections::BTreeMap; +use std::env; + +/// 支付宝配置（从环境变量读取） +pub struct AlipayConfig { + pub app_id: String, + pub private_key: String, + pub alipay_public_key: String, + pub gateway: String, +} + +impl AlipayConfig { + pub fn from_env() -> Option { + let app_id = env::var("ALIPAY_APP_ID").ok()?; + let private_key = env::var("ALIPAY_PRIVATE_KEY").ok()?; + let alipay_public_key = env::var("ALIPAY_ALIPAY_PUBLIC_KEY").ok()?; + let gateway = env::var("ALIPAY_GATEWAY") + .unwrap_or_else(|_| "https://openapi.alipay.com/gateway.do".to_string()); + Some(Self { + app_id, + private_key, + alipay_public_key, + gateway, + }) + } + + pub fn isConfigured() -> bool { + Self::from_env().is_some() + } +} + +/// 对 map 按 key 排序后构建 query string（用于签名） +fn build_query_string(params: &BTreeMap<&str, &str>) -> String { + params + .iter() + .filter(|(_, v)| !v.is_empty()) + .map(|(k, v)| format!("{}={}", k, urlencoding(v))) + .collect::>() + .join("&") +} + +/// URL 编码（简单实现） +fn urlencoding(s: &str) -> String { + let mut result = String::new(); + for c in s.chars() { + match c { + 'A'..='Z' | 'a'..='z' | '0'..='9' | '-' | '_' | '.' | '~' => result.push(c), + _ => { + for b in c.to_string().as_bytes() { + result.push_str(&format!("%{:02X}", b)); + } + } + } + } + result +} + +/// 使用 RSA2 (SHA256) 对内容签名 +pub fn rsa2_sign(content: &str, private_key_pem: &str) -> Result { + // 解析 PKCS8 格式的私钥 + let private_key = RsaPrivateKey::from_pkcs8_pem(private_key_pem) + .map_err(|e| format!("私钥解析失败: {}", e))?; + + let signing_input = content.as_bytes(); + let signature = private_key + .sign(Pkcs1v15Sign::new::(), signing_input) + .map_err(|e| format!("签名失败: {}", e))?; + + Ok(base64::Engine::encode( + &base64::engine::general_purpose::STANDARD, + &signature, + )) +} + +/// 验证 RSA2 签名 +pub fn rsa2_verify(content: &str, sign: &str, public_key_pem: &str) -> Result { + use rsa::pkcs8::DecodePublicKey; + + let public_key = rsa::RsaPublicKey::from_public_key_pem(public_key_pem) + .map_err(|e| format!("支付宝公钥解析失败: {}", e))?; + + let sig_bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, sign) + .map_err(|e| format!("签名 Base64 解码失败: {}", e))?; + + Ok(public_key + .verify( + Pkcs1v15Sign::new::(), + content.as_bytes(), + &sig_bytes, + ) + .is_ok()) +} + +/// 构建支付宝请求 URL（含签名） +/// 返回 (url, sign)，sign 已 URL 编码 +pub fn build_signed_request( + config: &AlipayConfig, + biz_content: &str, + other_params: Option>, +) -> Result<(String, String), String> { + let mut params: BTreeMap<&str, &str> = BTreeMap::new(); + + params.insert("app_id", &config.app_id); + params.insert("method", "alipay.trade.page.pay"); + params.insert("format", "JSON"); + params.insert("charset", "utf-8"); + params.insert("sign_type", "RSA2"); + params.insert( + "timestamp", + &chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string(), + ); + params.insert("version", "1.0"); + params.insert("biz_content", biz_content); + + // 加入其他参数（如 return_url, notify_url） + if let Some(ref extra) = other_params { + for (k, v) in extra { + params.insert(k, v); + } + } + + // 按 RFC 3986 编码后拼接待签名内容 + let sign_source: String = params + .iter() + .map(|(k, v)| format!("{}={}", k, urlencoding(v))) + .collect::>() + .join("&"); + + let sign = rsa2_sign(&sign_source, &config.private_key)?; + + // 构建最终 URL + let query = build_query_string(¶ms) + + "&sign=" + + &urlencoding(&sign); + + let url = config.gateway.clone() + "?" + &query; + + Ok((url, sign)) +} + +/// 调用支付宝接口并解析响应 +pub async fn call_alipay( + config: &AlipayConfig, + biz_content: &str, + other_params: Option>, +) -> Result { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + .map_err(|e| format!("HTTP 客户端创建失败: {}", e))?; + + let mut params: BTreeMap<&str, &str> = BTreeMap::new(); + params.insert("app_id", &config.app_id); + params.insert("method", "alipay.trade.page.pay"); + params.insert("format", "JSON"); + params.insert("charset", "utf-8"); + params.insert("sign_type", "RSA2"); + params.insert( + "timestamp", + &chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string(), + ); + params.insert("version", "1.0"); + params.insert("biz_content", biz_content); + + if let Some(ref extra) = other_params { + for (k, v) in extra { + params.insert(k, v); + } + } + + let sign_source: String = params + .iter() + .map(|(k, v)| format!("{}={}", k, urlencoding(v))) + .collect::>() + .join("&"); + + let sign = rsa2_sign(&sign_source, &config.private_key)?; + + let query = build_query_string(¶ms) + + "&sign=" + + &urlencoding(&sign); + + let url = config.gateway.clone(); + + let resp = client + .post(&url) + .header("Content-Type", "application/x-www-form-urlencoded") + .body(query) + .send() + .await + .map_err(|e| format!("请求支付宝失败: {}", e))?; + + let body = resp.text().await.map_err(|e| format!("读取响应失败: {}", e))?; + + // 支付宝返回格式: alipay_trade_page_pay_response={...}&sign=xxx + let parts: Vec<&str> = body.splitn(2, "&sign=").collect(); + if parts.len() != 2 { + return Err(format!("支付宝响应格式异常: {}", body)); + } + + let json_str = parts[0] + .strip_prefix("alipay_trade_page_pay_response=") + .unwrap_or(parts[0]); + + let sign_from_alipay = parts[1]; + + // 验签（确保响应来自支付宝） + if !rsa2_verify(json_str, sign_from_alipay, &config.alipay_public_key)? { + return Err("支付宝响应验签失败".to_string()); + } + + let json: serde_json::Value = + serde_json::from_str(json_str).map_err(|e| format!("JSON 解析失败: {}", e))?; + + if json.get("code").and_then(|v| v.as_str()) != Some("10000") { + let sub_msg = json.get("sub_msg").and_then(|v| v.as_str()).unwrap_or(""); + return Err(format!( + "支付宝接口错误: {} - {}", + json.get("msg").and_then(|v| v.as_str()).unwrap_or("未知"), + sub_msg + )); + } + + Ok(json) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_url_encoding() { + assert_eq!(urlencoding("hello"), "hello"); + assert_eq!(urlencoding("hello world"), "hello%20world"); + assert_eq!(urlencoding("中文"), "%E4%B8%AD%E6%96%87"); + } +} diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..954188c --- /dev/null +++ b/src/error.rs @@ -0,0 +1,180 @@ +use actix_web::{ResponseError, http::StatusCode, HttpResponse}; +use serde::Serialize; +use std::fmt; + +#[derive(Debug, Serialize)] +pub struct ErrorResponse { + pub success: bool, + pub data: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +impl ErrorResponse { + pub fn success(data: T) -> Self { + Self { + success: true, + data: Some(data), + error: None, + } + } + + pub fn error(msg: impl Into) -> Self { + Self { + success: false, + data: None, + error: Some(msg.into()), + } + } +} + +impl ErrorResponse { + pub fn to_json_response(self) -> HttpResponse { + HttpResponse::Ok().json(self) + } +} + +impl ErrorResponse<()> { + pub fn to_error_response(&self, status: StatusCode) -> HttpResponse { + let body = serde_json::json!({ + "success": false, + "error": self.error.clone().unwrap_or_else(|| "Unknown error".to_string()) + }); + HttpResponse::build(status).json(body) + } +} + +#[derive(Debug)] +pub enum AppError { + Unauthorized(String), + Forbidden(String), + NotFound(String), + BadRequest(String), + Internal(String), + Database(String), + TooManyRequests(String), +} + +impl fmt::Display for AppError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + AppError::Unauthorized(msg) => write!(f, "{}", msg), + AppError::Forbidden(msg) => write!(f, "{}", msg), + AppError::NotFound(msg) => write!(f, "{}", msg), + AppError::BadRequest(msg) => write!(f, "{}", msg), + AppError::Internal(msg) => write!(f, "{}", msg), + AppError::Database(msg) => write!(f, "{}", msg), + AppError::TooManyRequests(msg) => write!(f, "{}", msg), + } + } +} + +impl ResponseError for AppError { + fn status_code(&self) -> StatusCode { + match self { + AppError::Unauthorized(_) => StatusCode::UNAUTHORIZED, + AppError::Forbidden(_) => StatusCode::FORBIDDEN, + AppError::NotFound(_) => StatusCode::NOT_FOUND, + AppError::BadRequest(_) => StatusCode::BAD_REQUEST, + AppError::Internal(_) | AppError::Database(_) => StatusCode::INTERNAL_SERVER_ERROR, + AppError::TooManyRequests(_) => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> actix_web::HttpResponse { + let error_message = self.to_string(); + let status_code = self.status_code().as_u16(); + + let html = format!(r#" + + + + + 错误 {status_code} + + + +
+
{status_code}
+

{error_title}

+
{error_message}
+

请返回小程序重新生成新的下载链接

+
如果问题持续存在，请联系技术支持
+
+ +"#, + error_title = match self { + AppError::Unauthorized(_) => "认证失败", + AppError::Forbidden(_) => "权限不足", + AppError::NotFound(_) => "内容未找到", + AppError::BadRequest(_) => "请求无效", + AppError::Internal(_) | AppError::Database(_) => "服务器错误", + AppError::TooManyRequests(_) => "请求过于频繁", + } + ); + + HttpResponse::build(self.status_code()) + .content_type("text/html; charset=utf-8") + .body(html) + } +} + diff --git a/src/handlers/auth.rs b/src/handlers/auth.rs index dfa2d2f..ba86b6d 100644 --- a/src/handlers/auth.rs +++ b/src/handlers/auth.rs @@ -1,4 +1,4 @@ -use actix_web::{web, HttpResponse, Responder, post, get}; +use actix_web::{web, HttpResponse, HttpRequest, Responder, post, get}; use chrono::{Utc, Duration}; use serde::{Deserialize, Serialize}; use reqwest::Client; @@ -13,6 +13,8 @@ use crate::models::{ AppState, LoginResponse, RefreshTokenRequest, TokenRefreshResponse, WeChatApiResponse, WeChatLoginRequest, }; +use crate::rate_limiter::LOGIN_RATE_LIMITER; +use std::sync::Arc; #[post("/api/login")] pub async fn login( @@ -20,7 +22,27 @@ pub async fn login( req: web::Json, http_client: web::Data, app_state: web::Data, + http_req: HttpRequest, ) -> impl Responder { + let client_ip = http_req + .headers() + .get("X-Forwarded-For") + .and_then(|v| v.to_str().ok()) + .map(|s| s.split(',').next().unwrap_or(s).trim().to_string()) + .or_else(|| { + http_req + .headers() + .get("X-Real-IP") + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()) + }) + .unwrap_or_else(|| "unknown".to_string()); + + if let Err(e) = LOGIN_RATE_LIMITER.check_rate_limit(&client_ip).await { + warn!("登录请求被限流: client_ip={}", client_ip); + return HttpResponse::TooManyRequests().json(ErrorResponse::<()>::error(e.to_string())); + } + let url = format!( "https://api.weixin.qq.com/sns/jscode2session?appid={}&secret={}&js_code={}&grant_type=authorization_code", app_state.wechat_appid, app_state.wechat_secret, req.code diff --git a/src/handlers/meta.rs b/src/handlers/meta.rs new file mode 100644 index 0000000..2e9f340 --- /dev/null +++ b/src/handlers/meta.rs @@ -0,0 +1,112 @@ +use actix_web::{get, web, HttpResponse, Responder}; +use sqlx::postgres::PgPool; + +const HTML_TEMPLATE: &str = r#" + + + + + 大气稳定度判定系统 + + + +
+

大气稳定度判定系统

+

后端服务状态

+ +
+
{status_icon}
+
{status_text}
+
+
数据库: {database}
+
+
+ +
v{version}
+
+ +"#; + +#[get("/")] +pub async fn root(pool: web::Data) -> impl Responder { + let (status_class, status_icon, status_text, database) = match sqlx::query("SELECT 1") + .fetch_one(pool.get_ref()) + .await + { + Ok(_) => ("status-ok", "✅", "服务正常", "已连接"), + Err(_) => ("status-error", "❌", "服务异常", "未连接"), + }; + + let html = HTML_TEMPLATE + .replace("{status_class}", status_class) + .replace("{status_icon}", status_icon) + .replace("{status_text}", status_text) + .replace("{database}", database) + .replace("{version}", env!("CARGO_PKG_VERSION")); + + HttpResponse::Ok() + .content_type("text/html; charset=utf-8") + .body(html) +} diff --git a/src/main.rs b/src/main.rs index 260e881..06e0943 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,6 +12,7 @@ mod db; mod error; mod handlers; mod models; +mod rate_limiter; use auth::jwt_middleware; use config::AppConfig; diff --git a/src/rate_limiter.rs b/src/rate_limiter.rs new file mode 100644 index 0000000..9aa488e --- /dev/null +++ b/src/rate_limiter.rs @@ -0,0 +1,157 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::sync::RwLock; +use crate::error::AppError; + +pub struct RateLimiter { + requests: Arc>>>, + max_requests: usize, + window_secs: u64, +} + +impl RateLimiter { + pub fn new(max_requests: usize, window_secs: u64) -> Self { + Self { + requests: Arc::new(RwLock::new(HashMap::new())), + max_requests, + window_secs, + } + } + + pub async fn check_rate_limit(&self, client_ip: &str) -> Result<(), AppError> { + let now = Instant::now(); + let window = Duration::from_secs(self.window_secs); + + let mut requests = self.requests.write().await; + + let client_requests = requests.entry(client_ip.to_string()).or_insert_with(Vec::new); + + client_requests.retain(|&t| now.duration_since(t) < window); + + if client_requests.len() >= self.max_requests { + return Err(AppError::TooManyRequests( + format!("请求过于频繁，请{}秒后再试", self.window_secs) + )); + } + + client_requests.push(now); + Ok(()) + } +} + +pub fn extract_client_ip_from_header(headers: &actix_web::http::header::HeaderMap) -> String { + headers + .get("X-Forwarded-For") + .and_then(|v| v.to_str().ok()) + .map(|s| s.split(',').next().unwrap_or(s).trim().to_string()) + .or_else(|| { + headers + .get("X-Real-IP") + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()) + }) + .unwrap_or_else(|| "unknown".to_string()) +} + +pub fn create_login_rate_limiter() -> Arc { + Arc::new(RateLimiter::new(5, 60)) +} + +pub static LOGIN_RATE_LIMITER: std::sync::LazyLock> = + std::sync::LazyLock::new(create_login_rate_limiter); + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_rate_limiter_allows_requests_under_limit() { + let limiter = Arc::new(RateLimiter::new(5, 60)); + + for _ in 0..5 { + let result = limiter.check_rate_limit("192.168.1.1").await; + assert!(result.is_ok()); + } + } + + #[tokio::test] + async fn test_rate_limiter_blocks_excessive_requests() { + let limiter = Arc::new(RateLimiter::new(3, 60)); + + for _ in 0..3 { + assert!(limiter.check_rate_limit("192.168.1.2").await.is_ok()); + } + + let result = limiter.check_rate_limit("192.168.1.2").await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_rate_limiter_independent_per_client() { + let limiter = Arc::new(RateLimiter::new(2, 60)); + + assert!(limiter.check_rate_limit("192.168.1.100").await.is_ok()); + assert!(limiter.check_rate_limit("192.168.1.100").await.is_ok()); + assert!(limiter.check_rate_limit("192.168.1.100").await.is_err()); + + assert!(limiter.check_rate_limit("192.168.1.101").await.is_ok()); + assert!(limiter.check_rate_limit("192.168.1.101").await.is_ok()); + assert!(limiter.check_rate_limit("192.168.1.101").await.is_err()); + } + + #[test] + fn test_extract_client_ip_from_x_forwarded_for() { + use actix_web::http::header::{HeaderMap, HeaderName, HeaderValue}; + + let mut headers = HeaderMap::new(); + headers.insert( + HeaderName::from_static("x-forwarded-for"), + HeaderValue::from_static("203.0.113.195, 70.41.3.18"), + ); + + let ip = extract_client_ip_from_header(&headers); + assert_eq!(ip, "203.0.113.195"); + } + + #[test] + fn test_extract_client_ip_from_x_real_ip() { + use actix_web::http::header::{HeaderMap, HeaderName, HeaderValue}; + + let mut headers = HeaderMap::new(); + headers.insert( + HeaderName::from_static("x-real-ip"), + HeaderValue::from_static("203.0.113.195"), + ); + + let ip = extract_client_ip_from_header(&headers); + assert_eq!(ip, "203.0.113.195"); + } + + #[test] + fn test_extract_client_ip_prefers_x_forwarded_for() { + use actix_web::http::header::{HeaderMap, HeaderName, HeaderValue}; + + let mut headers = HeaderMap::new(); + headers.insert( + HeaderName::from_static("x-forwarded-for"), + HeaderValue::from_static("203.0.113.195"), + ); + headers.insert( + HeaderName::from_static("x-real-ip"), + HeaderValue::from_static("198.51.100.178"), + ); + + let ip = extract_client_ip_from_header(&headers); + assert_eq!(ip, "203.0.113.195"); + } + + #[test] + fn test_extract_client_ip_fallback_to_unknown() { + use actix_web::http::header::HeaderMap; + + let headers = HeaderMap::new(); + let ip = extract_client_ip_from_header(&headers); + assert_eq!(ip, "unknown"); + } +}