1#[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 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")]
60static 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")]
70static 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")]
84static 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")]
94static 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")]
109static 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#[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#[cfg(feature = "server")]
141fn ensure_limiter_gc() {
142 static SPAWNED: std::sync::Once = std::sync::Once::new();
143 SPAWNED.call_once(|| {
144 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")]
169pub 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")]
179pub 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")]
194pub 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#[cfg(feature = "server")]
245fn ip_from_x_forwarded_for(value: &str, trusted_proxy_count: usize) -> Option<String> {
246 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 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 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")]
317pub 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")]
328pub 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")]
337pub 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")]
359pub 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")]
369pub 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 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 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 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 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}