Merge pull request #17 from FxQnLr/0.3.2

0.3.2

12 files changed, 416 insertions, 331 deletions

diff --git a/src/auth.rs b/src/auth.rs
deleted file mode 100644
index feca652..0000000
--- a/src/auth.rs
+++ /dev/null
@@ -1,43 +0,0 @@
-use axum::http::{StatusCode, HeaderValue};
-use axum::http::header::ToStrError;
-use tracing::{debug, error, trace};
-use crate::auth::Error::{MissingSecret, WrongSecret};
-use crate::config::Config;
-
-pub fn auth(config: &Config, secret: Option<&HeaderValue>) -> Result<bool, Error> {
-    debug!("auth request with secret {:?}", secret);
-    if let Some(value) = secret {
-        trace!("value exists");
-        let key = &config.apikey;
-        if value.to_str().map_err(Error::HeaderToStr)? == key.as_str() {
-            debug!("successful auth");
-            Ok(true)
-        } else {
-            debug!("unsuccessful auth (wrong secret)");
-            Err(WrongSecret)
-        }
-    } else {
-        debug!("unsuccessful auth (no secret)");
-        Err(MissingSecret)
-    }
-}
-
-#[derive(Debug)]
-pub enum Error {
-    WrongSecret,
-    MissingSecret,
-    HeaderToStr(ToStrError)
-}
-
-impl Error {
-    pub fn get(self) -> (StatusCode, &'static str) {
-        match self {
-            Self::WrongSecret => (StatusCode::UNAUTHORIZED, "Wrong credentials"),
-            Self::MissingSecret => (StatusCode::BAD_REQUEST, "Missing credentials"),
-            Self::HeaderToStr(err) => {
-                error!("server error: {}", err.to_string());
-                (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
-            },
-        }
-    }
-}
diff --git a/src/db.rs b/src/db.rs
index 489a000..47e907d 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -1,13 +1,13 @@
 use serde::Serialize;
-use sqlx::{PgPool, postgres::PgPoolOptions};
+use sqlx::{PgPool, postgres::PgPoolOptions, types::{ipnetwork::IpNetwork, mac_address::MacAddress}};
 use tracing::{debug, info};
 
 #[derive(Serialize, Debug)]
 pub struct Device {
     pub id: String,
-    pub mac: String,
+    pub mac: MacAddress,
     pub broadcast_addr: String,
-    pub ip: String,
+    pub ip: IpNetwork,
     pub times: Option<Vec<i64>>
 }
 
diff --git a/src/error.rs b/src/error.rs
index 56d6c52..513b51b 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,44 +1,80 @@
-use std::io;
+use ::ipnetwork::IpNetworkError;
+use axum::http::header::ToStrError;
 use axum::http::StatusCode;
-use axum::Json;
 use axum::response::{IntoResponse, Response};
+use axum::Json;
+use mac_address::MacParseError;
 use serde_json::json;
+use std::io;
 use tracing::error;
-use crate::auth::Error as AuthError;
 
-#[derive(Debug)]
+#[derive(Debug, thiserror::Error)]
 pub enum Error {
-    Generic,
-    Auth(AuthError),
-    DB(sqlx::Error),
-    IpParse(<std::net::IpAddr as std::str::FromStr>::Err),
-    BufferParse(std::num::ParseIntError),
-    Broadcast(io::Error),
+    #[error("db: {source}")]
+    Db {
+        #[from]
+        source: sqlx::Error,
+    },
+
+    #[error("buffer parse: {source}")]
+    ParseInt {
+        #[from]
+        source: std::num::ParseIntError,
+    },
+
+    #[error("header parse: {source}")]
+    ParseHeader {
+        #[from]
+        source: ToStrError,
+    },
+
+    #[error("string parse: {source}")]
+    IpParse {
+        #[from]
+        source: IpNetworkError,
+    },
+
+    #[error("mac parse: {source}")]
+    MacParse {
+        #[from]
+        source: MacParseError,
+    },
+
+    #[error("io: {source}")]
+    Io {
+        #[from]
+        source: io::Error,
+    },
 }
 
 impl IntoResponse for Error {
     fn into_response(self) -> Response {
+        error!("{}", self.to_string());
         let (status, error_message) = match self {
-            Self::Auth(err) => {
-                err.get()
-            },
-            Self::Generic => (StatusCode::INTERNAL_SERVER_ERROR, ""),
-            Self::IpParse(err) => {
-                error!("server error: {}", err.to_string());
+            Self::Db { source } => {
+                error!("{source}");
+                (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
+            }
+            Self::Io { source } => {
+                error!("{source}");
+                (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
+            }
+            Self::ParseHeader { source } => {
+                error!("{source}");
                 (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
-            },
-            Self::DB(err) => {
-                error!("server error: {}", err.to_string());
+            }
+            Self::ParseInt { source } => {
+                error!("{source}");
                 (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
-            },
-            Self::Broadcast(err) => {
-                error!("server error: {}", err.to_string());
+            }
+            Self::MacParse { source } => {
+                error!("{source}");
                 (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
-            },
-            Self::BufferParse(err) => {
-                error!("server error: {}", err.to_string());
+            }
+            Self::IpParse { source } => {
+                error!("{source}");
                 (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
-            },
+            }
         };
         let body = Json(json!({
             "error": error_message,
diff --git a/src/extractors.rs b/src/extractors.rs
new file mode 100644
index 0000000..4d441e9
--- /dev/null
+++ b/src/extractors.rs
@@ -0,0 +1,24 @@
+use axum::{
+    extract::{Request, State},
+    http::{HeaderMap, StatusCode},
+    middleware::Next,
+    response::Response,
+};
+
+use crate::AppState;
+
+pub async fn auth(
+    State(state): State<AppState>,
+    headers: HeaderMap,
+    request: Request,
+    next: Next,
+) -> Result<Response, StatusCode> {
+    let secret = headers.get("authorization");
+    match secret {
+        Some(token) if token == state.config.apikey.as_str() => {
+            let response = next.run(request).await;
+            Ok(response)
+        }
+        _ => Err(StatusCode::UNAUTHORIZED),
+    }
+}
diff --git a/src/main.rs b/src/main.rs
index 4ef129b..d17984f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,42 +1,44 @@
-use std::env;
-use std::sync::Arc;
-use axum::{Router, routing::post};
-use axum::routing::{get, put};
-use dashmap::DashMap;
-use sqlx::PgPool;
-use time::util::local_offset;
-use tokio::sync::broadcast::{channel, Sender};
-use tracing::{info, level_filters::LevelFilter};
-use tracing_subscriber::{EnvFilter, fmt::{self, time::LocalTime}, prelude::*};
 use crate::config::Config;
 use crate::db::init_db_pool;
 use crate::routes::device;
 use crate::routes::start::start;
 use crate::routes::status::status;
-use crate::services::ping::{BroadcastCommands, StatusMap};
+use crate::services::ping::StatusMap;
+use axum::middleware::from_fn_with_state;
+use axum::routing::{get, put};
+use axum::{routing::post, Router};
+use dashmap::DashMap;
+use services::ping::BroadcastCommand;
+use sqlx::PgPool;
+use std::env;
+use std::sync::Arc;
+use tokio::sync::broadcast::{channel, Sender};
+use tracing::{info, level_filters::LevelFilter};
+use tracing_subscriber::fmt::time::UtcTime;
+use tracing_subscriber::{fmt, prelude::*, EnvFilter};
 
-mod auth;
 mod config;
-mod routes;
-mod wol;
 mod db;
 mod error;
+mod extractors;
+mod routes;
 mod services;
+mod wol;
 
 #[tokio::main]
 async fn main() -> color_eyre::eyre::Result<()> {
-
     color_eyre::install()?;
 
-    unsafe { local_offset::set_soundness(local_offset::Soundness::Unsound); }
     let time_format =
         time::macros::format_description!("[year]-[month]-[day] [hour]:[minute]:[second]");
-    let loc = LocalTime::new(time_format);
+    let loc = UtcTime::new(time_format);
+
+    let file_appender = tracing_appender::rolling::daily("logs", "webol.log");
+    let (non_blocking, _guard) = tracing_appender::non_blocking(file_appender);
 
     tracing_subscriber::registry()
-        .with(fmt::layer()
-            .with_timer(loc)
-        )
+        .with(fmt::layer().with_writer(non_blocking).with_ansi(false))
+        .with(fmt::layer().with_timer(loc))
         .with(
             EnvFilter::builder()
                 .with_default_directive(LevelFilter::INFO.into())
@@ -56,8 +58,13 @@ async fn main() -> color_eyre::eyre::Result<()> {
     let (tx, _) = channel(32);
 
     let ping_map: StatusMap = DashMap::new();
-    
-    let shared_state = Arc::new(AppState { db, config: config.clone(), ping_send: tx, ping_map });
+
+    let shared_state = AppState {
+        db,
+        config: config.clone(),
+        ping_send: tx,
+        ping_map,
+    };
 
     let app = Router::new()
         .route("/start", post(start))
@@ -65,20 +72,21 @@ async fn main() -> color_eyre::eyre::Result<()> {
         .route("/device", put(device::put))
         .route("/device", post(device::post))
         .route("/status", get(status))
-        .with_state(shared_state);
+        .route_layer(from_fn_with_state(shared_state.clone(), extractors::auth))
+        .with_state(Arc::new(shared_state));
 
     let addr = config.serveraddr;
     info!("start server on {}", addr);
-    let listener = tokio::net::TcpListener::bind(addr)
-        .await?;
+    let listener = tokio::net::TcpListener::bind(addr).await?;
     axum::serve(listener, app).await?;
 
     Ok(())
 }
 
+#[derive(Clone)]
 pub struct AppState {
     db: PgPool,
     config: Config,
-    ping_send: Sender<BroadcastCommands>,
+    ping_send: Sender<BroadcastCommand>,
     ping_map: StatusMap,
 }
diff --git a/src/routes/mod.rs b/src/routes.rs
index d5ab0d6..d5ab0d6 100644
--- a/src/routes/mod.rs
+++ b/src/routes.rs
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/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()
 }
diff --git a/src/services/mod.rs b/src/services.rs
index a766209..a766209 100644
--- a/src/services/mod.rs
+++ b/src/services.rs
diff --git a/src/services/ping.rs b/src/services/ping.rs
index 9b164c8..9191f86 100644
--- a/src/services/ping.rs
+++ b/src/services/ping.rs
@@ -1,59 +1,58 @@
-use std::str::FromStr;
-use std::net::IpAddr;
-use std::sync::Arc;
-
-use axum::extract::ws::WebSocket;
-use axum::extract::ws::Message;
+use crate::config::Config;
+use crate::db::Device;
 use dashmap::DashMap;
+use ipnetwork::IpNetwork;
 use sqlx::PgPool;
+use std::fmt::Display;
 use time::{Duration, Instant};
 use tokio::sync::broadcast::Sender;
 use tracing::{debug, error, trace};
-use crate::AppState;
-use crate::config::Config;
-use crate::db::Device;
 
 pub type StatusMap = DashMap<String, Value>;
 
 #[derive(Debug, Clone)]
 pub struct Value {
-    pub ip: String,
-    pub online: bool
+    pub ip: IpNetwork,
+    pub online: bool,
 }
 
-pub async fn spawn(tx: Sender<BroadcastCommands>, config: &Config, device: Device, uuid: String, ping_map: &StatusMap, db: &PgPool) {
+pub async fn spawn(
+    tx: Sender<BroadcastCommand>,
+    config: &Config,
+    device: Device,
+    uuid: String,
+    ping_map: &StatusMap,
+    db: &PgPool,
+) {
     let timer = Instant::now();
     let payload = [0; 8];
 
-    let ping_ip = IpAddr::from_str(&device.ip).expect("bad ip");
-
-    let mut msg: Option<BroadcastCommands> = None;
+    let mut msg: Option<BroadcastCommand> = None;
     while msg.is_none() {
-        let ping = surge_ping::ping(
-            ping_ip,
-            &payload
-        ).await;
+        let ping = surge_ping::ping(device.ip.ip(), &payload).await;
 
         if let Err(ping) = ping {
             let ping_timeout = matches!(ping, surge_ping::SurgeError::Timeout { .. });
             if !ping_timeout {
                 error!("{}", ping.to_string());
-                msg = Some(BroadcastCommands::Error(uuid.clone()));
+                msg = Some(BroadcastCommand::error(uuid.clone()));
             }
             if timer.elapsed() >= Duration::minutes(config.pingtimeout) {
-                msg = Some(BroadcastCommands::Timeout(uuid.clone()));
+                msg = Some(BroadcastCommand::timeout(uuid.clone()));
             }
         } else {
-            let (_, duration) = ping.map_err(|err| error!("{}", err.to_string())).expect("fatal error");
+            let (_, duration) = ping
+                .map_err(|err| error!("{}", err.to_string()))
+                .expect("fatal error");
             debug!("ping took {:?}", duration);
-            msg = Some(BroadcastCommands::Success(uuid.clone()));
+            msg = Some(BroadcastCommand::success(uuid.clone()));
         };
     }
 
     let msg = msg.expect("fatal error");
 
     let _ = tx.send(msg.clone());
-    if let BroadcastCommands::Success(..) = msg {
+    if let BroadcastCommands::Success = msg.command {
         sqlx::query!(
             r#"
             UPDATE devices
@@ -62,8 +61,17 @@ pub async fn spawn(tx: Sender<BroadcastCommands>, config: &Config, device: Devic
             "#,
             timer.elapsed().whole_seconds(),
             device.id
-        ).execute(db).await.unwrap();
-        ping_map.insert(uuid.clone(), Value { ip: device.ip.clone(), online: true });
+        )
+        .execute(db)
+        .await
+        .unwrap();
+        ping_map.insert(
+            uuid.clone(),
+            Value {
+                ip: device.ip,
+                online: true,
+            },
+        );
         tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
     }
     trace!("remove {} from ping_map", uuid);
@@ -72,74 +80,48 @@ pub async fn spawn(tx: Sender<BroadcastCommands>, config: &Config, device: Devic
 
 #[derive(Clone, Debug, PartialEq)]
 pub enum BroadcastCommands {
-    Success(String),
-    Timeout(String),
-    Error(String),
+    Success,
+    Timeout,
+    Error,
 }
 
-pub async fn status_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;
+#[derive(Clone, Debug, PartialEq)]
+pub struct BroadcastCommand {
+    pub uuid: String,
+    pub command: BroadcastCommands,
+}
 
-    let device_exists = state.ping_map.contains_key(&uuid);
-    if device_exists {
-        let _ = socket.send(process_device(state.clone(), uuid).await).await;
-    } else {
-        debug!("didn't find any device");
-        let _ = socket.send(Message::Text(format!("notfound_{uuid}"))).await;
-    };
+impl Display for BroadcastCommand {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let prefix = match self.command {
+            BroadcastCommands::Success => "start",
+            BroadcastCommands::Timeout => "timeout",
+            BroadcastCommands::Error => "error",
+        };
 
-    let _ = socket.close().await;
+        f.write_str(format!("{prefix}_{}", self.uuid).as_str())
+    }
 }
 
-async fn get_eta(db: &PgPool) -> i64 {
-    let query = sqlx::query!(
-        r#"SELECT times FROM devices;"#
-    ).fetch_one(db).await.unwrap();
-
-    let times = match query.times {
-        None => { vec![0] },
-        Some(t) => t,
-    };
-    times.iter().sum::<i64>() / i64::try_from(times.len()).unwrap()
+impl BroadcastCommand {
+    pub fn success(uuid: String) -> Self {
+        Self {
+            uuid,
+            command: BroadcastCommands::Success,
+        }
+    }
 
-}
+    pub fn timeout(uuid: String) -> Self {
+        Self {
+            uuid,
+            command: BroadcastCommands::Timeout,
+        }
+    }
 
-async fn process_device(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(format!("start_{uuid}"))
-    } else {
-        loop {
-            trace!("wait for tx message");
-            let message = state.ping_send.subscribe().recv().await.expect("fatal error");
-            trace!("got message {:?}", message);
-            return match message {
-                BroadcastCommands::Success(msg_uuid) => {
-                    if msg_uuid != uuid { continue; }
-                    trace!("message == uuid success");
-                    Message::Text(format!("start_{uuid}"))
-                },
-                BroadcastCommands::Timeout(msg_uuid) => {
-                    if msg_uuid != uuid { continue; }
-                    trace!("message == uuid timeout");
-                    Message::Text(format!("timeout_{uuid}"))
-                },
-                BroadcastCommands::Error(msg_uuid) => {
-                    if msg_uuid != uuid { continue; }
-                    trace!("message == uuid error");
-                    Message::Text(format!("error_{uuid}"))
-                }
-            }
+    pub fn error(uuid: String) -> Self {
+        Self {
+            uuid,
+            command: BroadcastCommands::Error,
         }
     }
 }
diff --git a/src/wol.rs b/src/wol.rs
index 83c0ee6..31cf350 100644
--- a/src/wol.rs
+++ b/src/wol.rs
@@ -1,4 +1,4 @@
-use std::net::{SocketAddr, UdpSocket};
+use std::net::{ToSocketAddrs, UdpSocket};
 
 use crate::error::Error;
 
@@ -11,8 +11,8 @@ pub fn create_buffer(mac_addr: &str) -> Result<Vec<u8>, Error> {
     let mut mac = Vec::new();
     let sp = mac_addr.split(':');
     for f in sp {
-        mac.push(u8::from_str_radix(f, 16).map_err(Error::BufferParse)?);
-    };
+        mac.push(u8::from_str_radix(f, 16)?);
+    }
     let mut buf = vec![255; 6];
     for _ in 0..16 {
         for i in &mac {
@@ -23,8 +23,12 @@ pub fn create_buffer(mac_addr: &str) -> Result<Vec<u8>, Error> {
 }
 
 /// Sends a buffer on UDP broadcast
-pub fn send_packet(bind_addr: &SocketAddr, broadcast_addr: &SocketAddr, buffer: &[u8]) -> Result<usize, Error> {
-    let socket = UdpSocket::bind(bind_addr).map_err(Error::Broadcast)?;
-    socket.set_broadcast(true).map_err(Error::Broadcast)?;
-    socket.send_to(buffer, broadcast_addr).map_err(Error::Broadcast)
+pub fn send_packet<A: ToSocketAddrs>(
+    bind_addr: A,
+    broadcast_addr: A,
+    buffer: &[u8],
+) -> Result<usize, Error> {
+    let socket = UdpSocket::bind(bind_addr)?;
+    socket.set_broadcast(true)?;
+    Ok(socket.send_to(buffer, broadcast_addr)?)
 }