From 3d7a8d17de57a5d794a5b2380be75da12169fdc6 Mon Sep 17 00:00:00 2001
From: Mike Dilger <mike@mikedilger.com>
Date: Tue, 20 Feb 2024 21:17:33 +1300
Subject: [PATCH] IpData, ban for longer if error bans repeat

---
 src/globals.rs | 29 ++++++++++++++++++++---------
 src/ip.rs      |  7 +++++++
 src/main.rs    | 21 ++++++++++++++++-----
 3 files changed, 43 insertions(+), 14 deletions(-)
 create mode 100644 src/ip.rs

diff --git a/src/globals.rs b/src/globals.rs
index 760108a..34780db 100644
--- a/src/globals.rs
+++ b/src/globals.rs
@@ -1,15 +1,15 @@
 use crate::config::Config;
+use crate::ip::IpData;
 use crate::store::Store;
 use crate::types::Time;
+use dashmap::DashMap;
 use hyper::server::conn::Http;
 use lazy_static::lazy_static;
-use std::collections::HashMap;
 use std::net::IpAddr;
 use std::sync::atomic::AtomicUsize;
 use std::sync::OnceLock;
 use tokio::sync::broadcast::Sender as BroadcastSender;
 use tokio::sync::watch::Sender as WatchSender;
-use tokio::sync::RwLock;
 
 pub struct Globals {
     pub config: OnceLock<Config>,
@@ -25,7 +25,8 @@ pub struct Globals {
 
     pub num_clients: AtomicUsize,
     pub shutting_down: WatchSender<bool>,
-    pub banlist: RwLock<HashMap<IpAddr, Time>>,
+
+    pub ip_data: DashMap<IpAddr, IpData>,
 }
 
 lazy_static! {
@@ -45,18 +46,28 @@ lazy_static! {
             new_events,
             num_clients: AtomicUsize::new(0),
             shutting_down,
-            banlist: RwLock::new(HashMap::new()),
+            ip_data: DashMap::new(),
         }
     };
 }
 
 impl Globals {
-    pub async fn ban(ipaddr: std::net::IpAddr, seconds: u64) {
+    pub async fn ban(ipaddr: std::net::IpAddr, seconds: u64, is_an_error_ban: bool) {
         let mut until = Time::now();
         until.0 += seconds;
-        if let Some(current_ban) = GLOBALS.banlist.read().await.get(&ipaddr) {
-            until.0 = current_ban.0.max(until.0);
-        }
-        GLOBALS.banlist.write().await.insert(ipaddr, until);
+
+        GLOBALS
+            .ip_data
+            .entry(ipaddr)
+            .and_modify(|ipdata| {
+                ipdata.ban_until = Time(ipdata.ban_until.0.max(until.0));
+                if is_an_error_ban {
+                    ipdata.number_of_error_bans += 1;
+                }
+            })
+            .or_insert(IpData {
+                ban_until: until,
+                number_of_error_bans: if is_an_error_ban { 1 } else { 0 },
+            });
     }
 }
diff --git a/src/ip.rs b/src/ip.rs
new file mode 100644
index 0000000..ef96f6c
--- /dev/null
+++ b/src/ip.rs
@@ -0,0 +1,7 @@
+use crate::types::Time;
+
+#[derive(Debug)]
+pub struct IpData {
+    pub ban_until: Time,
+    pub number_of_error_bans: usize,
+}
diff --git a/src/main.rs b/src/main.rs
index bc86a23..56db12b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,6 +3,7 @@ include!("macros.rs");
 pub mod config;
 pub mod error;
 pub mod globals;
+pub mod ip;
 pub mod nostr;
 pub mod reply;
 pub mod store;
@@ -105,10 +106,11 @@ async fn main() -> Result<(), Error> {
                 let (tcp_stream, peer_addr) = v?;
                 let ipaddr = peer_addr.ip();
 
-                if let Some(ban_until) = GLOBALS.banlist.read().await.get(&ipaddr) {
+                if let Some(ip_data) = GLOBALS.ip_data.get(&ipaddr) {
                     let now = Time::now();
-                    if *ban_until > now {
-                        log::debug!("{peer_addr}: Blocking reconnection until {ban_until}");
+                    if ip_data.ban_until > now {
+                        log::debug!("{peer_addr}: Blocking reconnection until {}",
+                                    ip_data.ban_until);
                         continue;
                     }
                 }
@@ -273,6 +275,7 @@ async fn handle_http_request(
                     // Everybody gets a 4-second ban on disconnect to prevent
                     // rapid reconnection
                     let mut ban_seconds: u64 = 4;
+                    let mut is_an_error_ban: bool = false;
 
                     // Handle the websocket
                     if let Err(e) = ws_service.handle_websocket_stream().await {
@@ -284,7 +287,15 @@ async fn handle_http_request(
                                 // No big deal, no extra ban for that.
                             }
                             ChorusError::TooManyErrors => {
-                                ban_seconds = 60;
+                                is_an_error_ban = true;
+
+                                let number_of_error_bans = match GLOBALS.ip_data.get(&peer.ip()) {
+                                    Some(ipdata) => ipdata.number_of_error_bans,
+                                    None => 0,
+                                };
+
+                                // Ban for longer if they've had error-based bans already
+                                ban_seconds = 60 + 60 * number_of_error_bans as u64;
                             }
                             _ => {
                                 log::error!("{}: {}", peer, e);
@@ -299,7 +310,7 @@ async fn handle_http_request(
                     log::info!("{}: TOTAL={}, Disconnection", peer, old_num_websockets - 1);
 
                     // Ban for the appropriate duration
-                    Globals::ban(peer.ip(), ban_seconds).await;
+                    Globals::ban(peer.ip(), ban_seconds, is_an_error_ban).await;
                 }
                 Err(e) => {
                     log::error!("{}: {}", peer, e);