use std::{
collections::{HashMap, VecDeque},
time::Duration,
};
const CONNECTION_FAILURE_LIMIT: u8 = 3;
const SOURCE_FAILURE_LIMIT: usize = 5;
const GLOBAL_FAILURE_LIMIT: usize = 30;
const FAILURE_WINDOW: Duration = Duration::from_secs(10 * 60);
const SOURCE_COOLDOWN: Duration = Duration::from_secs(10 * 60);
const GLOBAL_COOLDOWN: Duration = Duration::from_secs(5 * 60);
const SOURCE_IDLE_EXPIRY: Duration = Duration::from_secs(20 * 60);
const MAX_SOURCE_BUCKETS: usize = 10_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct ConnectionId(u64);
impl ConnectionId {
pub(crate) const fn new(value: u64) -> Self {
Self(value)
}
}
use super::types::SourceIdentity;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct FailureDisposition {
close_connection: bool,
}
impl FailureDisposition {
pub(crate) const fn close_connection(self) -> bool {
self.close_connection
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct RateLimit {
retry_after_seconds: u64,
close_connection: bool,
}
impl RateLimit {
pub(crate) const fn public_retry_after_seconds(self) -> u64 {
self.retry_after_seconds
}
pub(crate) const fn close_connection(self) -> bool {
self.close_connection
}
}
#[derive(Default)]
struct SourceBucket {
failures: VecDeque<Duration>,
cooldown_until: Option<Duration>,
last_seen: Duration,
insertion_order: u64,
}
#[derive(Default)]
pub(crate) struct AttemptLimiter {
connections: HashMap<ConnectionId, u8>,
sources: HashMap<SourceIdentity, SourceBucket>,
global_failures: VecDeque<Duration>,
global_cooldown_until: Option<Duration>,
next_insertion_order: u64,
}
impl AttemptLimiter {
#[cfg(test)]
pub(crate) fn new() -> Self {
Self::default()
}
pub(crate) fn check(
&mut self,
connection: ConnectionId,
source: SourceIdentity,
now: Duration,
) -> Result<(), RateLimit> {
self.prune(now);
if self.connections.get(&connection).copied().unwrap_or(0) >= CONNECTION_FAILURE_LIMIT {
return Err(RateLimit {
retry_after_seconds: 1,
close_connection: true,
});
}
if let Some(bucket) = self.sources.get_mut(&source) {
bucket.last_seen = now;
if let Some(until) = bucket.cooldown_until {
if now < until {
return Err(RateLimit {
retry_after_seconds: retry_after_seconds(until, now, SOURCE_COOLDOWN),
close_connection: false,
});
}
bucket.cooldown_until = None;
}
}
if let Some(until) = self.global_cooldown_until {
if now < until {
return Err(RateLimit {
retry_after_seconds: retry_after_seconds(until, now, GLOBAL_COOLDOWN),
close_connection: false,
});
}
self.global_cooldown_until = None;
}
Ok(())
}
pub(crate) fn record_failure(
&mut self,
connection: ConnectionId,
source: SourceIdentity,
now: Duration,
) -> FailureDisposition {
self.prune(now);
let connection_failures = self.connections.entry(connection).or_default();
*connection_failures = connection_failures.saturating_add(1);
if !self.sources.contains_key(&source) && self.sources.len() >= MAX_SOURCE_BUCKETS {
let oldest = self
.sources
.iter()
.min_by_key(|(identity, bucket)| {
(
bucket.last_seen,
bucket.insertion_order,
(*identity).clone(),
)
})
.map(|(identity, _)| identity.clone());
if let Some(oldest) = oldest {
self.sources.remove(&oldest);
}
}
let insertion_order = self.next_insertion_order;
let bucket = self.sources.entry(source).or_insert_with(|| {
self.next_insertion_order = self.next_insertion_order.saturating_add(1);
SourceBucket {
last_seen: now,
insertion_order,
..SourceBucket::default()
}
});
bucket.last_seen = now;
bucket.failures.push_back(now);
if bucket.failures.len() >= SOURCE_FAILURE_LIMIT {
bucket.failures.clear();
bucket.cooldown_until = Some(now.saturating_add(SOURCE_COOLDOWN));
}
self.global_failures.push_back(now);
if self.global_failures.len() >= GLOBAL_FAILURE_LIMIT {
self.global_failures.clear();
self.global_cooldown_until = Some(now.saturating_add(GLOBAL_COOLDOWN));
}
FailureDisposition {
close_connection: *connection_failures >= CONNECTION_FAILURE_LIMIT,
}
}
pub(crate) fn clear_authenticated_connection(&mut self, connection: ConnectionId) {
self.connections.remove(&connection);
}
pub(crate) fn dispose_connection(&mut self, connection: ConnectionId) {
self.connections.remove(&connection);
}
#[cfg(test)]
pub(crate) fn connection_failures(&self, connection: ConnectionId) -> u8 {
self.connections.get(&connection).copied().unwrap_or(0)
}
#[cfg(test)]
pub(crate) fn source_failures(&mut self, source: SourceIdentity, now: Duration) -> usize {
self.prune(now);
self.sources
.get(&source)
.map_or(0, |bucket| bucket.failures.len())
}
#[cfg(test)]
pub(crate) fn source_bucket_count(&self) -> usize {
self.sources.len()
}
#[cfg(test)]
pub(crate) fn contains_source(&self, source: SourceIdentity) -> bool {
self.sources.contains_key(&source)
}
fn prune(&mut self, now: Duration) {
prune_window(&mut self.global_failures, now);
self.sources
.retain(|_, bucket| now.saturating_sub(bucket.last_seen) < SOURCE_IDLE_EXPIRY);
for bucket in self.sources.values_mut() {
prune_window(&mut bucket.failures, now);
}
}
}
fn prune_window(failures: &mut VecDeque<Duration>, now: Duration) {
while failures
.front()
.is_some_and(|failure| now.saturating_sub(*failure) >= FAILURE_WINDOW)
{
failures.pop_front();
}
}
fn retry_after_seconds(until: Duration, now: Duration, bound: Duration) -> u64 {
let remaining = until.saturating_sub(now).min(bound);
let nanos = remaining.as_nanos();
let seconds = nanos.saturating_add(999_999_999) / 1_000_000_000;
u64::try_from(seconds)
.unwrap_or(bound.as_secs())
.clamp(1, bound.as_secs())
}