Add WS server for front-end

Cargo.lock

@@ -204,6 +204,22 @@ dependencies = [
  "syn 2.0.119",
 ]
 
+[[package]]
+name = "actix-ws"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "decf53c3cdd63dd6f289980b430238f9a2f6d19f8bce8e418272e08d3da43f0f"
+dependencies = [
+ "actix-codec",
+ "actix-http",
+ "actix-web",
+ "bytestring",
+ "futures-core",
+ "futures-sink",
+ "tokio",
+ "tokio-util",
+]
+
 [[package]]
 name = "adler2"
 version = "2.0.1"
@@ -710,6 +726,7 @@ version = "0.1.0"
 dependencies = [
  "actix-codec",
  "actix-web",
+ "actix-ws",
  "argon2",
  "awc",
  "dotenvy",
@@ -1497,6 +1514,7 @@ dependencies = [
  "bytes",
  "futures-core",
  "futures-sink",
+ "futures-util",
  "libc",
  "pin-project-lite",
  "tokio",

Cargo.toml

@@ -6,6 +6,7 @@ edition = "2024"
 [dependencies]
 actix-codec = "0.5.2"
 actix-web = "4.14.0"
+actix-ws = "0.4.0"
 argon2 = { version = "0.5.3", features = ["std"]}
 awc = "3.8.2"
 dotenvy = "0.15.7"

src/controllers/mod.rs

@@ -1,3 +1,4 @@
 pub mod home;
 pub mod first_admin;
 pub mod login;
+pub mod ws_connect;

src/controllers/ws_connect.rs

@@ -0,0 +1,65 @@
+use actix_web::{HttpRequest, HttpResponse, web, get};
+use serde_json::json;
+use tokio::sync::mpsc;
+use uuid::Uuid;
+use crate::{
+    users::User,
+    ws_registry::ClientRegistry
+};
+use std::sync::Mutex;
+
+#[get("/ws")]
+pub async fn route(
+    req: HttpRequest,
+    body: web::Payload,
+    users: web::Data<Mutex<Vec<User>>>,
+    registry: web::Data<ClientRegistry>
+) -> HttpResponse {
+    let cookie = match req.cookie("user") {
+        Some(c) => c,
+        None => { return HttpResponse::Unauthorized().json(json!({"msg": "Unauthorized"})) }
+    };
+    let user_id = match Uuid::parse_str(&cookie.value()) {
+        Ok(id) => id,
+        Err(_) => return HttpResponse::InternalServerError().finish()
+    };
+    let user = match users.lock() {
+        Ok(us) => match us.iter().find(|u| u.id == user_id).cloned() {
+            Some(a) => a,
+            None => { return HttpResponse::Unauthorized().json(json!({"msg": "Unauthorized"})) }
+        },
+        Err(_) => { return HttpResponse::InternalServerError().finish()}
+    };
+
+    let (response, mut session, mut msg_stream) =  match actix_ws::handle(&req, body) {
+        Ok(a) => a,
+        Err(e) => { return HttpResponse::InternalServerError().json(json!({"msg": e.to_string()}))}
+    };
+
+    let (tx, mut rx) = mpsc::unbounded_channel::<String>();
+    let client_id = registry.register(tx, user.id);
+
+    let mut push_session = session.clone();
+    actix_web::rt::spawn(async move {
+        while let Some(payload) = rx.recv().await {
+            if push_session.text(payload).await.is_err() {
+                break;
+            }
+        }
+    });
+
+    let registry = registry.clone();
+    actix_web::rt::spawn(async move {
+        use futures_util::StreamExt;
+
+        while let Some(Ok(msg)) = msg_stream.next().await {
+            if let actix_ws::Message::Text(text) = msg {
+                let _ = session.text(text).await;
+            }
+        }
+
+        registry.unregister(&client_id);
+    });
+
+    response
+}

src/main.rs

@@ -6,10 +6,11 @@ use crate::environment::Environment;
 
 mod routes;
 mod controllers;
-mod websocket;
+mod ha_websocket;
 mod users;
 mod logic;
 mod environment;
+mod ws_registry;
 
 #[actix_web::main]
 async fn main() -> Result<(), std::io::Error> {
@@ -29,7 +30,10 @@ async fn main() -> Result<(), std::io::Error> {
     //Connect to HA via websocket
     let (outbound_tx, outbound_rx) = mpsc::unbounded_channel::<Message>();
     let ha_sender = web::Data::new(outbound_tx);
-    rt::spawn(websocket::listen(ha_ip, ha_token, outbound_rx));
+    rt::spawn(ha_websocket::listen(ha_ip, ha_token, outbound_rx));
+
+    //Create websocket server
+    let registry = web::Data::new(ws_registry::ClientRegistry::new());
 
     //Create HTTP server
     HttpServer::new(move || {
@@ -37,6 +41,7 @@ async fn main() -> Result<(), std::io::Error> {
             .app_data(users_data.clone())
             .app_data(ha_sender.clone())
             .app_data(ae.clone())
+            .app_data(registry.clone())
             .wrap(middleware::Compress::default())
             .configure(routes::config)
     })

src/routes.rs

@@ -2,7 +2,8 @@ use actix_web::web::ServiceConfig;
 use crate::controllers::{
     home,
     first_admin,
-    login
+    login,
+    ws_connect
 };
 
 pub fn config(cfg: &mut ServiceConfig) {
@@ -10,4 +11,5 @@ pub fn config(cfg: &mut ServiceConfig) {
     
     cfg.service(first_admin::route);
     cfg.service(login::route);
+    cfg.service(ws_connect::route);
 }

src/ws_registry.rs

@@ -0,0 +1,34 @@
+use std::{collections::HashMap, sync::Mutex};
+use tokio::sync::mpsc::UnboundedSender;
+use uuid::Uuid;
+
+pub struct ClientRegistry {
+    pub clients: Mutex<HashMap<Uuid, UnboundedSender<String>>>
+}
+
+impl ClientRegistry {
+    pub fn new() -> Self {
+        Self { clients: Mutex::new(HashMap::new()) }
+    }
+
+    pub fn register(&self, tx: UnboundedSender<String>, id: Uuid) -> Uuid {
+        self.clients.lock().unwrap().insert(id, tx);
+        id
+    }
+
+    pub fn unregister(&self, id: &Uuid) {
+        self.clients.lock().unwrap().remove(id);
+    }
+
+    pub fn send_to(&self, id: &Uuid, payload: &str) {
+        if let Some(tx) = self.clients.lock().unwrap().get(id) {
+            let _ = tx.send(payload.to_string());
+        }
+    }
+
+    pub fn broadcast(&self, payload: &str) {
+        for tx in self.clients.lock().unwrap().values() {
+            let _ = tx.send(payload.to_string());
+        }
+    }
+}