-
Notifications
You must be signed in to change notification settings - Fork 20
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(server): proper rate limiting behind reverse-proxy
- Loading branch information
Showing
5 changed files
with
109 additions
and
63 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,63 +1,111 @@ | ||
| use std::{sync::Arc, time::Duration}; | ||
| use std::{net::IpAddr, sync::Arc, time::Duration}; | ||
|
|
||
| use axum::Router; | ||
| use governor::middleware::StateInformationMiddleware; | ||
| use serde::{Deserialize, Serialize}; | ||
|
|
||
| use governor::middleware::StateInformationMiddleware; | ||
| use tower_governor::{ | ||
| governor::GovernorConfigBuilder, key_extractor::PeerIpKeyExtractor, GovernorLayer, | ||
| governor::{GovernorConfig, GovernorConfigBuilder}, | ||
| key_extractor::{PeerIpKeyExtractor, SmartIpKeyExtractor}, | ||
| }; | ||
|
|
||
| pub use tower_governor::GovernorLayer; | ||
|
|
||
| #[derive(Serialize, Deserialize, Debug)] | ||
| pub struct RateLimiterConfig { | ||
| pub(crate) behind_proxy: bool, | ||
| pub(crate) per_second: u64, | ||
| pub(crate) burst_size: u32, | ||
| } | ||
|
|
||
| impl Default for RateLimiterConfig { | ||
| fn default() -> Self { | ||
| Self { | ||
| behind_proxy: false, | ||
| per_second: 2, | ||
| burst_size: 10, | ||
| } | ||
| } | ||
| } | ||
|
|
||
| pub type RateLimiterLayer = GovernorLayer<PeerIpKeyExtractor, StateInformationMiddleware>; | ||
|
|
||
| /// Create the default rate-limiting layer. | ||
| /// | ||
| /// This will be used by the [crate::http_server::HttpServer] to guard all endpoints (GET and PUT) | ||
| /// and in [crate::dht_server::DhtServer] before calling [pkarr::mainline::rpc::Rpc::get] | ||
| /// after a cache miss or if its cached packet is expired. | ||
| /// | ||
| /// The purpose is to limit DHT queries as much as possible, while serving honest clients still. | ||
| /// | ||
| /// This spawns a background thread to clean up the rate limiting cache. | ||
| /// | ||
| /// # Limits | ||
| /// | ||
| /// * allow a burst of `10 requests` per IP address | ||
| /// * replenish `1 request` every `2 seconds` | ||
| pub fn create(config: &RateLimiterConfig) -> RateLimiterLayer { | ||
| let governor_config = GovernorConfigBuilder::default() | ||
| .use_headers() | ||
| .per_second(config.per_second) | ||
| .burst_size(config.burst_size) | ||
| .finish() | ||
| .expect("failed to build rate-limiting governor"); | ||
|
|
||
| let governor_config = Arc::new(governor_config); | ||
|
|
||
| // The governor needs a background task for garbage collection (to clear expired records) | ||
| let gc_interval = Duration::from_secs(60); | ||
| let governor_limiter = governor_config.limiter().clone(); | ||
| std::thread::spawn(move || loop { | ||
| std::thread::sleep(gc_interval); | ||
| tracing::debug!("rate limiting storage size: {}", governor_limiter.len()); | ||
| governor_limiter.retain_recent(); | ||
| }); | ||
|
|
||
| GovernorLayer { | ||
| config: governor_config, | ||
| #[derive(Debug, Clone)] | ||
| /// A rate limiter that works for direct connections (Peer) or behind reverse-proxy (Proxy) | ||
| pub enum IpRateLimiter { | ||
| Peer(Arc<GovernorConfig<PeerIpKeyExtractor, StateInformationMiddleware>>), | ||
| Proxy(Arc<GovernorConfig<SmartIpKeyExtractor, StateInformationMiddleware>>), | ||
| } | ||
|
|
||
| impl IpRateLimiter { | ||
| /// Create an [IpRateLimiter] | ||
| /// | ||
| /// This spawns a background thread to clean up the rate limiting cache. | ||
| pub fn new(config: &RateLimiterConfig) -> Self { | ||
| match config.behind_proxy { | ||
| true => { | ||
| let config = Arc::new( | ||
| GovernorConfigBuilder::default() | ||
| .use_headers() | ||
| .per_second(config.per_second) | ||
| .burst_size(config.burst_size) | ||
| .key_extractor(SmartIpKeyExtractor) | ||
| .finish() | ||
| .expect("failed to build rate-limiting governor"), | ||
| ); | ||
|
|
||
| // The governor needs a background task for garbage collection (to clear expired records) | ||
| let gc_interval = Duration::from_secs(60); | ||
|
|
||
| let governor_limiter = config.limiter().clone(); | ||
| std::thread::spawn(move || loop { | ||
| std::thread::sleep(gc_interval); | ||
| tracing::debug!("rate limiting storage size: {}", governor_limiter.len()); | ||
| governor_limiter.retain_recent(); | ||
| }); | ||
|
|
||
| Self::Proxy(config) | ||
| } | ||
| false => { | ||
| let config = Arc::new( | ||
| GovernorConfigBuilder::default() | ||
| .use_headers() | ||
| .per_second(config.per_second) | ||
| .burst_size(config.burst_size) | ||
| .finish() | ||
| .expect("failed to build rate-limiting governor"), | ||
| ); | ||
|
|
||
| // The governor needs a background task for garbage collection (to clear expired records) | ||
| let gc_interval = Duration::from_secs(60); | ||
|
|
||
| let governor_limiter = config.limiter().clone(); | ||
| std::thread::spawn(move || loop { | ||
| std::thread::sleep(gc_interval); | ||
| tracing::debug!("rate limiting storage size: {}", governor_limiter.len()); | ||
| governor_limiter.retain_recent(); | ||
| }); | ||
|
|
||
| Self::Peer(config) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /// Check if the Ip is allowed to make more requests | ||
| pub fn is_limited(&self, ip: &IpAddr) -> bool { | ||
| match self { | ||
| IpRateLimiter::Peer(config) => config.limiter().check_key(ip).is_err(), | ||
| IpRateLimiter::Proxy(config) => config.limiter().check_key(ip).is_err(), | ||
| } | ||
| } | ||
|
|
||
| /// Add a [GovernorLayer] on the provided [Router] | ||
| pub fn layer(&self, router: &Router) { | ||
| let _ = match self { | ||
| IpRateLimiter::Peer(config) => router.clone().layer(GovernorLayer { | ||
| config: config.clone(), | ||
| }), | ||
| IpRateLimiter::Proxy(config) => router.clone().layer(GovernorLayer { | ||
| config: config.clone(), | ||
| }), | ||
| }; | ||
| } | ||
| } |