Re-implement the ws client

This commit is contained in:
2024-10-03 12:49:33 +02:00
parent 25df2642e9
commit fd4d43596f
6 changed files with 179 additions and 115 deletions

View File

@ -3,7 +3,7 @@ use chrono::{NaiveDateTime, Utc};
#[macro_use]
extern crate rocket;
use rocket::{fairing::AdHoc, Build, Rocket, State};
use rocket::{fairing::AdHoc, futures::SinkExt, trace::error, Build, Rocket, State};
use rocket_db_pools::{
rocket::{
figment::{
@ -20,15 +20,17 @@ use rocket_db_pools::{
use rocket_db_pools::diesel::mysql::{Mysql, MysqlValue};
use rocket_db_pools::diesel::serialize::IsNull;
use rocket_db_pools::diesel::sql_types::Text;
use rocket_db_pools::diesel::{deserialize, serialize};
use rocket_db_pools::diesel::MysqlPool;
use rocket_db_pools::diesel::{deserialize, serialize};
use rocket_db_pools::Database;
use rocket_ws::WebSocket;
use server::Worker;
use worker::detection::{detect_scanner, get_dns_client, Scanners};
use std::{io::Write, net::SocketAddr};
use std::path::PathBuf;
use std::{env, fmt};
use std::{io::Write, net::SocketAddr};
use uuid::Uuid;
use serde::{Deserialize, Deserializer, Serialize};
@ -460,19 +462,56 @@ 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() -> PlainText {
info!("establish_ws_connection");
PlainText("ok".to_string())
// Ok(HttpResponse::Unauthorized().json(e))
/*
match result {
Ok(response) => Ok(response.into()),
Err(e) => {
error!("ws connection error: {:?}", e);
Err(e)
},
}*/
pub async fn ws(ws: WebSocket) -> rocket_ws::Channel<'static> {
use rocket::futures::StreamExt;
use rocket::tokio;
use rocket_ws as ws;
use std::time::Duration;
ws.channel(move |mut stream: ws::stream::DuplexStream| {
Box::pin(async move {
let mut interval = tokio::time::interval(Duration::from_secs(10));
tokio::spawn(async move {
let mut worker = Worker::initial(&mut stream);
loop {
tokio::select! {
_ = interval.tick() => {
// Send message every 10 seconds
//let reading = get_latest_readings().await.unwrap();
//let _ = stream.send(ws::Message::Text(json!(reading).to_string())).await;
// info!("Sent message");
}
Ok(false) = worker.poll() => {
// Continue the loop
}
else => {
break;
}
}
}
});
tokio::signal::ctrl_c().await.unwrap();
Ok(())
})
})
}
struct AppConfigs {
@ -484,7 +523,8 @@ async fn report_counts(rocket: Rocket<Build>) -> Rocket<Build> {
let conn = SnowDb::fetch(&rocket)
.expect("database is attached")
.get().await
.get()
.await
.unwrap_or_else(|e| {
span_error!("failed to connect to MySQL database" => error!("{e}"));
panic!("aborting launch");
@ -499,9 +539,12 @@ async fn report_counts(rocket: Rocket<Build>) -> Rocket<Build> {
#[launch]
fn rocket() -> _ {
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)")
env.parse()
.expect("The ENV SERVER_ADDRESS should be a valid socket address (address:port)")
} else {
"127.0.0.1:8000".parse().expect("The default address should be valid")
"127.0.0.1:8000"
.parse()
.expect("The default address should be valid")
};
let static_data_dir: String = match env::var("STATIC_DATA_DIR") {
@ -544,32 +587,4 @@ fn rocket() -> _ {
ws,
],
)
/*
match server {
Ok(server) => {
match ws2::listen(worker_server_address.as_str()) {
Ok(mut ws_server) => {
std::thread::spawn(move || {
let pool = get_connection(db_url.as_str());
// note that obtaining a connection from the pool is also potentially blocking
let conn = &mut pool.get().unwrap();
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);
}
});
}
Err(err) => error!("Unable to listen on {worker_server_address}: {err}"),
};
}*/
}