130 lines
4.0 KiB
Rust
130 lines
4.0 KiB
Rust
use std::collections::hash_map::Entry;
|
|
use std::collections::HashMap;
|
|
use std::net::IpAddr;
|
|
|
|
use actix::{Actor, AsyncContext, Context, Handler, Message};
|
|
|
|
use crate::constants::{FAIL_LOGIN_ATTEMPT_CLEANUP_INTERVAL, KEEP_FAILED_LOGIN_ATTEMPTS_FOR};
|
|
use crate::utils::time::time;
|
|
|
|
#[derive(Message)]
|
|
#[rtype(result = "()")]
|
|
pub struct RecordFailedAttempt {
|
|
pub ip: IpAddr,
|
|
}
|
|
|
|
#[derive(Message)]
|
|
#[rtype(result = "usize")]
|
|
pub struct CountFailedAttempt {
|
|
pub ip: IpAddr,
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
pub struct BruteForceActor {
|
|
failed_attempts: HashMap<IpAddr, Vec<u64>>,
|
|
}
|
|
|
|
impl BruteForceActor {
|
|
pub fn clean_attempts(&mut self) {
|
|
#[allow(clippy::map_clone)]
|
|
let keys = self.failed_attempts.keys().map(|i| *i).collect::<Vec<_>>();
|
|
|
|
for ip in keys {
|
|
// Remove old attempts
|
|
let attempts = self.failed_attempts.get_mut(&ip).unwrap();
|
|
attempts.retain(|i| i + KEEP_FAILED_LOGIN_ATTEMPTS_FOR > time());
|
|
|
|
// Remove empty entry keys
|
|
if attempts.is_empty() {
|
|
self.failed_attempts.remove(&ip);
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn insert_failed_attempt(&mut self, ip: IpAddr) {
|
|
if let Entry::Vacant(e) = self.failed_attempts.entry(ip) {
|
|
e.insert(vec![time()]);
|
|
} else {
|
|
self.failed_attempts.get_mut(&ip).unwrap().push(time());
|
|
}
|
|
}
|
|
|
|
pub fn count_failed_attempts(&mut self, ip: &IpAddr) -> usize {
|
|
self.failed_attempts.get(ip).map(Vec::len).unwrap_or(0)
|
|
}
|
|
}
|
|
|
|
impl Actor for BruteForceActor {
|
|
type Context = Context<Self>;
|
|
|
|
fn started(&mut self, ctx: &mut Self::Context) {
|
|
// Clean up at a regular interval failed attempts
|
|
ctx.run_interval(FAIL_LOGIN_ATTEMPT_CLEANUP_INTERVAL, |act, _ctx| {
|
|
log::trace!("Cleaning up failed login attempts");
|
|
act.clean_attempts();
|
|
});
|
|
}
|
|
}
|
|
|
|
impl Handler<RecordFailedAttempt> for BruteForceActor {
|
|
type Result = ();
|
|
|
|
fn handle(&mut self, attempt: RecordFailedAttempt, _ctx: &mut Self::Context) -> Self::Result {
|
|
self.insert_failed_attempt(attempt.ip)
|
|
}
|
|
}
|
|
|
|
impl Handler<CountFailedAttempt> for BruteForceActor {
|
|
type Result = usize;
|
|
|
|
fn handle(&mut self, attempt: CountFailedAttempt, _ctx: &mut Self::Context) -> Self::Result {
|
|
self.count_failed_attempts(&attempt.ip)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod test {
|
|
use std::net::{IpAddr, Ipv4Addr};
|
|
|
|
use crate::actors::bruteforce_actor::BruteForceActor;
|
|
use crate::utils::time::time;
|
|
|
|
const IP_1: IpAddr = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1));
|
|
const IP_2: IpAddr = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 2));
|
|
const IP_3: IpAddr = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 5));
|
|
|
|
#[test]
|
|
fn test_clean() {
|
|
let mut actor = BruteForceActor::default();
|
|
actor.failed_attempts.insert(IP_1, vec![1, 10]);
|
|
actor.failed_attempts.insert(IP_2, vec![1, 10, time() + 10]);
|
|
actor
|
|
.failed_attempts
|
|
.insert(IP_3, vec![time() + 10, time() + 20]);
|
|
|
|
actor.clean_attempts();
|
|
|
|
let keys = actor.failed_attempts.keys().collect::<Vec<_>>();
|
|
assert_eq!(keys.len(), 2);
|
|
assert_eq!(actor.count_failed_attempts(&IP_1), 0);
|
|
assert_eq!(actor.count_failed_attempts(&IP_2), 1);
|
|
assert_eq!(actor.count_failed_attempts(&IP_3), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_insert_and_count() {
|
|
let mut actor = BruteForceActor::default();
|
|
assert_eq!(actor.count_failed_attempts(&IP_1), 0);
|
|
assert_eq!(actor.count_failed_attempts(&IP_2), 0);
|
|
actor.insert_failed_attempt(IP_1);
|
|
assert_eq!(actor.count_failed_attempts(&IP_1), 1);
|
|
assert_eq!(actor.count_failed_attempts(&IP_2), 0);
|
|
actor.insert_failed_attempt(IP_1);
|
|
assert_eq!(actor.count_failed_attempts(&IP_1), 2);
|
|
assert_eq!(actor.count_failed_attempts(&IP_2), 0);
|
|
actor.insert_failed_attempt(IP_2);
|
|
assert_eq!(actor.count_failed_attempts(&IP_1), 2);
|
|
assert_eq!(actor.count_failed_attempts(&IP_2), 1);
|
|
}
|
|
}
|