Switched WebSocket start motor to allow callbacks when client connects

addons/client/XEH_postInit.sqf

@@ -4,6 +4,8 @@ if (!hasInterface) exitWith {};
 
 _local_address = "armatak" callExtension ["local_ip", []] select 0;
 
+CALLEXT(websocket:start);
+
 SETVAR(player,GVAR(localAddress),_local_address);
 
 [{

src/lib.rs

@@ -31,10 +31,10 @@ pub fn init() -> Extension {
 
     log4rs::init_config(config).unwrap();
 
-    websocket::start();
-
     Extension::build()
         .group("websocket", Group::new()
+            .command("start", websocket::start)
+            .command("stop", websocket::stop)
             .command("message", websocket::message)
             .command("location", websocket::location)    
         )

src/websocket.rs

@@ -1,77 +1,128 @@
+use arma_rs::Context;
+use lazy_static::lazy_static;
 use log::info;
 use std::sync::mpsc::{self, Receiver, Sender};
 use std::sync::{Arc, Mutex};
 use std::thread;
-use ws::{listen, Message, Result as WsResult};
+use ws::{listen, CloseCode, Handler, Handshake, Message, Result as WsResult, Sender as WsSender};
+
 use crate::structs::{IntoMessage, LocationPayload};
-use lazy_static::lazy_static;
 
 pub enum WsCommand {
     SendMessage(String),
     Stop,
 }
 
+enum WsEvent {
+    FirstClientConnected,
+    LastClientDisconnected,
+}
+
 pub struct WsServer {
     pub(crate) tx: Sender<WsCommand>,
 }
 
+struct WsHandler {
+    out: WsSender,
+    clients: Arc<Mutex<Vec<WsSender>>>,
+    event_tx: Sender<WsEvent>,
+}
+
+impl Handler for WsHandler {
+    fn on_open(&mut self, _: Handshake) -> WsResult<()> {
+        let mut clients = self.clients.lock().unwrap();
+        clients.push(self.out.clone());
+
+        let count = clients.len();
+        info!("New client connected. Total clients: {}", count);
+
+        if count == 1 {
+            let _ = self.event_tx.send(WsEvent::FirstClientConnected);
+        }
+
+        Ok(())
+    }
+
+    fn on_close(&mut self, _: CloseCode, _: &str) {
+        let mut clients = self.clients.lock().unwrap();
+        clients.retain(|client| client.token() != self.out.token());
+
+        let count = clients.len();
+        info!("Client disconnected. Total clients: {}", count);
+
+        if count == 0 {
+            let _ = self.event_tx.send(WsEvent::LastClientDisconnected);
+        }
+    }
+
+    fn on_message(&mut self, msg: Message) -> WsResult<()> {
+        info!("Received: {}", msg);
+        Ok(())
+    }
+}
+
 impl WsServer {
-  pub fn start(&self, rx: Receiver<WsCommand>) {
+    pub fn start(&self, ctx: Context, rx: Receiver<WsCommand>) {
         if let Some(ref server) = *WEBSOCKET_SERVER.lock().unwrap() {
             server.stop();
         }
 
         let clients = Arc::new(Mutex::new(Vec::new()));
         let clients_clone = Arc::clone(&clients);
+        let (event_tx, event_rx): (Sender<WsEvent>, Receiver<WsEvent>) = mpsc::channel();
 
+        // Start WebSocket listener in a background thread
         thread::spawn(move || {
-          let mut running = true;
+            listen("0.0.0.0:4152", |out| WsHandler {
-
+                out,
-          let ws_thread = thread::spawn(move || {
+                clients: Arc::clone(&clients_clone),
-              listen("0.0.0.0:4152", |out| {
+                event_tx: event_tx.clone(),
-                  let clients_inner = Arc::clone(&clients_clone);
+            }).expect("WebSocket server failed to start");
-                  {
-                      let mut clients_guard = clients_inner.lock().unwrap();
-                      clients_guard.push(out.clone());
-                  }
-
-                  move |msg: Message| -> WsResult<()> {
-                      info!("Received: {}", msg);
-                      Ok(())
-                  }
-              })
-              .unwrap();
         });
 
-          while running {
+        // This thread owns ctx and reacts to commands and events
-              match rx.recv() {
+        thread::spawn(move || {
+            loop {
+                // WebSocket message commands
+                match rx.try_recv() {
                     Ok(WsCommand::SendMessage(message)) => {
                         let clients_guard = clients.lock().unwrap();
                         for client in clients_guard.iter() {
-                          client.send(message.clone()).unwrap();
+                            let _ = client.send(message.clone());
                         }
                     }
                     Ok(WsCommand::Stop) => {
-                      running = false;
                         info!("Stopping WebSocket server.");
+                        break;
                     }
-                  Err(error) => {
+                    Err(_) => {}
-                      info!("Error receiving command: {}", error.to_string());
-                  }
-              }
                 }
 
-          ws_thread.join().unwrap();
+                // Handle event callbacks with valid Context
+                match event_rx.try_recv() {
+                    Ok(WsEvent::FirstClientConnected) => {
+                        info!("Triggering callback: client_connected");
+                        let _ = ctx.callback_null("armatak_websocket", "client_connected");
+                    }
+                    Ok(WsEvent::LastClientDisconnected) => {
+                        info!("Triggering callback: client_disconnected");
+                        let _ = ctx.callback_null("armatak_websocket", "client_disconnected");
+                    }
+                    Err(_) => {}
+                }
+
+                thread::sleep(std::time::Duration::from_millis(10));
+            }
         });
     }
 
     pub fn send_message<T: IntoMessage>(&self, payload: T) {
         let message = payload.into_message();
-      self.tx.send(WsCommand::SendMessage(message)).unwrap();
+        let _ = self.tx.send(WsCommand::SendMessage(message));
     }
 
     pub fn stop(&self) {
-      self.tx.send(WsCommand::Stop).unwrap();
+        let _ = self.tx.send(WsCommand::Stop);
     }
 }
 
@@ -79,17 +130,16 @@ lazy_static! {
     static ref WEBSOCKET_SERVER: Arc<Mutex<Option<WsServer>>> = Arc::new(Mutex::new(None));
 }
 
-pub fn start() -> &'static str {
+pub fn start(ctx: Context) -> &'static str {
     let (tx, rx): (Sender<WsCommand>, Receiver<WsCommand>) = mpsc::channel();
 
     let server = WsServer { tx };
-  server.start(rx);
+    server.start(ctx, rx);
 
     let mut server_guard = WEBSOCKET_SERVER.lock().unwrap();
     *server_guard = Some(server);
 
     info!("WebSocket server started.");
-
     "Starting WebSocket Server"
 }
 
@@ -110,5 +160,16 @@ pub fn location(payload: LocationPayload) -> &'static str {
     } else {
         info!("WebSocket server is not running.");
     }
-  "sending location to all WebSocket clients"
+
+    "Sending location to all WebSocket clients"
+}
+
+pub fn stop() -> &'static str {
+    if let Some(ref server) = *WEBSOCKET_SERVER.lock().unwrap() {
+        server.stop();
+    } else {
+        info!("WebSocket server is not running.");
+    }
+
+    "Stopping WebSocket Server"
 }