Improve banning

src/globals.rs

@@ -1,7 +1,6 @@
 use crate::config::Config;
-use crate::ip::IpData;
+use crate::ip::{Ban, IpData};
 use crate::store::Store;
-use crate::types::Time;
 use dashmap::DashMap;
 use hyper::server::conn::Http;
 use lazy_static::lazy_static;
@@ -52,22 +51,11 @@ lazy_static! {
 }
 
 impl Globals {
-    pub async fn ban(ipaddr: std::net::IpAddr, seconds: u64, is_an_error_ban: bool) {
-        let mut until = Time::now();
-        until.0 += seconds;
-
+    pub async fn ban(ipaddr: std::net::IpAddr, bankind: Ban) {
         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 },
-            });
+            .and_modify(|ipdata| ipdata.ban(bankind))
+            .or_insert(IpData::new(bankind));
     }
 }

src/ip.rs

@@ -1,7 +1,57 @@
 use crate::types::Time;
 
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum Ban {
+    General,
+    ErrorExit,
+    TooManyErrors,
+    Timeout,
+}
+
 #[derive(Debug)]
 pub struct IpData {
     pub ban_until: Time,
-    pub number_of_error_bans: usize,
+    pub number_of_error_exits: u64,
+    pub number_of_too_many_error_bans: u64,
+    pub number_of_timeouts: u64,
+}
+
+impl IpData {
+    pub fn new(ban: Ban) -> IpData {
+        let mut ipdata = IpData {
+            ban_until: Time::now(),
+            number_of_error_exits: 0,
+            number_of_too_many_error_bans: 0,
+            number_of_timeouts: 0,
+        };
+
+        ipdata.ban(ban);
+
+        ipdata
+    }
+
+    pub fn ban(&mut self, ban: Ban) {
+        // Update numbers
+        match ban {
+            Ban::ErrorExit => self.number_of_error_exits += 1,
+            Ban::TooManyErrors => self.number_of_too_many_error_bans += 1,
+            Ban::Timeout => self.number_of_timeouts += 1,
+            _ => (),
+        };
+
+        // Compute ban_until
+        let mut until = Time::now();
+        until.0 += self.ban_seconds(ban);
+
+        self.ban_until = Time(self.ban_until.0.max(until.0));
+    }
+
+    fn ban_seconds(&self, thisban: Ban) -> u64 {
+        match thisban {
+            Ban::General => 3,
+            Ban::ErrorExit => 3 + self.number_of_error_exits * 10,
+            Ban::TooManyErrors => 3 + self.number_of_too_many_error_bans * 15,
+            Ban::Timeout => 3 + self.number_of_timeouts * 5,
+        }
+    }
 }

src/main.rs

@@ -14,6 +14,7 @@ pub mod web;
 use crate::config::{Config, FriendlyConfig};
 use crate::error::{ChorusError, Error};
 use crate::globals::{Globals, GLOBALS};
+use crate::ip::Ban;
 use crate::reply::NostrReply;
 use crate::store::Store;
 use crate::tls::MaybeTlsStream;
@@ -274,8 +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;
+                    let mut bankind: Ban = Ban::General;
                     let mut msg = "Closed";
 
                     // Handle the websocket
@@ -289,35 +289,17 @@ async fn handle_http_request(
                                 msg = "Reset";
                             }
                             ChorusError::TooManyErrors => {
-                                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;
-
-                                msg = "Errored Out, temporarily banned (long and growing)";
+                                bankind = Ban::TooManyErrors;
+                                msg = "Errored Out (too many)";
                             }
                             ChorusError::TimedOut => {
-                                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;
-
-                                msg = "Timed Out, temporarily banned (long and growing)";
+                                bankind = Ban::Timeout;
+                                msg = "Timed Out (with no subscriptions)";
                             }
                             _ => {
                                 log::error!("{}: {}", peer, e);
-                                ban_seconds = 15;
-                                msg = "Errored, temporarily banned (short, fixed)";
+                                bankind = Ban::ErrorExit;
+                                msg = "Error Exited";
                             }
                         }
                     }
@@ -328,7 +310,7 @@ async fn handle_http_request(
                     log::info!("{}: TOTAL={}, {}", peer, old_num_websockets - 1, msg);
 
                     // Ban for the appropriate duration
-                    Globals::ban(peer.ip(), ban_seconds, is_an_error_ban).await;
+                    Globals::ban(peer.ip(), bankind).await;
                 }
                 Err(e) => {
                     log::error!("{}: {}", peer, e);