use std::net::IpAddr;
use std::num::NonZeroU32;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use governor::{DefaultDirectRateLimiter, DefaultKeyedRateLimiter, Quota};
use iroh::TransportAddr;
use iroh::endpoint::Connection;
use crate::dht::routing::NodeId;
use crate::dispatch::source_key;
use crate::prune_guard::PruneGuard;
#[derive(Debug, Clone)]
pub struct RateLimitConfig {
pub per_peer_rate_per_sec: f64,
pub per_peer_burst: u32,
pub per_ip_rate_per_sec: f64,
pub per_ip_burst: u32,
pub global_rate_per_sec: f64,
pub global_burst: u32,
pub max_tracked_per_ip: usize,
pub max_tracked_per_peer: usize,
}
impl Default for RateLimitConfig {
fn default() -> Self {
Self {
per_peer_rate_per_sec: 20.0,
per_peer_burst: 40,
per_ip_rate_per_sec: 100.0,
per_ip_burst: 200,
global_rate_per_sec: 1000.0,
global_burst: 2000,
max_tracked_per_ip: 4096,
max_tracked_per_peer: 4096,
}
}
}
impl From<&decdn_common::config::ResolvedDht> for RateLimitConfig {
fn from(r: &decdn_common::config::ResolvedDht) -> Self {
Self {
per_peer_rate_per_sec: r.per_peer_rate_per_sec,
per_peer_burst: r.per_peer_burst,
per_ip_rate_per_sec: r.per_ip_rate_per_sec,
per_ip_burst: r.per_ip_burst,
global_rate_per_sec: r.global_rate_per_sec,
global_burst: r.global_burst,
max_tracked_per_ip: r.max_tracked_per_ip,
max_tracked_per_peer: r.max_tracked_per_peer,
}
}
}
impl From<&decdn_common::config::ResolvedProbe> for RateLimitConfig {
fn from(r: &decdn_common::config::ResolvedProbe) -> Self {
Self {
per_peer_rate_per_sec: r.per_peer_rate_per_sec,
per_peer_burst: r.per_peer_burst,
per_ip_rate_per_sec: r.per_ip_rate_per_sec,
per_ip_burst: r.per_ip_burst,
global_rate_per_sec: r.global_rate_per_sec,
global_burst: r.global_burst,
max_tracked_per_ip: r.max_tracked_per_ip,
max_tracked_per_peer: r.max_tracked_per_peer,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RejectLayer {
PerPeer,
PerIp,
Global,
}
impl RejectLayer {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::PerPeer => "per_peer",
Self::PerIp => "per_ip",
Self::Global => "global",
}
}
}
pub trait RateLimitMetricsSink: Send + Sync {
fn rejected(&self, layer: RejectLayer);
fn prune_sweep_per_ip(&self);
fn prune_sweep_per_peer(&self);
fn set_tracked_per_ip(&self, n: usize);
fn set_tracked_per_peer(&self, n: usize);
}
#[allow(missing_debug_implementations)]
pub struct ThreeLayerRateLimiter {
global: Option<Arc<DefaultDirectRateLimiter>>,
per_ip: Option<Arc<DefaultKeyedRateLimiter<IpAddr>>>,
per_peer: Option<Arc<DefaultKeyedRateLimiter<NodeId>>>,
cap_per_ip: usize,
cap_per_peer: usize,
pruning_per_ip: AtomicBool,
pruning_per_peer: AtomicBool,
metrics: Arc<dyn RateLimitMetricsSink>,
}
impl ThreeLayerRateLimiter {
#[must_use]
pub fn new(cfg: &RateLimitConfig, metrics: Arc<dyn RateLimitMetricsSink>) -> Self {
Self {
global: build_direct_limiter(cfg.global_rate_per_sec, cfg.global_burst),
per_ip: build_keyed_limiter(cfg.per_ip_rate_per_sec, cfg.per_ip_burst),
per_peer: build_keyed_limiter(cfg.per_peer_rate_per_sec, cfg.per_peer_burst),
cap_per_ip: cfg.max_tracked_per_ip,
cap_per_peer: cfg.max_tracked_per_peer,
pruning_per_ip: AtomicBool::new(false),
pruning_per_peer: AtomicBool::new(false),
metrics,
}
}
pub fn check(&self, peer_node_id: &NodeId, peer_ip: Option<IpAddr>) -> Result<(), RejectLayer> {
match self.try_admit_one(peer_node_id, peer_ip, true) {
Ok(()) => Ok(()),
Err(layer) => {
self.metrics.rejected(layer);
Err(layer)
}
}
}
fn try_admit_one(
&self,
peer_node_id: &NodeId,
peer_ip: Option<IpAddr>,
prune: bool,
) -> Result<(), RejectLayer> {
if let Some(g) = self.global.as_ref()
&& g.check().is_err()
{
return Err(RejectLayer::Global);
}
if let (Some(ip), Some(limiter)) = (peer_ip, self.per_ip.as_ref()) {
let ip = source_key(ip);
let result = limiter.check_key(&ip);
if prune {
self.maybe_prune_per_ip(limiter);
}
if result.is_err() {
return Err(RejectLayer::PerIp);
}
}
if let Some(limiter) = self.per_peer.as_ref() {
let result = limiter.check_key(peer_node_id);
if prune {
self.maybe_prune_per_peer(limiter);
}
if result.is_err() {
return Err(RejectLayer::PerPeer);
}
}
Ok(())
}
#[must_use]
pub fn admit_batch_extra(
&self,
peer_node_id: &NodeId,
peer_ip: Option<IpAddr>,
extra: usize,
) -> usize {
let mut granted = 0usize;
while granted < extra {
if self.try_admit_one(peer_node_id, peer_ip, false).is_err() {
break;
}
granted = granted.saturating_add(1);
}
granted
}
fn maybe_prune_per_ip(&self, limiter: &Arc<DefaultKeyedRateLimiter<IpAddr>>) {
if self.cap_per_ip > 0
&& limiter.len() > self.cap_per_ip.saturating_add(self.cap_per_ip / 10)
&& self
.pruning_per_ip
.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_ok()
{
let _guard = PruneGuard(&self.pruning_per_ip);
limiter.retain_recent();
self.metrics.prune_sweep_per_ip();
self.metrics.set_tracked_per_ip(limiter.len());
}
}
fn maybe_prune_per_peer(&self, limiter: &Arc<DefaultKeyedRateLimiter<NodeId>>) {
if self.cap_per_peer > 0
&& limiter.len() > self.cap_per_peer.saturating_add(self.cap_per_peer / 10)
&& self
.pruning_per_peer
.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_ok()
{
let _guard = PruneGuard(&self.pruning_per_peer);
limiter.retain_recent();
self.metrics.prune_sweep_per_peer();
self.metrics.set_tracked_per_peer(limiter.len());
}
}
#[must_use]
pub fn gc_per_ip(&self) -> Option<(usize, usize)> {
let limiter = self.per_ip.as_ref()?;
if self
.pruning_per_ip
.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_ok()
{
let _guard = PruneGuard(&self.pruning_per_ip);
let before = limiter.len();
limiter.retain_recent();
let after = limiter.len();
self.metrics.prune_sweep_per_ip();
self.metrics.set_tracked_per_ip(after);
Some((before, after))
} else {
None
}
}
#[must_use]
pub fn gc_per_peer(&self) -> Option<(usize, usize)> {
let limiter = self.per_peer.as_ref()?;
if self
.pruning_per_peer
.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_ok()
{
let _guard = PruneGuard(&self.pruning_per_peer);
let before = limiter.len();
limiter.retain_recent();
let after = limiter.len();
self.metrics.prune_sweep_per_peer();
self.metrics.set_tracked_per_peer(after);
Some((before, after))
} else {
None
}
}
#[must_use]
pub fn per_ip_tracked(&self) -> usize {
self.per_ip.as_ref().map_or(0, |l| l.len())
}
#[must_use]
pub fn per_peer_tracked(&self) -> usize {
self.per_peer.as_ref().map_or(0, |l| l.len())
}
}
pub(crate) fn peer_ip(conn: &Connection) -> Option<IpAddr> {
let paths = conn.paths();
paths
.iter()
.find(iroh::endpoint::Path::is_selected)
.and_then(|p| match p.remote_addr() {
TransportAddr::Ip(addr) => Some(addr.ip()),
_ => None,
})
.or_else(|| {
paths.iter().find_map(|p| match p.remote_addr() {
TransportAddr::Ip(addr) => Some(addr.ip()),
_ => None,
})
})
}
fn build_direct_limiter(rate_per_sec: f64, burst: u32) -> Option<Arc<DefaultDirectRateLimiter>> {
let quota = make_quota(rate_per_sec, burst)?;
Some(Arc::new(governor::RateLimiter::direct(quota)))
}
fn build_keyed_limiter<K>(rate_per_sec: f64, burst: u32) -> Option<Arc<DefaultKeyedRateLimiter<K>>>
where
K: std::hash::Hash + Eq + Clone + Send + Sync + 'static,
{
let quota = make_quota(rate_per_sec, burst)?;
Some(Arc::new(DefaultKeyedRateLimiter::keyed(quota)))
}
fn make_quota(rate_per_sec: f64, burst: u32) -> Option<Quota> {
if rate_per_sec <= 0.0 || burst == 0 || !rate_per_sec.is_finite() {
return None;
}
let period_secs = 1.0_f64 / rate_per_sec;
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
let period_ns = (period_secs * 1_000_000_000.0).max(1.0) as u64;
let period = Duration::from_nanos(period_ns);
let burst_nz = NonZeroU32::new(burst)?;
Some(Quota::with_period(period)?.allow_burst(burst_nz))
}
#[cfg(test)]
#[allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::panic,
// float_cmp: the mapped `f64` slots are copied literals, never computed, so
// exact equality is the property being pinned — an epsilon compare here
// would be strictly laxer than the mapping guarantee.
clippy::float_cmp
)]
mod tests;