From d8ad684b0ecf0422f4c1547daa90b9ed47f4ebcf Mon Sep 17 00:00:00 2001 From: shamoo53 Date: Sun, 22 Feb 2026 12:00:58 +0100 Subject: [PATCH] Security-Prevent-Admin-Rug-Pull Security-Prevent-Admin-Rug-Pull --- contracts/vesting_contracts/src/lib.rs | 50 +++++++++++ contracts/vesting_contracts/src/test.rs | 115 ++++++++++++++++++++++++ 2 files changed, 165 insertions(+) diff --git a/contracts/vesting_contracts/src/lib.rs b/contracts/vesting_contracts/src/lib.rs index 49ca59f..72cfa16 100644 --- a/contracts/vesting_contracts/src/lib.rs +++ b/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); diff --git a/contracts/vesting_contracts/src/test.rs b/contracts/vesting_contracts/src/test.rs index 4c4090c..4f14465 100644 --- a/contracts/vesting_contracts/src/test.rs +++ b/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 // -------------------------------------------------------------------------