Skip to main content

yggdrasil/api/
rate_limit.rs

1//! 基于 governor 的多级限流。
2//!
3//! 提供 strict、upload、image、comment、comment_upload(双层)、code_exec(双层)
4//! 等限流器,支持从 `X-Forwarded-For` / `X-Real-IP` 中提取客户端 IP,
5//! 并可通过 `TRUSTED_PROXY_COUNT` 配置信任代理层数。
6//!
7//! 当未配置可信代理时,Axum handler 可回退到 TCP 连接的对端地址;
8//! Dioxus server function 无法获取对端地址,会退回到 `"unknown"` key,
9//! 此时所有请求共享一个限流桶。生产环境应在反向代理后部署并正确配置
10//! `TRUSTED_PROXY_COUNT`。
11//!
12//! 仅在 `feature = "server"` 时生效。
13
14#[cfg(feature = "server")]
15use axum::http::StatusCode;
16#[cfg(feature = "server")]
17use axum::response::{IntoResponse, Response};
18#[cfg(feature = "server")]
19use governor::{clock::Clock, DefaultKeyedRateLimiter, Quota, RateLimiter};
20#[cfg(feature = "server")]
21use std::num::NonZeroU32;
22#[cfg(feature = "server")]
23use std::sync::LazyLock;
24#[cfg(feature = "server")]
25use std::sync::Once;
26#[cfg(feature = "server")]
27use std::time::Duration;
28
29#[cfg(feature = "server")]
30fn nz(val: u32) -> NonZeroU32 {
31    // val.max(1) 保证 ≥ 1,NonZeroU32::new 必然 Some;expect 说明该不变量。
32    NonZeroU32::new(val.max(1)).expect("val.max(1) 保证非零,NonZeroU32::new 不可能失败")
33}
34
35#[cfg(feature = "server")]
36static STRICT_LIMITER: LazyLock<DefaultKeyedRateLimiter<String>> = LazyLock::new(|| {
37    let c = crate::config::rate_limit();
38    RateLimiter::keyed(Quota::per_second(nz(c.strict_per_sec)).allow_burst(nz(c.strict_burst)))
39});
40
41#[cfg(feature = "server")]
42static UPLOAD_LIMITER: LazyLock<DefaultKeyedRateLimiter<String>> = LazyLock::new(|| {
43    let c = crate::config::rate_limit();
44    RateLimiter::keyed(Quota::per_second(nz(c.upload_per_sec)).allow_burst(nz(c.upload_burst)))
45});
46
47#[cfg(feature = "server")]
48static IMAGE_LIMITER: LazyLock<DefaultKeyedRateLimiter<String>> = LazyLock::new(|| {
49    let c = crate::config::rate_limit();
50    RateLimiter::keyed(Quota::per_second(nz(c.image_per_sec)).allow_burst(nz(c.image_burst)))
51});
52
53#[cfg(feature = "server")]
54static COMMENT_LIMITER: LazyLock<DefaultKeyedRateLimiter<String>> = LazyLock::new(|| {
55    let c = crate::config::rate_limit();
56    RateLimiter::keyed(Quota::per_second(nz(c.comment_per_sec)).allow_burst(nz(c.comment_burst)))
57});
58
59#[cfg(feature = "server")]
60/// 评论图片上传单 IP 每秒限流(默认 1 req/s,突发 5)。
61/// 评论区允许匿名上传图片,与 admin 的 UPLOAD_LIMITER 隔离,避免互占额度。
62static COMMENT_UPLOAD_LIMITER: LazyLock<DefaultKeyedRateLimiter<String>> = LazyLock::new(|| {
63    let c = crate::config::rate_limit();
64    RateLimiter::keyed(
65        Quota::per_second(nz(c.comment_upload_per_sec)).allow_burst(nz(c.comment_upload_burst)),
66    )
67});
68
69#[cfg(feature = "server")]
70/// 评论图片上传单 IP 每日限额(默认 30 次/天)。
71/// 匿名可传图意味着任何人都能消耗磁盘,需要硬性日上限(模拟方式同
72/// CODE_EXEC_DAILY_LIMITER:24h 补充 1 token、突发上限即日额度)。
73static COMMENT_UPLOAD_DAILY_LIMITER: LazyLock<DefaultKeyedRateLimiter<String>> =
74    LazyLock::new(|| {
75        let c = crate::config::rate_limit();
76        RateLimiter::keyed(
77            Quota::with_period(Duration::from_secs(86_400))
78                .expect("with_period 仅在 Duration 为 0 时返回 None;86_400s 必然 Some")
79                .allow_burst(nz(c.comment_upload_daily)),
80        )
81    });
82
83#[cfg(feature = "server")]
84/// 代码执行单 IP 每秒限流(默认 1 req/s,突发 3)。
85/// 防止单个客户端高频提交容器任务,与下方日限额共同构成双层速率限制。
86static CODE_EXEC_LIMITER: LazyLock<DefaultKeyedRateLimiter<String>> = LazyLock::new(|| {
87    let c = crate::config::rate_limit();
88    RateLimiter::keyed(
89        Quota::per_second(nz(c.code_exec_per_sec)).allow_burst(nz(c.code_exec_burst)),
90    )
91});
92
93#[cfg(feature = "server")]
94/// 代码执行单 IP 每日限额(默认 50 次/天)。
95/// 容器执行成本高(CPU/内存/启动延迟),需要硬性日上限防止资源耗尽。
96///
97/// governor 0.8 无 `Quota::per_day`,用 `with_period(24h)` + `allow_burst(daily)`
98/// 模拟:每 24h 补充 1 token、突发上限即日额度,等效于「每日最多 daily 次」。
99static CODE_EXEC_DAILY_LIMITER: LazyLock<DefaultKeyedRateLimiter<String>> = LazyLock::new(|| {
100    let c = crate::config::rate_limit();
101    RateLimiter::keyed(
102        Quota::with_period(Duration::from_secs(86_400))
103            .expect("with_period 仅在 Duration 为 0 时返回 None;86_400s 必然 Some")
104            .allow_burst(nz(c.code_exec_daily)),
105    )
106});
107
108#[cfg(feature = "server")]
109/// 当无法识别真实客户端 IP("unknown")时使用的宽松限流桶。
110///
111/// TRUSTED_PROXY_COUNT=0(默认)时,Dioxus server function 拿不到 TCP 对端地址,
112/// get_client_ip 会返回 "unknown",导致所有匿名请求共享同一个严格桶
113/// (1 req/s, burst 5),正常用户的高频请求被误杀。此桶阈值更高,
114/// 可通过 settings 表的 ratelimit_unknown_* 配置项调整(重启生效)。
115static UNKNOWN_BUCKET_LIMITER: LazyLock<DefaultKeyedRateLimiter<String>> = LazyLock::new(|| {
116    let c = crate::config::rate_limit();
117    RateLimiter::keyed(Quota::per_second(nz(c.unknown_per_sec)).allow_burst(nz(c.unknown_burst)))
118});
119
120/// 限流桶 GC 间隔(秒):周期性调用 governor 的 `retain_recent`,回收已恢复为
121/// 「初始态」的 IP 键,防止 IP 轮换攻击下键空间无限膨胀。
122///
123/// 默认 300 秒。governor 的 `retain_recent` 只丢弃与「新桶」不可区分的键(即限流
124/// 窗口早已冷却、保留与否都不影响后续请求),因此即便间隔较长,内存占用也只反映
125/// 「最近活跃过且仍在限流窗口内」的 IP 集合,而非历史全集。由 settings 表的
126/// ratelimit_gc_interval_secs 配置项控制(重启生效),值越小越激进,回收越勤。
127#[cfg(feature = "server")]
128fn limiter_gc_interval() -> Duration {
129    let secs = crate::config::rate_limit().gc_interval_secs;
130    Duration::from_secs(u64::from(secs.max(1)))
131}
132
133/// 启动后台限流桶 GC 任务(全进程仅生效一次)。
134///
135/// 九个 IP 键控限流器均为独立的 `DefaultKeyedRateLimiter`,没有集中状态表,故采用
136/// 惰性启动:首个请求到达任一 `check_*` 时,经 `Once` 派生一个常驻 tokio 任务,按
137/// [`limiter_gc_interval`] 周期对全部限流器调用 `retain_recent`,回收长时间未命中的
138/// 键。这直接缓解 IP 轮换攻击下的内存膨胀——攻击者不断换 IP 制造新键,GC 周期性剔除
139/// 已冷却的旧键,使键集合大小收敛到「限流窗口内的活跃 IP」。
140#[cfg(feature = "server")]
141fn ensure_limiter_gc() {
142    static SPAWNED: std::sync::Once = std::sync::Once::new();
143    SPAWNED.call_once(|| {
144        // 仅在 tokio 运行时上下文中派生 GC 任务。测试(非 #[tokio::test])无运行时,
145        // tokio::spawn 会 panic 并毒化 Once,使后续调用全部 panic。try_current 在无
146        // 运行时时返回 Err,静默跳过——测试不依赖 GC,生产必有运行时。
147        if tokio::runtime::Handle::try_current().is_err() {
148            return;
149        }
150        let interval = limiter_gc_interval();
151        tokio::spawn(async move {
152            loop {
153                tokio::time::sleep(interval).await;
154                STRICT_LIMITER.retain_recent();
155                UPLOAD_LIMITER.retain_recent();
156                IMAGE_LIMITER.retain_recent();
157                COMMENT_LIMITER.retain_recent();
158                COMMENT_UPLOAD_LIMITER.retain_recent();
159                COMMENT_UPLOAD_DAILY_LIMITER.retain_recent();
160                CODE_EXEC_LIMITER.retain_recent();
161                CODE_EXEC_DAILY_LIMITER.retain_recent();
162                UNKNOWN_BUCKET_LIMITER.retain_recent();
163            }
164        });
165    });
166}
167
168#[cfg(feature = "server")]
169/// 检查评论请求是否超出限流阈值。
170pub fn check_comment_limit(ip: &str) -> Result<(), String> {
171    ensure_limiter_gc();
172    COMMENT_LIMITER
173        .check_key(&ip.to_string())
174        .map(|_| ())
175        .map_err(|_| "评论过于频繁,请稍后再试".to_string())
176}
177
178#[cfg(feature = "server")]
179/// 检查评论图片上传的双层速率限制(每秒突发 + 每日总额)。
180///
181/// 两层任一被限即拒绝。评论区允许匿名上传,日限额是防磁盘耗尽的关键防线。
182pub fn check_comment_upload_limit(ip: &str) -> Result<(), String> {
183    ensure_limiter_gc();
184    COMMENT_UPLOAD_LIMITER
185        .check_key(&ip.to_string())
186        .map_err(|_| "上传过于频繁,请稍后再试".to_string())?;
187    COMMENT_UPLOAD_DAILY_LIMITER
188        .check_key(&ip.to_string())
189        .map_err(|_| "今日上传次数已达上限,请明天再试".to_string())?;
190    Ok(())
191}
192
193#[cfg(feature = "server")]
194/// 检查图片访问请求是否超出限流阈值。
195///
196/// 超限时返回带 `Retry-After` 头的 429 响应:等待秒数取自 governor 的
197/// `NotUntil`(到下一个可用令牌的精确时长,向上取整、至少 1s),
198/// 客户端可据此安排重试而非盲等。
199pub fn check_image_limit(ip: &str) -> Result<(), Box<Response>> {
200    ensure_limiter_gc();
201    IMAGE_LIMITER
202        .check_key(&ip.to_string())
203        .map(|_| ())
204        .map_err(|not_until| {
205            let wait = not_until.wait_time_from(governor::clock::DefaultClock::default().now());
206            let secs = (wait.as_millis() as u64).div_ceil(1000).max(1);
207            Box::new(
208                (
209                    StatusCode::TOO_MANY_REQUESTS,
210                    [(axum::http::header::RETRY_AFTER, secs.to_string())],
211                )
212                    .into_response(),
213            )
214        })
215}
216
217#[cfg(feature = "server")]
218async fn trusted_proxy_count() -> usize {
219    crate::api::settings::runtime_security_settings()
220        .await
221        .trusted_proxy_count as usize
222}
223
224#[cfg(feature = "server")]
225fn is_valid_ip(ip: &str) -> bool {
226    ip.parse::<std::net::IpAddr>().is_ok()
227}
228
229/// 从 `X-Forwarded-For` 头按信任代理层数提取真实客户端 IP。
230///
231/// # 伪造风险(务必正确配置 `TRUSTED_PROXY_COUNT`)
232///
233/// XFF 头由客户端**可写**,本函数以 `parts[len-1-trusted_proxy_count]` 选取
234/// 「可信代理链最左侧之外」的地址。该取值**完全依赖 `TRUSTED_PROXY_COUNT` 与真实
235/// 代理跳数精确相等**,一旦不符即可被滥用:
236/// - **配得偏大**:会选中客户端伪造的地址——攻击者在请求里塞入多个伪 XFF 段,
237///   使选中的「客户端 IP」落为其伪造值,从而绕过按 IP 的限流(每次伪造一个新 IP),
238///   或令多个真实用户被错误归并到同一代理 IP 的桶里。
239/// - **配得偏小**:会选中某一跳中间代理的 IP,导致所有用户共享同一个限流桶。
240///
241/// 因此生产部署**必须**由最外层反向代理覆盖/重写 XFF(而非原样转发客户端 XFF),
242/// 并令 `TRUSTED_PROXY_COUNT` 等于「客户端 → 服务端」之间的真实代理跳数。
243/// 详见 `AGENTS.md` 中 `TRUSTED_PROXY_COUNT` 的部署说明。
244#[cfg(feature = "server")]
245fn ip_from_x_forwarded_for(value: &str, trusted_proxy_count: usize) -> Option<String> {
246    // X-Forwarded-For 格式:client, proxy1, proxy2, ..., proxyN
247    // 越靠右的地址离服务端越近。
248    let parts: Vec<&str> = value
249        .split(',')
250        .map(str::trim)
251        .filter(|s| !s.is_empty())
252        .collect();
253
254    if trusted_proxy_count == 0 || parts.len() <= trusted_proxy_count {
255        return None;
256    }
257
258    // 真实客户端 IP 位于右侧第 trusted_proxy_count + 1 个。
259    let idx = parts.len() - 1 - trusted_proxy_count;
260    let ip = parts[idx].to_string();
261    if is_valid_ip(&ip) {
262        Some(ip)
263    } else {
264        None
265    }
266}
267
268#[cfg(feature = "server")]
269fn ip_from_x_real_ip(value: &str) -> Option<String> {
270    let ip = value.trim().to_string();
271    if is_valid_ip(&ip) {
272        Some(ip)
273    } else {
274        None
275    }
276}
277
278#[cfg(feature = "server")]
279fn get_client_ip_internal(
280    headers: &http::HeaderMap,
281    trusted: usize,
282    peer: Option<std::net::SocketAddr>,
283) -> String {
284    if trusted > 0 {
285        if let Some(value) = headers.get("x-forwarded-for").and_then(|v| v.to_str().ok()) {
286            if let Some(ip) = ip_from_x_forwarded_for(value, trusted) {
287                return ip;
288            }
289        }
290
291        if let Some(ip) = headers
292            .get("x-real-ip")
293            .and_then(|v| v.to_str().ok())
294            .and_then(ip_from_x_real_ip)
295        {
296            return ip;
297        }
298    }
299
300    if let Some(addr) = peer {
301        return addr.ip().to_string();
302    }
303
304    // 落到 unknown 是预期的常见路径(TRUSTED_PROXY_COUNT=0 + server function 无 TCP
305    // 对端地址),每个请求都打 WARN 会淹没日志,故进程级仅警告一次。
306    static UNKNOWN_IP_WARNED: Once = Once::new();
307    UNKNOWN_IP_WARNED.call_once(|| {
308        tracing::warn!(
309            "无法获取客户端真实 IP(TRUSTED_PROXY_COUNT=0 不解析代理头,\
310             且当前上下文无 TCP 对端地址,通常为 server function),限流将按 'unknown' 键聚合"
311        );
312    });
313    "unknown".to_string()
314}
315
316#[cfg(feature = "server")]
317/// 根据信任代理层数从请求头中提取客户端 IP,并校验 IP 合法性。
318///
319/// 当未配置可信代理时,不会信任任何 `X-Forwarded-For` / `X-Real-IP` 头,
320/// 而是直接返回 `peer` 中的 TCP 对端地址(如果提供)。
321pub async fn get_client_ip_with_peer(
322    headers: &http::HeaderMap,
323    peer: Option<std::net::SocketAddr>,
324) -> String {
325    get_client_ip_internal(headers, trusted_proxy_count().await, peer)
326}
327#[cfg(feature = "server")]
328/// 使用「站点配置 → 安全」面板的代理层数提取客户端 IP。
329///
330/// 适用于 Dioxus server function 等无法获取 `ConnectInfo` 的场景。
331/// 生产环境建议配合反向代理与设置面板的 TRUSTED_PROXY_COUNT 使用。
332pub async fn get_client_ip(headers: &http::HeaderMap) -> String {
333    get_client_ip_internal(headers, trusted_proxy_count().await, None)
334}
335
336#[cfg(feature = "server")]
337/// 检查严格限流(注册、登录等敏感接口)。
338///
339/// 当 IP 为 "unknown"(无法识别真实客户端,通常是 TRUSTED_PROXY_COUNT=0
340/// 且调用方为 Dioxus server function 时)改用宽松桶,避免所有匿名请求共享
341/// 严格桶导致正常用户被误杀。生产环境配好 TRUSTED_PROXY_COUNT 后走真实 IP,
342/// 始终命中严格桶。
343pub fn check_strict_limit(ip: &str) -> Result<(), String> {
344    ensure_limiter_gc();
345    if ip == "unknown" {
346        UNKNOWN_BUCKET_LIMITER
347            .check_key(&ip.to_string())
348            .map(|_| ())
349            .map_err(|_| "服务繁忙,请稍后再试".to_string())
350    } else {
351        STRICT_LIMITER
352            .check_key(&ip.to_string())
353            .map(|_| ())
354            .map_err(|_| "请求过于频繁,请稍后再试".to_string())
355    }
356}
357
358#[cfg(feature = "server")]
359/// 检查上传请求是否超出限流阈值。
360pub fn check_upload_limit(ip: &str) -> Result<(), String> {
361    ensure_limiter_gc();
362    UPLOAD_LIMITER
363        .check_key(&ip.to_string())
364        .map(|_| ())
365        .map_err(|_| "上传过于频繁,请稍后再试".to_string())
366}
367
368#[cfg(feature = "server")]
369/// 检查代码执行请求的双层速率限制(每秒突发 + 每日总额)。
370///
371/// 两层任一被限即拒绝。返回中文错误消息,供 server function 直接透传给前端。
372pub fn check_code_exec_limit(ip: &str) -> Result<(), String> {
373    ensure_limiter_gc();
374    CODE_EXEC_LIMITER
375        .check_key(&ip.to_string())
376        .map_err(|_| "请求过于频繁,请稍后再试".to_string())?;
377    CODE_EXEC_DAILY_LIMITER
378        .check_key(&ip.to_string())
379        .map_err(|_| "今日运行次数已达上限,请明天再试".to_string())?;
380    Ok(())
381}
382
383#[cfg(all(test, feature = "server"))]
384mod tests {
385    use super::*;
386    use http::HeaderMap;
387    use std::net::{IpAddr, Ipv4Addr, SocketAddr};
388
389    #[test]
390    fn get_client_ip_from_x_forwarded_for_with_one_trusted_proxy() {
391        let mut headers = HeaderMap::new();
392        headers.insert("x-forwarded-for", "1.2.3.4, 5.6.7.8".parse().unwrap());
393        assert_eq!(
394            get_client_ip_with_trusted_and_peer(&headers, 1, None),
395            "1.2.3.4"
396        );
397    }
398
399    #[test]
400    fn get_client_ip_from_x_forwarded_for_with_two_trusted_proxies() {
401        let mut headers = HeaderMap::new();
402        headers.insert(
403            "x-forwarded-for",
404            "1.2.3.4, 5.6.7.8, 9.10.11.12".parse().unwrap(),
405        );
406        assert_eq!(
407            get_client_ip_with_trusted_and_peer(&headers, 2, None),
408            "1.2.3.4"
409        );
410    }
411
412    #[test]
413    fn get_client_ip_ignores_x_forwarded_for_when_no_trusted_proxies() {
414        let mut headers = HeaderMap::new();
415        headers.insert("x-forwarded-for", "1.2.3.4, 5.6.7.8".parse().unwrap());
416        assert_eq!(
417            get_client_ip_with_trusted_and_peer(&headers, 0, None),
418            "unknown"
419        );
420    }
421
422    #[test]
423    fn get_client_ip_falls_back_to_peer_when_no_trusted_proxies() {
424        let mut headers = HeaderMap::new();
425        headers.insert("x-forwarded-for", "1.2.3.4, 5.6.7.8".parse().unwrap());
426        let peer = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 12345);
427        assert_eq!(
428            get_client_ip_with_trusted_and_peer(&headers, 0, Some(peer)),
429            "127.0.0.1"
430        );
431    }
432
433    #[test]
434    fn get_client_ip_from_x_real_ip_when_trusted() {
435        let mut headers = HeaderMap::new();
436        headers.insert("x-real-ip", "9.8.7.6".parse().unwrap());
437        assert_eq!(
438            get_client_ip_with_trusted_and_peer(&headers, 1, None),
439            "9.8.7.6"
440        );
441    }
442
443    #[test]
444    fn get_client_ip_x_real_ip_ignored_when_not_trusted() {
445        let mut headers = HeaderMap::new();
446        headers.insert("x-real-ip", "9.8.7.6".parse().unwrap());
447        let peer = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)), 12345);
448        assert_eq!(
449            get_client_ip_with_trusted_and_peer(&headers, 0, Some(peer)),
450            "192.168.1.1"
451        );
452    }
453
454    #[test]
455    fn get_client_ip_x_forwarded_for_takes_priority_over_x_real_ip() {
456        let mut headers = HeaderMap::new();
457        headers.insert("x-forwarded-for", "1.1.1.1, 2.2.2.2".parse().unwrap());
458        headers.insert("x-real-ip", "3.3.3.3".parse().unwrap());
459        assert_eq!(
460            get_client_ip_with_trusted_and_peer(&headers, 1, None),
461            "1.1.1.1"
462        );
463    }
464
465    #[test]
466    fn get_client_ip_no_headers_returns_unknown() {
467        let headers = HeaderMap::new();
468        assert_eq!(
469            get_client_ip_with_trusted_and_peer(&headers, 1, None),
470            "unknown"
471        );
472    }
473
474    #[test]
475    fn get_client_ip_ignores_short_x_forwarded_for_list() {
476        let mut headers = HeaderMap::new();
477        headers.insert("x-forwarded-for", "1.2.3.4".parse().unwrap());
478        assert_eq!(
479            get_client_ip_with_trusted_and_peer(&headers, 2, None),
480            "unknown"
481        );
482    }
483
484    #[test]
485    fn get_client_ip_ignores_x_forwarded_for_equal_to_proxy_count() {
486        let mut headers = HeaderMap::new();
487        headers.insert("x-forwarded-for", "1.2.3.4, 5.6.7.8".parse().unwrap());
488        assert_eq!(
489            get_client_ip_with_trusted_and_peer(&headers, 2, None),
490            "unknown"
491        );
492    }
493
494    #[test]
495    fn get_client_ip_ignores_empty_x_forwarded_for_entries() {
496        let mut headers = HeaderMap::new();
497        headers.insert(
498            "x-forwarded-for",
499            " , 1.2.3.4 , 5.6.7.8 , ".parse().unwrap(),
500        );
501        assert_eq!(
502            get_client_ip_with_trusted_and_peer(&headers, 1, None),
503            "1.2.3.4"
504        );
505    }
506
507    #[test]
508    fn get_client_ip_rejects_invalid_x_forwarded_for_value() {
509        let mut headers = HeaderMap::new();
510        headers.insert("x-forwarded-for", "not-an-ip, 5.6.7.8".parse().unwrap());
511        assert_eq!(
512            get_client_ip_with_trusted_and_peer(&headers, 1, None),
513            "unknown"
514        );
515    }
516
517    #[test]
518    fn get_client_ip_rejects_invalid_x_real_ip_value() {
519        let mut headers = HeaderMap::new();
520        headers.insert("x-real-ip", "not-an-ip".parse().unwrap());
521        assert_eq!(
522            get_client_ip_with_trusted_and_peer(&headers, 1, None),
523            "unknown"
524        );
525    }
526
527    #[test]
528    fn get_client_ip_prefers_xff_over_peer() {
529        let mut headers = HeaderMap::new();
530        headers.insert("x-forwarded-for", "1.2.3.4, 5.6.7.8".parse().unwrap());
531        let peer = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 12345);
532        assert_eq!(
533            get_client_ip_with_trusted_and_peer(&headers, 1, Some(peer)),
534            "1.2.3.4"
535        );
536    }
537
538    #[tokio::test]
539    async fn get_client_ip_defaults_to_unknown_without_db() {
540        // 无 DB 连接的单元测试环境,trusted_proxy_count 回退默认 0,
541        // 不信任任何 XFF 头,返回 "unknown"。
542        let mut headers = HeaderMap::new();
543        headers.insert("x-forwarded-for", "1.2.3.4, 5.6.7.8".parse().unwrap());
544        assert_eq!(get_client_ip(&headers).await, "unknown");
545    }
546
547    #[test]
548    #[serial_test::serial]
549    fn check_strict_unknown_ip_uses_lenient_bucket() {
550        // "unknown" 桶 burst 为 100,少量请求应全部放行,不被严格桶误杀。
551        // 用 serial 隔离,因为 UNKNOWN_BUCKET_LIMITER 是全局状态。
552        for _ in 0..20 {
553            assert!(
554                super::check_strict_limit("unknown").is_ok(),
555                "unknown bucket should allow small bursts, not hit strict 1 req/s limit"
556            );
557        }
558    }
559
560    #[test]
561    #[serial_test::serial]
562    fn check_strict_real_ip_uses_strict_bucket() {
563        // 真实 IP 命中严格桶(1 req/s, burst 5)。连发超过 burst 应被限流。
564        // 用一个唯一的 IP 避免与其他测试状态冲突。
565        let unique_ip = "198.51.100.42";
566        let mut allowed = 0;
567        let mut blocked = false;
568        for _ in 0..50 {
569            match super::check_strict_limit(unique_ip) {
570                Ok(()) => allowed += 1,
571                Err(_) => blocked = true,
572            }
573            if blocked {
574                break;
575            }
576        }
577        assert!(
578            blocked,
579            "strict bucket should eventually block real IP burst"
580        );
581        assert!(
582            allowed <= 6,
583            "strict burst is 5, allowed should be <= 6, got {allowed}"
584        );
585    }
586
587    // 测试辅助函数:绕过环境变量读取,直接指定 trusted_proxy_count。
588    fn get_client_ip_with_trusted_and_peer(
589        headers: &HeaderMap,
590        trusted: usize,
591        peer: Option<SocketAddr>,
592    ) -> String {
593        get_client_ip_internal(headers, trusted, peer)
594    }
595}