feat(iroh-dns-server)!: Make http rate limit configurable (#2772)

Hello,

We are currently testing some new cloud features at
[Spacedrive](https://github.com/spacedriveapp/spacedrive), and our
implementation relies heavily on iroh. As part of this, we are deploying
our own iroh-dns-server. However, since all of our backend services
operate behind a reverse proxy, we noticed that the iroh-dns-server was
frequently hitting its rate limit because it wasn’t aware of the proxy
setup. To address this, I decided to implement a configurable rate limit
for the iroh-dns-server.

## Description

This PR adds a new entry to the `iroh-dns-server` TOML file for
configuring the HTTP rate limit. The new configuration allows for
disabling the rate limit and also supports configuring it to use the
[SmartIPKeyExtract](https://github.com/benwis/tower-governor/blob/v0.4.2/src/key_extractor.rs#L85-L119),
making it compatible with reverse proxies.

## Breaking Changes

- `iroh-dns-server`'s configuration structure now has a new field
allowing to choose the rate limiting algorithms.

## Notes & open questions

:)

## Change checklist

- [x] Self-review.
- [x] Documentation updates following the [style
guide](https://rust-lang.github.io/rfcs/1574-more-api-documentation-conventions.html#appendix-a-full-conventions-text),
if relevant.
- [x] Tests if relevant.
- [x] All breaking changes documented.

Author: HeavenVolkoff

iroh-dns-server/config.dev.toml

@@ -1,3 +1,5 @@
+pkarr_put_rate_limit = "disabled"
+
 [http]
 port = 8080
 bind_addr = "127.0.0.1"

iroh-dns-server/config.prod.toml

@@ -1,3 +1,5 @@
+pkarr_put_rate_limit = "smart"
+
 [https]
 port = 443
 domains = ["irohdns.example.org"]

iroh-dns-server/src/config.rs

@@ -12,7 +12,7 @@ use tracing::info;
 
 use crate::{
     dns::DnsConfig,
-    http::{CertMode, HttpConfig, HttpsConfig},
+    http::{CertMode, HttpConfig, HttpsConfig, RateLimitConfig},
 };
 
 const DEFAULT_METRICS_ADDR: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 9117);
@@ -43,6 +43,10 @@ pub struct Config {
 
     /// Config for the mainline lookup.
     pub mainline: Option<MainlineConfig>,
+
+    /// Config for pkarr rate limit
+    #[serde(default)]
+    pub pkarr_put_rate_limit: RateLimitConfig,
 }
 
 /// The config for the metrics server.
@@ -185,6 +189,7 @@ impl Default for Config {
             },
             metrics: None,
             mainline: None,
+            pkarr_put_rate_limit: RateLimitConfig::default(),
         }
     }
 }

iroh-dns-server/src/http.rs

@@ -30,7 +30,7 @@ mod pkarr;
 mod rate_limiting;
 mod tls;
 
-pub use self::tls::CertMode;
+pub use self::{rate_limiting::RateLimitConfig, tls::CertMode};
 use crate::{config::Config, metrics::Metrics, state::AppState};
 
 /// Config for the HTTP server
