Merge pull request #41 from shamoo53/Security-Prevent-Admin-Rug-Pull

Security-Prevent-Admin-Rug-Pull

Author: JerryIdoko

contracts/vesting_contracts/src/lib.rs

@@ -28,6 +28,7 @@ pub struct Vault {
     pub start_time: u64,
     pub end_time: u64,
     pub is_initialized: bool, // Lazy initialization flag
+    pub is_irrevocable: bool, // Security flag to prevent admin withdrawal
 }
 
 #[contracttype]
@@ -135,6 +136,7 @@ impl VestingContract {
             start_time,
             end_time,
             is_initialized: true, // Mark as fully initialized
+            is_irrevocable: false, // Default to revocable
         };
 
         // Store vault data immediately (expensive gas usage)
@@ -187,6 +189,7 @@ impl VestingContract {
             start_time,
             end_time,
             is_initialized: false, // Mark as lazy initialized
+            is_irrevocable: false, // Default to revocable
         };
 
         // Store only essential data initially (cheaper gas)
@@ -226,6 +229,7 @@ impl VestingContract {
                     start_time: 0,
                     end_time: 0,
                     is_initialized: false,
+                    is_irrevocable: false,
                 }
             });
 
@@ -341,6 +345,7 @@ impl VestingContract {
                 start_time: batch_data.start_times.get(i).unwrap(),
                 end_time: batch_data.end_times.get(i).unwrap(),
                 is_initialized: false, // Lazy initialization
+                is_irrevocable: false, // Default to revocable
             };
 
             // Store vault data (minimal writes)
@@ -392,6 +397,7 @@ impl VestingContract {
                 start_time: batch_data.start_times.get(i).unwrap(),
                 end_time: batch_data.end_times.get(i).unwrap(),
                 is_initialized: true, // Full initialization
+                is_irrevocable: false, // Default to revocable
             };
 
             // Store vault data (expensive writes)
@@ -438,6 +444,7 @@ impl VestingContract {
                     start_time: 0,
                     end_time: 0,
                     is_initialized: false,
+                    is_irrevocable: false,
                 }
             });
 
@@ -469,6 +476,7 @@ impl VestingContract {
                         start_time: 0,
                         end_time: 0,
                         is_initialized: false,
+                        is_irrevocable: false,
                     }
                 });
 
@@ -490,6 +498,9 @@ impl VestingContract {
                 panic!("Vault not found");
             });
 
+        // Security check: Cannot revoke from irrevocable vaults
+        require!(!vault.is_irrevocable, "Vault is irrevocable");
+        
         // Calculate amount to return (unreleased tokens)
         let unreleased_amount = vault.total_amount - vault.released_amount;
         require!(unreleased_amount > 0, "No tokens available to revoke");
@@ -525,6 +536,9 @@ impl VestingContract {
                 panic!("Vault not found");
             });
 
+        // Security check: Cannot revoke from irrevocable vaults
+        require!(!vault.is_irrevocable, "Vault is irrevocable");
+        
         // Calculate unvested balance (tokens not yet released)
         let unvested_balance = vault.total_amount - vault.released_amount;
         require!(amount > 0, "Amount to revoke must be positive");
@@ -551,6 +565,42 @@ impl VestingContract {
         amount
     }
 
