refactor: [torrust#1096] extract BanService

Author: josecelano

src/servers/udp/handlers.rs

@@ -11,14 +11,14 @@ use aquatic_udp_protocol::{
     ResponsePeer, ScrapeRequest, ScrapeResponse, TorrentScrapeStatistics, TransactionId,
 };
 use bittorrent_primitives::info_hash::InfoHash;
-use bloom::{CountingBloomFilter, ASMS};
 use tokio::sync::RwLock;
 use torrust_tracker_clock::clock::Time as _;
 use tracing::{instrument, Level};
 use uuid::Uuid;
 use zerocopy::network_endian::I32;
 
 use super::connection_cookie::{check, make};
+use super::server::banning::BanService;
 use super::RawRequest;
 use crate::core::{statistics, PeersWanted, Tracker};
 use crate::servers::udp::error::Error;
@@ -53,13 +53,13 @@ impl CookieTimeValues {
 /// - Delegating the request to the correct handler depending on the request type.
 ///
 /// It will return an `Error` response if the request is invalid.
-#[instrument(fields(request_id), skip(udp_request, tracker, cookie_time_values, connection_id_errors_per_ip), ret(level = Level::TRACE))]
+#[instrument(fields(request_id), skip(udp_request, tracker, cookie_time_values, ban_service), ret(level = Level::TRACE))]
 pub(crate) async fn handle_packet(
     udp_request: RawRequest,
     tracker: &Tracker,
     local_addr: SocketAddr,
     cookie_time_values: CookieTimeValues,
-    connection_id_errors_per_ip: Arc<RwLock<CountingBloomFilter>>,
+    ban_service: Arc<RwLock<BanService>>,
 ) -> Response {
     tracing::Span::current().record("request_id", Uuid::new_v4().to_string());
     tracing::debug!("Handling Packets: {udp_request:?}");
@@ -76,10 +76,8 @@ pub(crate) async fn handle_packet(
                         | Error::CookieValueExpired { .. }
                         | Error::CookieValueFromFuture { .. } => {
                             // code-review: should we include `RequestParseError` and `BadRequest`?
-                            connection_id_errors_per_ip
-                                .write()
-                                .await
-                                .insert(&udp_request.from.ip().to_string());
+                            let mut ban_service = ban_service.write().await;
+                            ban_service.increase_counter(&udp_request.from.ip());
                         }
                         _ => {}
                     }

src/servers/udp/server/banning.rs

@@ -0,0 +1,149 @@
+use std::net::IpAddr;
+use std::time::Duration;
+
+use bloom::{CountingBloomFilter, ASMS};
+use tokio::time::Instant;
+use url::Url;
+
+use crate::servers::udp::UDP_TRACKER_LOG_TARGET;
+
+/// The maximum number of connection id errors per ip. Clients will be banned if
+/// they exceed this limit.
+pub const MAX_CONNECTION_ID_ERRORS_PER_IP: u32 = 10;
+pub const RESET_CONNECTION_ID_ERRORS_COUNTER_FREQUENCY_IN_SECS: u64 = 3600;
+
+pub struct BanService {
+    max_connection_id_errors_per_ip: u32,
+    ban_duration: Duration,
+    connection_id_errors_per_ip: CountingBloomFilter,
+    local_addr: Url,
+    last_connection_id_errors_reset: Instant,
+}
+
+impl BanService {
+    #[must_use]
+    pub fn new(max_connection_id_errors_per_ip: u32, duration_in_seconds: u64, local_addr: Url) -> Self {
+        Self {
+            max_connection_id_errors_per_ip,
+            ban_duration: Duration::from_secs(duration_in_seconds),
+            local_addr,
+            connection_id_errors_per_ip: CountingBloomFilter::with_rate(4, 0.01, 100),
+            last_connection_id_errors_reset: tokio::time::Instant::now(),
+        }
+    }
+
+    pub fn increase_counter(&mut self, ip: &IpAddr) {
+        self.connection_id_errors_per_ip.insert(&ip.to_string());
+    }
+
+    pub fn get_counter(&mut self, ip: &IpAddr) -> u32 {
+        self.connection_id_errors_per_ip.estimate_count(&ip.to_string())
+    }
+
+    /// Returns true if the given ip address is banned.
+    #[must_use]
+    pub fn is_banned(&self, ip: &IpAddr) -> bool {
+        let connection_id_errors_from_ip = self.connection_id_errors_per_ip.estimate_count(&ip.to_string());
+
+        connection_id_errors_from_ip > self.max_connection_id_errors_per_ip
+    }
+
+    pub fn run_bans_cleaner(&mut self) {
+        if self.last_connection_id_errors_reset.elapsed() >= self.ban_duration {
+            self.reset_filter();
+        }
+    }
+
+    /// Resets the filter and updates the reset timestamp.
+    pub fn reset_filter(&mut self) {
+        self.connection_id_errors_per_ip.clear();
+
+        self.last_connection_id_errors_reset = Instant::now();
+
+        let local_addr = self.local_addr.to_string();
+        tracing::info!(target: UDP_TRACKER_LOG_TARGET, local_addr, "Udp::run_udp_server::loop (connection id errors filter cleared)");
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::net::IpAddr;
+    use std::time::Duration;
+
+    use tokio::time::sleep;
+
+    use super::BanService;
+
+    /// Sample service with one day ban duration.
+    fn service_with_one_day_ban(counter_limit: u32) -> BanService {
+        let one_day_in_seconds = 86400;
+        let udp_tracker_url = "udp://127.0.0.1".parse().unwrap();
+
+        BanService::new(counter_limit, one_day_in_seconds, udp_tracker_url)
+    }
+
+    #[test]
+    fn it_should_increase_the_ip_counter() {
+        let mut ban_service = service_with_one_day_ban(1);
+
+        let ip: IpAddr = "127.0.0.2".parse().unwrap();
+
+        ban_service.increase_counter(&ip);
+
+        assert_eq!(ban_service.get_counter(&ip), 1);
+    }
+
+    #[test]
+    fn it_should_ban_ips_with_counters_exceeding_a_predefined_limit() {
+        let mut ban_service = service_with_one_day_ban(1);
+
+        let ip: IpAddr = "127.0.0.2".parse().unwrap();
+
+        ban_service.increase_counter(&ip); // Counter = 1
+        ban_service.increase_counter(&ip); // Counter = 2
+
+        assert!(ban_service.is_banned(&ip));
+    }
+
+    #[test]
+    fn it_should_not_ban_ips_whose_counters_do_not_exceed_the_predefined_limit() {
+        let mut ban_service = service_with_one_day_ban(1);
+
+        let ip: IpAddr = "127.0.0.2".parse().unwrap();
+
+        ban_service.increase_counter(&ip);
+
+        assert!(!ban_service.is_banned(&ip));
+    }
+
+    #[test]
+    fn it_should_allow_resetting_all_the_counters() {
+        let mut ban_service = service_with_one_day_ban(1);
+
+        let ip: IpAddr = "127.0.0.2".parse().unwrap();
+
+        ban_service.increase_counter(&ip); // Counter = 1
+
+        ban_service.reset_filter();
+
+        assert_eq!(ban_service.get_counter(&ip), 0);
+    }
+
+    #[tokio::test]
+    async fn it_should_allow_run_a_bans_cleaner_to_reset_the_counters_periodically() {
+        let udp_tracker_url = "udp://127.0.0.1".parse().unwrap();
+        let ban_duration_in_secs = 1;
+
+        let mut ban_service = BanService::new(1, ban_duration_in_secs, udp_tracker_url);
+
+        let ip: IpAddr = "127.0.0.2".parse().unwrap();
+
+        ban_service.increase_counter(&ip); // Counter = 1
+
+        sleep(Duration::from_secs(2)).await;
+
+        ban_service.run_bans_cleaner();
+
+        assert_eq!(ban_service.get_counter(&ip), 0);
+    }
+}

src/servers/udp/server/launcher.rs

@@ -3,13 +3,13 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use bittorrent_tracker_client::udp::client::check;
-use bloom::{CountingBloomFilter, ASMS};
 use derive_more::Constructor;
 use futures_util::StreamExt;
 use tokio::select;
 use tokio::sync::{oneshot, RwLock};
 use tracing::instrument;
 
+use super::banning::{BanService, MAX_CONNECTION_ID_ERRORS_PER_IP, RESET_CONNECTION_ID_ERRORS_COUNTER_FREQUENCY_IN_SECS};
 use super::request_buffer::ActiveRequests;
 use crate::bootstrap::jobs::Started;
 use crate::core::{statistics, Tracker};
@@ -21,11 +21,6 @@ use crate::servers::udp::server::processor::Processor;
 use crate::servers::udp::server::receiver::Receiver;
 use crate::servers::udp::UDP_TRACKER_LOG_TARGET;
 
-/// The maximum number of connection id errors per ip. Clients will be banned if
-/// they exceed this limit.
-const MAX_CONNECTION_ID_ERRORS_PER_IP: u32 = 10;
-const RESET_CONNECTION_ID_ERRORS_COUNTER_FREQUENCY_IN_SECS: u64 = 3600;
-
 /// A UDP server instance launcher.
 #[derive(Constructor)]
 pub struct Launcher;
@@ -120,27 +115,21 @@ impl Launcher {
     async fn run_udp_server_main(mut receiver: Receiver, tracker: Arc<Tracker>, cookie_lifetime: Duration) {
         let active_requests = &mut ActiveRequests::default();
 
-        // Create a counting bloom filter that uses 4 bits per element and has a
-        // false positive rate of 0.01 when 100 items have been inserted
-        let connection_id_errors_per_ip = Arc::new(RwLock::new(CountingBloomFilter::with_rate(4, 0.01, 100)));
-
-        // Timer to track when to clear the filter
-        let mut last_connection_id_errors_reset = tokio::time::Instant::now();
-
         let addr = receiver.bound_socket_address();
 
         let local_addr = format!("udp://{addr}");
 
         let cookie_lifetime = cookie_lifetime.as_secs_f64();
 
+        let ban_service = Arc::new(RwLock::new(BanService::new(
+            MAX_CONNECTION_ID_ERRORS_PER_IP,
+            RESET_CONNECTION_ID_ERRORS_COUNTER_FREQUENCY_IN_SECS,
+            local_addr.parse().unwrap(),
+        )));
+
         loop {
-            if last_connection_id_errors_reset.elapsed()
-                >= Duration::from_secs(RESET_CONNECTION_ID_ERRORS_COUNTER_FREQUENCY_IN_SECS)
-            {
-                connection_id_errors_per_ip.write().await.clear();
-                tracing::info!(target: UDP_TRACKER_LOG_TARGET, local_addr, "Udp::run_udp_server::loop (connection id errors filter cleared)");
-                last_connection_id_errors_reset = tokio::time::Instant::now();
-            }
+            // code-review: the ban service could spawn a task to clear the bans.
+            ban_service.write().await.run_bans_cleaner();
 
             let processor = Processor::new(receiver.socket.clone(), tracker.clone(), cookie_lifetime);
 
@@ -171,12 +160,7 @@ impl Launcher {
                     }
                 }
 
-                let connection_id_errors_from_ip = connection_id_errors_per_ip
-                    .read()
-                    .await
-                    .estimate_count(&req.from.ip().to_string());
-
-                if connection_id_errors_from_ip > MAX_CONNECTION_ID_ERRORS_PER_IP {
+                if ban_service.read().await.is_banned(&req.from.ip()) {
                     tracing::debug!(target: UDP_TRACKER_LOG_TARGET, local_addr,  "Udp::run_udp_server::loop continue: (banned ip)");
                     continue;
                 }
@@ -193,7 +177,7 @@ impl Launcher {
                 // chance to finish. However, the buffer is yielding before
                 // aborting one tasks, giving it the chance to finish.
                 let abort_handle: tokio::task::AbortHandle =
-                    tokio::task::spawn(processor.process_request(req, connection_id_errors_per_ip.clone())).abort_handle();
+                    tokio::task::spawn(processor.process_request(req, ban_service.clone())).abort_handle();
 
                 if abort_handle.is_finished() {
                     continue;

src/servers/udp/server/mod.rs

@@ -6,6 +6,7 @@ use thiserror::Error;
 
 use super::RawRequest;
 
+pub mod banning;
 pub mod bound_socket;
 pub mod launcher;
 pub mod processor;

src/servers/udp/server/processor.rs

@@ -3,10 +3,10 @@ use std::net::{IpAddr, SocketAddr};
 use std::sync::Arc;
 
 use aquatic_udp_protocol::Response;
-use bloom::CountingBloomFilter;
 use tokio::sync::RwLock;
 use tracing::{instrument, Level};
 
+use super::banning::BanService;
 use super::bound_socket::BoundSocket;
 use crate::core::{statistics, Tracker};
 use crate::servers::udp::handlers::CookieTimeValues;
@@ -27,15 +27,15 @@ impl Processor {
         }
     }
 
-    #[instrument(skip(self, request, connection_id_errors_per_ip))]
-    pub async fn process_request(self, request: RawRequest, connection_id_errors_per_ip: Arc<RwLock<CountingBloomFilter>>) {
+    #[instrument(skip(self, request, ban_service))]
+    pub async fn process_request(self, request: RawRequest, ban_service: Arc<RwLock<BanService>>) {
         let from = request.from;
         let response = handlers::handle_packet(
             request,
             &self.tracker,
             self.socket.address(),
             CookieTimeValues::new(self.cookie_lifetime),
-            connection_id_errors_per_ip,
+            ban_service,
         )
         .await;