const BAD_THRESHOLD: u16 = 1000;
const BAD_STRIKE_MUL: u16 = 4;
pub struct BadnessFilter {
addresses: std::collections::HashMap<core::net::IpAddr, u16>,
is_enabled: bool,
}
impl BadnessFilter {
pub fn new(enable: bool) -> Self {
Self {
addresses: std::collections::HashMap::new(),
is_enabled: enable,
}
}
pub fn is_bad(&self, addr: core::net::SocketAddr) -> bool {
if !self.is_enabled { return false; }
if let Some(badness) = self.addresses.get(&addr.ip()) {
*badness > BAD_THRESHOLD
} else {
false
}
}
pub fn tick(&mut self) {
if !self.is_enabled { return; }
for val in self.addresses.values_mut() {
*val = val.saturating_sub(1);
}
}
pub fn strike(&mut self, addr: core::net::SocketAddr) {
if !self.is_enabled { return; }
if let Some(badness) = self.addresses.get_mut(&addr.ip()) {
*badness = badness.saturating_mul(BAD_STRIKE_MUL);
} else {
self.addresses.insert(addr.ip(), BAD_STRIKE_MUL);
}
}
pub fn reset(&mut self, addr: core::net::SocketAddr) {
if !self.is_enabled { return; }
if let Some(badness) = self.addresses.get_mut(&addr.ip()) {
*badness = 1;
}
}
}