Merge pull request #17 from FxQnLr/0.3.2

0.3.2

4 files changed, 214 insertions, 143 deletions

diff --git a/src/routes/device.rs b/src/routes/device.rs
index c85df1b..d39d98e 100644
--- a/src/routes/device.rs
+++ b/src/routes/device.rs
@@ -1,34 +1,34 @@
-use std::sync::Arc;
+use crate::db::Device;
+use crate::error::Error;
 use axum::extract::State;
 use axum::Json;
-use axum::http::HeaderMap;
+use mac_address::MacAddress;
 use serde::{Deserialize, Serialize};
 use serde_json::{json, Value};
+use sqlx::types::ipnetwork::IpNetwork;
+use std::{sync::Arc, str::FromStr};
 use tracing::{debug, info};
-use crate::auth::auth;
-use crate::db::Device;
-use crate::error::Error;
 
-pub async fn get(State(state): State<Arc<crate::AppState>>, headers: HeaderMap, Json(payload): Json<GetDevicePayload>) -> Result<Json<Value>, Error> {
-    info!("add device {}", payload.id);
-    let secret = headers.get("authorization");
-    if auth(&state.config, secret).map_err(Error::Auth)? {
-        let device = sqlx::query_as!(
-            Device,
-            r#"
-            SELECT id, mac, broadcast_addr, ip, times
-            FROM devices
-            WHERE id = $1;
-            "#,
-            payload.id
-        ).fetch_one(&state.db).await.map_err(Error::DB)?;
+pub async fn get(
+    State(state): State<Arc<crate::AppState>>,
+    Json(payload): Json<GetDevicePayload>,
+) -> Result<Json<Value>, Error> {
+    info!("get device {}", payload.id);
+    let device = sqlx::query_as!(
+        Device,
+        r#"
+        SELECT id, mac, broadcast_addr, ip, times
+        FROM devices
+        WHERE id = $1;
+        "#,
+        payload.id
+    )
+    .fetch_one(&state.db)
+    .await?;
 
-        debug!("got device {:?}", device);
+    debug!("got device {:?}", device);
 
-        Ok(Json(json!(device)))
-    } else {
-        Err(Error::Generic)
-    }
+    Ok(Json(json!(device)))
 }
 
 #[derive(Deserialize)]
@@ -36,25 +36,31 @@ pub struct GetDevicePayload {
     id: String,
 }
 
-pub async fn put(State(state): State<Arc<crate::AppState>>, headers: HeaderMap, Json(payload): Json<PutDevicePayload>) -> Result<Json<Value>, Error> {
-    info!("add device {} ({}, {}, {})", payload.id, payload.mac, payload.broadcast_addr, payload.ip);
-    let secret = headers.get("authorization");
-    if auth(&state.config, secret).map_err(Error::Auth)? {
-        sqlx::query!(
-            r#"
-            INSERT INTO devices (id, mac, broadcast_addr, ip)
-            VALUES ($1, $2, $3, $4);
-            "#,
-            payload.id,
-            payload.mac,
-            payload.broadcast_addr,
-            payload.ip
-        ).execute(&state.db).await.map_err(Error::DB)?;
+pub async fn put(
+    State(state): State<Arc<crate::AppState>>,
+    Json(payload): Json<PutDevicePayload>,
+) -> Result<Json<Value>, Error> {
+    info!(
+        "add device {} ({}, {}, {})",
+        payload.id, payload.mac, payload.broadcast_addr, payload.ip
+    );
+    
+    let ip = IpNetwork::from_str(&payload.ip)?;
+    let mac = MacAddress::from_str(&payload.mac)?;
+    sqlx::query!(
+        r#"
+        INSERT INTO devices (id, mac, broadcast_addr, ip)
+        VALUES ($1, $2, $3, $4);
+        "#,
+        payload.id,
+        mac,
+        payload.broadcast_addr,
+        ip
+    )
+    .execute(&state.db)
+    .await?;
 
-        Ok(Json(json!(PutDeviceResponse { success: true })))
-    } else {
-        Err(Error::Generic)
-    }
+    Ok(Json(json!(PutDeviceResponse { success: true })))
 }
 
 #[derive(Deserialize)]
@@ -62,35 +68,40 @@ pub struct PutDevicePayload {
     id: String,
     mac: String,
     broadcast_addr: String,
-    ip: String
+    ip: String,
 }
 
 #[derive(Serialize)]
 pub struct PutDeviceResponse {
-    success: bool
+    success: bool,
 }
 
-pub async fn post(State(state): State<Arc<crate::AppState>>, headers: HeaderMap, Json(payload): Json<PostDevicePayload>) -> Result<Json<Value>, Error> {
-    info!("edit device {} ({}, {}, {})", payload.id, payload.mac, payload.broadcast_addr, payload.ip);
-    let secret = headers.get("authorization");
-    if auth(&state.config, secret).map_err(Error::Auth)? {
-        let device = sqlx::query_as!(
-            Device,
-            r#"
-            UPDATE devices
-            SET mac = $1, broadcast_addr = $2, ip = $3 WHERE id = $4
-            RETURNING id, mac, broadcast_addr, ip, times;
-            "#,
-            payload.mac,
-            payload.broadcast_addr,
-            payload.ip,
-            payload.id
-        ).fetch_one(&state.db).await.map_err(Error::DB)?;
+pub async fn post(
+    State(state): State<Arc<crate::AppState>>,
+    Json(payload): Json<PostDevicePayload>,
+) -> Result<Json<Value>, Error> {
+    info!(
+        "edit device {} ({}, {}, {})",
+        payload.id, payload.mac, payload.broadcast_addr, payload.ip
+    );
+    let ip = IpNetwork::from_str(&payload.ip)?;
+    let mac = MacAddress::from_str(&payload.mac)?;
+    let device = sqlx::query_as!(
+        Device,
+        r#"
+        UPDATE devices
+        SET mac = $1, broadcast_addr = $2, ip = $3 WHERE id = $4
+        RETURNING id, mac, broadcast_addr, ip, times;
+        "#,
+        mac,
+        payload.broadcast_addr,
+        ip,
+        payload.id
+    )
+    .fetch_one(&state.db)
+    .await?;
 
-        Ok(Json(json!(device)))
-    } else {
-        Err(Error::Generic)
-    }
+    Ok(Json(json!(device)))
 }
 
 #[derive(Deserialize)]
diff --git a/src/routes/mod.rs b/src/routes/mod.rs
deleted file mode 100644
index d5ab0d6..0000000
--- a/src/routes/mod.rs
+++ /dev/null
@@ -1,3 +0,0 @@
-pub mod start;
-pub mod device;
-pub mod status;
\ No newline at end of file
diff --git a/src/routes/start.rs b/src/routes/start.rs
index ce95bf3..d4c0802 100644
--- a/src/routes/start.rs
+++ b/src/routes/start.rs
@@ -1,10 +1,8 @@
-use crate::auth::auth;
 use crate::db::Device;
 use crate::error::Error;
 use crate::services::ping::Value as PingValue;
 use crate::wol::{create_buffer, send_packet};
 use axum::extract::State;
-use axum::http::HeaderMap;
 use axum::Json;
 use serde::{Deserialize, Serialize};
 use serde_json::{json, Value};
@@ -12,86 +10,82 @@ use std::sync::Arc;
 use tracing::{debug, info};
 use uuid::Uuid;
 
-#[axum_macros::debug_handler]
 pub async fn start(
     State(state): State<Arc<crate::AppState>>,
-    headers: HeaderMap,
     Json(payload): Json<Payload>,
 ) -> Result<Json<Value>, Error> {
     info!("POST request");
-    let secret = headers.get("authorization");
-    let authorized = auth(&state.config, secret).map_err(Error::Auth)?;
-    if authorized {
-        let device = sqlx::query_as!(
-            Device,
-            r#"
-            SELECT id, mac, broadcast_addr, ip, times
-            FROM devices
-            WHERE id = $1;
-            "#,
-            payload.id
-        )
-        .fetch_one(&state.db)
-        .await
-        .map_err(Error::DB)?;
-
-        info!("starting {}", device.id);
-
-        let bind_addr = "0.0.0.0:0";
+    let device = sqlx::query_as!(
+        Device,
+        r#"
+        SELECT id, mac, broadcast_addr, ip, times
+        FROM devices
+        WHERE id = $1;
+        "#,
+        payload.id
+    )
+    .fetch_one(&state.db)
+    .await?;
 
-        let _ = send_packet(
-            &bind_addr.parse().map_err(Error::IpParse)?,
-            &device.broadcast_addr.parse().map_err(Error::IpParse)?,
-            &create_buffer(&device.mac)?,
-        )?;
-        let dev_id = device.id.clone();
-        let uuid = if payload.ping.is_some_and(|ping| ping) {
-            let mut uuid: Option<String> = None;
-            for (key, value) in state.ping_map.clone() {
-                if value.ip == device.ip {
-                    debug!("service already exists");
-                    uuid = Some(key);
-                    break;
-                }
-            }
-            let uuid_gen = match uuid {
-                Some(u) => u,
-                None => Uuid::new_v4().to_string(),
-            };
-            let uuid_genc = uuid_gen.clone();
+    info!("starting {}", device.id);
 
-            tokio::spawn(async move {
-                debug!("init ping service");
-                state.ping_map.insert(
-                    uuid_gen.clone(),
-                    PingValue {
-                        ip: device.ip.clone(),
-                        online: false,
-                    },
-                );
+    let bind_addr = "0.0.0.0:0";
 
-                crate::services::ping::spawn(
-                    state.ping_send.clone(),
-                    &state.config,
-                    device,
-                    uuid_gen.clone(),
-                    &state.ping_map,
-                    &state.db,
-                )
-                .await;
-            });
-            Some(uuid_genc)
-        } else {
-            None
-        };
-        Ok(Json(json!(Response {
-            id: dev_id,
-            boot: true,
-            uuid
-        })))
+    let _ = send_packet(
+        bind_addr,
+        &device.broadcast_addr,
+        &create_buffer(&device.mac.to_string())?,
+    )?;
+    let dev_id = device.id.clone();
+    let uuid = if payload.ping.is_some_and(|ping| ping) {
+        Some(setup_ping(state, device))
     } else {
-        Err(Error::Generic)
+        None
+    };
+    Ok(Json(json!(Response {
+        id: dev_id,
+        boot: true,
+        uuid
+    })))
+}
+
+fn setup_ping(state: Arc<crate::AppState>, device: Device) -> String {
+    let mut uuid: Option<String> = None;
+    for (key, value) in state.ping_map.clone() {
+        if value.ip == device.ip {
+            debug!("service already exists");
+            uuid = Some(key);
+            break;
+        }
     }
+    let uuid_gen = match uuid {
+        Some(u) => u,
+        None => Uuid::new_v4().to_string(),
+    };
+    let uuid_ret = uuid_gen.clone();
+
+    debug!("init ping service");
+    state.ping_map.insert(
+        uuid_gen.clone(),
+        PingValue {
+            ip: device.ip,
+            online: false,
+        },
+    );
+
+    tokio::spawn(async move {
+        crate::services::ping::spawn(
+            state.ping_send.clone(),
+            &state.config,
+            device,
+            uuid_gen,
+            &state.ping_map,
+            &state.db,
+        )
+        .await;
+    });
+
+    uuid_ret
 }
 
 #[derive(Deserialize)]
diff --git a/src/routes/status.rs b/src/routes/status.rs
index 31ef996..0e25f7d 100644
--- a/src/routes/status.rs
+++ b/src/routes/status.rs
@@ -1,10 +1,79 @@
-use std::sync::Arc;
+use crate::services::ping::BroadcastCommand;
+use crate::AppState;
+use axum::extract::ws::{Message, WebSocket};
 use axum::extract::{State, WebSocketUpgrade};
 use axum::response::Response;
-use crate::AppState;
-use crate::services::ping::status_websocket;
+use sqlx::PgPool;
+use std::sync::Arc;
+use tracing::{debug, trace};
 
-#[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))
+    ws.on_upgrade(move |socket| websocket(socket, state))
+}
+
+pub async fn websocket(mut socket: WebSocket, state: Arc<AppState>) {
+    trace!("wait for ws message (uuid)");
+    let msg = socket.recv().await;
+    let uuid = msg.unwrap().unwrap().into_text().unwrap();
+
+    trace!("Search for uuid: {}", uuid);
+
+    let eta = get_eta(&state.db).await;
+    let _ = socket
+        .send(Message::Text(format!("eta_{eta}_{uuid}")))
+        .await;
+
+    let device_exists = state.ping_map.contains_key(&uuid);
+    if device_exists {
+        let _ = socket
+            .send(receive_ping_broadcast(state.clone(), uuid).await)
+            .await;
+    } else {
+        debug!("didn't find any device");
+        let _ = socket.send(Message::Text(format!("notfound_{uuid}"))).await;
+    };
+
+    let _ = socket.close().await;
+}
+
+async fn receive_ping_broadcast(state: Arc<AppState>, uuid: String) -> Message {
+    let pm = state.ping_map.clone().into_read_only();
+    let device = pm.get(&uuid).expect("fatal error");
+    debug!("got device: {} (online: {})", device.ip, device.online);
+    if device.online {
+        debug!("already started");
+        Message::Text(BroadcastCommand::success(uuid).to_string())
+    } else {
+        loop {
+            trace!("wait for tx message");
+            let message = state
+                .ping_send
+                .subscribe()
+                .recv()
+                .await
+                .expect("fatal error");
+            trace!("got message {:?}", message);
+
+            if message.uuid != uuid {
+                continue;
+            }
+            trace!("message == uuid success");
+            return Message::Text(message.to_string());
+        }
+    }
+}
+
+async fn get_eta(db: &PgPool) -> i64 {
+    let query = sqlx::query!(r#"SELECT times FROM devices;"#)
+        .fetch_one(db)
+        .await
+        .unwrap();
+
+    let times = if let Some(times) = query.times {
+        times
+    } else {
+        vec![0]
+    };
+
+    times.iter().sum::<i64>() / i64::try_from(times.len()).unwrap()
 }