From fe684c23e60fc8c35d3ec02bf088f36e1e248c50 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADtor=20Vasconcellos?= Date: Tue, 29 Oct 2024 11:35:38 -0300 Subject: [PATCH] feat(iroh-dns-server)!: Make http rate limit configurable (#2772) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- iroh-dns-server/config.dev.toml | 2 + iroh-dns-server/config.prod.toml | 2 + iroh-dns-server/src/config.rs | 7 ++- iroh-dns-server/src/http.rs | 15 ++++-- iroh-dns-server/src/http/rate_limiting.rs | 59 ++++++++++++++++++++--- iroh-dns-server/src/server.rs | 8 ++- 6 files changed, 78 insertions(+), 15 deletions(-) diff --git a/iroh-dns-server/config.dev.toml b/iroh-dns-server/config.dev.toml index a43b10364a..37569002bf 100644 --- a/iroh-dns-server/config.dev.toml +++ b/iroh-dns-server/config.dev.toml @@ -1,3 +1,5 @@ +pkarr_put_rate_limit = "disabled" + [http] port = 8080 bind_addr = "127.0.0.1" diff --git a/iroh-dns-server/config.prod.toml b/iroh-dns-server/config.prod.toml index 64b3f88f67..9868f67fc9 100644 --- a/iroh-dns-server/config.prod.toml +++ b/iroh-dns-server/config.prod.toml @@ -1,3 +1,5 @@ +pkarr_put_rate_limit = "smart" + [https] port = 443 domains = ["irohdns.example.org"] diff --git a/iroh-dns-server/src/config.rs b/iroh-dns-server/src/config.rs index 89222f9daf..732d65e4e8 100644 --- a/iroh-dns-server/src/config.rs +++ b/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, + + /// 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(), } } } diff --git a/iroh-dns-server/src/http.rs b/iroh-dns-server/src/http.rs index d7c8499396..bdd4766d11 100644 --- a/iroh-dns-server/src/http.rs +++ b/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, https_config: Option, + rate_limit_config: RateLimitConfig, state: AppState, ) -> Result { 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!" })) diff --git a/iroh-dns-server/src/http/rate_limiting.rs b/iroh-dns-server/src/http/rate_limiting.rs index 991ff6e3b0..53952f339e 100644 --- a/iroh-dns-server/src/http/rate_limiting.rs +++ b/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> { +pub fn create( + rate_limit_config: &RateLimitConfig, +) -> Option>> { + 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