Merge pull request #6 from FxQnLr/ping

Ping

11 files changed, 243 insertions, 55 deletions

diff --git a/src/auth.rs b/src/auth.rs
index 0fffa60..e4b1c2f 100644
--- a/src/auth.rs
+++ b/src/auth.rs
@@ -1,8 +1,8 @@
-use std::error::Error;
 use axum::headers::HeaderValue;
 use axum::http::StatusCode;
+use axum::http::header::ToStrError;
 use tracing::{debug, error, trace};
-use crate::auth::AuthError::{MissingSecret, ServerError, WrongSecret};
+use crate::auth::AuthError::{MissingSecret, WrongSecret};
 use crate::config::SETTINGS;
 
 pub fn auth(secret: Option<&HeaderValue>) -> Result<bool, AuthError> {
@@ -11,8 +11,8 @@ pub fn auth(secret: Option<&HeaderValue>) -> Result<bool, AuthError> {
         trace!("value exists");
         let key = SETTINGS
             .get_string("apikey")
-            .map_err(|err| ServerError(Box::new(err)))?;
-        if value.to_str().map_err(|err| ServerError(Box::new(err)))? == key.as_str() {
+            .map_err(AuthError::Config)?;
+        if value.to_str().map_err(AuthError::HeaderToStr)? == key.as_str() {
             debug!("successful auth");
             Ok(true)
         } else {
@@ -29,15 +29,20 @@ pub fn auth(secret: Option<&HeaderValue>) -> Result<bool, AuthError> {
 pub enum AuthError {
     WrongSecret,
     MissingSecret,
-    ServerError(Box<dyn Error>),
+    Config(config::ConfigError),
+    HeaderToStr(ToStrError)
 }
 
 impl AuthError {
     pub fn get(self) -> (StatusCode, &'static str) {
         match self {
-            AuthError::WrongSecret => (StatusCode::UNAUTHORIZED, "Wrong credentials"),
-            AuthError::MissingSecret => (StatusCode::BAD_REQUEST, "Missing credentials"),
-            AuthError::ServerError(err) => {
+            Self::WrongSecret => (StatusCode::UNAUTHORIZED, "Wrong credentials"),
+            Self::MissingSecret => (StatusCode::BAD_REQUEST, "Missing credentials"),
+            Self::Config(err) => {
+                error!("server error: {}", err.to_string());
+                (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
+            },
+            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 3c51e2b..c012b47 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -8,11 +8,12 @@ use tracing::{debug, info};
 #[cfg(not(debug_assertions))]
 use crate::config::SETTINGS;
 
-#[derive(Serialize)]
+#[derive(Serialize, Debug)]
 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/error.rs b/src/error.rs
index db2fc86..5b82534 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,4 +1,4 @@
-use std::error::Error;
+use std::io;
 use axum::http::StatusCode;
 use axum::Json;
 use axum::response::{IntoResponse, Response};
@@ -8,25 +8,41 @@ use crate::auth::AuthError;
 
 #[derive(Debug)]
 pub enum WebolError {
-    Auth(AuthError),
     Generic,
-    Server(Box<dyn Error>),
+    Auth(AuthError),
+    DB(sqlx::Error),
+    IpParse(<std::net::IpAddr as std::str::FromStr>::Err),
+    BufferParse(std::num::ParseIntError),
+    Broadcast(io::Error),
 }
 
 impl IntoResponse for WebolError {
     fn into_response(self) -> Response {
         let (status, error_message) = match self {
-            WebolError::Auth(err) => err.get(),
-            WebolError::Generic => (StatusCode::INTERNAL_SERVER_ERROR, ""),
-            WebolError::Server(err) => {
+            Self::Auth(err) => {
+                err.get()
+            },
+            Self::Generic => (StatusCode::INTERNAL_SERVER_ERROR, ""),
+            Self::IpParse(err) => {
+                error!("server error: {}", err.to_string());
+                (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
+            },
+            Self::DB(err) => {
+                error!("server error: {}", err.to_string());
+                (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
+            },
+            Self::Broadcast(err) => {
+                error!("server error: {}", err.to_string());
+                (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
+            },
+            Self::BufferParse(err) => {
                 error!("server error: {}", err.to_string());
                 (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
             },
-
         };
         let body = Json(json!({
             "error": error_message,
         }));
         (status, body).into_response()
     }
-}
+}
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index ce12cf6..e96b736 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,14 +2,18 @@ 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::SETTINGS;
 use crate::db::init_db_pool;
 use crate::routes::device::{get_device, post_device, put_device};
 use crate::routes::start::start;
+use crate::routes::status::status;
+use crate::services::ping::{BroadcastCommands, PingMap};
 
 mod auth;
 mod config;
@@ -17,6 +21,7 @@ mod routes;
 mod wol;
 mod db;
 mod error;
+mod services;
 
 #[tokio::main]
 async fn main() {
@@ -43,13 +48,18 @@ async fn main() {
     let db = init_db_pool().await;
     sqlx::migrate!().run(&db).await.unwrap();
 
-    let shared_state = Arc::new(AppState { db });
+    let (tx, _) = channel(32);
+
+    let ping_map: PingMap = DashMap::new();
+    
+    let shared_state = Arc::new(AppState { db, ping_send: tx, ping_map });
 
     let app = Router::new()
         .route("/start", post(start))
         .route("/device", get(get_device))
         .route("/device", put(put_device))
         .route("/device", post(post_device))
+        .route("/status", get(status))
         .with_state(shared_state);
 
     let addr = SETTINGS.get_string("serveraddr").unwrap_or("0.0.0.0:7229".to_string());
@@ -61,5 +71,7 @@ async fn main() {
 }
 
 pub struct AppState {
-    db: PgPool
-}
+    db: PgPool,
+    ping_send: Sender<BroadcastCommands>,
+    ping_map: PingMap,
+}
\ No newline at end of file
diff --git a/src/routes/device.rs b/src/routes/device.rs
index 025c7d0..1eeff0b 100644
--- a/src/routes/device.rs
+++ b/src/routes/device.rs
@@ -4,24 +4,26 @@ use axum::headers::HeaderMap;
 use axum::Json;
 use serde::{Deserialize, Serialize};
 use serde_json::{json, Value};
-use tracing::info;
+use tracing::{debug, info};
 use crate::auth::auth;
 use crate::db::Device;
 use crate::error::WebolError;
 
 pub async fn get_device(State(state): State<Arc<crate::AppState>>, headers: HeaderMap, Json(payload): Json<GetDevicePayload>) -> Result<Json<Value>, WebolError> {
-    info!("GET request");
+    info!("add device {}", payload.id);
     let secret = headers.get("authorization");
     if auth(secret).map_err(WebolError::Auth)? {
         let device = sqlx::query_as!(
             Device,
             r#"
-            SELECT id, mac, broadcast_addr
+            SELECT id, mac, broadcast_addr, ip
             FROM devices
             WHERE id = $1;
             "#,
             payload.id
-        ).fetch_one(&state.db).await.map_err(|err| WebolError::Server(Box::new(err)))?;
+        ).fetch_one(&state.db).await.map_err(WebolError::DB)?;
+
+        debug!("got device {:?}", device);
 
         Ok(Json(json!(device)))
     } else {
@@ -35,18 +37,19 @@ pub struct GetDevicePayload {
 }
 
 pub async fn put_device(State(state): State<Arc<crate::AppState>>, headers: HeaderMap, Json(payload): Json<PutDevicePayload>) -> Result<Json<Value>, WebolError> {
-    info!("PUT request");
+    info!("add device {} ({}, {}, {})", payload.id, payload.mac, payload.broadcast_addr, payload.ip);
     let secret = headers.get("authorization");
     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
-        ).execute(&state.db).await.map_err(|err| WebolError::Server(Box::new(err)))?;
+            payload.broadcast_addr,
+            payload.ip
+        ).execute(&state.db).await.map_err(WebolError::DB)?;
 
         Ok(Json(json!(PutDeviceResponse { success: true })))
     } else {
@@ -59,6 +62,7 @@ pub struct PutDevicePayload {
     id: String,
     mac: String,
     broadcast_addr: String,
+    ip: String
 }
 
 #[derive(Serialize)]
@@ -67,20 +71,21 @@ pub struct PutDeviceResponse {
 }
 
 pub async fn post_device(State(state): State<Arc<crate::AppState>>, headers: HeaderMap, Json(payload): Json<PostDevicePayload>) -> Result<Json<Value>, WebolError> {
-    info!("POST request");
+    info!("edit device {} ({}, {}, {})", payload.id, payload.mac, payload.broadcast_addr, payload.ip);
     let secret = headers.get("authorization");
     if auth(secret).map_err(WebolError::Auth)? {
         let device = sqlx::query_as!(
             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(|err| WebolError::Server(Box::new(err)))?;
+        ).fetch_one(&state.db).await.map_err(WebolError::DB)?;
 
         Ok(Json(json!(device)))
     } else {
@@ -93,4 +98,5 @@ pub struct PostDevicePayload {
     id: String,
     mac: String,
     broadcast_addr: String,
-}
\ No newline at end of file
+    ip: String,
+}
diff --git a/src/routes/mod.rs b/src/routes/mod.rs
index 12fbfab..d5ab0d6 100644
--- a/src/routes/mod.rs
+++ b/src/routes/mod.rs
@@ -1,2 +1,3 @@
 pub mod start;
-pub mod device;
\ No newline at end of file
+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 163d58c..271f924 100644
--- a/src/routes/start.rs
+++ b/src/routes/start.rs
@@ -4,26 +4,30 @@ use serde::{Deserialize, Serialize};
 use std::sync::Arc;
 use axum::extract::State;
 use serde_json::{json, Value};
-use tracing::info;
+use tracing::{debug, info};
+use uuid::Uuid;
 use crate::auth::auth;
 use crate::config::SETTINGS;
 use crate::wol::{create_buffer, send_packet};
 use crate::db::Device;
 use crate::error::WebolError;
+use crate::services::ping::PingValue;
 
+#[axum_macros::debug_handler]
 pub async fn start(State(state): State<Arc<crate::AppState>>, headers: HeaderMap, Json(payload): Json<StartPayload>) -> Result<Json<Value>, WebolError> {
     info!("POST request");
     let secret = headers.get("authorization");
-    if auth(secret).map_err(WebolError::Auth)? {
+    let authorized = auth(secret).map_err(WebolError::Auth)?;
+    if authorized {
         let device = sqlx::query_as!(
             Device,
             r#"
-            SELECT id, mac, broadcast_addr
+            SELECT id, mac, broadcast_addr, ip
             FROM devices
             WHERE id = $1;
             "#,
             payload.id
-        ).fetch_one(&state.db).await.map_err(|err| WebolError::Server(Box::new(err)))?;
+        ).fetch_one(&state.db).await.map_err(WebolError::DB)?;
 
         info!("starting {}", device.id);
 
@@ -32,11 +36,23 @@ pub async fn start(State(state): State<Arc<crate::AppState>>, headers: HeaderMap
             .unwrap_or("0.0.0.0:1111".to_string());
 
         let _ = send_packet(
-            &bind_addr.parse().map_err(|err| WebolError::Server(Box::new(err)))?,
-            &device.broadcast_addr.parse().map_err(|err| WebolError::Server(Box::new(err)))?,
-            create_buffer(&device.mac).map_err(|err| WebolError::Server(Box::new(err)))?
-        ).map_err(|err| WebolError::Server(Box::new(err)));
-        Ok(Json(json!(StartResponse { id: device.id, boot: true })))
+            &bind_addr.parse().map_err(WebolError::IpParse)?,
+            &device.broadcast_addr.parse().map_err(WebolError::IpParse)?,
+            create_buffer(&device.mac)?
+        )?;
+
+        let uuid = if payload.ping.is_some_and(|ping| ping) {
+            let uuid_gen = Uuid::new_v4().to_string();
+            let uuid_genc = uuid_gen.clone();
+            tokio::spawn(async move {
+                debug!("init ping service");
+                state.ping_map.insert(uuid_gen.clone(), PingValue { ip: device.ip.clone(), online: false });
+
+                crate::services::ping::spawn(state.ping_send.clone(), device.ip, uuid_gen.clone(), &state.ping_map).await
+            });
+            Some(uuid_genc)
+        } else { None };
+        Ok(Json(json!(StartResponse { id: device.id, boot: true, uuid })))
     } else {
         Err(WebolError::Generic)
     }
@@ -45,11 +61,12 @@ pub async fn start(State(state): State<Arc<crate::AppState>>, headers: HeaderMap
 #[derive(Deserialize)]
 pub struct StartPayload {
     id: String,
-    _test: Option<bool>,
+    ping: Option<bool>,
 }
 
 #[derive(Serialize)]
 struct StartResponse {
     id: String,
     boot: bool,
-}
\ No newline at end of file
+    uuid: Option<String>,
+}
diff --git a/src/routes/status.rs b/src/routes/status.rs
new file mode 100644
index 0000000..45f3e51
--- /dev/null
+++ b/src/routes/status.rs
@@ -0,0 +1,10 @@
+use std::sync::Arc;
+use axum::extract::{State, WebSocketUpgrade};
+use axum::response::Response;
+use crate::AppState;
+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))
+}
\ No newline at end of file
diff --git a/src/services/mod.rs b/src/services/mod.rs
new file mode 100644
index 0000000..a766209
--- /dev/null
+++ b/src/services/mod.rs
@@ -0,0 +1 @@
+pub mod ping;
diff --git a/src/services/ping.rs b/src/services/ping.rs
new file mode 100644
index 0000000..d900acb
--- /dev/null
+++ b/src/services/ping.rs
@@ -0,0 +1,118 @@
+use std::sync::Arc;
+
+use axum::extract::{ws::WebSocket};
+use axum::extract::ws::Message;
+use dashmap::DashMap;
+use time::{Duration, Instant};
+use tokio::sync::broadcast::{Sender};
+use tracing::{debug, error, trace};
+use crate::AppState;
+use crate::config::SETTINGS;
+
+pub type PingMap = DashMap<String, PingValue>;
+
+#[derive(Debug, Clone)]
+pub struct PingValue {
+    pub ip: String,
+    pub online: bool
+}
+
+pub async fn spawn(tx: Sender<BroadcastCommands>, ip: String, uuid: String, ping_map: &PingMap) {
+    let timer = Instant::now();
+    let payload = [0; 8];
+
+    let mut cont = true;
+    while cont {
+        let ping = surge_ping::ping(
+            ip.parse().expect("bad ip"),
+            &payload
+        ).await;
+
+        if let Err(ping) = ping {
+            cont = matches!(ping, surge_ping::SurgeError::Timeout { .. });
+            if !cont {
+                error!("{}", ping.to_string());
+            }
+            if timer.elapsed() >= Duration::minutes(SETTINGS.get_int("pingtimeout").unwrap_or(10)) {
+                let _ = tx.send(BroadcastCommands::PingTimeout(uuid.clone()));
+                trace!("remove {} from ping_map after timeout", uuid);
+                ping_map.remove(&uuid);
+                cont = false;
+            }
+        } else {
+            let (_, duration) = ping.map_err(|err| error!("{}", err.to_string())).expect("fatal error");
+            debug!("ping took {:?}", duration);
+            cont = false;
+            handle_broadcast_send(&tx, ip.clone(), ping_map, uuid.clone()).await;
+        };
+    }
+}
+
+async fn handle_broadcast_send(tx: &Sender<BroadcastCommands>, ip: String, ping_map: &PingMap, uuid: String) {
+    debug!("send pingsuccess message");
+    let _ = tx.send(BroadcastCommands::PingSuccess(uuid.clone()));
+    trace!("sent message");
+    ping_map.insert(uuid.clone(), PingValue { ip: ip.clone(), online: true });
+    trace!("updated ping_map");
+    tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
+    debug!("remove {} from ping_map after success", uuid);
+    ping_map.remove(&uuid);
+}
+
+#[derive(Clone, Debug)]
+pub enum BroadcastCommands {
+    PingSuccess(String),
+    PingTimeout(String)
+}
+
+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 device_exists = state.ping_map.contains_key(&uuid);
+    match device_exists {
+        true => {
+            let _ = socket.send(process_device(state.clone(), uuid).await).await;
+        },
+        false => {
+            debug!("didn't find any device");
+            let _ = socket.send(Message::Text(format!("notfound_{}", uuid))).await;
+        },
+    };
+
+    let _ = socket.close().await;
+}
+
+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);
+    match device.online {
+        true => {
+            debug!("already started");
+            Message::Text(format!("start_{}", uuid))
+        },
+        false => {
+            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::PingSuccess(msg_uuid) => {
+                        if msg_uuid != uuid { continue; }
+                        trace!("message == uuid success");
+                        Message::Text(format!("start_{}", uuid))
+                    },
+                    BroadcastCommands::PingTimeout(msg_uuid) => {
+                        if msg_uuid != uuid { continue; }
+                        trace!("message == uuid timeout");
+                        Message::Text(format!("timeout_{}", uuid))
+                    }
+                }
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/wol.rs b/src/wol.rs
index 80b66cd..0cdcae3 100644
--- a/src/wol.rs
+++ b/src/wol.rs
@@ -1,16 +1,17 @@
 use std::net::{SocketAddr, UdpSocket};
-use std::num::ParseIntError;
+
+use crate::error::WebolError;
 
 /// Creates the magic packet from a mac address
 ///
 /// # Panics
 ///
 /// Panics if `mac_addr` is an invalid mac
-pub fn create_buffer(mac_addr: &str) -> Result<Vec<u8>, ParseIntError> {
+pub fn create_buffer(mac_addr: &str) -> Result<Vec<u8>, WebolError> {
     let mut mac = Vec::new();
     let sp = mac_addr.split(':');
     for f in sp {
-        mac.push(u8::from_str_radix(f, 16)?);
+        mac.push(u8::from_str_radix(f, 16).map_err(WebolError::BufferParse)?)
     };
     let mut buf = vec![255; 6];
     for _ in 0..16 {
@@ -22,8 +23,8 @@ pub fn create_buffer(mac_addr: &str) -> Result<Vec<u8>, ParseIntError> {
 }
 
 /// Sends a buffer on UDP broadcast
-pub fn send_packet(bind_addr: &SocketAddr, broadcast_addr: &SocketAddr, buffer: Vec<u8>) -> Result<usize, std::io::Error> {
-    let socket = UdpSocket::bind(bind_addr)?;
-    socket.set_broadcast(true)?;
-    socket.send_to(&buffer, broadcast_addr)
-}
\ No newline at end of file
+pub fn send_packet(bind_addr: &SocketAddr, broadcast_addr: &SocketAddr, buffer: Vec<u8>) -> Result<usize, WebolError> {
+    let socket = UdpSocket::bind(bind_addr).map_err(WebolError::Broadcast)?;
+    socket.set_broadcast(true).map_err(WebolError::Broadcast)?;
+    socket.send_to(&buffer, broadcast_addr).map_err(WebolError::Broadcast)
+}