Re-implement the ws client
This commit is contained in:
@ -1,14 +1,21 @@
|
||||
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 crate::{worker::detection::detect_scanner_from_name, DbConn, Scanner};
|
||||
use crate::worker::{
|
||||
detection::detect_scanner_from_name,
|
||||
modules::{Network, WorkerMessages},
|
||||
};
|
||||
use crate::{DbConn, Scanner};
|
||||
|
||||
pub struct Server {
|
||||
pub clients: HashMap<u32, Worker>,
|
||||
pub struct Server<'a> {
|
||||
pub clients: HashMap<u32, Worker<'a>>,
|
||||
pub new_scanners: HashMap<String, IpAddr>,
|
||||
}
|
||||
|
||||
impl Server {
|
||||
impl<'a> Server<'a> {
|
||||
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();
|
||||
@ -44,18 +51,81 @@ impl Server {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Worker {
|
||||
pub authenticated: bool,
|
||||
pub login: Option<String>,
|
||||
pub struct Worker<'a> {
|
||||
authenticated: bool,
|
||||
login: Option<String>,
|
||||
stream: &'a mut rocket_ws::stream::DuplexStream,
|
||||
}
|
||||
|
||||
impl Worker {
|
||||
pub fn initial() -> Worker {
|
||||
impl<'a> Worker<'a> {
|
||||
pub fn initial(stream: &mut rocket_ws::stream::DuplexStream) -> Worker {
|
||||
info!("New worker");
|
||||
Worker {
|
||||
authenticated: false,
|
||||
login: None,
|
||||
stream,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn send(&mut self, msg: Message) -> Result<(), rocket_ws::result::Error> {
|
||||
self.stream.send(msg).await
|
||||
}
|
||||
|
||||
pub fn next(&mut self) -> Next<'_, rocket_ws::stream::DuplexStream> {
|
||||
self.stream.next()
|
||||
}
|
||||
|
||||
pub async fn poll(&mut self) -> Result<bool, ()> {
|
||||
let message = self.next();
|
||||
|
||||
match message.await {
|
||||
Some(Ok(message)) => {
|
||||
match message {
|
||||
rocket_ws::Message::Text(_) => match self.on_message(&message).await {
|
||||
Ok(_) => {}
|
||||
Err(err) => error!("Processing error: {err}"),
|
||||
},
|
||||
rocket_ws::Message::Binary(data) => {
|
||||
// Handle Binary message
|
||||
info!("Received Binary message: {:?}", data);
|
||||
}
|
||||
rocket_ws::Message::Close(close_frame) => {
|
||||
// Handle Close message
|
||||
info!("Received Close message: {:?}", close_frame);
|
||||
let close_frame = rocket_ws::frame::CloseFrame {
|
||||
code: rocket_ws::frame::CloseCode::Normal,
|
||||
reason: "Client disconected".to_string().into(),
|
||||
};
|
||||
let _ = self.stream.close(Some(close_frame)).await;
|
||||
return Ok(true);
|
||||
}
|
||||
rocket_ws::Message::Ping(ping_data) => {
|
||||
match self.send(rocket_ws::Message::Pong(ping_data)).await {
|
||||
Ok(_) => {}
|
||||
Err(err) => error!("Processing error: {err}"),
|
||||
}
|
||||
}
|
||||
rocket_ws::Message::Pong(pong_data) => {
|
||||
// Handle Pong message
|
||||
info!("Received Pong message: {:?}", pong_data);
|
||||
}
|
||||
_ => {
|
||||
info!("Received other message: {:?}", message);
|
||||
}
|
||||
};
|
||||
Ok(false)
|
||||
}
|
||||
Some(Err(err)) => {
|
||||
info!("Connection closed");
|
||||
let close_frame = rocket_ws::frame::CloseFrame {
|
||||
code: rocket_ws::frame::CloseCode::Normal,
|
||||
reason: "Client disconected".to_string().into(),
|
||||
};
|
||||
let _ = self.stream.close(Some(close_frame)).await;
|
||||
// The connection is closed by the client
|
||||
Ok(true)
|
||||
}
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
@ -77,63 +147,43 @@ impl Worker {
|
||||
self.authenticated = true;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
impl ws2::Handler for Server {
|
||||
fn on_open(&mut self, ws: &WebSocket) -> Pod {
|
||||
info!("New client: {ws}");
|
||||
let worker = Worker::initial();
|
||||
// Add the client
|
||||
self.clients.insert(ws.id(), worker);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn on_close(&mut self, ws: &WebSocket) -> Pod {
|
||||
info!("Client /quit: {ws}");
|
||||
// Drop the client
|
||||
self.clients.remove(&ws.id());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn on_message(&mut self, ws: &WebSocket, msg: String) -> Pod {
|
||||
let client = self.clients.get_mut(&ws.id());
|
||||
if client.is_none() {
|
||||
// Impossible, close in case
|
||||
return ws.close();
|
||||
}
|
||||
let worker: &mut Worker = client.unwrap();
|
||||
|
||||
info!("on message: {msg}, {ws}");
|
||||
pub async fn on_message(&mut self, msg: &Message) -> Result<(), String> {
|
||||
info!("on message: {msg}");
|
||||
|
||||
let mut worker_reply: Option<WorkerMessages> = None;
|
||||
let worker_request: WorkerMessages = msg.clone().into();
|
||||
let worker_request: WorkerMessages = match msg.clone().try_into() {
|
||||
Ok(data) => data,
|
||||
Err(err) => return Err(err),
|
||||
};
|
||||
|
||||
let result = match worker_request {
|
||||
WorkerMessages::AuthenticateRequest { login } => {
|
||||
if !worker.is_authenticated() {
|
||||
worker.authenticate(login);
|
||||
if !self.is_authenticated() {
|
||||
self.authenticate(login);
|
||||
return Ok(());
|
||||
} else {
|
||||
error!("Already authenticated: {ws}");
|
||||
error!("Already authenticated");
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
WorkerMessages::ScannerFoundResponse { name, address } => {
|
||||
info!("Detected {name} for {address}");
|
||||
self.new_scanners.insert(name, address);
|
||||
//self.new_scanners.insert(name, address);
|
||||
Ok(())
|
||||
}
|
||||
WorkerMessages::GetWorkRequest {} => {
|
||||
let net = IpCidr::from_str("52.189.78.0/24").unwrap();
|
||||
worker_reply = Some(WorkerMessages::DoWorkRequest {
|
||||
neworks: vec![Network(IpCidr::from_str("52.189.78.0/24")?)],
|
||||
neworks: vec![Network(net)],
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
WorkerMessages::DoWorkRequest { .. } | WorkerMessages::Invalid { .. } => {
|
||||
error!("Unable to understand: {msg}, {ws}");
|
||||
error!("Unable to understand: {msg}");
|
||||
// Unable to understand, close the connection
|
||||
return ws.close();
|
||||
//return ws.close();
|
||||
Err("Unable to understand: {msg}}")
|
||||
} /*msg => {
|
||||
error!("No implemented: {:#?}", msg);
|
||||
Ok(())
|
||||
@ -143,14 +193,14 @@ impl ws2::Handler for Server {
|
||||
// it has a request to send
|
||||
if let Some(worker_reply) = worker_reply {
|
||||
let msg_string: String = worker_reply.to_string();
|
||||
match ws.send(msg_string) {
|
||||
match self.send(rocket_ws::Message::Text(msg_string)).await {
|
||||
Ok(_) => match worker_reply {
|
||||
WorkerMessages::DoWorkRequest { .. } => {}
|
||||
msg => error!("No implemented: {:#?}", msg),
|
||||
},
|
||||
Err(err) => error!("Error sending reply to {ws}: {err}"),
|
||||
Err(err) => error!("Error sending reply: {err}"),
|
||||
}
|
||||
}
|
||||
result
|
||||
Ok(result?)
|
||||
}
|
||||
}*/
|
||||
}
|
||||
|
Reference in New Issue
Block a user