Fix bugs and add controller action for controller ws

emgauwa-controller/src/ws/mod.rs

@@ -0,0 +1,109 @@
+use emgauwa_lib::db::DbController;
+use emgauwa_lib::errors::{DatabaseError, EmgauwaError};
+use emgauwa_lib::models::Controller;
+use emgauwa_lib::types::ControllerWsAction;
+use futures::{SinkExt, StreamExt};
+use sqlx::pool::PoolConnection;
+use sqlx::{Pool, Sqlite};
+use tokio_tungstenite::tungstenite::Message;
+use tokio_tungstenite::{connect_async, tungstenite};
+
+pub async fn run_websocket(
+	pool: Pool<Sqlite>,
+	this: Controller,
+	url: &str,
+) -> Result<(), EmgauwaError> {
+	match connect_async(url).await {
+		Ok(connection) => {
+			let (ws_stream, _) = connection;
+
+			let (mut write, read) = ws_stream.split();
+
+			let ws_action = ControllerWsAction::Register(this.clone());
+
+			let ws_action_json = serde_json::to_string(&ws_action)?;
+			if let Err(err) = write.send(Message::text(ws_action_json)).await {
+				log::error!("Failed to register at websocket: {}", err);
+				return Ok(());
+			}
+
+			let read_handler = read.for_each(|msg| handle_message(pool.clone(), this.clone(), msg));
+
+			read_handler.await;
+
+			log::warn!("Lost connection to websocket");
+		}
+		Err(err) => {
+			log::warn!("Failed to connect to websocket: {}", err,);
+		}
+	}
+	Ok(())
+}
+
+async fn handle_message(
+	pool: Pool<Sqlite>,
+	this: Controller,
+	message_result: Result<Message, tungstenite::Error>,
+) {
+	let msg = match message_result {
+		Ok(msg) => msg,
+		Err(err) => {
+			log::error!("Error reading message: {}", err);
+			return;
+		}
+	};
+	match msg {
+		Message::Text(text) => match serde_json::from_str(&text) {
+			Ok(action) => {
+				log::debug!("Received action: {:?}", action);
+				let mut pool_conn = match pool.acquire().await {
+					Ok(conn) => conn,
+					Err(err) => {
+						log::error!("Failed to acquire database connection: {:?}", err);
+						return;
+					}
+				};
+				let action_res = handle_action(&mut pool_conn, this, action).await;
+				if let Err(e) = action_res {
+					log::error!("Error handling action: {:?}", e);
+				}
+			}
+			Err(e) => {
+				log::error!("Error deserializing action: {:?}", e);
+			}
+		},
+		_ => (),
+	}
+}
+
+pub async fn handle_action(
+	conn: &mut PoolConnection<Sqlite>,
+	this: Controller,
+	action: ControllerWsAction,
+) -> Result<(), EmgauwaError> {
+	match action {
+		ControllerWsAction::Controller(controller) => {
+			handle_controller(conn, this, controller).await
+		}
+		_ => Ok(()),
+	}
+}
+
+pub async fn handle_controller(
+	conn: &mut PoolConnection<Sqlite>,
+	this: Controller,
+	controller: Controller,
+) -> Result<(), EmgauwaError> {
+	if controller.c.uid != this.c.uid {
+		return Err(EmgauwaError::Other(String::from(
+			"Controller UID mismatch during update",
+		)));
+	}
+	DbController::get_by_uid(conn, &controller.c.uid)
+		.await?
+		.ok_or(DatabaseError::NotFound)?
+		.update(conn, controller.c.name.as_str(), this.c.relay_count)
+		.await?;
+
+	Ok(())
+}