From b76ea99b60f9198be8d872549d63809694844156 Mon Sep 17 00:00:00 2001
From: Mike Dilger <mike@mikedilger.com>
Date: Mon, 19 Feb 2024 12:36:07 +1300
Subject: [PATCH] IP banning

---
 src/globals.rs |  5 +++++
 src/main.rs    | 20 +++++++++++++++++++-
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/globals.rs b/src/globals.rs
index 77bcb68..832a68e 100644
--- a/src/globals.rs
+++ b/src/globals.rs
@@ -1,7 +1,10 @@
 use crate::config::{Config, FriendlyConfig};
 use crate::store::Store;
+use crate::types::Time;
 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;
@@ -22,6 +25,7 @@ pub struct Globals {
 
     pub num_clients: AtomicUsize,
     pub shutting_down: WatchSender<bool>,
+    pub banlist: RwLock<HashMap<IpAddr, Time>>,
 }
 
 lazy_static! {
@@ -41,6 +45,7 @@ lazy_static! {
             new_events,
             num_clients: AtomicUsize::new(0),
             shutting_down,
+            banlist: RwLock::new(HashMap::new()),
         }
     };
 }
diff --git a/src/main.rs b/src/main.rs
index df0f51c..0a66b05 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -16,7 +16,7 @@ use crate::globals::GLOBALS;
 use crate::reply::NostrReply;
 use crate::store::Store;
 use crate::tls::MaybeTlsStream;
-use crate::types::{OwnedFilter, Pubkey};
+use crate::types::{OwnedFilter, Pubkey, Time};
 use futures::{sink::SinkExt, stream::StreamExt};
 use hyper::service::Service;
 use hyper::upgrade::Upgraded;
@@ -102,6 +102,15 @@ async fn main() -> Result<(), Error> {
             // Accepts network connections and spawn a task to serve each one
             v = listener.accept() => {
                 let (tcp_stream, peer_addr) = v?;
+                let ipaddr = peer_addr.ip();
+
+                if let Some(ban_until) = GLOBALS.banlist.read().await.get(&ipaddr) {
+                    let now = Time::now();
+                    if *ban_until > now {
+                        log::debug!("{peer_addr}: Blocking reconnection until {ban_until}");
+                        continue;
+                    }
+                }
 
                 if let Some(tls_acceptor) = &maybe_tls_acceptor {
                     let tls_acceptor_clone = tls_acceptor.clone();
@@ -275,6 +284,15 @@ async fn handle_http_request(
                         peer,
                         old_num_websockets - 1
                     );
+
+                    // Everybody gets a 4 second ban on disconnect, to prevent rapid reconnection
+                    let ipaddr = peer.ip();
+                    let mut until = Time::now();
+                    until.0 += 4;
+                    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);
                 }
                 Err(e) => {
                     log::error!("{}", e);