diff --git a/snow-scanner/src/main.rs b/snow-scanner/src/main.rs index a2d32b3..8ad5d14 100644 --- a/snow-scanner/src/main.rs +++ b/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) -> 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::(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) -> Rocket { +async fn report_counts<'a>(rocket: Rocket) -> Rocket { 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::shutdown_to_all(&clients.workers); + } + }) + })) + .manage(SharedData::init()) .manage(AppConfigs { static_data_dir }) .mount( "/", diff --git a/snow-scanner/src/server.rs b/snow-scanner/src/server.rs index a1e37b8..62f0e7c 100644 --- a/snow-scanner/src/server.rs +++ b/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 std::{collections::HashMap, net::IpAddr, str::FromStr}; +use rocket_ws::{frame::CloseFrame, Message}; +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 clients: HashMap>, +pub type WorkersList = Vec>; + +pub struct SharedData { + pub workers: Mutex, +} + +impl SharedData { + pub fn init() -> SharedData { + SharedData { + workers: Mutex::new(vec![]), + } + } + + pub fn add_worker(tx: Sender, workers: &Mutex) -> () { + 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) -> () { + 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, 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, pub new_scanners: HashMap, } -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(); diff --git a/snow-scanner/src/worker/worker.rs b/snow-scanner/src/worker/worker.rs index 7d83a89..34ea9b5 100644 --- a/snow-scanner/src/worker/worker.rs +++ b/snow-scanner/src/worker/worker.rs @@ -1,4 +1,3 @@ -use std::any::Any; use std::{env, net::IpAddr}; use chrono::{Duration, NaiveDateTime, Utc};