From dd303dc41e4d500e48760f79351720cc2c1c9ffe Mon Sep 17 00:00:00 2001 From: FxQnLr Date: Sun, 29 Oct 2023 21:09:46 +0100 Subject: add ip to database and use for ping, remove arc from pingmap --- src/db.rs | 3 ++- src/main.rs | 2 +- src/routes/device.rs | 16 ++++++++++------ src/routes/start.rs | 11 +++++------ src/routes/status.rs | 2 +- src/services/ping.rs | 20 +++++++++++--------- 6 files changed, 30 insertions(+), 24 deletions(-) (limited to 'src') diff --git a/src/db.rs b/src/db.rs index 3c51e2b..51ea469 100644 --- a/src/db.rs +++ b/src/db.rs @@ -12,7 +12,8 @@ use crate::config::SETTINGS; pub struct Device { pub id: String, pub mac: String, - pub broadcast_addr: String + pub broadcast_addr: String, + pub ip: String } pub async fn init_db_pool() -> PgPool { diff --git a/src/main.rs b/src/main.rs index 762a817..ee540af 100644 --- a/src/main.rs +++ b/src/main.rs @@ -52,7 +52,7 @@ async fn main() { let ping_map: DashMap = DashMap::new(); - let shared_state = Arc::new(AppState { db, ping_send: tx, ping_map: Arc::new(ping_map) }); + let shared_state = Arc::new(AppState { db, ping_send: tx, ping_map }); let app = Router::new() .route("/start", post(start)) diff --git a/src/routes/device.rs b/src/routes/device.rs index 248d1e0..7353733 100644 --- a/src/routes/device.rs +++ b/src/routes/device.rs @@ -16,7 +16,7 @@ pub async fn get_device(State(state): State>, headers: Head let device = sqlx::query_as!( Device, r#" - SELECT id, mac, broadcast_addr + SELECT id, mac, broadcast_addr, ip FROM devices WHERE id = $1; "#, @@ -40,12 +40,13 @@ pub async fn put_device(State(state): State>, headers: Head if auth(secret).map_err(WebolError::Auth)? { sqlx::query!( r#" - INSERT INTO devices (id, mac, broadcast_addr) - VALUES ($1, $2, $3); + INSERT INTO devices (id, mac, broadcast_addr, ip) + VALUES ($1, $2, $3, $4); "#, payload.id, payload.mac, - payload.broadcast_addr + payload.broadcast_addr, + payload.ip ).execute(&state.db).await.map_err(WebolError::DB)?; Ok(Json(json!(PutDeviceResponse { success: true }))) @@ -59,6 +60,7 @@ pub struct PutDevicePayload { id: String, mac: String, broadcast_addr: String, + ip: String } #[derive(Serialize)] @@ -74,11 +76,12 @@ pub async fn post_device(State(state): State>, headers: Hea Device, r#" UPDATE devices - SET mac = $1, broadcast_addr = $2 WHERE id = $3 - RETURNING id, mac, broadcast_addr; + SET mac = $1, broadcast_addr = $2, ip = $3 WHERE id = $4 + RETURNING id, mac, broadcast_addr, ip; "#, payload.mac, payload.broadcast_addr, + payload.ip, payload.id ).fetch_one(&state.db).await.map_err(WebolError::DB)?; @@ -93,4 +96,5 @@ pub struct PostDevicePayload { id: String, mac: String, broadcast_addr: String, + ip: String, } diff --git a/src/routes/start.rs b/src/routes/start.rs index 5b73281..3bccb0f 100644 --- a/src/routes/start.rs +++ b/src/routes/start.rs @@ -22,7 +22,7 @@ pub async fn start(State(state): State>, headers: HeaderMap let device = sqlx::query_as!( Device, r#" - SELECT id, mac, broadcast_addr + SELECT id, mac, broadcast_addr, ip FROM devices WHERE id = $1; "#, @@ -44,16 +44,15 @@ pub async fn start(State(state): State>, headers: HeaderMap let uuid = if payload.ping.is_some_and(|ping| ping) { let uuid_gen = Uuid::new_v4().to_string(); let uuid_genc = uuid_gen.clone(); - let uuid_gencc = uuid_gen.clone(); - tokio::spawn(async move{ + tokio::spawn(async move { debug!("Init ping service"); - state.ping_map.insert(uuid_gen, ("192.168.178.94".to_string(), false)); + state.ping_map.insert(uuid_gen.clone(), (device.ip.clone(), false)); warn!("{:?}", state.ping_map); - crate::services::ping::spawn(state.ping_send.clone(), "192.168.178.94".to_string(), uuid_genc.clone(), state.ping_map.clone()).await + crate::services::ping::spawn(state.ping_send.clone(), device.ip, uuid_gen.clone(), &state.ping_map).await }); - Some(uuid_gencc) + Some(uuid_genc) } else { None }; Ok(Json(json!(StartResponse { id: device.id, boot: true, uuid }))) } else { diff --git a/src/routes/status.rs b/src/routes/status.rs index 4a5ec67..45f3e51 100644 --- a/src/routes/status.rs +++ b/src/routes/status.rs @@ -6,5 +6,5 @@ use crate::services::ping::status_websocket; #[axum_macros::debug_handler] pub async fn status(State(state): State>, ws: WebSocketUpgrade) -> Response { - ws.on_upgrade(move |socket| status_websocket(socket, state.ping_send.clone(), state.ping_map.clone())) + ws.on_upgrade(move |socket| status_websocket(socket, state)) } \ No newline at end of file diff --git a/src/services/ping.rs b/src/services/ping.rs index ed848fc..04ad511 100644 --- a/src/services/ping.rs +++ b/src/services/ping.rs @@ -6,12 +6,13 @@ use axum::extract::ws::{CloseFrame, Message}; use dashmap::DashMap; use tokio::sync::broadcast::{Sender}; use tracing::{debug, trace, warn}; +use crate::AppState; use crate::error::WebolError; -pub type PingMap = Arc>; +pub type PingMap = DashMap; -pub async fn spawn(tx: Sender, ip: String, uuid: String, ping_map: PingMap) -> Result<(), WebolError> { +pub async fn spawn(tx: Sender, ip: String, uuid: String, ping_map: &PingMap) -> Result<(), WebolError> { let payload = [0; 8]; // TODO: Better while @@ -31,14 +32,14 @@ pub async fn spawn(tx: Sender, ip: String, uuid: String, ping let (_, duration) = ping.unwrap(); debug!("Ping took {:?}", duration); cont = false; - handle_broadcast_send(&tx, ip.clone(), ping_map.clone(), uuid.clone()).await; + handle_broadcast_send(&tx, ip.clone(), &ping_map, uuid.clone()).await; }; } Ok(()) } -async fn handle_broadcast_send(tx: &Sender, ip: String, ping_map: PingMap, uuid: String) { +async fn handle_broadcast_send(tx: &Sender, ip: String, ping_map: &PingMap, uuid: String) { debug!("sending pingsuccess message"); ping_map.insert(uuid.clone(), (ip.clone(), true)); let _ = tx.send(BroadcastCommands::PingSuccess(ip)); @@ -52,8 +53,8 @@ pub enum BroadcastCommands { PingSuccess(String) } -pub async fn status_websocket(mut socket: WebSocket, tx: Sender, ping_map: PingMap) { - warn!("{:?}", ping_map); +pub async fn status_websocket(mut socket: WebSocket, state: Arc) { + warn!("{:?}", state.ping_map); trace!("wait for ws message (uuid)"); let msg = socket.recv().await; @@ -62,13 +63,14 @@ pub async fn status_websocket(mut socket: WebSocket, tx: Sender { debug!("already started"); + // TODO: What's better? // socket.send(Message::Text(format!("start_{}", uuid))).await.unwrap(); // socket.close().await.unwrap(); socket.send(Message::Close(Some(CloseFrame { code: 4001, reason: Cow::from(format!("start_{}", uuid)) }))).await.unwrap(); @@ -77,7 +79,7 @@ pub async fn status_websocket(mut socket: WebSocket, tx: Sender