Finally make it capable of talking to all hosts

snow-scanner/src/event_bus.rs

@@ -0,0 +1,93 @@
+use rocket::futures::channel::mpsc as rocket_mpsc;
+use rocket::futures::StreamExt;
+use rocket::tokio;
+
+/// Handles all the raw events being streamed from balancers and parses and filters them into only the events we care about.
+pub struct EventBus {
+    events_rx: rocket_mpsc::Receiver<rocket_ws::Message>,
+    events_tx: rocket_mpsc::Sender<rocket_ws::Message>,
+    bus_tx: tokio::sync::broadcast::Sender<EventBusEvent>,
+}
+
+impl EventBus {
+    pub fn new() -> Self {
+        let (events_tx, events_rx) = rocket_mpsc::channel(100);
+        let (bus_tx, _) = tokio::sync::broadcast::channel(100);
+        Self {
+            events_rx,
+            events_tx,
+            bus_tx,
+        }
+    }
+
+    pub async fn run(&mut self) {
+        info!("EventBus started");
+        loop {
+            tokio::select! {
+                Some(event) = self.events_rx.next() => {
+                    info!("EventBus received: {event}");
+                    self.handle_event(event);
+                }
+                else => {
+                    warn!("EventBus stopped");
+                    break;
+                }
+            }
+        }
+    }
+
+    fn handle_event(&self, event: rocket_ws::Message) {
+        info!("Received event: {}", event);
+        if self.bus_tx.receiver_count() == 0 {
+            return;
+        }
+        match self.bus_tx.send(event) {
+            Ok(count) => {
+                info!("Event sent to {count} subscribers");
+            }
+            Err(err) => {
+                error!("Error sending event to subscribers: {}", err);
+            }
+        }
+    }
+
+    pub fn subscriber(&self) -> EventBusSubscriber {
+        EventBusSubscriber::new(self.bus_tx.clone())
+    }
+
+    pub fn writer(&self) -> EventBusWriter {
+        EventBusWriter::new(self.events_tx.clone())
+    }
+}
+
+pub type EventBusEvent = rocket_ws::Message;
+
+/// Enables subscriptions to the event bus
+pub struct EventBusSubscriber {
+    bus_tx: tokio::sync::broadcast::Sender<EventBusEvent>,
+}
+
+/// Enables subscriptions to the event bus
+pub struct EventBusWriter {
+    bus_tx: rocket_mpsc::Sender<EventBusEvent>,
+}
+
+impl EventBusWriter {
+    pub fn new(bus_tx: rocket_mpsc::Sender<EventBusEvent>) -> Self {
+        Self { bus_tx }
+    }
+
+    pub fn write(&self) -> rocket_mpsc::Sender<EventBusEvent> {
+        self.bus_tx.clone()
+    }
+}
+
+impl EventBusSubscriber {
+    pub fn new(bus_tx: tokio::sync::broadcast::Sender<EventBusEvent>) -> Self {
+        Self { bus_tx }
+    }
+
+    pub fn subscribe(&self) -> tokio::sync::broadcast::Receiver<EventBusEvent> {
+        self.bus_tx.subscribe()
+    }
+}

snow-scanner/src/main.rs

@@ -3,8 +3,8 @@ use chrono::{NaiveDateTime, Utc};
 #[macro_use]
 extern crate rocket;
 
