Implement broadcast to all nodes

snow-scanner/src/main.rs

@@ -3,7 +3,8 @@ use chrono::{NaiveDateTime, Utc};
 #[macro_use]
 extern crate rocket;
 
-use rocket::{fairing::AdHoc, trace::error, Build, Rocket, State};
+use rocket::futures::channel::mpsc::channel;
+use rocket::{fairing::AdHoc, trace::error, Rocket, State};
 use rocket_db_pools::{
     rocket::{
         figment::{
@@ -24,8 +25,9 @@ use rocket_db_pools::diesel::MysqlPool;
 use rocket_db_pools::diesel::{deserialize, serialize};
 use rocket_db_pools::Database;
 
+use ::ws::Message;
 use rocket_ws::WebSocket;
-use server::Worker;
+use server::{SharedData, Worker};
 use worker::detection::{detect_scanner, get_dns_client, Scanners};
 
 use std::path::PathBuf;
@@ -462,28 +464,20 @@ async fn pong() -> PlainText {
     PlainText("pong".to_string())
 }
 
-/*
-    let mut ws_server_handles = Server {
-        clients: HashMap::new(),
-        new_scanners: HashMap::new(),
-    };
-    info!("Worker server is listening on: {worker_server_address}");
-    loop {
-        match ws_server.process(&mut ws_server_handles, 0.5) {
-            Ok(_) => {}
-            Err(err) => error!("Processing error: {err}"),
-        }
-        ws_server_handles.cleanup(&ws_server);
-        ws_server_handles.commit(conn);
-}*/
-
 #[get("/ws")]
-pub async fn ws(ws: WebSocket) -> rocket_ws::Channel<'static> {
+pub async fn ws(ws: WebSocket, shared: &State<SharedData>) -> rocket_ws::Channel<'static> {
+    use crate::rocket::futures::StreamExt;
     use rocket::tokio;
     use rocket_ws as ws;
     use std::time::Duration;
 
-    ws.channel(move |mut stream: ws::stream::DuplexStream| {
+    let (tx, mut rx) = channel::<ws::Message>(1);
+
+    SharedData::add_worker(tx, &shared.workers);
+
+    SharedData::send_to_all(&shared.workers, "I am new !");
+
+    let channel = ws.channel(move |mut stream: ws::stream::DuplexStream| {
         Box::pin(async move {
             let mut interval = tokio::time::interval(Duration::from_secs(60));
 
@@ -497,6 +491,10 @@ pub async fn ws(ws: WebSocket) -> rocket_ws::Channel<'static> {
                                 break;
                             }
                         }
+                        Some(message) = rx.next() => {
+                            println!("Received message from other client: {:?}", message);
+                            let _ = worker.send(message).await;
+                        },
                         Ok(false) = worker.poll() => {
                             // Continue the loop
                         }
@@ -510,14 +508,16 @@ pub async fn ws(ws: WebSocket) -> rocket_ws::Channel<'static> {
             tokio::signal::ctrl_c().await.unwrap();
             Ok(())
         })
-    })
+    });
+
+    channel
 }
 
 struct AppConfigs {
     static_data_dir: String,
 }
 
-async fn report_counts(rocket: Rocket<Build>) -> Rocket<Build> {
+async fn report_counts<'a>(rocket: Rocket<rocket::Build>) -> Rocket<rocket::Build> {
     use rocket_db_pools::diesel::AsyncConnectionWrapper;
 
     let conn = SnowDb::fetch(&rocket)
@@ -572,6 +572,14 @@ fn rocket() -> _ {
     rocket::custom(config_figment)
         .attach(SnowDb::init())
         .attach(AdHoc::on_ignite("Report counts", report_counts))
+        .attach(AdHoc::on_shutdown("Close Websockets", |r| {
+            Box::pin(async move {
+                if let Some(clients) = r.state::<SharedData>() {
+                    SharedData::shutdown_to_all(&clients.workers);
+                }
+            })
+        }))
+        .manage(SharedData::init())
         .manage(AppConfigs { static_data_dir })
         .mount(
             "/",

snow-scanner/src/server.rs

@@ -1,21 +1,92 @@
 use cidr::IpCidr;
 use hickory_resolver::Name;
 use rocket::futures::{stream::Next, SinkExt, StreamExt};
-use rocket_ws::Message;
+use rocket_ws::{frame::CloseFrame, Message};
-use std::{collections::HashMap, net::IpAddr, str::FromStr};
+use std::{collections::HashMap, net::IpAddr, ops::Deref, str::FromStr, sync::Mutex};
 
 use crate::worker::{
     detection::detect_scanner_from_name,
     modules::{Network, WorkerMessages},
 };
 use crate::{DbConn, Scanner};
+use rocket::futures::channel::mpsc::Sender;
 
-pub struct Server<'a> {
+pub type WorkersList = Vec<Sender<Message>>;
-    pub clients: HashMap<u32, Worker<'a>>,
+
+pub struct SharedData {
+    pub workers: Mutex<WorkersList>,
+}
+
+impl SharedData {
+    pub fn init() -> SharedData {
+        SharedData {
+            workers: Mutex::new(vec![]),
+        }
+    }
+
+    pub fn add_worker(tx: Sender<Message>, workers: &Mutex<WorkersList>) -> () {
+        let workers_lock = workers.try_lock();
+        if let Ok(mut workers) = workers_lock {
+            workers.push(tx);
+            info!("Clients: {}", workers.len());
+            std::mem::drop(workers);
+        } else {
+            error!("Unable to lock workers");
+        }
+    }
+
+    pub fn shutdown_to_all(workers: &Mutex<WorkersList>) -> () {
+        let workers_lock = workers.try_lock();
+        if let Ok(ref workers) = workers_lock {
+            workers.iter().for_each(|tx| {
+                let res = tx.clone().try_send(Message::Close(Some(CloseFrame {
+                    code: rocket_ws::frame::CloseCode::Away,
+                    reason: "Server stop".into(),
+                })));
+                match res {
+                    Ok(_) => {
+                        info!("Worker did receive stop signal.");
+                    }
+                    Err(err) => {
+                        error!("Send error: {err}");
+                    }
+                };
+            });
+            info!("Currently {} workers online.", workers.len());
+            std::mem::drop(workers_lock);
+        } else {
+            error!("Unable to lock workers");
+        }
+    }
+
+    pub fn send_to_all(workers: &Mutex<WorkersList>, message: &str) -> () {
+        let workers_lock = workers.try_lock();
+        if let Ok(ref workers) = workers_lock {
+            workers.iter().for_each(|tx| {
+                let res = tx.clone().try_send(Message::Text(message.to_string()));
+                match res {
+                    Ok(_) => {
+                        info!("Message sent to worker !");
+                    }
+                    Err(err) => {
+                        error!("Send error: {err}");
+                    }
+                };
+            });
+            info!("Currently {} workers online.", workers.len());
+            std::mem::drop(workers_lock);
+        } else {
+            error!("Unable to lock workers");
+        }
+    }
+}
+
+pub struct Server {
+    pub clients: HashMap<u32, String>,
     pub new_scanners: HashMap<String, IpAddr>,
 }
 
-impl<'a> Server<'a> {
+impl Server {
     pub async fn commit(&mut self, conn: &mut DbConn) -> &Server {
         for (name, query_address) in self.new_scanners.clone() {
             let scanner_name = Name::from_str(name.as_str()).unwrap();

snow-scanner/src/worker/worker.rs

@@ -1,4 +1,3 @@
-use std::any::Any;
 use std::{env, net::IpAddr};
 
 use chrono::{Duration, NaiveDateTime, Utc};