+    // Mark a vault as irrevocable to prevent admin withdrawal
+    pub fn mark_irrevocable(env: Env, vault_id: u64) {
+        Self::require_admin(&env);
+        
+        let mut vault: Vault = env.storage().instance()
+            .get(&VAULT_DATA, &vault_id)
+            .unwrap_or_else(|| {
+                panic!("Vault not found");
+            });
+        
+        // Cannot mark already irrevocable vaults
+        require!(!vault.is_irrevocable, "Vault is already irrevocable");
+        
+        // Mark vault as irrevocable
+        vault.is_irrevocable = true;
+        env.storage().instance().set(&VAULT_DATA, &vault_id, &vault);
+        
+        // Emit IrrevocableMarked event
+        let timestamp = env.ledger().timestamp();
+        env.events().publish(
+            (Symbol::new(&env, "IrrevocableMarked"), vault_id),
+            (timestamp),
+        );
+    }
+    
+    // Check if a vault is irrevocable
+    pub fn is_vault_irrevocable(env: Env, vault_id: u64) -> bool {
+        let vault: Vault = env.storage().instance()
+            .get(&VAULT_DATA, &vault_id)
+            .unwrap_or_else(|| {
+                panic!("Vault not found");
+            });
+        
+        vault.is_irrevocable
+    }
+    
     // Get contract state for invariant checking
     pub fn get_contract_state(env: Env) -> (i128, i128, i128) {
         let initial_supply: i128 = env.storage().instance().get(&INITIAL_SUPPLY).unwrap_or(0);

contracts/vesting_contracts/src/test.rs

@@ -549,6 +549,121 @@ fn test_usdc_integration_mock_token() {
     assert_eq!(updated_vault.released_amount, claimable);
 }
 
+#[test]
+fn test_irrevocable_vault_security() {
+    let env = Env::default();
+    let contract_id = env.register(VestingContract, ());
+    let client = VestingContractClient::new(&env, &contract_id);
+    
+    // Create addresses for testing
+    let admin = Address::generate(&env);
+    let vault_owner = Address::generate(&env);
+    let unauthorized_user = Address::generate(&env);
+    
+    // Initialize contract with admin
+    let initial_supply = 1000000i128;
+    client.initialize(&admin, &initial_supply);
+    
+    // Create a vault
+    env.as_contract(&contract_id, || {
+        env.current_contract_address().set(&admin);
+    });
+    
+    let vault_amount = 1000i128;
+    let vault_id = client.create_vault_full(&vault_owner, &vault_amount, &100u64, &200u64);
+    
+    // Verify vault is initially revocable
+    assert_eq!(client.is_vault_irrevocable(&vault_id), false);
+    
+    // Test: Unauthorized user cannot mark vault as irrevocable
+    env.as_contract(&contract_id, || {
+        env.current_contract_address().set(&unauthorized_user);
+    });
+    
+    let result = std::panic::catch_unwind(|| {
+        client.mark_irrevocable(&vault_id);
+    });
+    assert!(result.is_err());
+    
+    // Test: Admin can mark vault as irrevocable
+    env.as_contract(&contract_id, || {
+        env.current_contract_address().set(&admin);
+    });
+    
+    client.mark_irrevocable(&vault_id);
+    
+    // Verify vault is now irrevocable
+    assert_eq!(client.is_vault_irrevocable(&vault_id), true);
+    
+    // Test: Cannot mark already irrevocable vault
+    let result = std::panic::catch_unwind(|| {
+        client.mark_irrevocable(&vault_id);
+    });
+    assert!(result.is_err());
+    
+    // Test: Admin cannot revoke tokens from irrevocable vault (full revocation)
+    let result = std::panic::catch_unwind(|| {
+        client.revoke_tokens(&vault_id);
+    });
+    assert!(result.is_err());
+    
+    // Test: Admin cannot revoke partial tokens from irrevocable vault
+    let result = std::panic::catch_unwind(|| {
+        client.revoke_partial(&vault_id, &100i128);
+    });
+    assert!(result.is_err());
+    
+    // Verify the vault state remains unchanged
+    let vault = client.get_vault(&vault_id);
+    assert_eq!(vault.released_amount, 0);
+    assert_eq!(vault.total_amount, vault_amount);
+    assert_eq!(vault.is_irrevocable, true);
+}
+
+#[test]
+fn test_irrevocable_vault_with_claims() {
+    let env = Env::default();
+    let contract_id = env.register(VestingContract, ());
+    let client = VestingContractClient::new(&env, &contract_id);
+    
+    // Create addresses for testing
+    let admin = Address::generate(&env);
+    let vault_owner = Address::generate(&env);
+    
+    // Initialize contract with admin
+    let initial_supply = 1000000i128;
+    client.initialize(&admin, &initial_supply);
+    
+    // Create a vault
+    env.as_contract(&contract_id, || {
+        env.current_contract_address().set(&admin);
+    });
+    
+    let vault_amount = 1000i128;
+    let vault_id = client.create_vault_full(&vault_owner, &vault_amount, &100u64, &200u64);
+    
+    // Mark vault as irrevocable
+    client.mark_irrevocable(&vault_id);
+    
+    // Beneficiary can still claim tokens from irrevocable vault
+    env.ledger().set_timestamp(150); // Halfway through vesting period
+    
+    let claimable_amount = vault_amount / 2; // 500 tokens should be claimable
+    let claimed = client.claim_tokens(&vault_id, &claimable_amount);
+    assert_eq!(claimed, claimable_amount);
+    
+    // Verify vault state after claim
+    let vault = client.get_vault(&vault_id);
+    assert_eq!(vault.released_amount, claimable_amount);
+    assert_eq!(vault.is_irrevocable, true);
+    
+    // Admin still cannot revoke even after claims
+    let result = std::panic::catch_unwind(|| {
+        client.revoke_partial(&vault_id, &100i128);
+    });
+    assert!(result.is_err());
+}
+
 // -------------------------------------------------------------------------
 // Additional beneficiary-transfer tests added for coverage
 // -------------------------------------------------------------------------