-use rocket::futures::channel::mpsc::channel;
+use event_bus::{EventBusSubscriber, EventBusWriter};
-use rocket::{fairing::AdHoc, trace::error, Rocket, State};
+use rocket::{fairing::AdHoc, futures::SinkExt, trace::error, Rocket, State};
 use rocket_db_pools::{
     rocket::{
         figment::{
@@ -13,7 +13,7 @@ use rocket_db_pools::{
         },
         form::Form,
         fs::NamedFile,
-        launch, Responder,
+        Responder,
     },
     Connection,
 };
@@ -25,9 +25,8 @@ 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::{SharedData, Worker};
+use server::Server;
 use worker::detection::{detect_scanner, get_dns_client, Scanners};
 
 use std::path::PathBuf;
@@ -39,6 +38,7 @@ use serde::{Deserialize, Deserializer, Serialize};
 
 use dns_ptr_resolver::{get_ptr, ResolvedResult};
 
+pub mod event_bus;
 pub mod models;
 pub mod schema;
 pub mod server;
@@ -460,55 +460,27 @@ async fn index() -> HtmlContents {
 }
 
 #[get("/ping")]
-async fn pong() -> PlainText {
+async fn pong(event_bus_writer: &State<EventBusWriter>) -> PlainText {
+    let mut bus_tx = event_bus_writer.write();
+    let _ = bus_tx
+        .send(rocket_ws::Message::Text("Groumpf!".to_string()))
+        .await;
     PlainText("pong".to_string())
 }
 
 #[get("/ws")]
-pub async fn ws(ws: WebSocket, shared: &State<SharedData>) -> rocket_ws::Channel<'static> {
+pub async fn ws(
-    use crate::rocket::futures::StreamExt;
+    ws: WebSocket,
-    use rocket::tokio;
+    event_bus: &State<EventBusSubscriber>,
-    use rocket_ws as ws;
+    event_bus_writer: &State<EventBusWriter>,
-    use std::time::Duration;
+) -> rocket_ws::Channel<'static> {
+    use rocket::futures::channel::mpsc as rocket_mpsc;
 
-    let (tx, mut rx) = channel::<ws::Message>(1);
+    let (_, ws_receiver) = rocket_mpsc::channel::<rocket_ws::Message>(1);
-
+    let bus_rx = event_bus.subscribe();
-    SharedData::add_worker(tx, &shared.workers);
+    let bus_tx = event_bus_writer.write();
-
+    let channel: rocket_ws::Channel =
-    SharedData::send_to_all(&shared.workers, "I am new !");
+        ws.channel(|stream| Server::handle(stream, bus_rx, bus_tx, ws_receiver));
-
-    let channel = ws.channel(move |mut stream: ws::stream::DuplexStream| {
-        Box::pin(async move {
-            let mut interval = tokio::time::interval(Duration::from_secs(60));
-
-            tokio::spawn(async move {
-                let mut worker = Worker::initial(&mut stream);
-                loop {
-                    tokio::select! {
-                        _ = interval.tick() => {
-                            // Send message every X seconds
-                            if let Ok(true) = worker.tick().await {
-                                break;
-                            }
-                        }
-                        Some(message) = rx.next() => {
-                            println!("Received message from other client: {:?}", message);
-                            let _ = worker.send(message).await;
-                        },
-                        Ok(false) = worker.poll() => {
-                            // Continue the loop
-                        }
-                        else => {
-                            break;
-                        }
-                    }
-                }
-            });
-
-            tokio::signal::ctrl_c().await.unwrap();
-            Ok(())
-        })
-    });
 
     channel
 }
@@ -535,8 +507,8 @@ async fn report_counts<'a>(rocket: Rocket<rocket::Build>) -> Rocket<rocket::Buil
     rocket
 }
 
-#[launch]
+#[rocket::main]
-fn rocket() -> _ {
+async fn main() -> Result<(), rocket::Error> {
     let server_address: SocketAddr = if let Ok(env) = env::var("SERVER_ADDRESS") {
         env.parse()
             .expect("The ENV SERVER_ADDRESS should be a valid socket address (address:port)")
@@ -569,18 +541,30 @@ fn rocket() -> _ {
         .merge(("port", server_address.port()))
         .merge(("databases", map!["snow_scanner_db" => db]));
 
-    rocket::custom(config_figment)
+    let mut event_bus = event_bus::EventBus::new();
+    let event_subscriber = event_bus.subscriber();
+    let event_writer = event_bus.writer();
+
+    let _ = 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>() {
+                if let Some(writer) = r.state::<EventBusWriter>() {
-                    SharedData::shutdown_to_all(&clients.workers);
+                    Server::shutdown_to_all(writer);
                 }
             })
         }))
-        .manage(SharedData::init())
+        .attach(AdHoc::on_liftoff("Run websocket client manager", |_| {
+            Box::pin(async move {
+                rocket::tokio::spawn(async move {
+                    event_bus.run().await;
+                });
+            })
+        }))
         .manage(AppConfigs { static_data_dir })
+        .manage(event_subscriber)
+        .manage(event_writer)
         .mount(
             "/",
             routes![
@@ -594,4 +578,7 @@ fn rocket() -> _ {
                 ws,
             ],
         )
+        .launch()
+        .await;
+    Ok(())
 }

snow-scanner/src/server.rs

@@ -1,30 +1,105 @@
 use cidr::IpCidr;
-use hickory_resolver::Name;
 use rocket::futures::{stream::Next, SinkExt, StreamExt};
 use rocket_ws::{frame::CloseFrame, Message};
-use std::{collections::HashMap, net::IpAddr, ops::Deref, str::FromStr, sync::Mutex};
+use std::{pin::Pin, str::FromStr};
 
