Notify system

src/encryption.rs

@@ -12,19 +12,13 @@ use openssl::{
 
 use crate::response::Error;
 
-fn aes_encrypt(
+fn aes_encrypt(key: &[u8], plaintext: &[u8]) -> Result<Vec<u8>, Error> {
-    key: &[u8],
-    plaintext: &[u8],
-) -> Result<Vec<u8>, Error> {
     let cipher = Cipher::aes_256_ecb();
     encrypt(cipher, key, None, plaintext)
         .map_err(|e| Error::EncryptionError(format!("Failed to encrypt: {}", e)))
 }
 
-fn aes_decrypt(
+fn aes_decrypt(key: &[u8], ciphertext: &[u8]) -> Result<Vec<u8>, Error> {
-    key: &[u8],
-    ciphertext: &[u8],
-) -> Result<Vec<u8>, Error> {
     let cipher = Cipher::aes_256_ecb();
     decrypt(cipher, key, None, ciphertext)
         .map_err(|e| Error::EncryptionError(format!("Failed to decrypt: {}", e)))
@@ -32,31 +26,38 @@ fn aes_decrypt(
 
 #[allow(deprecated)]
 pub fn handshake(mut stream: TcpStream, rsa: Rsa<Private>) -> Result<ETcp, Error> {
-    let stream_clone = stream.try_clone()
+    let stream_clone = stream
+        .try_clone()
         .map_err(|e| Error::NetworkError(format!("Failed to clone stream: {}", e)))?;
     let mut reader = BufReader::new(stream_clone);
 
-    let public_key = rsa.public_key_to_pem()
+    let public_key = rsa
+        .public_key_to_pem()
         .map_err(|e| Error::EncryptionError(format!("Failed to get public key: {}", e)))?;
-    stream.write_all(&public_key)
+    stream
+        .write_all(&public_key)
         .map_err(|e| Error::NetworkError(format!("Failed to write public key: {}", e)))?;
 
-    stream.flush()
+    stream
+        .flush()
         .map_err(|e| Error::NetworkError(format!("Failed to flush stream: {}", e)))?;
 
     let mut aes_encoded = String::new();
-    reader.read_line(&mut aes_encoded)
+    reader
+        .read_line(&mut aes_encoded)
         .map_err(|e| Error::NetworkError(format!("Failed to read AES key: {}", e)))?;
 
     let encrypted_data = base64::decode(aes_encoded.trim())
         .map_err(|e| Error::EncryptionError(format!("Failed to decode base64: {}", e)))?;
 
     let mut aes_decrypted = vec![0; 256];
-    let aes_len = rsa.private_decrypt(
+    let aes_len = rsa
-        &encrypted_data,
+        .private_decrypt(
-        &mut aes_decrypted,
+            &encrypted_data,
-        openssl::rsa::Padding::PKCS1,
+            &mut aes_decrypted,
-    ).map_err(|e| Error::EncryptionError(format!("Failed to decrypt AES key: {}", e)))?;
+            openssl::rsa::Padding::PKCS1,
+        )
+        .map_err(|e| Error::EncryptionError(format!("Failed to decrypt AES key: {}", e)))?;
 
     println!("len: {}", aes_len);
     aes_decrypted.truncate(aes_len);
@@ -86,10 +87,12 @@ impl ETcp {
         let encrypted = aes_encrypt(&self.aes, d.to_string().as_bytes())?;
         let encoded = base64::encode(encrypted) + "\n";
 
-        self.stream.write_all(encoded.as_bytes())
+        self.stream
+            .write_all(encoded.as_bytes())
             .map_err(|e| Error::NetworkError(format!("Failed to write to stream: {}", e)))?;
 
-        self.stream.flush()
+        self.stream
+            .flush()
             .map_err(|e| Error::NetworkError(format!("Failed to flush stream: {}", e)))
     }
 
@@ -97,10 +100,14 @@ impl ETcp {
     pub fn read(&mut self) -> Result<Option<String>, Error> {
         let mut v_encoded = String::new();
 
-        if self.reader.read_line(&mut v_encoded)
+        if self
-            .map_err(|e| Error::NetworkError(format!("Failed to read line: {}", e)))? == 0 {
+            .reader
-                return Ok(None);
+            .read_line(&mut v_encoded)
-            }
+            .map_err(|e| Error::NetworkError(format!("Failed to read line: {}", e)))?
+            == 0
+        {
+            return Ok(None);
+        }
 
         let encrypted_data = base64::decode(v_encoded.trim())
             .map_err(|e| Error::EncryptionError(format!("Failed to decode base64: {}", e)))?;
@@ -113,14 +120,18 @@ impl ETcp {
     }
 
     pub fn close(&self) -> Result<(), Error> {
-        self.stream.shutdown(std::net::Shutdown::Both)
+        self.stream
+            .shutdown(std::net::Shutdown::Both)
             .map_err(|e| Error::NetworkError(format!("Failed to shutdown stream: {}", e)))
     }
 
-    pub fn clone(&mut self) -> Result<Self, Error> {
+    pub fn try_clone(&mut self) -> Result<Self, Error> {
-        let stream = self.stream.try_clone()
+        let stream = self
+            .stream
+            .try_clone()
             .map_err(|e| Error::NetworkError(format!("Failed to clone stream: {}", e)))?;
-        let stream_clone = stream.try_clone()
+        let stream_clone = stream
+            .try_clone()
             .map_err(|e| Error::NetworkError(format!("Failed to clone stream: {}", e)))?;
 
         Ok(ETcp {
@@ -130,3 +141,23 @@ impl ETcp {
         })
     }
 }
+
+impl Clone for ETcp {
+    fn clone(&self) -> Self {
+        let stream = self
+            .stream
+            .try_clone()
+            .map_err(|e| Error::NetworkError(format!("Failed to clone stream: {}", e)))
+            .unwrap();
+        let stream_clone = stream
+            .try_clone()
+            .map_err(|e| Error::NetworkError(format!("Failed to clone stream: {}", e)))
+            .unwrap();
+
+        ETcp {
+            reader: BufReader::new(stream_clone),
+            stream,
+            aes: self.aes.clone(),
+        }
+    }
+}

src/main.rs

@@ -4,25 +4,32 @@ use std::{
     time::{Duration, Instant},
 };
 
-use openssl::{rsa::Rsa, pkey::Private};
+use openssl::{pkey::Private, rsa::Rsa};
 
+use crate::methods::SocketMap;
+
+pub mod cosmetics;
 pub mod database;
 pub mod encryption;
 pub mod methods;
 pub mod parser;
 pub mod response;
-pub mod cosmetics;
 
 fn main() -> Result<(), response::Error> {
-    let rsa = Rsa::generate(2048)
+    let rsa = Rsa::generate(2048).map_err(|e| {
-        .map_err(|e| response::Error::EncryptionError(format!("Failed to generate RSA keys: {}", e)))?;
+        response::Error::EncryptionError(format!("Failed to generate RSA keys: {}", e))
+    })?;
 
-    let client = mongodb::sync::Client::with_uri_str("mongodb://admin:admin@localhost/")
+    let client =
-        .map_err(|e| response::Error::DatabaseError(format!("Failed to connect to MongoDB: {}", e)))?;
+        mongodb::sync::Client::with_uri_str("mongodb://admin:[email protected]/").map_err(|e| {
+            response::Error::DatabaseError(format!("Failed to connect to MongoDB: {}", e))
+        })?;
     let database = Arc::new(database::Database::new(&client));
+    let sockets = SocketMap::default();
 
-    let listener = TcpListener::bind("0.0.0.0:8080")
+    let listener = TcpListener::bind("0.0.0.0:8080").map_err(|e| {
-        .map_err(|e| response::Error::NetworkError(format!("Failed to bind to port 8080: {}", e)))?;
+        response::Error::NetworkError(format!("Failed to bind to port 8080: {}", e))
+    })?;
 
     println!("Server listening on port 8080");
 
@@ -37,10 +44,11 @@ fn main() -> Result<(), response::Error> {
 
         println!("New connection");
         let database = Arc::clone(&database);
+        let sockets = Arc::clone(&sockets);
         let rsa = rsa.clone();
 
         std::thread::spawn(move || {
-            if let Err(e) = handle_client(stream, rsa, database) {
+            if let Err(e) = handle_client(stream, rsa, database, sockets) {
                 eprintln!("Client error: {}", e);
             }
         });
@@ -49,11 +57,16 @@ fn main() -> Result<(), response::Error> {
     Ok(())
 }
 
-fn handle_client(stream: TcpStream, rsa: Rsa<Private>, database: Arc<database::Database>) -> Result<(), response::Error> {
+fn handle_client(
+    stream: TcpStream,
+    rsa: Rsa<Private>,
+    database: Arc<database::Database>,
+    sockets: SocketMap,
+) -> Result<(), response::Error> {
     let mut stream = encryption::handshake(stream, rsa)?;
     let mut last_activity = Instant::now();
 
-    match methods::Session::new(stream.clone()?, database) {
+    match methods::Session::new(stream.try_clone()?, database, sockets) {
         Ok((session, res)) => {
             stream.send(res)?;
 

src/methods/emote.rs

@@ -0,0 +1,14 @@
+use crate::response::{Response, Result};
+
+use super::Session;
+
+pub fn emote(session: &Session, emote: String, notify: Vec<&str>) -> Result {
+    if notify.len() > 0 {
+        session.notify(
+            &notify,
+            &format!("emote@uuid={}@name={emote}", session.local_player.id),
+        )?;
+    }
+
+    Ok(Response::Success)
+}

src/methods/mod.rs

@@ -1,11 +1,12 @@
 mod buy;
+mod emote;
 pub mod player;
 mod set_cloak;
 mod set_hat;
 
 use std::{
     collections::HashMap,
-    sync::{mpsc, Arc},
+    sync::{mpsc, Arc, Mutex},
 };
 
 use serde::{Deserialize, Serialize};
@@ -22,20 +23,26 @@ pub struct LocalPlayer {
     pub name: String,
 }
 
+pub type SocketMap = Arc<Mutex<HashMap<String, ETcp>>>;
+
 pub struct Session {
     pub session_token: String,
     pub database: Arc<crate::database::Database>,
     pub local_player: LocalPlayer,
+    sockets: SocketMap,
 }
 
 impl Session {
     pub fn new(
         mut stream: ETcp,
         database: Arc<crate::database::Database>,
+        sockets: SocketMap,
     ) -> Result<(Self, Response)> {
         let (token_send, token_recv) =
             mpsc::channel::<std::result::Result<String, crate::response::Error>>();
 
+        let stream_2 = stream.try_clone()?;
+
         std::thread::spawn(move || match stream.read() {
             Ok(Some(session_token)) => token_send.send(Ok(session_token)),
             _ => token_send.send(Err(crate::response::Error::InvalidHandshake(
@@ -76,11 +83,15 @@ impl Session {
                     session_token,
                     database,
                     local_player,
+                    sockets,
                 };
 
                 // Capture the player
                 match player::login(&session) {
-                    Ok(player) => Ok((session, player)),
+                    Ok(player) => {
+                        session.add_socket(stream_2);
+                        Ok((session, player))
+                    }
                     Err(e) => Err(e),
                 }
             }
@@ -97,9 +108,35 @@ impl Session {
         match method {
             "ping" => Ok(Response::Pong),
 
-            "set_cloak" => set_cloak::set_cloak(self, params.parse_param("cloak")?),
+            "set_cloak" => set_cloak::set_cloak(
+                self,
+                params.parse_param("cloak")?,
+                params
+                    .parse_param::<String>("notify")
+                    .unwrap_or_default()
+                    .split("$")
+                    .collect(),
+            ),
 
-            "set_hat" => set_hat::set_hat(self, params.parse_param("hat")?),
+            "set_hat" => set_hat::set_hat(
+                self,
+                params.parse_param("hat")?,
+                params
+                    .parse_param::<String>("notify")
+                    .unwrap_or_default()
+                    .split("$")
+                    .collect(),
+            ),
+
+            "emote" => emote::emote(
+                self,
+                params.parse_param("name")?,
+                params
+                    .parse_param::<String>("notify")
+                    .unwrap_or_default()
+                    .split("$")
+                    .collect(),
+            ),
 
             "player" => player::player(self, params.parse_param("uuid")?),
 
@@ -134,3 +171,27 @@ impl Session {
         }
     }
 }
+
+impl Session {
+    pub fn add_socket(&self, sock: ETcp) {
+        self.sockets
+            .lock()
+            .unwrap()
+            .insert(self.local_player.id.clone(), sock);
+    }
+
+    pub fn remove_from_sockets(&self) {
+        self.sockets.lock().unwrap().remove(&self.local_player.id);
+    }
+
+    pub fn notify(&self, players: &[&str], message: &str) -> Result<()> {
+        let mut sockets = self.sockets.lock().unwrap();
+        for player in players {
+            if let Some(i) = sockets.get_mut(&player.to_string()) {
+                i.send(message)
+                    .map_err(|e| crate::response::Error::DatabaseError(format!("{e}")))?;
+            }
+        }
+        Ok(())
+    }
+}

src/methods/player.rs

@@ -103,6 +103,7 @@ pub fn login(session: &Session) -> Result {
 }
 
 pub fn logout(session: &Session) -> Result {
+    session.remove_from_sockets();
     let uuid = session.local_player.id.clone();
 
     // Update player's online status to false

src/methods/set_cloak.rs

@@ -4,13 +4,14 @@ use crate::response::{Error, Response, Result};
 
 use super::{player, Session};
 
-pub fn set_cloak(session: &Session, cloak: String) -> Result {
+pub fn set_cloak(session: &Session, cloak: String, notify: Vec<&str>) -> Result {
     match player::player(session, session.local_player.id.clone())? {
         Response::Player(crate::response::PlayerResponse::Player { cloaks, .. }) => {
             if !cloaks.contains(&cloak) && !cloak.is_empty() {
-                return Err(Error::ValidationError(
+                return Err(Error::ValidationError(format!(
-                    format!("Player does not own cloak: {}", cloak)
+                    "Player does not own cloak: {}",
-                ));
+                    cloak
+                )));
             }
 
             session
@@ -27,12 +28,22 @@ pub fn set_cloak(session: &Session, cloak: String) -> Result {
                     },
                 )
                 .run()
-                .map_err(|e| Error::DatabaseError(
+                .map_err(|e| Error::DatabaseError(format!("Failed to update cloak: {}", e)))?;
-                    format!("Failed to update cloak: {}", e)
+
-                ))?;
+            if notify.len() > 0 {
+                session.notify(
+                    &notify,
+                    &format!(
+                        "update_cloak@uuid={}@cloak={cloak}",
+                        session.local_player.id
+                    ),
+                )?;
+            }
 
             Ok(Response::Success)
         }
-        _ => Err(Error::ValidationError("Invalid player data received".to_string())),
+        _ => Err(Error::ValidationError(
+            "Invalid player data received".to_string(),
+        )),
     }
 }

src/methods/set_hat.rs

@@ -4,13 +4,14 @@ use crate::response::{Error, Response, Result};
 
 use super::{player, Session};
 
-pub fn set_hat(session: &Session, hat: String) -> Result {
+pub fn set_hat(session: &Session, hat: String, notify: Vec<&str>) -> Result {
     match player::player(session, session.local_player.id.clone())? {
         Response::Player(crate::response::PlayerResponse::Player { hats, .. }) => {
             if !hats.contains(&hat) && !hat.is_empty() {
-                return Err(Error::ValidationError(
+                return Err(Error::ValidationError(format!(
-                    format!("Player does not own hat: {}", hat)
+                    "Player does not own hat: {}",
-                ));
+                    hat
+                )));
             }
 
             session
@@ -27,12 +28,19 @@ pub fn set_hat(session: &Session, hat: String) -> Result {
                     },
                 )
                 .run()
-                .map_err(|e| Error::DatabaseError(
+                .map_err(|e| Error::DatabaseError(format!("Failed to update hat: {}", e)))?;
-                    format!("Failed to update hat: {}", e)
+
-                ))?;
+            if notify.len() > 0 {
+                session.notify(
+                    &notify,
+                    &format!("update_hat@uuid={}@hat={hat}", session.local_player.id),
+                )?;
+            }
 
             Ok(Response::Success)
         }
-        _ => Err(Error::ValidationError("Invalid player data received".to_string())),
+        _ => Err(Error::ValidationError(
+            "Invalid player data received".to_string(),
+        )),
     }
 }