ping timeout and cleanup

5 files changed, 59 insertions, 74 deletions

diff --git a/src/db.rs b/src/db.rs
index 51ea469..c012b47 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -8,7 +8,7 @@ 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,
diff --git a/src/error.rs b/src/error.rs
index 1592a78..5b82534 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,4 +1,3 @@
-use std::error::Error;
 use std::io;
 use axum::http::StatusCode;
 use axum::Json;
@@ -10,9 +9,7 @@ use crate::auth::AuthError;
 #[derive(Debug)]
 pub enum WebolError {
     Generic,
-    // User(UserError),
     Auth(AuthError),
-    Ping(surge_ping::SurgeError),
     DB(sqlx::Error),
     IpParse(<std::net::IpAddr as std::str::FromStr>::Err),
     BufferParse(std::num::ParseIntError),
@@ -22,13 +19,10 @@ pub enum WebolError {
 impl IntoResponse for WebolError {
     fn into_response(self) -> Response {
         let (status, error_message) = match self {
-            Self::Auth(err) => err.get(),
-            // Self::User(err) => err.get(),
-            Self::Generic => (StatusCode::INTERNAL_SERVER_ERROR, ""),
-            Self::Ping(err) => {
-                error!("Ping: {}", err.source().unwrap());
-                (StatusCode::INTERNAL_SERVER_ERROR, "Server Error")
+            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")
@@ -51,17 +45,4 @@ impl IntoResponse for WebolError {
         }));
         (status, body).into_response()
     }
-}
-
-// #[derive(Debug)]
-// pub enum UserError {
-//     UnknownUUID,
-// }
-//
-// impl UserError {
-//     pub fn get(self) -> (StatusCode, &'static str) {
-//         match self {
-//             Self::UnknownUUID => (StatusCode::UNPROCESSABLE_ENTITY, "Unknown UUID"),
-//         }
-//     }
-// }
+}
\ No newline at end of file
diff --git a/src/routes/device.rs b/src/routes/device.rs
index 7353733..1eeff0b 100644
--- a/src/routes/device.rs
+++ b/src/routes/device.rs
@@ -4,13 +4,13 @@ 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!(
@@ -23,6 +23,8 @@ pub async fn get_device(State(state): State<Arc<crate::AppState>>, headers: Head
             payload.id
         ).fetch_one(&state.db).await.map_err(WebolError::DB)?;
 
+        debug!("got device {:?}", device);
+
         Ok(Json(json!(device)))
     } else {
         Err(WebolError::Generic)
@@ -35,7 +37,7 @@ 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!(
@@ -69,7 +71,7 @@ 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!(
diff --git a/src/routes/start.rs b/src/routes/start.rs
index c2c9378..9cd358b 100644
--- a/src/routes/start.rs
+++ b/src/routes/start.rs
@@ -49,8 +49,6 @@ pub async fn start(State(state): State<Arc<crate::AppState>>, headers: HeaderMap
                 debug!("Init ping service");
                 state.ping_map.insert(uuid_gen.clone(), PingValue { ip: device.ip.clone(), online: false });
 
-                warn!("{:?}", state.ping_map);
-
                 crate::services::ping::spawn(state.ping_send.clone(), device.ip, uuid_gen.clone(), &state.ping_map).await
             });
             Some(uuid_genc)
diff --git a/src/services/ping.rs b/src/services/ping.rs
index f0cc4a3..a26dacc 100644
--- a/src/services/ping.rs
+++ b/src/services/ping.rs
@@ -1,14 +1,13 @@
-use std::borrow::Cow;
 use std::sync::Arc;
 
 use axum::extract::{ws::WebSocket};
-use axum::extract::ws::{CloseFrame, Message};
+use axum::extract::ws::Message;
 use dashmap::DashMap;
+use time::{Duration, Instant};
 use tokio::sync::broadcast::{Sender};
-use tracing::{debug, trace, warn};
+use tracing::{debug, error, trace};
 use crate::AppState;
-
-use crate::error::WebolError;
+use crate::config::SETTINGS;
 
 pub type PingMap = DashMap<String, PingValue>;
 
@@ -18,92 +17,97 @@ pub struct PingValue {
     pub online: 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) {
+    let timer = Instant::now();
     let payload = [0; 8];
 
-    // TODO: Better while
     let mut cont = true;
     while cont {
         let ping = surge_ping::ping(
-            ip.parse().map_err(WebolError::IpParse)?,
+            ip.parse().expect("bad ip"),
             &payload
         ).await;
 
         if let Err(ping) = ping {
             cont = matches!(ping, surge_ping::SurgeError::Timeout { .. });
             if !cont {
-                return Err(ping).map_err(WebolError::Ping)
+                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.unwrap();
+            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;
         };
     }
-
-    Ok(())
 }
 
 async fn handle_broadcast_send(tx: &Sender<BroadcastCommands>, ip: String, ping_map: &PingMap, uuid: String) {
-    debug!("sending pingsuccess message");
+    debug!("send pingsuccess message");
     ping_map.insert(uuid.clone(), PingValue { ip: ip.clone(), online: true });
-    let _ = tx.send(BroadcastCommands::PingSuccess(ip));
+    let _ = tx.send(BroadcastCommands::PingSuccess(uuid.clone()));
     tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
-    trace!("remove {} from ping_map", uuid);
+    trace!("remove {} from ping_map after success", uuid);
     ping_map.remove(&uuid);
 }
 
 #[derive(Clone, Debug)]
 pub enum BroadcastCommands {
-    PingSuccess(String)
+    PingSuccess(String),
+    PingTimeout(String)
 }
 
 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;
     let uuid = msg.unwrap().unwrap().into_text().unwrap();
 
     trace!("Search for uuid: {:?}", uuid);
 
-    // TODO: Handle Error
-    let device = state.ping_map.get(&uuid).unwrap().to_owned();
+    match state.ping_map.get(&uuid) {
+        Some(device) => {
+            debug!("got device: {} (online: {})", device.ip, device.online);
+            let _ = socket.send(process_device(state.clone(), uuid, device.to_owned()).await).await;
+        },
+        None => {
+            debug!("didn't find any device");
+            let _ = socket.send(Message::Text(format!("notfound_{}", uuid))).await;
+        },
+    };
 
-    trace!("got device: {:?}", device);
+    let _ = socket.close().await;
+}
 
+async fn process_device(state: Arc<AppState>, uuid: String, device: PingValue) -> Message {
     match device.online {
         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();
+            Message::Text(format!("start_{}", uuid))
         },
         false => {
-            let ip = device.ip.to_owned();
             loop{
                 trace!("wait for tx message");
-                let message = state.ping_send.subscribe().recv().await.unwrap();
-                trace!("GOT = {:?}", message);
-                // if let BroadcastCommands::PingSuccess(msg_ip) = message {
-                //     if msg_ip == ip {
-                //         trace!("message == ip");
-                //         break;
-                //     }
-                // }
-                let BroadcastCommands::PingSuccess(msg_ip) = message;
-                if msg_ip == ip {
-                    trace!("message == ip");
-                    break;
+                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))
+                    }
                 }
-            };
-
-            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!("{:?}", state.ping_map);
+            }
         }
     }
 }
\ No newline at end of file