-use crate::worker::{
+use crate::{
-    detection::detect_scanner_from_name,
+    event_bus::{EventBusEvent, EventBusWriter},
-    modules::{Network, WorkerMessages},
+    worker::modules::{Network, WorkerMessages},
 };
-use crate::{DbConn, Scanner};
+use rocket::futures::channel::mpsc as rocket_mpsc;
-use rocket::futures::channel::mpsc::Sender;
 
-pub type WorkersList = Vec<Sender<Message>>;
+pub struct WsChat {}
 
-pub struct SharedData {
+impl WsChat {
-    pub workers: Mutex<WorkersList>,
+    pub async fn work(
-}
+        mut stream: rocket_ws::stream::DuplexStream,
+        mut bus_rx: rocket::tokio::sync::broadcast::Receiver<EventBusEvent>,
+        mut bus_tx: rocket_mpsc::Sender<EventBusEvent>,
+        mut ws_receiver: rocket_mpsc::Receiver<rocket_ws::Message>,
+    ) {
+        use crate::rocket::futures::StreamExt;
+        use rocket::tokio;
 
-impl SharedData {
+        let _ = bus_tx
-    pub fn init() -> SharedData {
+            .send(rocket_ws::Message::Text("I am new !".to_string()))
-        SharedData {
+            .await;
-            workers: Mutex::new(vec![]),
+        //SharedData::send_to_all(&workers, "I am new !");
+        let mut worker = Worker::initial(&mut stream);
+        let mut interval = rocket::tokio::time::interval(std::time::Duration::from_secs(60));
+        loop {
+            tokio::select! {
+                _ = interval.tick() => {
+                    // Send message every X seconds
+                    if let Ok(true) = worker.tick().await {
+                        break;
+                    }
+                }
+                result = bus_rx.recv() => {
+                    let message = match result {
+                        Ok(message) => message,
+                        Err(err) => {
+                            error!("Bus error: {err}");
+                            continue;
+                        }
+                    };
+                    if let Err(err) = worker.send(message).await {
+                        error!("Error sending event to Event bus WebSocket: {}", err);
+                        break;
+                    }
+                }
+                Some(message) = ws_receiver.next() => {
+                    info!("Received message from other client: {:?}", message);
+                    let _ = worker.send(message).await;
+                },
+                Ok(false) = worker.poll() => {
+                    // Continue the loop
+                }
+                else => {
+                    break;
+                }
+            }
         }
     }
+}
 
-    pub fn add_worker(tx: Sender<Message>, workers: &Mutex<WorkersList>) -> () {
+pub struct Server {}
+
+type HandleBox = Pin<
+    Box<dyn std::future::Future<Output = Result<(), rocket_ws::result::Error>> + std::marker::Send>,
+>;
+
+impl Server {
+    pub fn handle(
+        stream: rocket_ws::stream::DuplexStream,
+        bus_rx: rocket::tokio::sync::broadcast::Receiver<EventBusEvent>,
+        bus_tx: rocket_mpsc::Sender<EventBusEvent>,
+        ws_receiver: rocket_mpsc::Receiver<rocket_ws::Message>,
+    ) -> HandleBox {
+        use rocket::tokio;
+
+        //SharedData::add_worker(tx.clone(), &shared.workers);
+        //move |mut stream: ws::stream::DuplexStream| {
+        Box::pin(async move {
+            let work_fn = WsChat::work(
+                stream,
+                bus_rx,
+                bus_tx,
+                ws_receiver,
+                //workers
+            );
+            tokio::spawn(work_fn);
+
+            tokio::signal::ctrl_c().await.unwrap();
+            Ok(())
+        })
+    }
+
+    pub fn new() -> Server {
+        Server {}
+    }
+
+    /*pub fn add_worker(tx: rocket_mpsc::Sender<Message>, workers: &Mutex<WorkersList>) -> () {
         let workers_lock = workers.try_lock();
         if let Ok(mut workers) = workers_lock {
             workers.push(tx);
@@ -33,33 +108,24 @@ impl SharedData {
         } else {
             error!("Unable to lock workers");
         }
+    }*/
+
+    pub fn shutdown_to_all(server: &EventBusWriter) -> () {
+        let res = server.write().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}");
+            }
+        };
     }
 
-    pub fn shutdown_to_all(workers: &Mutex<WorkersList>) -> () {
+    /*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::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| {
@@ -78,48 +144,7 @@ impl SharedData {
         } else {
             error!("Unable to lock workers");
         }
-    }
+    }*/
-}
-
-pub struct Server {
-    pub clients: HashMap<u32, String>,
-    pub new_scanners: HashMap<String, IpAddr>,
-}
-
-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();
-
-            match detect_scanner_from_name(&scanner_name) {
-                Ok(Some(scanner_type)) => {
-                    match Scanner::find_or_new(
-                        query_address,
-                        scanner_type,
-                        Some(scanner_name),
-                        conn,
-                    )
-                    .await
-                    {
-                        Ok(scanner) => {
-                            // Got saved
-                            self.new_scanners.remove(&name);
-                            info!(
-                                "Saved {scanner_type}: {name} for {query_address}: {:?}",
-                                scanner.ip_ptr
-                            );
-                        }
-                        Err(err) => {
-                            error!("Unable to find or new {:?}", err);
-                        }
-                    };
-                }
-                Ok(None) => {}
-                Err(_) => {}
-            }
-        }
-        self
-    }
 }
 
 pub struct Worker<'a> {
@@ -196,7 +221,7 @@ impl<'a> Worker<'a> {
                 };
                 Ok(false)
             }
-            Some(Err(err)) => {
+            Some(Err(_)) => {
                 info!("Connection closed");
                 let close_frame = rocket_ws::frame::CloseFrame {
                     code: rocket_ws::frame::CloseCode::Normal,