use std::{net::IpAddr, time::Duration};
use crate::{
DetectorKind, OwnedAnomaly, Timestamp,
anomaly_fields::KeyFields,
correlate::{KeyIndexed, TimeBucketedCounter},
detect::registry::SrcHost,
event::Severity,
};
pub struct ConnectionFloodDetector {
starts: TimeBucketedCounter<IpAddr>,
threshold: u64,
window: Duration,
cooldown: Duration,
last_emitted: KeyIndexed<IpAddr, ()>,
}
#[derive(Debug, Clone, Copy)]
#[non_exhaustive]
pub struct FloodScore {
pub src: IpAddr,
pub count: u64,
pub window_secs: f64,
}
impl ConnectionFloodDetector {
pub fn new() -> Self {
Self::with_params(Duration::from_secs(10), Duration::from_secs(1), 100)
}
pub fn with_params(window: Duration, bucket_width: Duration, threshold: u64) -> Self {
Self {
starts: TimeBucketedCounter::new(window, bucket_width, 10_000),
threshold,
window,
cooldown: Duration::from_secs(60),
last_emitted: KeyIndexed::new(Duration::from_secs(60), 10_000),
}
}
pub fn with_cooldown(mut self, cooldown: Duration) -> Self {
self.cooldown = cooldown;
self.last_emitted = KeyIndexed::new(cooldown, 10_000);
self
}
pub fn observe(&mut self, src: IpAddr, now: Timestamp) -> Option<FloodScore> {
self.starts.bump(src, now);
let count = self.starts.count(&src, now);
if count < self.threshold {
return None;
}
if self.last_emitted.get(&src, now).is_some() {
return None;
}
self.last_emitted.insert(src, (), now);
Some(FloodScore {
src,
count,
window_secs: self.window.as_secs_f64(),
})
}
pub fn evict_expired(&mut self, now: Timestamp) {
self.starts.evict_expired(now);
self.last_emitted.evict_expired(now);
}
pub fn tracked(&self) -> usize {
self.starts.len()
}
}
impl Default for ConnectionFloodDetector {
fn default() -> Self {
Self::new()
}
}
impl FloodScore {
pub fn into_anomaly(self, ts: Timestamp) -> OwnedAnomaly {
OwnedAnomaly::new(DetectorKind::ConnectionFlood, Severity::Warning, ts)
.with_key(&SrcHost(self.src))
.with_metric("count", self.count as f64)
.with_metric("window_secs", self.window_secs)
}
}
impl crate::DetectorScore for FloodScore {
fn kind(&self) -> DetectorKind {
DetectorKind::ConnectionFlood
}
fn into_anomaly(self, ts: Timestamp) -> OwnedAnomaly {
self.into_anomaly(ts)
}
}
impl<K: KeyFields + Send> crate::detect::Detector<K> for ConnectionFloodDetector {
fn kind(&self) -> DetectorKind {
DetectorKind::ConnectionFlood
}
fn on_flow_start(&mut self, key: &K, ts: Timestamp, out: &mut Vec<OwnedAnomaly>) {
let Some(src) = key.src_ip() else {
return;
};
if let Some(score) = self.observe(src, ts) {
out.push(score.into_anomaly(ts));
}
}
fn tracked(&self) -> usize {
ConnectionFloodDetector::tracked(self)
}
fn evict_expired(&mut self, now: Timestamp) {
ConnectionFloodDetector::evict_expired(self, now);
}
}
#[cfg(test)]
mod tests {
use std::net::Ipv4Addr;
use super::*;
fn ip(last: u8) -> IpAddr {
IpAddr::from(Ipv4Addr::new(10, 0, 0, last))
}
#[test]
fn burst_over_threshold_fires_once() {
let mut d = ConnectionFloodDetector::with_params(
Duration::from_secs(10),
Duration::from_secs(1),
100,
);
let mut alerts = 0;
for i in 0..150 {
if d.observe(ip(1), Timestamp::new(0, i)).is_some() {
alerts += 1;
}
}
assert_eq!(alerts, 1, "one alert per cooldown");
}
#[test]
fn slow_rate_stays_quiet() {
let mut d = ConnectionFloodDetector::with_params(
Duration::from_secs(10),
Duration::from_secs(1),
100,
);
for s in 0..60 {
assert!(d.observe(ip(1), Timestamp::new(s, 0)).is_none());
}
}
#[test]
fn per_source_isolation() {
let mut d = ConnectionFloodDetector::with_params(
Duration::from_secs(10),
Duration::from_secs(1),
50,
);
for i in 0..60 {
d.observe(ip(1), Timestamp::new(0, i));
}
assert!(d.observe(ip(2), Timestamp::new(0, 0)).is_none());
}
#[test]
fn cooldown_then_realert_after_window() {
let mut d = ConnectionFloodDetector::with_params(
Duration::from_secs(2),
Duration::from_secs(1),
10,
)
.with_cooldown(Duration::from_secs(5));
let mut alerts = 0;
for i in 0..20 {
if d.observe(ip(1), Timestamp::new(0, i)).is_some() {
alerts += 1;
}
}
for i in 0..20 {
if d.observe(ip(1), Timestamp::new(10, i)).is_some() {
alerts += 1;
}
}
assert_eq!(alerts, 2, "re-alert after cooldown elapses");
}
}