ipblc/src/websocket.rs

101 lines
3.0 KiB
Rust

use crate::config::{Context, WebSocketCfg};
use crate::ip::IpEvent;
use crate::utils::gethostname;
use serde_json::json;
use std::net::TcpStream;
use std::sync::Arc;
use tokio::sync::mpsc::Sender;
use tokio::sync::RwLock;
use tungstenite::stream::*;
use tungstenite::*;
pub async fn websocketconnect<'a>(
wscfg: &WebSocketCfg,
hostname: String,
) -> Result<WebSocket<MaybeTlsStream<TcpStream>>, Error> {
let (mut socket, _response) = connect(&wscfg.endpoint).expect("Can't connect");
let msg = json!({ "hostname": hostname });
socket
.write_message(Message::Text(msg.to_string()))
.unwrap();
Ok(socket)
}
pub async fn websocketinit(
ctxarc: &Arc<RwLock<Context>>,
ipeventtx: Arc<RwLock<Sender<IpEvent>>>,
) -> WebSocket<MaybeTlsStream<TcpStream>> {
let (wssocketps, mut wssocketrr, bootstrap_event);
{
let ctx = ctxarc.read().await;
bootstrap_event = ctxarc.read().await.cfg.bootstrap_event().clone();
wssocketps = Arc::new(RwLock::new(
websocketconnect(&ctx.cfg.ws.get("pubsub").unwrap(), ctx.hostname.clone())
.await
.unwrap(),
));
wssocketrr = websocketconnect(&ctx.cfg.ws.get("reqrep").unwrap(), ctx.hostname.clone())
.await
.unwrap();
send_to_ipbl_websocket(&mut wssocketrr, &bootstrap_event).await;
}
wslistenpubsub(wssocketps, ipeventtx).await;
return wssocketrr;
}
async fn wslistenpubsub(
websocket: Arc<RwLock<WebSocket<MaybeTlsStream<TcpStream>>>>,
txpubsub: Arc<RwLock<Sender<IpEvent>>>,
) {
tokio::spawn(async move {
loop {
{
let mut ws = websocket.write().await;
match ws.read_message() {
Ok(msg) => {
let tosend: IpEvent =
serde_json::from_str(msg.to_string().as_str()).unwrap();
if tosend.ipdata.hostname != gethostname(true)
|| tosend.msgtype == "init".to_string()
{
let txps = txpubsub.write().await;
txps.send(tosend).await.unwrap();
}
}
Err(e) => {
println!("error: {e:?}");
ws.close(None).unwrap();
return;
}
};
}
}
});
}
pub async fn send_to_ipbl_websocket(
ws: &mut WebSocket<MaybeTlsStream<TcpStream>>,
ip: &IpEvent,
) -> bool {
let msg = format!("{val}", val = serde_json::to_string(&ip).unwrap());
match ws.write_message(Message::Text(msg)) {
Ok(_) => {}
Err(e) => {
println!("err 1: {e:?}");
return false;
}
};
match ws.read_message() {
Ok(_) => {}
Err(e) => {
println!("err 2: {e:?}");
return false;
}
};
true
}