10 files changed, 260 insertions, 37 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 09f5744..bf813bf 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -89,6 +89,7 @@ checksum = "3b829e4e32b91e643de6eafe82b1d90675f5874230191a4ffbc1b336dec4d6bf"
 dependencies = [
  "async-trait",
  "axum-core",
+ "base64 0.21.4",
  "bitflags 1.3.2",
  "bytes",
  "futures-util",
@@ -107,8 +108,10 @@ dependencies = [
  "serde_json",
  "serde_path_to_error",
  "serde_urlencoded",
+ "sha1",
  "sync_wrapper",
  "tokio",
+ "tokio-tungstenite",
  "tower",
  "tower-layer",
  "tower-service",
@@ -294,6 +297,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "data-encoding"
+version = "2.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c2e66c9d817f1720209181c316d28635c050fa304f9c79e47a520882661b7308"
+
+[[package]]
 name = "der"
 version = "0.7.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -517,6 +526,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6fb8d784f27acf97159b40fc4db5ecd8aa23b9ad5ef69cdd136d3bc80665f0c0"
 
 [[package]]
+name = "glob"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b"
+
+[[package]]
 name = "hashbrown"
 version = "0.12.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -843,6 +858,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "no-std-net"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "43794a0ace135be66a25d3ae77d41b91615fb68ae937f904090203e81f755b65"
+
+[[package]]
 name = "nom"
 version = "7.1.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1115,6 +1136,48 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964"
 
 [[package]]
+name = "pnet_base"
+version = "0.33.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "872e46346144ebf35219ccaa64b1dffacd9c6f188cd7d012bd6977a2a838f42e"
+dependencies = [
+ "no-std-net",
+]
+
+[[package]]
+name = "pnet_macros"
+version = "0.33.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2a780e80005c2e463ec25a6e9f928630049a10b43945fea83207207d4a7606f4"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "regex",
+ "syn 1.0.109",
+]
+
+[[package]]
+name = "pnet_macros_support"
+version = "0.33.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6d932134f32efd7834eb8b16d42418dac87086347d1bc7d142370ef078582bc"
+dependencies = [
+ "pnet_base",
+]
+
+[[package]]
+name = "pnet_packet"
+version = "0.33.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8bde678bbd85cb1c2d99dc9fc596e57f03aa725f84f3168b0eaf33eeccb41706"
+dependencies = [
+ "glob",
+ "pnet_base",
+ "pnet_macros",
+ "pnet_macros_support",
+]
+
+[[package]]
 name = "powerfmt"
 version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1694,6 +1757,22 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "81cdd64d312baedb58e21336b31bc043b77e01cc99033ce76ef539f78e965ebc"
 
 [[package]]
+name = "surge-ping"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "af341b2be485d647b5dc4cfb2da99efac35b5c95748a08fb7233480fedc5ead3"
+dependencies = [
+ "hex",
+ "parking_lot",
+ "pnet_packet",
+ "rand",
+ "socket2 0.5.5",
+ "thiserror",
+ "tokio",
+ "tracing",
+]
+
+[[package]]
 name = "syn"
 version = "1.0.109"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1850,6 +1929,18 @@ dependencies = [
 ]
 
 [[package]]
+name = "tokio-tungstenite"
+version = "0.20.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "212d5dcb2a1ce06d81107c3d0ffa3121fe974b73f068c8282cb1c32328113b6c"
+dependencies = [
+ "futures-util",
+ "log",
+ "tokio",
+ "tungstenite",
+]
+
+[[package]]
 name = "toml"
 version = "0.5.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1956,6 +2047,25 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3528ecfd12c466c6f163363caf2d02a71161dd5e1cc6ae7b34207ea2d42d81ed"
 
 [[package]]
+name = "tungstenite"
+version = "0.20.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9e3dac10fd62eaf6617d3a904ae222845979aec67c615d1c842b4002c7666fb9"
+dependencies = [
+ "byteorder",
+ "bytes",
+ "data-encoding",
+ "http",
+ "httparse",
+ "log",
+ "rand",
+ "sha1",
+ "thiserror",
+ "url",
+ "utf-8",
+]
+
+[[package]]
 name = "typenum"
 version = "1.17.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -2012,6 +2122,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "utf-8"
+version = "0.7.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9"
+
+[[package]]
 name = "valuable"
 version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -2054,6 +2170,7 @@ dependencies = [
  "serde",
  "serde_json",
  "sqlx",
+ "surge-ping",
  "time",
  "tokio",
  "tracing",
diff --git a/Cargo.toml b/Cargo.toml
index e772810..9bdc4da 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,7 +6,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-axum = { version = "0.6.20", features = ["headers"] }
+axum = { version = "0.6.20", features = ["headers", "ws"] }
 tokio = { version = "1.32.0", features = ["macros", "rt-multi-thread"] }
 tracing = "0.1.37"
 tracing-subscriber = { version = "0.3.17", features = ["env-filter", "local-time", "time"] }
@@ -16,3 +16,4 @@ serde_json = "1.0.107"
 config = "0.13.3"
 once_cell = "1.18.0"
 sqlx = { version = "0.7.1", features = ["postgres", "runtime-tokio"]}
+surge-ping = "0.8.0"
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/error.rs b/src/error.rs
index db2fc86..f143ee9 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,4 +1,5 @@
 use std::error::Error;
+use std::io;
 use axum::http::StatusCode;
 use axum::Json;
 use axum::response::{IntoResponse, Response};
@@ -8,21 +9,45 @@ use crate::auth::AuthError;
 
 #[derive(Debug)]
 pub enum WebolError {
-    Auth(AuthError),
     Generic,
-    Server(Box<dyn Error>),
+    Auth(AuthError),
+    Ping(surge_ping::SurgeError),
+    DB(sqlx::Error),
+    IpParse(<std::net::IpAddr as std::str::FromStr>::Err),
+    BufferParse(std::num::ParseIntError),
+    Broadcast(io::Error),
+    Axum(axum::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::Ping(err) => {
+                error!("Ping: {}", err.source().unwrap());
+                (StatusCode::INTERNAL_SERVER_ERROR, "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")
+            },
+            Self::Axum(err) => {
                 error!("server error: {}", err.to_string());
                 (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
             },
-
         };
         let body = Json(json!({
             "error": error_message,
diff --git a/src/main.rs b/src/main.rs
index ce12cf6..9c31ec8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,12 +4,14 @@ use axum::{Router, routing::post};
 use axum::routing::{get, put};
 use sqlx::PgPool;
 use time::util::local_offset;
+use tokio::sync::mpsc::{self, 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::services::ping::ws_ping;
 
 mod auth;
 mod config;
@@ -17,6 +19,7 @@ mod routes;
 mod wol;
 mod db;
 mod error;
+mod services;
 
 #[tokio::main]
 async fn main() {
@@ -43,13 +46,23 @@ async fn main() {
     let db = init_db_pool().await;
     sqlx::migrate!().run(&db).await.unwrap();
 
-    let shared_state = Arc::new(AppState { db });
+    let (tx, mut rx) = mpsc::channel(32);
+    
+    // FIXME: once_cell? or just static mutable
+    tokio::spawn( async move {
+        while let Some(message) = rx.recv().await {
+            println!("GOT = {}", message);
+        }
+    });
+
+    let shared_state = Arc::new(AppState { db, ping_send: tx });
 
     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(ws_ping))
         .with_state(shared_state);
 
     let addr = SETTINGS.get_string("serveraddr").unwrap_or("0.0.0.0:7229".to_string());
@@ -61,5 +74,7 @@ async fn main() {
 }
 
 pub struct AppState {
-    db: PgPool
+    db: PgPool,
+    ping_send: Sender<String>,
+    // ping_receive: Receiver<String>
 }
diff --git a/src/routes/device.rs b/src/routes/device.rs
index 025c7d0..248d1e0 100644
--- a/src/routes/device.rs
+++ b/src/routes/device.rs
@@ -21,7 +21,7 @@ pub async fn get_device(State(state): State<Arc<crate::AppState>>, headers: Head
             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)?;
 
         Ok(Json(json!(device)))
     } else {
@@ -46,7 +46,7 @@ pub async fn put_device(State(state): State<Arc<crate::AppState>>, headers: Head
             payload.id,
             payload.mac,
             payload.broadcast_addr
-        ).execute(&state.db).await.map_err(|err| WebolError::Server(Box::new(err)))?;
+        ).execute(&state.db).await.map_err(WebolError::DB)?;
 
         Ok(Json(json!(PutDeviceResponse { success: true })))
     } else {
@@ -80,7 +80,7 @@ pub async fn post_device(State(state): State<Arc<crate::AppState>>, headers: Hea
             payload.mac,
             payload.broadcast_addr,
             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 +93,4 @@ pub struct PostDevicePayload {
     id: String,
     mac: String,
     broadcast_addr: String,
-}
\ No newline at end of file
+}
diff --git a/src/routes/start.rs b/src/routes/start.rs
index 163d58c..b45fe5b 100644
--- a/src/routes/start.rs
+++ b/src/routes/start.rs
@@ -14,7 +14,8 @@ use crate::error::WebolError;
 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#"
@@ -23,7 +24,7 @@ pub async fn start(State(state): State<Arc<crate::AppState>>, headers: HeaderMap
             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,10 +33,14 @@ 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)));
+            &bind_addr.parse().map_err(WebolError::IpParse)?,
+            &device.broadcast_addr.parse().map_err(WebolError::IpParse)?,
+            create_buffer(&device.mac)?
+        )?;
+
+        if payload.ping.is_some_and(|ping| ping) {
+            tokio::spawn(async move {crate::services::ping::spawn(state.ping_send.clone()).await});
+        }
         Ok(Json(json!(StartResponse { id: device.id, boot: true })))
     } else {
         Err(WebolError::Generic)
@@ -45,11 +50,11 @@ 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
+}
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..6e710ec
--- /dev/null
+++ b/src/services/ping.rs
@@ -0,0 +1,53 @@
+use std::sync::Arc;
+
+use axum::{extract::{WebSocketUpgrade, ws::WebSocket, State}, response::Response};
+use tokio::sync::mpsc::Sender;
+use tracing::{debug, error};
+
+use crate::{error::WebolError, AppState};
+
+pub async fn spawn(tx: Sender<String>) -> Result<(), WebolError> {
+    let payload = [0; 8];
+
+    let mut cont = true;
+    while cont {
+        let ping = surge_ping::ping(
+            "192.168.178.28".parse().map_err(WebolError::IpParse)?,
+            &payload
+        ).await;
+
+        if let Err(ping) = ping {
+            cont = matches!(ping, surge_ping::SurgeError::Timeout { .. });
+
+            debug!("{}", cont);
+            
+            if !cont {
+                return Err(ping).map_err(WebolError::Ping)
+            }
+
+        } else {
+            let (_, duration) = ping.unwrap();
+            debug!("Ping took {:?}", duration);
+            cont = false;
+            // FIXME: remove unwrap
+            tx.send("Got ping".to_string()).await.unwrap();
+        };
+    }
+
+    Ok(())
+}
+
+pub async fn ws_ping(ws: WebSocketUpgrade, State(_state): State<Arc<AppState>>) -> Response {
+    ws.on_upgrade(handle_socket)
+}
+
+// FIXME: Handle commands through enum
+async fn handle_socket(mut socket: WebSocket) {
+    // TODO: Understand Cow
+
+    // match socket.send(axum::extract::ws::Message::Close(Some(CloseFrame { code: 4000, reason: Cow::Owned("started".to_owned()) }))).await.map_err(WebolError::Axum) {
+    match socket.send(axum::extract::ws::Message::Text("started".to_string())).await.map_err(WebolError::Axum) {
+        Ok(..) => (),
+        Err(err) => { error!("Server Error: {:?}", err) }
+    };
+}
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)
+}