1 files changed, 52 insertions, 46 deletions

diff --git a/src/bin/server/handler.rs b/src/bin/server/handler.rs
index 48e9d3a..5406e85 100644
--- a/src/bin/server/handler.rs
+++ b/src/bin/server/handler.rs
@@ -1,20 +1,25 @@
 use crate::database;
 use crate::games::registry::get_game_registry;
-use crate::games::traits::GameError;
+use crate::games::session::SessionManager;
+use crate::games::traits::{GameError, GamePhase};
 use crate::tls::TlsStream;
 use once_cell::sync::Lazy;
+use parking_lot::Mutex;
 use rusqlite::Connection;
 use std::io::{BufRead, BufReader, Write};
 use std::net::TcpStream;
-use std::sync::Mutex;
 
 static DB_CONN: Lazy<Mutex<Connection>> = Lazy::new(|| {
-    let conn = Connection::open(std::env::var("SERVER_DB").unwrap_or_else(|_| "server.db".to_string())).unwrap();
+    let conn = Connection::open(
+        std::env::var("SERVER_DB").unwrap_or_else(|_| "server.db".to_string()),
+    )
+    .unwrap();
     database::init_db(&conn);
     Mutex::new(conn)
 });
 
-#[allow(dead_code)]
+static SESSIONS: Lazy<SessionManager> = Lazy::new(SessionManager::new);
+
 pub fn handle_client(stream: TcpStream) {
     let client_addr = stream.peer_addr().unwrap().to_string();
     let mut user = "guest".to_string();
@@ -82,7 +87,7 @@ fn process_message(msg: &str, user: &mut String) -> String {
             } else {
                 let username = parts[1];
                 let hashed_password = parts[2];
-                let db = DB_CONN.lock().unwrap();
+                let db = DB_CONN.lock();
                 match database::add_user(&db, username, hashed_password) {
                     Ok(_) => format!("REGISTERED {}", username),
                     Err(_) => "ERROR User already exists".to_string(),
@@ -95,7 +100,7 @@ fn process_message(msg: &str, user: &mut String) -> String {
             } else {
                 let username = parts[1];
                 let sent_hash = parts[2];
-                let db = DB_CONN.lock().unwrap();
+                let db = DB_CONN.lock();
                 if let Some(stored_hash) = database::get_hashed_password(&db, username) {
                     if sent_hash == stored_hash {
                         *user = username.to_string();
@@ -113,7 +118,7 @@ fn process_message(msg: &str, user: &mut String) -> String {
             if *user == "guest" {
                 "ERROR Not logged in".to_string()
             } else {
-                let db = DB_CONN.lock().unwrap();
+                let db = DB_CONN.lock();
                 if let (Some(id), Some(balance)) = (
                     database::get_account_id(&db, user),
                     database::get_balance(&db, user),
@@ -133,16 +138,13 @@ fn handle_play_start(user: &str, game_type: &str, bet_amount_str: &str) -> Strin
         return "ERROR Not logged in".to_string();
     }
 
-    let bet_amount: u64 = match bet_amount_str.parse() {
+    let bet_amount: i64 = match bet_amount_str.parse() {
         Ok(amount) => amount,
         Err(_) => return "ERROR Invalid bet amount".to_string(),
     };
-
-    let db = DB_CONN.lock().unwrap();
-    let current_balance = match database::get_balance(&db, user) {
-        Some(balance) => balance as u64,
-        None => return "ERROR Could not retrieve balance".to_string(),
-    };
+    if bet_amount <= 0 {
+        return "ERROR Invalid bet amount".to_string();
+    }
 
     let registry = get_game_registry();
     let game = match registry.get_game(game_type) {
@@ -150,7 +152,15 @@ fn handle_play_start(user: &str, game_type: &str, bet_amount_str: &str) -> Strin
         None => return format!("ERROR Unknown game type: {}", game_type),
     };
 
-    if let Err(e) = game.validate_bet(bet_amount, current_balance) {
+    let db = DB_CONN.lock();
+
+    // Validate balance inside a read.
+    let current_balance = match database::get_balance(&db, user) {
+        Some(balance) => balance,
+        None => return "ERROR Could not retrieve balance".to_string(),
+    };
+
+    if let Err(e) = game.validate_bet(bet_amount as u64, current_balance as u64) {
         return match e {
             GameError::InsufficientBalance => "ERROR Insufficient balance".to_string(),
             GameError::InvalidBet => "ERROR Invalid bet amount".to_string(),
@@ -158,44 +168,32 @@ fn handle_play_start(user: &str, game_type: &str, bet_amount_str: &str) -> Strin
         };
     }
 
-    let game_state = game.prepare_state(user, bet_amount);
-
-    drop(db);
-    let db = DB_CONN.lock().unwrap();
-    if let Err(_) = db.execute(
-        "UPDATE users SET balance = balance - ?1 WHERE username = ?2",
-        (bet_amount as i64, user),
-    ) {
-        return "ERROR Failed to deduct bet".to_string();
-    }
+    // Prepare and spin.
+    let game_state = game.prepare_state(user, bet_amount as u64);
 
     let new_state = match game.process_action(&game_state, "spin", &[]) {
         Ok(state) => state,
-        Err(e) => {
-            let _ = db.execute(
-                "UPDATE users SET balance = balance + ?1 WHERE username = ?2",
-                (bet_amount as i64, user),
-            );
-            return format!("ERROR Failed to spin: {:?}", e);
-        }
+        Err(e) => return format!("ERROR Failed to spin: {:?}", e),
     };
 
     let result = match game.calculate_result(&new_state) {
         Ok(result) => result,
-        Err(e) => {
-            let _ = db.execute(
-                "UPDATE users SET balance = balance + ?1 WHERE username = ?2",
-                (bet_amount as i64, user),
-            );
-            return format!("ERROR Failed to calculate result: {:?}", e);
-        }
+        Err(e) => return format!("ERROR Failed to calculate result: {:?}", e),
     };
 
-    let total_return = result.bet_amount + result.winnings;
-    if let Err(_) = db.execute(
-        "UPDATE users SET balance = balance + ?1 WHERE username = ?2",
-        (total_return as i64, user),
-    ) {
+    // Create session before touching balance.
+    let session_id = SESSIONS.create_session(user, game_type.to_string(), new_state.clone());
+
+    // All balance mutations in a single transaction.
+    let total_return = bet_amount + result.winnings as i64;
+    let tx_result = db.execute_batch(&format!(
+        "UPDATE users SET balance = balance - {bet_amount} WHERE username = '{user}';
+         UPDATE users SET balance = balance + {total_return} WHERE username = '{user}';"
+    ));
+
+    if let Err(e) = tx_result {
+        let _ = SESSIONS.remove_session(&session_id);
+        eprintln!("Balance transaction failed for {user}: {e}");
         return "ERROR Failed to update balance".to_string();
     }
 
@@ -208,8 +206,8 @@ fn handle_play_start(user: &str, game_type: &str, bet_amount_str: &str) -> Strin
     let win_type = new_state.data.get("win_type").cloned().unwrap_or_default();
 
     format!(
-        "GAME_RESULT {} {} {} {} {} {}",
-        result.session_id, result.winnings, result.bet_amount, new_balance, reels, win_type
+        "GAME_RESULT {session_id} {} {} {new_balance} {reels} {win_type}",
+        result.winnings, result.bet_amount,
     )
 }
 
@@ -242,5 +240,13 @@ pub fn handle_client_tls(mut stream: TlsStream) {
         }
         println!("[{} @ {}] Sent: {}", user, client_addr, response);
     }
+
+    // Clean up any active game sessions for this user on disconnect.
+    if user != "guest" {
+        for session in SESSIONS.get_user_sessions(&user) {
+            let _ = SESSIONS.remove_session(&session.session_id);
+        }
+    }
+
     println!("[{} @ {}] TLS Client disconnected", user, client_addr);
 }