From ce6ca78087790b987d99d38f2cd8b71617a8d6b0 Mon Sep 17 00:00:00 2001
From: Paul Lecuq <paul@paulbsd.com>
Date: Tue, 26 Dec 2023 13:13:30 +0100
Subject: [PATCH] added safety in ipblc

---
 src/config.rs    | 50 +++++++++++++++++++++++++-----------------------
 src/ipblc.rs     | 17 ++++++++--------
 src/utils.rs     |  5 +++++
 src/websocket.rs | 12 +++++-------
 4 files changed, 45 insertions(+), 39 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 886b90e..1772b8f 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -169,39 +169,41 @@ impl Context {
     }
 
     pub async fn update_blocklist(&mut self, ipevent: &IpEvent) -> Option<IpEvent> {
-        let ipdata = &ipevent.ipdata.clone().unwrap();
-        match self.cfg.sets.get(&ipdata.src) {
-            Some(set) => {
-                let starttime = DateTime::parse_from_rfc3339(ipdata.date.as_str())
-                    .unwrap()
-                    .with_timezone(&chrono::Local);
-                let blocktime = set.blocktime;
-                if ipevent.mode == "file".to_string() && gethostname(true) == ipevent.hostname {
-                    let block =
+        match &ipevent.ipdata {
+            Some(ipdata) => match self.cfg.sets.get(&ipdata.src) {
+                Some(set) => {
+                    let starttime = DateTime::parse_from_rfc3339(ipdata.date.as_str())
+                        .unwrap()
+                        .with_timezone(&chrono::Local);
+                    let blocktime = set.blocktime;
+                    if ipevent.mode == "file".to_string() && gethostname(true) == ipevent.hostname {
+                        let block =
+                            self.blocklist
+                                .entry(ipdata.ip.to_string())
+                                .or_insert(BlockIpData {
+                                    ipdata: ipdata.clone(),
+                                    tryfail: 0,
+                                    starttime,
+                                    blocktime,
+                                });
+                        block.tryfail += 1;
+                        block.blocktime = blocktime;
+                        if block.tryfail >= set.tryfail {
+                            return Some(ipevent.clone());
+                        }
+                    } else {
                         self.blocklist
                             .entry(ipdata.ip.to_string())
                             .or_insert(BlockIpData {
                                 ipdata: ipdata.clone(),
-                                tryfail: 0,
+                                tryfail: set.tryfail,
                                 starttime,
                                 blocktime,
                             });
-                    block.tryfail += 1;
-                    block.blocktime = blocktime;
-                    if block.tryfail >= set.tryfail {
-                        return Some(ipevent.clone());
                     }
-                } else {
-                    self.blocklist
-                        .entry(ipdata.ip.to_string())
-                        .or_insert(BlockIpData {
-                            ipdata: ipdata.clone(),
-                            tryfail: set.tryfail,
-                            starttime,
-                            blocktime,
-                        });
                 }
-            }
+                None => {}
+            },
             None => {}
         }
         None
diff --git a/src/ipblc.rs b/src/ipblc.rs
index 0be1cd1..b5de02d 100644
--- a/src/ipblc.rs
+++ b/src/ipblc.rs
@@ -65,13 +65,6 @@ pub async fn run() {
 
         let ctxclone = Arc::clone(&ctxarc);
 
-        let ipe = ipevent!("ping", "ws", gethostname(true));
-        if !send_to_ipbl_websocket(&mut wssocketrr, &ipe).await {
-            wssocketrr.close(None).unwrap();
-            wssocketrr = websocketreqrep(&ctxwsrr).await;
-            continue;
-        }
-
         tokio::select! {
             ipevent = ipeventrx.recv() => {
                 let received_ip = ipevent.unwrap();
@@ -85,6 +78,7 @@ pub async fn run() {
                     for ip_to_send in toblock {
                         let ipe = ipevent!("init","ws",gethostname(true),Some(ip_to_send));
                         if !send_to_ipbl_websocket(&mut wssocketrr, &ipe).await {
+                            wssocketrr.close(None).unwrap();
                             wssocketrr = websocketreqrep(&ctxwsrr).await;
                             break;
                         }
@@ -104,13 +98,20 @@ pub async fn run() {
                         let ipe = ipevent!("add","ws",gethostname(true),ipevent.ipdata);
                         send_to_ipbl_api(&server.clone(), &ipe).await;
                         if !send_to_ipbl_websocket(&mut wssocketrr, &ipe).await {
+                            wssocketrr.close(None).unwrap();
                             wssocketrr = websocketreqrep(&ctxwsrr).await;
                             continue;
                         }
                     }
                 }
             }
-            _val = sleep_s(LOOP_MAX_WAIT) => {}
+            _val = sleep_s(LOOP_MAX_WAIT) => {
+                let ipe = ipevent!("ping", "ws", gethostname(true));
+                if !send_to_ipbl_websocket(&mut wssocketrr, &ipe).await {
+                    wssocketrr.close(None).unwrap();
+                    wssocketrr = websocketreqrep(&ctxwsrr).await;
+                }
+            }
         };
 
         let ctxclone = Arc::clone(&ctxarc);
diff --git a/src/utils.rs b/src/utils.rs
index fd04fb2..3c4bb40 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -21,6 +21,11 @@ pub async fn sleep_s(s: u64) {
     sleep(Duration::from_secs(s)).await;
 }
 
+#[allow(dead_code)]
+pub async fn sleep_ms(m: u64) {
+    sleep(Duration::from_millis(m)).await;
+}
+
 pub fn gethostname(show_fqdn: bool) -> String {
     let hostname_cstr = unistd::gethostname().expect("Failed getting hostname");
     let fqdn = hostname_cstr
diff --git a/src/websocket.rs b/src/websocket.rs
index d23aab7..436239d 100644
--- a/src/websocket.rs
+++ b/src/websocket.rs
@@ -59,7 +59,10 @@ pub async fn websocketpubsub(
                                 txps.send(tosend).await.unwrap();
                             }
                         }
-                        None => {}
+                        None => {
+                            let txps = txpubsub.read().await;
+                            txps.send(tosend.clone()).await.unwrap();
+                        }
                     }
                 }
                 Err(e) => {
@@ -111,14 +114,11 @@ pub async fn send_to_ipbl_websocket(
             Ok(_) => {}
             Err(e) => {
                 println!("err send read: {e:?}");
-                ws.close(None).unwrap_or(());
                 return false;
             }
         };
     } else {
         println!("can't write to socket");
-        ws.close(None).unwrap_or(());
-        sleep_s(1);
         return false;
     };
 
@@ -127,15 +127,13 @@ pub async fn send_to_ipbl_websocket(
             Ok(_) => {}
             Err(e) => {
                 println!("err send read: {e:?}");
-                ws.close(None).unwrap_or(());
                 return false;
             }
         };
     } else {
         println!("can't read from socket");
-        sleep_s(1);
-        ws.close(None).unwrap_or(());
         return false;
     };
+
     true
 }