Re-implement the ws client

snow-scanner/src/server.rs

@@ -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?)
     }
-}*/
+}