feat(auth): 实现登录限流和性能优化

性能优化:
- 添加数据库索引优化查询性能 (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/
This commit is contained in:
2026-04-24 16:13:49 +08:00
parent b3dc966576
commit e6048ea010
11 changed files with 852 additions and 9 deletions

41
.gitignore vendored
View File

@@ -1,9 +1,34 @@
# Ignore all files
*
# Build output
target/
# But track these
!AGENTS.md
!*.service
!*.sh
!deploy.sh
!test_deployment.sh
# 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

1
Cargo.lock generated
View File

@@ -1976,6 +1976,7 @@ dependencies = [
"log",
"openssl",
"pkcs8",
"rand 0.8.5",
"reqwest",
"rsa",
"serde",

View File

@@ -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 '过期时间';

View File

@@ -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();

View File

@@ -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';

241
src/alipay.rs Normal file
View File

@@ -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<Self> {
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::<Vec<_>>()
.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<String, String> {
// 解析 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::<Sha256>(), 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<bool, String> {
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::<Sha256>(),
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<BTreeMap<&str, &str>>,
) -> 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::<Vec<_>>()
.join("&");
let sign = rsa2_sign(&sign_source, &config.private_key)?;
// 构建最终 URL
let query = build_query_string(&params)
+ "&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<BTreeMap<&str, &str>>,
) -> Result<serde_json::Value, String> {
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::<Vec<_>>()
.join("&");
let sign = rsa2_sign(&sign_source, &config.private_key)?;
let query = build_query_string(&params)
+ "&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");
}
}

180
src/error.rs Normal file
View File

@@ -0,0 +1,180 @@
use actix_web::{ResponseError, http::StatusCode, HttpResponse};
use serde::Serialize;
use std::fmt;
#[derive(Debug, Serialize)]
pub struct ErrorResponse<T = ()> {
pub success: bool,
pub data: Option<T>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
impl<T> ErrorResponse<T> {
pub fn success(data: T) -> Self {
Self {
success: true,
data: Some(data),
error: None,
}
}
pub fn error(msg: impl Into<String>) -> Self {
Self {
success: false,
data: None,
error: Some(msg.into()),
}
}
}
impl<T: Serialize> ErrorResponse<T> {
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#"<!DOCTYPE html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>错误 {status_code}</title>
<style>
* {{ margin: 0; padding: 0; box-sizing: border-box; }}
body {{
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
background: linear-gradient(135deg, #e74c3c 0%, #c0392b 100%);
min-height: 100vh;
display: flex;
align-items: center;
justify-content: center;
padding: 20px;
}}
.container {{
background: white;
border-radius: 16px;
padding: 48px;
max-width: 480px;
width: 100%;
box-shadow: 0 20px 60px rgba(0,0,0,0.3);
text-align: center;
}}
.error-code {{
font-size: 72px;
font-weight: bold;
color: #e74c3c;
margin-bottom: 16px;
}}
h1 {{
color: #333;
margin-bottom: 16px;
font-size: 24px;
}}
.message {{
background: #f8f9fa;
border-radius: 8px;
padding: 16px;
margin: 24px 0;
color: #555;
font-size: 14px;
line-height: 1.6;
}}
.hint {{
color: #888;
font-size: 13px;
margin-top: 24px;
}}
.back-link {{
display: inline-block;
margin-top: 24px;
padding: 12px 24px;
background: #3498db;
color: white;
text-decoration: none;
border-radius: 8px;
font-size: 14px;
transition: background 0.3s;
}}
.back-link:hover {{
background: #2980b9;
}}
</style>
</head>
<body>
<div class="container">
<div class="error-code">{status_code}</div>
<h1>{error_title}</h1>
<div class="message">{error_message}</div>
<p style="margin-top: 24px; color: #3498db; font-size: 14px;">请返回小程序重新生成新的下载链接</p>
<div class="hint">如果问题持续存在,请联系技术支持</div>
</div>
</body>
</html>"#,
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)
}
}

View File

@@ -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<WeChatLoginRequest>,
http_client: web::Data<Client>,
app_state: web::Data<AppState>,
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

112
src/handlers/meta.rs Normal file
View File

@@ -0,0 +1,112 @@
use actix_web::{get, web, HttpResponse, Responder};
use sqlx::postgres::PgPool;
const HTML_TEMPLATE: &str = r#"<!DOCTYPE html>
<html lang="zh-CN">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>大气稳定度判定系统</title>
<style>
* { margin: 0; padding: 0; box-sizing: border-box; }
body {
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
min-height: 100vh;
display: flex;
align-items: center;
justify-content: center;
padding: 20px;
}
.container {
background: white;
border-radius: 16px;
padding: 48px;
max-width: 480px;
width: 100%;
box-shadow: 0 20px 60px rgba(0,0,0,0.3);
text-align: center;
}
h1 {
color: #333;
margin-bottom: 8px;
font-size: 28px;
}
.subtitle {
color: #666;
margin-bottom: 32px;
font-size: 14px;
}
.status-box {
padding: 24px;
border-radius: 12px;
margin-bottom: 24px;
}
.status-ok {
background: #d4edda;
border: 1px solid #28a745;
}
.status-error {
background: #f8d7da;
border: 1px solid #dc3545;
}
.status-icon {
font-size: 48px;
margin-bottom: 16px;
}
.status-text {
font-size: 24px;
font-weight: bold;
margin-bottom: 8px;
}
.status-ok .status-text { color: #155724; }
.status-error .status-text { color: #721c24; }
.status-detail {
color: #495057;
font-size: 14px;
}
.version {
color: #999;
font-size: 12px;
}
</style>
</head>
<body>
<div class="container">
<h1>大气稳定度判定系统</h1>
<p class="subtitle">后端服务状态</p>
<div class="status-box {status_class}">
<div class="status-icon">{status_icon}</div>
<div class="status-text">{status_text}</div>
<div class="status-detail">
<div>数据库: {database}</div>
</div>
</div>
<div class="version">v{version}</div>
</div>
</body>
</html>"#;
#[get("/")]
pub async fn root(pool: web::Data<PgPool>) -> 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)
}

View File

@@ -12,6 +12,7 @@ mod db;
mod error;
mod handlers;
mod models;
mod rate_limiter;
use auth::jwt_middleware;
use config::AppConfig;

157
src/rate_limiter.rs Normal file
View File

@@ -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<RwLock<HashMap<String, Vec<Instant>>>>,
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<RateLimiter> {
Arc::new(RateLimiter::new(5, 60))
}
pub static LOGIN_RATE_LIMITER: std::sync::LazyLock<Arc<RateLimiter>> =
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");
}
}