6 files changed, 30 insertions, 24 deletions

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<String, (String, bool)> = 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<Arc<crate::AppState>>, 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<Arc<crate::AppState>>, 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<Arc<crate::AppState>>, 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<Arc<crate::AppState>>, 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<Arc<crate::AppState>>, 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<Arc<AppState>>, 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<DashMap<String, (String, bool)>>;
+pub type PingMap = DashMap<String, (String, bool)>;
 
-pub async fn spawn(tx: Sender<BroadcastCommands>, ip: String, uuid: String, ping_map: PingMap) -> Result<(), WebolError> {
+pub async fn spawn(tx: Sender<BroadcastCommands>, 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<BroadcastCommands>, 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<BroadcastCommands>, ip: String, ping_map: PingMap, uuid: String) {
+async fn handle_broadcast_send(tx: &Sender<BroadcastCommands>, 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<BroadcastCommands>, ping_map: PingMap) {
-    warn!("{:?}", ping_map);
+pub async fn status_websocket(mut socket: WebSocket, state: Arc<AppState>) {
+    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<BroadcastCommand
     trace!("Search for uuid: {:?}", uuid);
 
     // TODO: Handle Error
-    let device = ping_map.get(&uuid).unwrap().to_owned();
+    let device = state.ping_map.get(&uuid).unwrap().to_owned();
 
     trace!("got device: {:?}", device);
 
     match device.1 {
         true => {
             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<BroadcastCommand
             let ip = device.0.to_owned();
             loop{
                 trace!("wait for tx message");
-                let message = tx.subscribe().recv().await.unwrap();
+                let message = state.ping_send.subscribe().recv().await.unwrap();
                 trace!("GOT = {:?}", message);
                 // if let BroadcastCommands::PingSuccess(msg_ip) = message {
                 //     if msg_ip == ip {
@@ -95,7 +97,7 @@ pub async fn status_websocket(mut socket: WebSocket, tx: Sender<BroadcastCommand
             socket.send(Message::Close(Some(CloseFrame { code: 4000, reason: Cow::from(format!("start_{}", uuid)) }))).await.unwrap();
             // socket.send(Message::Text(format!("start_{}", uuid))).await.unwrap();
             // socket.close().await.unwrap();
-            warn!("{:?}", ping_map);
+            warn!("{:?}", state.ping_map);
         }
     }
 }
\ No newline at end of file