333 lines
No EOL
12 KiB
Rust
333 lines
No EOL
12 KiB
Rust
use std::collections::HashMap;
|
|
use std::io;
|
|
use futures::executor::block_on;
|
|
use rand::random;
|
|
use sqlx::{Pool, Sqlite};
|
|
use tokio::signal;
|
|
use tokio::sync::mpsc;
|
|
use tokio::sync::mpsc::UnboundedSender;
|
|
use emgauwa_common::db::{DbController, DbJunctionRelaySchedule, DbRelay, DbSchedule};
|
|
use emgauwa_common::errors::{DatabaseError, EmgauwaError};
|
|
use emgauwa_common::models::{convert_db_list, Controller, FromDbModel, Relay};
|
|
use emgauwa_common::types::{ControllersWsAction, EmgauwaUid, RelayStates};
|
|
use emgauwa_common::utils;
|
|
use crate::server::{WsConnId, WsControllerAction, WsRelayAction, WsServerAction};
|
|
|
|
#[derive(Debug)]
|
|
pub struct WsServer {
|
|
/// Map of connection IDs to their message receivers.
|
|
controller_sessions: HashMap<WsConnId, (mpsc::UnboundedSender<WsControllerAction>, Option<EmgauwaUid>)>,
|
|
controllers: HashMap<EmgauwaUid, (Controller, WsConnId)>,
|
|
|
|
relay_sessions: HashMap<WsConnId, mpsc::UnboundedSender<WsRelayAction>>,
|
|
|
|
/// Command receiver.
|
|
cmd_rx: mpsc::UnboundedReceiver<WsServerAction>,
|
|
|
|
pub pool: Pool<Sqlite>,
|
|
}
|
|
|
|
impl WsServer {
|
|
pub fn new(pool: Pool<Sqlite>) -> (Self, mpsc::UnboundedSender<WsServerAction>) {
|
|
let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
|
|
|
|
let new_self = Self {
|
|
controller_sessions: HashMap::new(),
|
|
controllers: HashMap::new(),
|
|
relay_sessions: HashMap::new(),
|
|
cmd_rx,
|
|
pool,
|
|
};
|
|
(new_self, cmd_tx)
|
|
}
|
|
|
|
/// Register new session and assign unique ID to this session
|
|
async fn connect_controller(&mut self, tx: mpsc::UnboundedSender<WsControllerAction>) -> WsConnId {
|
|
log::info!("Controller connected");
|
|
|
|
// register a session with random connection ID
|
|
let id = random::<WsConnId>();
|
|
self.controller_sessions.insert(id, (tx, None));
|
|
|
|
// send id back
|
|
id
|
|
}
|
|
|
|
async fn connect_relay(&mut self, tx: mpsc::UnboundedSender<WsRelayAction>) -> Result<WsConnId, EmgauwaError> {
|
|
log::debug!("Relay client connected");
|
|
|
|
let relays_json = serde_json::to_string(&self.get_relays()?)?;
|
|
tx.send(WsRelayAction::Forward(relays_json.clone()))?;
|
|
|
|
// register a session with random connection ID
|
|
let id = random::<WsConnId>();
|
|
self.relay_sessions.insert(id, tx);
|
|
|
|
// send id back
|
|
Ok(id)
|
|
}
|
|
|
|
async fn disconnect_controller(&mut self, sender: UnboundedSender<WsControllerAction>, emgauwa_uid: EmgauwaUid, update_db: bool) -> Result<(), EmgauwaError> {
|
|
if let Some((controller, _)) = self.controllers.remove(&emgauwa_uid) {
|
|
if update_db {
|
|
let mut pool_conn = block_on(self.pool.acquire())?;
|
|
log::debug!("Marking controller {} as inactive", controller.c.uid);
|
|
if let Err(err) = block_on(controller.c.update_active(&mut pool_conn, false)) {
|
|
log::error!(
|
|
"Failed to mark controller {} as inactive: {:?}",
|
|
controller.c.uid,
|
|
err
|
|
);
|
|
return Err(EmgauwaError::from(err));
|
|
}
|
|
}
|
|
log::debug!("Telling websocket server to disconnect a controller");
|
|
sender.send(WsControllerAction::Disconnect)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Unregister the connection from the controller map and broadcast a disconnection message.
|
|
async fn disconnect_controller_conn(&mut self, conn_id: WsConnId, update_db: bool) -> Result<(), EmgauwaError> {
|
|
log::debug!("Trying to find controller connection for disconnect: {}", conn_id);
|
|
if let Some((sender, Some(emgauwa_uid))) = self.controller_sessions.remove(&conn_id) {
|
|
self.disconnect_controller(sender, emgauwa_uid, update_db).await?;
|
|
}
|
|
self.notify_relay_clients()?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn disconnect_relay_conn(&mut self, conn_id: WsConnId) -> Result<(), EmgauwaError> {
|
|
if let Some(relay_tx) = self.relay_sessions.remove(&conn_id) {
|
|
relay_tx.send(WsRelayAction::Disconnect)?;
|
|
}
|
|
self.notify_relay_clients()?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn get_relays(&mut self) -> Result<Vec<Relay>, EmgauwaError> {
|
|
let mut pool_conn = block_on(self.pool.acquire())?;
|
|
let db_controllers = block_on(DbController::get_all(&mut pool_conn))?;
|
|
let mut controllers: Vec<Controller> = convert_db_list(&mut pool_conn, db_controllers)?;
|
|
|
|
controllers.iter_mut().for_each(|c| {
|
|
if let Some((cc, _)) = self.controllers.get(&c.c.uid) {
|
|
c.apply_relay_states(&cc.get_relay_states())
|
|
}
|
|
});
|
|
|
|
let mut relays: Vec<Relay> = Vec::new();
|
|
controllers.iter().for_each(|c| {
|
|
relays.extend(c.relays.clone());
|
|
});
|
|
|
|
Ok(relays)
|
|
}
|
|
|
|
fn notify_relay_clients(&mut self) -> Result<(), EmgauwaError> {
|
|
let relays_json = serde_json::to_string(&self.get_relays()?)?;
|
|
|
|
self.relay_sessions.retain(|_, tx| {
|
|
tx.send(WsRelayAction::Forward(relays_json.clone())).is_ok()
|
|
});
|
|
Ok(())
|
|
}
|
|
|
|
// You may want to log errors rather than return them in this function.
|
|
pub async fn run(self) -> io::Result<()> {
|
|
tokio::select! {
|
|
output = self.real_run() => output,
|
|
_ = signal::ctrl_c() => Ok(())
|
|
}
|
|
}
|
|
|
|
pub async fn real_run(mut self) -> io::Result<()> {
|
|
while let Some(cmd) = self.cmd_rx.recv().await {
|
|
match cmd {
|
|
WsServerAction::ConnectController(conn_tx, res_tx) => {
|
|
let conn_id = self.connect_controller(conn_tx).await;
|
|
let _ = res_tx.send(conn_id);
|
|
}
|
|
|
|
WsServerAction::DisconnectControllerConn(conn, update_db) => {
|
|
if let Err(err) = self.disconnect_controller_conn(conn, update_db).await {
|
|
log::error!("Error disconnecting controller: {:?}", err);
|
|
}
|
|
}
|
|
WsServerAction::DisconnectController(emgauwa_uid, update_db) => {
|
|
if let Some((_, conn_id)) = self.controllers.get(&emgauwa_uid) {
|
|
if let Err(err ) = self.disconnect_controller_conn(*conn_id, update_db).await {
|
|
log::error!("Error disconnecting controller: {:?}", err);
|
|
}
|
|
}
|
|
}
|
|
|
|
WsServerAction::ControllerMessage(conn, msg) => {
|
|
match serde_json::from_str(&msg) {
|
|
Ok(action) => {
|
|
self.handle_action(conn, action);
|
|
}
|
|
Err(e) => {
|
|
log::error!("Error deserializing action: {:?}", e);
|
|
self.text_to_controller_conn(
|
|
conn,
|
|
serde_json::to_string(&EmgauwaError::Serialization(e))
|
|
.unwrap_or(String::from("Error in deserializing action")),
|
|
);
|
|
}
|
|
}
|
|
}
|
|
WsServerAction::ForwardToController(uid, action) => {
|
|
log::debug!("Forwarding action: {:?}", action);
|
|
if let Some((_, address)) = self.controllers.get(&uid) {
|
|
self.text_to_controller_conn(*address, serde_json::to_string(&action)?);
|
|
}
|
|
}
|
|
WsServerAction::GetRelays(tx) => {
|
|
tx.send(self.get_relays()).map_err(|_| io::Error::other("Error sending relays"))?;
|
|
}
|
|
WsServerAction::ConnectRelay(conn_tx, res_tx) => {
|
|
if let Ok(conn_id) = self.connect_relay(conn_tx).await {
|
|
if let Err(e) = res_tx.send(conn_id) {
|
|
log::error!("Error sending relay conn id: {}", e);
|
|
}
|
|
}
|
|
}
|
|
WsServerAction::DisconnectRelayConn(conn) => {
|
|
if let Err(err) = self.disconnect_relay_conn(conn).await {
|
|
log::error!("Error disconnecting relay conn: {:?}", err);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub fn handle_action(
|
|
&mut self,
|
|
conn: WsConnId,
|
|
action: ControllersWsAction,
|
|
) {
|
|
let action_res = match action {
|
|
ControllersWsAction::Register(controller) => self.handle_register(conn, controller),
|
|
ControllersWsAction::RelayStates((controller_uid, relay_states)) => {
|
|
self.handle_relay_states(controller_uid, relay_states)
|
|
}
|
|
_ => Ok(()),
|
|
};
|
|
if let Err(e) = action_res {
|
|
log::error!("Error handling action: {:?}", e);
|
|
self.text_to_controller_conn(
|
|
conn,
|
|
serde_json::to_string(&e).unwrap_or(format!("Error in handling action: {:?}", e)),
|
|
);
|
|
}
|
|
}
|
|
|
|
pub fn text_to_controller_conn(&self, target: WsConnId, msg: String) {
|
|
if let Some((tx, _)) = self.controller_sessions.get(&target) {
|
|
log::debug!("Sending message to controller conn {}", target);
|
|
// errors if a client has disconnected abruptly and hasn't been timed-out yet
|
|
let _ = tx.send(WsControllerAction::Forward(msg.clone()));
|
|
}
|
|
}
|
|
|
|
pub fn handle_register(
|
|
&mut self,
|
|
conn: WsConnId,
|
|
controller: Controller,
|
|
) -> Result<(), EmgauwaError> {
|
|
log::info!(
|
|
"Registering controller: {} ({})",
|
|
controller.c.name,
|
|
controller.c.uid
|
|
);
|
|
let mut pool_conn = block_on(self.pool.acquire())?;
|
|
|
|
let c = &controller.c;
|
|
let controller_db = block_on(DbController::get_by_uid_or_create(
|
|
&mut pool_conn,
|
|
&c.uid,
|
|
&c.name,
|
|
c.relay_count,
|
|
))?;
|
|
block_on(controller_db.update_active(&mut pool_conn, true))?;
|
|
// update only the relay count
|
|
block_on(controller_db.update(&mut pool_conn, &controller_db.name, c.relay_count))?;
|
|
|
|
for relay in &controller.relays {
|
|
log::debug!(
|
|
"Registering relay: {} ({})",
|
|
relay.r.name,
|
|
match relay.is_on {
|
|
Some(true) => "+",
|
|
Some(false) => "-",
|
|
None => "?",
|
|
}
|
|
);
|
|
let (new_relay, created) = block_on(DbRelay::get_by_controller_and_num_or_create(
|
|
&mut pool_conn,
|
|
&controller_db,
|
|
relay.r.number,
|
|
&relay.r.name,
|
|
))?;
|
|
if created {
|
|
let mut relay_schedules = Vec::new();
|
|
for schedule in &relay.schedules {
|
|
let (new_schedule, _) = block_on(DbSchedule::get_by_uid_or_create(
|
|
&mut pool_conn,
|
|
schedule.uid.clone(),
|
|
&schedule.name,
|
|
&schedule.periods,
|
|
))?;
|
|
relay_schedules.push(new_schedule);
|
|
}
|
|
|
|
block_on(DbJunctionRelaySchedule::set_schedules(
|
|
&mut pool_conn,
|
|
&new_relay,
|
|
relay_schedules.iter().collect(),
|
|
))?;
|
|
}
|
|
}
|
|
|
|
let relay_states = controller.get_relay_states();
|
|
|
|
let controller_uid = &controller.c.uid;
|
|
let controller_db = block_on(DbController::get_by_uid(&mut pool_conn, controller_uid))?
|
|
.ok_or(DatabaseError::InsertGetError)?;
|
|
let mut controller = Controller::from_db_model(&mut pool_conn, controller_db)?;
|
|
controller.apply_relay_states(&relay_states);
|
|
|
|
self.controllers.insert(controller_uid.clone(), (controller, conn));
|
|
|
|
// Update uid in controller session
|
|
if let Some((tx, _)) = self.controller_sessions.remove(&conn) {
|
|
log::debug!("Updating controller uid for conn {}", conn);
|
|
self.controller_sessions.insert(conn, (tx, Some(controller_uid.clone())));
|
|
}
|
|
|
|
log::debug!("Done registering controller");
|
|
Ok(())
|
|
}
|
|
|
|
pub fn handle_relay_states(
|
|
&mut self,
|
|
controller_uid: EmgauwaUid,
|
|
relay_states: RelayStates,
|
|
) -> Result<(), EmgauwaError> {
|
|
log::debug!(
|
|
"Received relay states: {} for {}",
|
|
utils::printable_relay_states(&relay_states),
|
|
controller_uid
|
|
);
|
|
|
|
if let Some((controller, _)) = self.controllers.get_mut(&controller_uid) {
|
|
controller.apply_relay_states(&relay_states);
|
|
}
|
|
self.notify_relay_clients()
|
|
}
|
|
} |