better error handling

src/encryption.rs

@@ -10,52 +10,55 @@ use openssl::{
     symm::{decrypt, encrypt, Cipher},
 };
 
+use crate::response::Error;
+
 fn aes_encrypt(
     key: &[u8],
-    // iv: &[u8],
     plaintext: &[u8],
-) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
+) -> Result<Vec<u8>, Error> {
     let cipher = Cipher::aes_256_ecb();
-    let ciphertext = encrypt(cipher, key, None, plaintext)?;
+    encrypt(cipher, key, None, plaintext)
-    Ok(ciphertext)
+        .map_err(|e| Error::EncryptionError(format!("Failed to encrypt: {}", e)))
 }
 
 fn aes_decrypt(
     key: &[u8],
-    // iv: &[u8],
     ciphertext: &[u8],
-) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
+) -> Result<Vec<u8>, Error> {
     let cipher = Cipher::aes_256_ecb();
-    let plaintext = decrypt(cipher, key, None, ciphertext)?;
+    decrypt(cipher, key, None, ciphertext)
-    Ok(plaintext)
+        .map_err(|e| Error::EncryptionError(format!("Failed to decrypt: {}", e)))
 }
 
 #[allow(deprecated)]
-pub fn handshake(mut stream: TcpStream, rsa: Rsa<Private>) -> ETcp {
+pub fn handshake(mut stream: TcpStream, rsa: Rsa<Private>) -> Result<ETcp, Error> {
-    let mut reader = BufReader::new(stream.try_clone().unwrap());
+    let stream_clone = stream.try_clone()
+        .map_err(|e| Error::NetworkError(format!("Failed to clone stream: {}", e)))?;
+    let mut reader = BufReader::new(stream_clone);
 
-    stream.write_all(&rsa.public_key_to_pem().unwrap()).unwrap();
+    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)
+        .map_err(|e| Error::NetworkError(format!("Failed to write public key: {}", e)))?;
 
-    stream.flush().unwrap();
+    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)
+        .map_err(|e| Error::NetworkError(format!("Failed to read AES key: {}", e)))?;
 
-    reader.read_line(&mut aes_encoded).unwrap();
+    let encrypted_data = base64::decode(aes_encoded.trim())
-
+        .map_err(|e| Error::EncryptionError(format!("Failed to decode base64: {}", e)))?;
-    let encrypted_data = base64::decode(aes_encoded.trim()).expect("Failed to decode base64");
 
     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,
-        )
-        .expect("Decryption failed");
 
     println!("len: {}", aes_len);
