diff --git a/src/lib.rs b/src/lib.rs index cfe1d79..c45f31c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -12,7 +12,8 @@ pub mod plugin; pub mod vfs; pub use anyhow::Result; -use tungstenite::accept; +pub use tungstenite; +use tungstenite::{WebSocket, accept}; use crate::plugin::DynPlugin; pub use once_cell; @@ -23,9 +24,10 @@ pub struct ServerConfig { } pub struct Server { - plugins: Mutex>, root: PathBuf, config: ServerConfig, + plugins: Mutex>, + clients: Mutex>>>>, } impl Default for ServerConfig { @@ -48,6 +50,7 @@ impl Server { plugins: Mutex::new(Vec::new()), root: root.to_path_buf(), config: ServerConfig::default(), + clients: Mutex::new(Vec::new()), }) } @@ -56,6 +59,7 @@ impl Server { plugins: Mutex::new(Vec::new()), root: root.to_path_buf(), config, + clients: Mutex::new(Vec::new()), }) } @@ -102,10 +106,46 @@ impl Server { fn handle_client(self: &Arc, stream: TcpStream) -> anyhow::Result<()> { Self::LOGGER.info(format!("New connection: {}", stream.peer_addr()?)); - let mut ws = accept(stream)?; + let ws = Arc::new(Mutex::new(accept(stream)?)); + + self.clients.lock().unwrap().push(ws.clone()); + + loop { + let req = ws.lock().unwrap().read()?; + + for plugin in self.plugins.lock().unwrap().iter_mut() { + plugin.on_request(&req, self); + } + + if req.is_close() { + let mut clients = self.clients.lock().unwrap(); + + if let Some(i) = clients.iter().position(|v| Arc::ptr_eq(v, &ws)) { + clients.remove(i); + } + } + + for c in self.clients.lock().unwrap().iter_mut() { + self.wrap_err(&ws, c.lock().unwrap().send(req.clone()))?; + } + } - let msg = ws.read()?; - ws.send(msg)?; Ok(()) } + + /// When there is a error it removes the client + fn wrap_err( + self: &Arc, + ws: &Arc>>, + res: std::result::Result, + ) -> std::result::Result { + if res.is_err() { + let mut clients = self.clients.lock().unwrap(); + if let Some(i) = clients.iter().position(|v| Arc::ptr_eq(v, &ws)) { + clients.remove(i); + } + } + + res + } } diff --git a/src/plugin.rs b/src/plugin.rs index 3e27c42..3766b22 100644 --- a/src/plugin.rs +++ b/src/plugin.rs @@ -1,9 +1,15 @@ use std::sync::Arc; +use tungstenite::Message; + use crate::Server; pub type DynPlugin = Box; pub trait Plugin { fn init(&mut self, server: &Arc); + #[allow(unused_variables)] + fn on_request(&mut self, msg: &Message, server: &Arc) -> bool { + false + } }