@@ -71,13 +71,14 @@ impl HttpServer {
     pub async fn spawn(
         http_config: Option<HttpConfig>,
         https_config: Option<HttpsConfig>,
+        rate_limit_config: RateLimitConfig,
         state: AppState,
     ) -> Result<HttpServer> {
         if http_config.is_none() && https_config.is_none() {
             bail!("Either http or https config is required");
         }
 
-        let app = create_app(state);
+        let app = create_app(state, &rate_limit_config);
 
         let mut tasks = JoinSet::new();
 
@@ -184,7 +185,7 @@ impl HttpServer {
     }
 }
 
-pub(crate) fn create_app(state: AppState) -> Router {
+pub(crate) fn create_app(state: AppState, rate_limit_config: &RateLimitConfig) -> Router {
     // configure cors middleware
     let cors = CorsLayer::new()
         // allow `GET` and `POST` when accessing the resource
@@ -209,7 +210,7 @@ pub(crate) fn create_app(state: AppState) -> Router {
     });
 
     // configure rate limiting middleware
-    let rate_limit = rate_limiting::create();
+    let rate_limit = rate_limiting::create(rate_limit_config);
 
     // configure routes
     //
@@ -218,7 +219,11 @@ pub(crate) fn create_app(state: AppState) -> Router {
         .route("/dns-query", get(doh::get).post(doh::post))
         .route(
             "/pkarr/:key",
-            get(pkarr::get).put(pkarr::put.layer(rate_limit)),
+            if let Some(rate_limit) = rate_limit {
+                get(pkarr::get).put(pkarr::put.layer(rate_limit))
+            } else {
+                get(pkarr::get).put(pkarr::put)
+            },
         )
         .route("/healthcheck", get(|| async { "OK" }))
         .route("/", get(|| async { "Hi!" }))

iroh-dns-server/src/http/rate_limiting.rs

@@ -1,21 +1,64 @@
 use std::time::Duration;
 
 use governor::{clock::QuantaInstant, middleware::NoOpMiddleware};
+use serde::{Deserialize, Serialize};
 use tower_governor::{
-    governor::GovernorConfigBuilder, key_extractor::PeerIpKeyExtractor, GovernorLayer,
+    governor::GovernorConfigBuilder,
+    key_extractor::{PeerIpKeyExtractor, SmartIpKeyExtractor},
+    GovernorLayer,
 };
 
+/// Config for http rate limit.
+#[derive(Debug, Deserialize, Default, Serialize, Clone)]
+#[serde(rename_all = "lowercase")]
+pub enum RateLimitConfig {
+    /// Disable rate limit for http server.
+    Disabled,
+    /// Enable rate limit for http server based on the connection peer IP address.
+    /// https://docs.rs/tower_governor/latest/tower_governor/key_extractor/struct.PeerIpKeyExtractor.html
+    #[default]
+    Simple,
+    /// Enable rate limit for http server based on a smart logic for extracting the connection original IP address, useful for reverse proxies.
+    /// https://docs.rs/tower_governor/latest/tower_governor/key_extractor/struct.SmartIpKeyExtractor.html
+    Smart,
+}
+
+impl Default for &RateLimitConfig {
+    fn default() -> Self {
+        &RateLimitConfig::Simple
+    }
+}
+
 /// Create the default rate-limiting layer.
 ///
 /// This spawns a background thread to clean up the rate limiting cache.
-pub fn create() -> GovernorLayer<'static, PeerIpKeyExtractor, NoOpMiddleware<QuantaInstant>> {
+pub fn create(
+    rate_limit_config: &RateLimitConfig,
+) -> Option<GovernorLayer<'static, PeerIpKeyExtractor, NoOpMiddleware<QuantaInstant>>> {
+    let use_smart_extractor = match rate_limit_config {
+        RateLimitConfig::Disabled => {
+            tracing::info!("Rate limiting disabled");
+            return None;
+        }
+        RateLimitConfig::Simple => false,
+        RateLimitConfig::Smart => true,
+    };
+
+    tracing::info!("Rate limiting enabled ({rate_limit_config:?})");
+
     // Configure rate limiting:
     // * allow bursts with up to five requests per IP address
     // * replenish one element every two seconds
-    let governor_conf = GovernorConfigBuilder::default()
-        // .use_headers()
-        .per_second(4)
-        .burst_size(2)
+    let mut governor_conf_builder = GovernorConfigBuilder::default();
+    // governor_conf_builder.use_headers()
+    governor_conf_builder.per_second(4);
+    governor_conf_builder.burst_size(2);
+
+    if use_smart_extractor {
+        governor_conf_builder.key_extractor(SmartIpKeyExtractor);
+    }
+
+    let governor_conf = governor_conf_builder
         .finish()
         .expect("failed to build rate-limiting governor");
 
@@ -34,7 +77,7 @@ pub fn create() -> GovernorLayer<'static, PeerIpKeyExtractor, NoOpMiddleware<Qua
         governor_limiter.retain_recent();
     });
 
-    GovernorLayer {
+    Some(GovernorLayer {
         config: &*governor_conf,
-    }
+    })
 }

iroh-dns-server/src/server.rs

@@ -52,7 +52,13 @@ impl Server {
             }
             Ok(())
         });
-        let http_server = HttpServer::spawn(config.http, config.https, state.clone()).await?;
+        let http_server = HttpServer::spawn(
+            config.http,
+            config.https,
+            config.pkarr_put_rate_limit,
+            state.clone(),
+        )
+        .await?;
         let dns_server = DnsServer::spawn(config.dns, state.dns_handler.clone()).await?;
         Ok(Self {
             http_server,