From a720562c3c3cca9c2b2d5cbe748e68c3b4e2c9c0 Mon Sep 17 00:00:00 2001 From: Paul Lecuq Date: Wed, 10 May 2023 21:32:27 +0200 Subject: [PATCH] some code refactor --- src/config.rs | 5 +---- src/ip.rs | 4 ++-- src/ipblc.rs | 20 +++++++++++++------- src/main.rs | 1 - src/webservice.rs | 6 +++--- src/websocket.rs | 36 ++++++++++++++++++------------------ 6 files changed, 37 insertions(+), 35 deletions(-) diff --git a/src/config.rs b/src/config.rs index 2a9d758..3f58aff 100644 --- a/src/config.rs +++ b/src/config.rs @@ -26,7 +26,6 @@ pub struct Context { pub cfg: Config, pub discovery: Discovery, pub flags: Flags, - pub hostname: String, pub instance: Box, pub sas: HashMap, pub hashwd: HashMap, @@ -52,7 +51,6 @@ impl Context { pub async fn new() -> Self { // Get flags let argp: ArgMatches = Context::argparse(); - //let debug: bool = argp.contains_id("debug"); let debug: bool = argp.get_one::("debug").unwrap().to_owned(); let server: String = argp.get_one::("server").unwrap().to_string(); @@ -60,7 +58,6 @@ impl Context { let mut ctx = Context { cfg: Config::new(), flags: Flags { debug, server }, - hostname: gethostname(true), discovery: Discovery { version: "1.0".to_string(), urls: HashMap::new(), @@ -178,7 +175,7 @@ impl Context { .unwrap() .with_timezone(&chrono::Local); let blocktime = set.blocktime; - if ipevent.mode == "file".to_string() && self.hostname == ipevent.hostname { + if ipevent.mode == "file".to_string() && gethostname(true) == ipevent.hostname { let block = self .blocklist .entry(ipevent.ipdata.ip.to_string()) diff --git a/src/ip.rs b/src/ip.rs index 3febd99..acfe6c3 100644 --- a/src/ip.rs +++ b/src/ip.rs @@ -78,7 +78,7 @@ impl Display for IpData { pub fn filter( lines: Box, - list: &mut Vec, + iplist: &mut Vec, trustnets: &Vec, regex: &Regex, src: &String, @@ -129,7 +129,7 @@ pub fn filter( }; if !is_trusted(&ipaddr, &trustnets) { - list.push(IpData { + iplist.push(IpData { ip: s_ipaddr, src: src.to_owned(), date: s_date.to_rfc3339().to_owned(), diff --git a/src/ipblc.rs b/src/ipblc.rs index 0fe28a4..45deb9f 100644 --- a/src/ipblc.rs +++ b/src/ipblc.rs @@ -64,10 +64,15 @@ pub async fn run() { ipevent = ipeventrx.recv() => { let received_ip = ipevent.unwrap(); - let mut ctx = ctxclone.write().await; + let (toblock,server); + { + let mut ctx = ctxclone.write().await; + toblock = ctx.get_blocklist_toblock().await; + server = ctx.flags.server.clone(); + } if received_ip.msgtype == "bootstrap".to_string() { - for ip_to_send in ctx.get_blocklist_toblock().await { + for ip_to_send in toblock { let ipe = IpEvent{ msgtype: String::from("init"), mode: String::from("ws"), @@ -75,17 +80,19 @@ pub async fn run() { ipdata: ip_to_send, }; if !send_to_ipbl_websocket(&mut wssocketrr, &ipe).await { - drop(ctx); wssocketrr = websocketreqrep(&ctxwsrr).await; break; } } - continue } // refresh context blocklist - let filtered_ipevent = ctx.update_blocklist(&received_ip).await; + let filtered_ipevent; + { + let mut ctx = ctxarc.write().await; + filtered_ipevent = ctx.update_blocklist(&received_ip).await; + } // send ip list to api and ws sockets if let Some(ipevent) = filtered_ipevent { @@ -97,9 +104,8 @@ pub async fn run() { hostname: gethostname(true), ipdata: ipevent.ipdata, }; - send_to_ipbl_api(&ctx.flags.server, &ipe).await; + send_to_ipbl_api(&server.clone(), &ipe).await; let status = send_to_ipbl_websocket(&mut wssocketrr, &ipe).await; - drop(ctx); if !status { wssocketrr = websocketreqrep(&ctxwsrr).await; continue; diff --git a/src/main.rs b/src/main.rs index 9d2566f..9b36929 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9,6 +9,5 @@ mod websocket; #[tokio::main] pub async fn main() { - // Create a new context ipblc::run().await; } diff --git a/src/webservice.rs b/src/webservice.rs index 5299078..48891f8 100644 --- a/src/webservice.rs +++ b/src/webservice.rs @@ -8,7 +8,7 @@ use reqwest::Error as ReqError; const MAX_FAILED_API_RATE: u64 = 10; pub async fn send_to_ipbl_api(server: &str, ip: &IpEvent) { - let mut i = 0; + let mut try_req = 0; let client = httpclient(); loop { match push_ip(&client, &server, &ip.ipdata).await { @@ -18,10 +18,10 @@ pub async fn send_to_ipbl_api(server: &str, ip: &IpEvent) { Err(err) => { println!("{err}"); sleep_s(1).await; - if i == MAX_FAILED_API_RATE { + if try_req == MAX_FAILED_API_RATE { break; } - i += 1; + try_req += 1; } }; } diff --git a/src/websocket.rs b/src/websocket.rs index 98dae57..596b3f3 100644 --- a/src/websocket.rs +++ b/src/websocket.rs @@ -14,16 +14,14 @@ use tungstenite::*; pub async fn websocketreqrep( ctxarc: &Arc>, ) -> WebSocket> { - let (mut wssocketrr, bootstrap_event); + let (mut wssocketrr, bootstrap_event, cfg); { let ctx = ctxarc.read().await; - bootstrap_event = ctxarc.read().await.cfg.bootstrap_event().clone(); - - wssocketrr = websocketconnect(&ctx.cfg.ws.get("reqrep").unwrap(), ctx.hostname.clone()) - .await - .unwrap(); - send_to_ipbl_websocket(&mut wssocketrr, &bootstrap_event).await; + bootstrap_event = ctx.cfg.bootstrap_event().clone(); + cfg = ctx.cfg.ws.get("reqrep").unwrap().clone(); } + wssocketrr = websocketconnect(&cfg, &gethostname(true)).await.unwrap(); + send_to_ipbl_websocket(&mut wssocketrr, &bootstrap_event).await; return wssocketrr; } @@ -32,12 +30,13 @@ pub async fn websocketpubsub( ctxarc: &Arc>, txpubsub: Arc>>, ) { - let ctx = ctxarc.read().await; - let cfg = ctx.cfg.ws.get("pubsub").unwrap().clone(); - let hostname = ctx.hostname.clone(); - drop(ctx); + let cfg; + { + let ctx = ctxarc.read().await; + cfg = ctx.cfg.ws.get("pubsub").unwrap().clone(); + } let mut websocket = Arc::new(RwLock::new( - websocketconnect(&cfg, hostname.clone()).await.unwrap(), + websocketconnect(&cfg, &gethostname(true)).await.unwrap(), )); tokio::spawn(async move { loop { @@ -62,7 +61,7 @@ pub async fn websocketpubsub( ws.close(None).unwrap(); drop(ws); websocket = Arc::new(RwLock::new( - websocketconnect(&cfg, hostname.clone()).await.unwrap(), + websocketconnect(&cfg, &gethostname(true)).await.unwrap(), )); } }; @@ -72,23 +71,24 @@ pub async fn websocketpubsub( pub async fn websocketconnect<'a>( wscfg: &WebSocketCfg, - hostname: String, + hostname: &String, ) -> Result>, Error> { - print!("connecting to {} ... ", &wscfg.endpoint); + let endpoint = &wscfg.endpoint; + print!("connecting to {} ... ", endpoint); io::stdout().flush().unwrap(); let mut socket; loop { - (socket, _) = match connect(&wscfg.endpoint) { + (socket, _) = match connect(endpoint) { Ok((o, e)) => (o, e), _ => { - println!("error connecting to {}, retrying", &wscfg.endpoint); + println!("error connecting to {endpoint}, retrying"); sleep_s(1).await; continue; } }; break; } - println!("connected to {}", &wscfg.endpoint); + println!("connected to {endpoint}"); let msg = json!({ "hostname": hostname }); socket .write_message(Message::Text(msg.to_string()))