-
     aes_decrypted.truncate(aes_len);
 
     println!(
@@ -64,11 +67,11 @@ pub fn handshake(mut stream: TcpStream, rsa: Rsa<Private>) -> ETcp {
         aes_decrypted
     );
 
-    ETcp {
+    Ok(ETcp {
         stream,
         aes: aes_decrypted,
         reader,
-    }
+    })
 }
 
 pub struct ETcp {
@@ -79,39 +82,51 @@ pub struct ETcp {
 
 impl ETcp {
     #[allow(deprecated)]
-    pub fn send<T: Display>(&mut self, d: T) {
+    pub fn send<T: Display>(&mut self, d: T) -> Result<(), Error> {
-        let a = aes_encrypt(&self.aes, d.to_string().as_bytes()).unwrap();
+        let encrypted = aes_encrypt(&self.aes, d.to_string().as_bytes())?;
+        let encoded = base64::encode(encrypted) + "\n";
 
-        self.stream
+        self.stream.write_all(encoded.as_bytes())
-            .write_all((base64::encode(a) + "\n").as_bytes())
+            .map_err(|e| Error::NetworkError(format!("Failed to write to stream: {}", e)))?;
-            .unwrap();
 
-        self.stream.flush().unwrap();
+        self.stream.flush()
+            .map_err(|e| Error::NetworkError(format!("Failed to flush stream: {}", e)))
     }
 
     #[allow(deprecated)]
-    pub fn read(&mut self) -> Option<String> {
+    pub fn read(&mut self) -> Result<Option<String>, Error> {
         let mut v_encoded = String::new();
 
-        self.reader.read_line(&mut v_encoded).unwrap();
+        if self.reader.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()).expect("Failed to decode base64");
+        let encrypted_data = base64::decode(v_encoded.trim())
+            .map_err(|e| Error::EncryptionError(format!("Failed to decode base64: {}", e)))?;
 
-        Some(String::from_utf8(aes_decrypt(&self.aes, &encrypted_data).unwrap()).unwrap())
+        let decrypted = aes_decrypt(&self.aes, &encrypted_data)?;
+        let text = String::from_utf8(decrypted)
+            .map_err(|e| Error::EncryptionError(format!("Failed to decode UTF-8: {}", e)))?;
+
+        Ok(Some(text))
     }
 
-    pub fn close(&self) {
+    pub fn close(&self) -> Result<(), Error> {
-        self.stream
+        self.stream.shutdown(std::net::Shutdown::Both)
-            .shutdown(std::net::Shutdown::Both)
+            .map_err(|e| Error::NetworkError(format!("Failed to shutdown stream: {}", e)))
-            .expect("shutdown call failed");
     }
 
-    pub fn clone(&mut self) -> Self {
+    pub fn clone(&mut self) -> Result<Self, Error> {
-        let stream = self.stream.try_clone().unwrap();
+        let stream = self.stream.try_clone()
-        ETcp {
+            .map_err(|e| Error::NetworkError(format!("Failed to clone stream: {}", e)))?;
-            reader: BufReader::new(stream.try_clone().unwrap()),
+        let stream_clone = stream.try_clone()
+            .map_err(|e| Error::NetworkError(format!("Failed to clone stream: {}", e)))?;
+
+        Ok(ETcp {
+            reader: BufReader::new(stream_clone),
             stream,
             aes: self.aes.clone(),
-        }
+        })
     }
 }

src/main.rs

@@ -1,10 +1,10 @@
 use std::{
-    net::TcpListener,
+    net::{TcpListener, TcpStream},
     sync::Arc,
     time::{Duration, Instant},
 };
 
-use openssl::rsa::Rsa;
+use openssl::{rsa::Rsa, pkey::Private};
 
 pub mod database;
 pub mod encryption;
@@ -12,70 +12,89 @@ pub mod methods;
 pub mod parser;
 pub mod response;
 
-fn main() -> std::io::Result<()> {
+fn main() -> Result<(), response::Error> {
-    let rsa = Rsa::generate(2048).expect("Failed to generate RSA keys");
+    let rsa = Rsa::generate(2048)
+        .map_err(|e| response::Error::EncryptionError(format!("Failed to generate RSA keys: {}", e)))?;
 
-    let client = mongodb::sync::Client::with_uri_str("mongodb://admin:admin@localhost/").unwrap();
+    let client = mongodb::sync::Client::with_uri_str("mongodb://admin:admin@localhost/")
+        .map_err(|e| response::Error::DatabaseError(format!("Failed to connect to MongoDB: {}", e)))?;
     let database = Arc::new(database::Database::new(&client));
 
-    let listener = TcpListener::bind("0.0.0.0:8080").unwrap();
+    let listener = TcpListener::bind("0.0.0.0:8080")
+        .map_err(|e| response::Error::NetworkError(format!("Failed to bind to port 8080: {}", e)))?;
+
+    println!("Server listening on port 8080");
+
+    for stream_result in listener.incoming() {
+        let stream = match stream_result {
+            Ok(stream) => stream,
+            Err(e) => {
+                eprintln!("Failed to accept connection: {}", e);
+                continue;
+            }
+        };
 
-    for stream in listener.incoming() {
         println!("New connection");
         let database = Arc::clone(&database);
         let rsa = rsa.clone();
 
-        std::thread::spawn({
+        std::thread::spawn(move || {
-            move || {
+            if let Err(e) = handle_client(stream, rsa, database) {
-                let mut stream = encryption::handshake(stream.unwrap(), rsa);
+                eprintln!("Client error: {}", e);
-                let mut last_activity = Instant::now();
-
-                match methods::Session::new(stream.clone(), database) {
-                    Ok((session, res)) => {
-                        stream.send(res);
-
-                        loop {
-                            // Check for inactivity
-                            if last_activity.elapsed() > Duration::from_secs(60) {
-                                println!(
-                                    "[MOJANG] {} inactive for too long",
-                                    session.local_player.name
-                                );
-                                methods::player::logout(&session).unwrap();
-                                println!("[MOJANG] {} went offline", session.local_player.name);
-                                break;
-                            }
-
-                            match stream.read() {
-                                None => {
-                                    // Client disconnected
-                                    println!("[MOJANG] {} disconnected", session.local_player.name);
-                                    methods::player::logout(&session).unwrap();
-                                    println!("[MOJANG] {} went offline", session.local_player.name);
-                                    break;
-                                }
-                                Some(request_string) => {
-                                    last_activity = Instant::now();
-                                    let (method, params) = parser::parse(&request_string).unwrap();
-                                    let response = session.handle_request(&method, &params);
-
-                                    stream.send(match response {
-                                        Ok(response) => response.to_string(),
-                                        Err(e) => format!("!{e}"),
-                                    });
-                                }
-                            };
-                        }
-                    }
-                    Err(e) => {
-                        stream.send(format!("!{e}"));
-                        println!("Disconnected");
-                        stream.close();
-                    }
-                }
             }
         });
     }
 
     Ok(())
 }
+
+fn handle_client(stream: TcpStream, rsa: Rsa<Private>, database: Arc<database::Database>) -> Result<(), response::Error> {
+    let mut stream = encryption::handshake(stream, rsa)?;
+    let mut last_activity = Instant::now();
+
+    match methods::Session::new(stream.clone()?, database) {
+        Ok((session, res)) => {
+            stream.send(res)?;
+
+            loop {
+                // Check for inactivity
+                if last_activity.elapsed() > Duration::from_secs(60) {
+                    println!(
+                        "[MOJANG] {} inactive for too long",
+                        session.local_player.name
+                    );
+                    methods::player::logout(&session)?;
+                    println!("[MOJANG] {} went offline", session.local_player.name);
+                    break;
+                }
+
+                match stream.read()? {
+                    None => {
+                        // Client disconnected
+                        println!("[MOJANG] {} disconnected", session.local_player.name);
+                        methods::player::logout(&session)?;
+                        println!("[MOJANG] {} went offline", session.local_player.name);
+                        break;
+                    }
+                    Some(request_string) => {
+                        last_activity = Instant::now();
+                        let (method, params) = parser::parse(&request_string)?;
+                        let response = session.handle_request(&method, &params);
+
+                        stream.send(match response {
+                            Ok(response) => response.to_string(),
+                            Err(e) => format!("!{e}"),
+                        })?;
+                    }
+                }
+            }
+        }
+        Err(e) => {
+            stream.send(format!("!{e}"))?;
+            println!("Disconnected");
+            stream.close()?;
+        }
+    }
+
+    Ok(())
+}

src/methods/mod.rs

@@ -36,8 +36,10 @@ impl Session {
             mpsc::channel::<std::result::Result<String, crate::response::Error>>();
 
         std::thread::spawn(move || match stream.read() {
-            Some(session_token) => token_send.send(Ok(session_token)),
+            Ok(Some(session_token)) => token_send.send(Ok(session_token)),
-            _ => token_send.send(Err(crate::response::Error::InvalidHandshake("Failed to read session token".to_string()))),
+            _ => token_send.send(Err(crate::response::Error::InvalidHandshake(
+                "Failed to read session token".to_string(),
+            ))),
         });
 
         match token_recv.recv_timeout(std::time::Duration::from_secs(20)) {
@@ -46,17 +48,26 @@ impl Session {
                 let response = minreq::get("https://api.minecraftservices.com/minecraft/profile")
                     .with_header("Authorization", &format!("Bearer {session_token}"))
                     .send()
-                    .map_err(|_| crate::response::Error::InvalidSession("Failed to validate session".to_string()))?;
+                    .map_err(|_| {
+                        crate::response::Error::InvalidSession(
+                            "Failed to validate session".to_string(),
+                        )
+                    })?;
 
                 // If the session is invalid, return an error
                 if response.status_code != 200 {
-                    return Err(crate::response::Error::InvalidSession(format!("Invalid session status code: {}", response.status_code)));
+                    return Err(crate::response::Error::InvalidSession(format!(
+                        "Invalid session status code: {}",
+                        response.status_code
+                    )));
                 }
 
                 // Parse the player data
-                let local_player: LocalPlayer = response
+                let local_player: LocalPlayer = response.json().map_err(|_| {
-                    .json()
+                    crate::response::Error::InvalidSession(
-                    .map_err(|_| crate::response::Error::InvalidSession("Failed to parse player data".to_string()))?;
+                        "Failed to parse player data".to_string(),
+                    )
+                })?;
 
                 println!("[MOJANG] {} successfully logged on", &local_player.name);
 
@@ -72,8 +83,12 @@ impl Session {
                     Err(e) => Err(e),
                 }
             }
-            Err(mpsc::RecvTimeoutError::Timeout) => Err(crate::response::Error::Timeout("Session handshake timed out".to_string())),
+            Err(mpsc::RecvTimeoutError::Timeout) => Err(crate::response::Error::Timeout(
-            _ => Err(crate::response::Error::InvalidHandshake("Failed to receive handshake response".to_string())),
+                "Session handshake timed out".to_string(),
+            )),
+            _ => Err(crate::response::Error::InvalidHandshake(
+                "Failed to receive handshake response".to_string(),
+            )),
         }
     }
 

src/response.rs

@@ -73,6 +73,7 @@ pub enum Error {
     Timeout(String),
     AuthenticationError(String),
     ValidationError(String),
+    EncryptionError(String),
 }
 
 impl std::fmt::Display for Error {
@@ -89,6 +90,7 @@ impl std::fmt::Display for Error {
             Error::Timeout(operation) => write!(f, "Operation timed out: {}", operation),
             Error::AuthenticationError(details) => write!(f, "Authentication failed: {}", details),
             Error::ValidationError(details) => write!(f, "Validation failed: {}", details),
+            Error::EncryptionError(details) => write!(f, "Encryption error: {}", details),
         }
     }
 }