iroh_dns_server/http/
rate_limiting.rs

1use std::{sync::Arc, time::Duration};
2
3use governor::{clock::QuantaInstant, middleware::NoOpMiddleware};
4use serde::{Deserialize, Serialize};
5use tower_governor::{
6    GovernorLayer,
7    governor::GovernorConfigBuilder,
8    key_extractor::{PeerIpKeyExtractor, SmartIpKeyExtractor},
9};
10
11/// Rate limiting strategy for the HTTP server.
12#[derive(Debug, Deserialize, Default, Serialize, Clone)]
13#[serde(rename_all = "lowercase")]
14#[non_exhaustive]
15pub enum RateLimitConfig {
16    /// Disables rate limiting entirely.
17    Disabled,
18    /// Rate limits by the connection's peer IP address.
19    ///
20    /// See [`PeerIpKeyExtractor`].
21    ///
22    /// [`PeerIpKeyExtractor`]: https://docs.rs/tower_governor/latest/tower_governor/key_extractor/struct.PeerIpKeyExtractor.html
23    #[default]
24    Simple,
25    /// Rate limits by the client IP, extracted from reverse-proxy headers.
26    ///
27    /// Uses the headers commonly set by reverse proxies (for example
28    /// `X-Forwarded-For`) to extract the original client IP, falling back to the
29    /// connection's peer IP address. See [`SmartIpKeyExtractor`].
30    ///
31    /// [`SmartIpKeyExtractor`]: https://docs.rs/tower_governor/latest/tower_governor/key_extractor/struct.SmartIpKeyExtractor.html
32    Smart,
33}
34
35impl Default for &RateLimitConfig {
36    fn default() -> Self {
37        &RateLimitConfig::Simple
38    }
39}
40
41/// Create the default rate-limiting layer.
42///
43/// This spawns a background thread to clean up the rate limiting cache.
44pub(super) fn create<RespBody>(
45    rate_limit_config: &RateLimitConfig,
46) -> Option<GovernorLayer<PeerIpKeyExtractor, NoOpMiddleware<QuantaInstant>, RespBody>> {
47    let use_smart_extractor = match rate_limit_config {
48        RateLimitConfig::Disabled => {
49            tracing::info!("Rate limiting disabled");
50            return None;
51        }
52        RateLimitConfig::Simple => false,
53        RateLimitConfig::Smart => true,
54    };
55
56    tracing::info!("Rate limiting enabled ({rate_limit_config:?})");
57
58    // Configure rate limiting:
59    // * allow bursts with up to five requests per IP address
60    // * replenish one element every two seconds
61    let mut governor_conf_builder = GovernorConfigBuilder::default();
62    // governor_conf_builder.use_headers()
63    governor_conf_builder.per_second(4);
64    governor_conf_builder.burst_size(2);
65
66    if use_smart_extractor {
67        governor_conf_builder.key_extractor(SmartIpKeyExtractor);
68    }
69
70    let governor_conf = governor_conf_builder
71        .finish()
72        .expect("failed to build rate-limiting governor");
73
74    let governor_conf = Arc::new(governor_conf);
75
76    // The governor needs a background task for garbage collection (to clear expired records)
77    let gc_interval = Duration::from_secs(60);
78    let governor_limiter = governor_conf.limiter().clone();
79    std::thread::spawn(move || {
80        loop {
81            std::thread::sleep(gc_interval);
82            tracing::debug!("rate limiting storage size: {}", governor_limiter.len());
83            governor_limiter.retain_recent();
84        }
85    });
86
87    Some(GovernorLayer::new(governor_conf))
88}
iroh_dns_server/http/rate_limiting.rs

iroh_dns_server/http/
rate_limiting.rs
