From 792ee3a7081ba1bb6e16c3fb35f382f3d85e9b53 Mon Sep 17 00:00:00 2001 From: Leo dev Date: Sun, 13 Apr 2025 06:24:48 +0200 Subject: [PATCH] better error handling --- src/encryption.rs | 105 ++++++++++++++++++++---------------- src/main.rs | 129 ++++++++++++++++++++++++++------------------- src/methods/mod.rs | 33 ++++++++---- src/response.rs | 2 + 4 files changed, 160 insertions(+), 109 deletions(-) diff --git a/src/encryption.rs b/src/encryption.rs index 9662529..b0e7d4c 100644 --- a/src/encryption.rs +++ b/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, Box> { +) -> Result, Error> { let cipher = Cipher::aes_256_ecb(); - let ciphertext = encrypt(cipher, key, None, plaintext)?; - Ok(ciphertext) + encrypt(cipher, key, None, plaintext) + .map_err(|e| Error::EncryptionError(format!("Failed to encrypt: {}", e))) } fn aes_decrypt( key: &[u8], - // iv: &[u8], ciphertext: &[u8], -) -> Result, Box> { +) -> Result, Error> { let cipher = Cipher::aes_256_ecb(); - let plaintext = decrypt(cipher, key, None, ciphertext)?; - Ok(plaintext) + decrypt(cipher, key, None, ciphertext) + .map_err(|e| Error::EncryptionError(format!("Failed to decrypt: {}", e))) } #[allow(deprecated)] -pub fn handshake(mut stream: TcpStream, rsa: Rsa) -> ETcp { - let mut reader = BufReader::new(stream.try_clone().unwrap()); +pub fn handshake(mut stream: TcpStream, rsa: Rsa) -> Result { + 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()).expect("Failed to decode base64"); + 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( - &encrypted_data, - &mut aes_decrypted, - openssl::rsa::Padding::PKCS1, - ) - .expect("Decryption failed"); + let aes_len = rsa.private_decrypt( + &encrypted_data, + &mut aes_decrypted, + 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); println!( @@ -64,11 +67,11 @@ pub fn handshake(mut stream: TcpStream, rsa: Rsa) -> 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(&mut self, d: T) { - let a = aes_encrypt(&self.aes, d.to_string().as_bytes()).unwrap(); + pub fn send(&mut self, d: T) -> Result<(), Error> { + let encrypted = aes_encrypt(&self.aes, d.to_string().as_bytes())?; + let encoded = base64::encode(encrypted) + "\n"; - self.stream - .write_all((base64::encode(a) + "\n").as_bytes()) - .unwrap(); + self.stream.write_all(encoded.as_bytes()) + .map_err(|e| Error::NetworkError(format!("Failed to write to stream: {}", e)))?; - 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 { + pub fn read(&mut self) -> Result, 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) { - self.stream - .shutdown(std::net::Shutdown::Both) - .expect("shutdown call failed"); + pub fn close(&self) -> Result<(), Error> { + self.stream.shutdown(std::net::Shutdown::Both) + .map_err(|e| Error::NetworkError(format!("Failed to shutdown stream: {}", e))) } - pub fn clone(&mut self) -> Self { - let stream = self.stream.try_clone().unwrap(); - ETcp { - reader: BufReader::new(stream.try_clone().unwrap()), + pub fn clone(&mut self) -> Result { + let stream = self.stream.try_clone() + .map_err(|e| Error::NetworkError(format!("Failed to clone stream: {}", e)))?; + 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(), - } + }) } } diff --git a/src/main.rs b/src/main.rs index 03ed25f..e8098db 100644 --- a/src/main.rs +++ b/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<()> { - let rsa = Rsa::generate(2048).expect("Failed to generate RSA keys"); +fn main() -> Result<(), response::Error> { + 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({ - move || { - let mut stream = encryption::handshake(stream.unwrap(), 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).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, ¶ms); - - stream.send(match response { - Ok(response) => response.to_string(), - Err(e) => format!("!{e}"), - }); - } - }; - } - } - Err(e) => { - stream.send(format!("!{e}")); - println!("Disconnected"); - stream.close(); - } - } + std::thread::spawn(move || { + if let Err(e) = handle_client(stream, rsa, database) { + eprintln!("Client error: {}", e); } }); } Ok(()) } + +fn handle_client(stream: TcpStream, rsa: Rsa, database: Arc) -> 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, ¶ms); + + stream.send(match response { + Ok(response) => response.to_string(), + Err(e) => format!("!{e}"), + })?; + } + } + } + } + Err(e) => { + stream.send(format!("!{e}"))?; + println!("Disconnected"); + stream.close()?; + } + } + + Ok(()) +} diff --git a/src/methods/mod.rs b/src/methods/mod.rs index 7294a35..50fa7b1 100644 --- a/src/methods/mod.rs +++ b/src/methods/mod.rs @@ -36,8 +36,10 @@ impl Session { mpsc::channel::>(); std::thread::spawn(move || match stream.read() { - Some(session_token) => token_send.send(Ok(session_token)), - _ => token_send.send(Err(crate::response::Error::InvalidHandshake("Failed to read session token".to_string()))), + 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(), + ))), }); 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 - .json() - .map_err(|_| crate::response::Error::InvalidSession("Failed to parse player data".to_string()))?; + let local_player: LocalPlayer = response.json().map_err(|_| { + crate::response::Error::InvalidSession( + "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(crate::response::Error::InvalidHandshake("Failed to receive handshake response".to_string())), + Err(mpsc::RecvTimeoutError::Timeout) => Err(crate::response::Error::Timeout( + "Session handshake timed out".to_string(), + )), + _ => Err(crate::response::Error::InvalidHandshake( + "Failed to receive handshake response".to_string(), + )), } } diff --git a/src/response.rs b/src/response.rs index 1891df8..65469af 100644 --- a/src/response.rs +++ b/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), } } }