From f4441376847612da28d842ad21ea8bd29519a8a2 Mon Sep 17 00:00:00 2001
From: Paul Lecuq <paul@paulbsd.com>
Date: Sat, 22 Apr 2023 19:28:38 +0200
Subject: [PATCH] misc changes

---
 src/config.rs    | 13 ++++++++++++-
 src/ipblc.rs     |  4 ++--
 src/websocket.rs | 39 +++++++++++++++++++++++++--------------
 3 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 82b4047..c9d3b41 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -17,7 +17,7 @@ use std::path::Path;
 pub const GIT_VERSION: &str = git_version!();
 const MASTERSERVER: &str = "ipbl.paulbsd.com";
 const WSSUBSCRIPTION: &str = "ipbl";
-const CONFIG_RETRY: u64 = 10;
+const CONFIG_RETRY: u64 = 5;
 
 #[derive(Debug, Clone)]
 pub struct Context {
@@ -131,18 +131,29 @@ impl Context {
         if cfg!(test) {
             return Ok(());
         }
+        let mut last_in_err = false;
         loop {
             match self.cfg.load(self.to_owned()).await {
                 Ok(()) => {
+                    if last_in_err {
+                        println!("loaded config");
+                    }
                     break;
                 }
                 Err(err) => {
                     println!("error loading config: {err}, retrying in {CONFIG_RETRY} secs");
+                    last_in_err = true;
                     sleep_s(CONFIG_RETRY).await;
                 }
             };
         }
+        if last_in_err {
+            println!("creating sas");
+        }
         self.create_sas().await?;
+        if last_in_err {
+            println!("created sas");
+        }
         Ok(())
     }
 
diff --git a/src/ipblc.rs b/src/ipblc.rs
index 5d41582..1a66860 100644
--- a/src/ipblc.rs
+++ b/src/ipblc.rs
@@ -74,11 +74,11 @@ pub async fn run() {
                             ipdata: ip_to_send,
                         };
                         if !send_to_ipbl_websocket(&mut wssocketrr, &ipe).await {
-                            drop(ctx);
                             wssocketrr = websocketreqrep(&ctxwsrr).await;
                             break;
                         }
                     }
+                    drop(ctx);
 
                     continue
                 }
@@ -98,8 +98,8 @@ pub async fn run() {
                         };
                         send_to_ipbl_api(&ctx.client, &ctx.flags.server, &ipe).await;
                         let status = send_to_ipbl_websocket(&mut wssocketrr, &ipe).await;
+                        drop(ctx);
                         if !status {
-                            drop(ctx);
                             wssocketrr = websocketreqrep(&ctxwsrr).await;
                             continue;
                         }
diff --git a/src/websocket.rs b/src/websocket.rs
index 3a9876f..98dae57 100644
--- a/src/websocket.rs
+++ b/src/websocket.rs
@@ -102,22 +102,33 @@ pub async fn send_to_ipbl_websocket(
 ) -> bool {
     let msg = format!("{val}", val = serde_json::to_string(&ip).unwrap());
 
-    match ws.write_message(Message::Text(msg)) {
-        Ok(_) => {}
-        Err(e) => {
-            println!("err send write: {e:?}");
-            ws.close(None).unwrap();
-            return false;
-        }
+    if ws.can_write() {
+        match ws.write_message(Message::Text(msg)) {
+            Ok(_) => {}
+            Err(e) => {
+                println!("err send read: {e:?}");
+                return handle_websocket_error(ws);
+            }
+        };
+    } else {
+        return handle_websocket_error(ws);
     };
 
-    match ws.read_message() {
-        Ok(_) => {}
-        Err(e) => {
-            println!("err send read: {e:?}");
-            ws.close(None).unwrap();
-            return false;
-        }
+    if ws.can_read() {
+        match ws.read_message() {
+            Ok(_) => {}
+            Err(e) => {
+                println!("err send read: {e:?}");
+                return handle_websocket_error(ws);
+            }
+        };
+    } else {
+        return handle_websocket_error(ws);
     };
     true
 }
+
+fn handle_websocket_error(ws: &mut WebSocket<MaybeTlsStream<TcpStream>>) -> bool {
+    ws.close(None).unwrap();
+    return false;
+}