feat: Add Reentrancy Protection (#366)

Author: tali-creator

contract/contracts/predifi-contract/src/lib.rs

@@ -155,6 +155,7 @@ pub enum DataKey {
     UserPredictionIndex(Address, u32),
     Config,
     Paused,
+    ReentrancyGuard,
     CategoryPoolCount(Symbol),
     CategoryPoolIndex(Symbol, u32),
     /// Token whitelist: TokenWhitelist(token_address) -> true if allowed for betting.
@@ -632,6 +633,18 @@ impl PredifiContract {
         }
     }
 
+    fn enter_reentrancy_guard(env: &Env) {
+        let key = DataKey::ReentrancyGuard;
+        if env.storage().temporary().has(&key) {
+            panic!("Reentrancy detected");
+        }
+        env.storage().temporary().set(&key, &true);
+    }
+
+    fn exit_reentrancy_guard(env: &Env) {
+        env.storage().temporary().remove(&DataKey::ReentrancyGuard);
+    }
+
     /// Returns true if the token is on the allowed betting whitelist.
     fn is_token_whitelisted(env: &Env, token: &Address) -> bool {
         let key = DataKey::TokenWhitelist(token.clone());
@@ -1294,6 +1307,8 @@ impl PredifiContract {
         user.require_auth();
         assert!(amount > 0, "amount must be positive");
 
+        Self::enter_reentrancy_guard(&env);
+
         let pool_key = DataKey::Pool(pool_id);
         let mut pool: Pool = env
             .storage()
@@ -1312,6 +1327,7 @@ impl PredifiContract {
             "outcome exceeds options_count"
         );
 
+        // --- INTERNAL CHECKS & EFFECTS ---
         // Validate: per-pool stake limits
         assert!(
             amount >= pool.min_stake,
@@ -1324,9 +1340,6 @@ impl PredifiContract {
             );
         }
 
-        let token_client = token::Client::new(&env, &pool.token);
-        token_client.transfer(&user, &env.current_contract_address(), &amount);
-
         let pred_key = DataKey::Prediction(user.clone(), pool_id);
         if !env.storage().persistent().has(&pred_key) {
             let pc_key = DataKey::ParticipantsCount(pool_id);
@@ -1358,6 +1371,13 @@ impl PredifiContract {
         env.storage().persistent().set(&count_key, &(count + 1));
         Self::extend_persistent(&env, &count_key);
 
+        // --- INTERACTIONS ---
+
+        let token_client = token::Client::new(&env, &pool.token);
+        token_client.transfer(&user, &env.current_contract_address(), &amount);
+
+        Self::exit_reentrancy_guard(&env);
+
         PredictionPlacedEvent {
             pool_id,
             user: user.clone(),
@@ -1399,6 +1419,8 @@ impl PredifiContract {
         Self::require_not_paused(&env);
         user.require_auth();
 
+        Self::enter_reentrancy_guard(&env);
+
         let pool_key = DataKey::Pool(pool_id);
         let pool: Pool = env
             .storage()
@@ -1408,6 +1430,7 @@ impl PredifiContract {
         Self::extend_persistent(&env, &pool_key);
 
         if pool.state == MarketState::Active {
+            Self::exit_reentrancy_guard(&env);
             return Err(PredifiError::PoolNotResolved);
         }
 
@@ -1420,12 +1443,11 @@ impl PredifiContract {
                 timestamp: env.ledger().timestamp(),
             }
             .publish(&env);
+            Self::exit_reentrancy_guard(&env);
             return Err(PredifiError::AlreadyClaimed);
         }
 
-        // Mark as claimed immediately to prevent re-entrancy (INV-3)
-        env.storage().persistent().set(&claimed_key, &true);
-        Self::extend_persistent(&env, &claimed_key);
+        // --- CHECKS ---
 
         let pred_key = DataKey::Prediction(user.clone(), pool_id);
         let prediction: Option<Prediction> = env.storage().persistent().get(&pred_key);
@@ -1436,14 +1458,25 @@ impl PredifiContract {
 
         let prediction = match prediction {
             Some(p) => p,
-            None => return Ok(0),
+            None => {
+                Self::exit_reentrancy_guard(&env);
+                return Ok(0);
+            }
         };
 
+        // --- EFFECTS ---
+
+        // Mark as claimed immediately to prevent re-entrancy (INV-3)
+        env.storage().persistent().set(&claimed_key, &true);
+        Self::extend_persistent(&env, &claimed_key);
+
         if pool.state == MarketState::Canceled {
-            // Refunds: user gets exactly what they put in.
+            // --- INTERACTIONS (Refund) ---
             let token_client = token::Client::new(&env, &pool.token);
             token_client.transfer(&env.current_contract_address(), &user, &prediction.amount);
 
+            Self::exit_reentrancy_guard(&env);
+
             WinningsClaimedEvent {
                 pool_id,
                 user: user.clone(),
@@ -1455,6 +1488,7 @@ impl PredifiContract {
         }
 
         if prediction.outcome != pool.outcome {
+            Self::exit_reentrancy_guard(&env);
             return Ok(0);
         }
 
@@ -1463,6 +1497,7 @@ impl PredifiContract {
         let winning_stake: i128 = stakes.get(pool.outcome).unwrap_or(0);
 
         if winning_stake == 0 {
+            Self::exit_reentrancy_guard(&env);
             return Ok(0);
         }
 
@@ -1472,9 +1507,12 @@ impl PredifiContract {
         // Verify invariant: winnings ≤ total_stake (INV-4)
         assert!(winnings <= pool.total_stake, "Winnings exceed total stake");
 
+        // --- INTERACTIONS (Winnings Payout) ---
         let token_client = token::Client::new(&env, &pool.token);
         token_client.transfer(&env.current_contract_address(), &user, &winnings);
 
+        Self::exit_reentrancy_guard(&env);
+
         WinningsClaimedEvent {
             pool_id,
             user,