From 6f56d9b728d233d16b1866752dd3fa386b669209 Mon Sep 17 00:00:00 2001 From: lifewithbigdamz Date: Wed, 25 Feb 2026 10:53:14 +0100 Subject: [PATCH] feat: Implement global pause switch for security (Issue #12) - Add IsPaused to DataKey enum for instance storage - Initialize pause state to false in initialize() function - Implement toggle_pause() function (admin-only) - Add pause checks to claim_tokens() and claim_as_delegate() functions - Add is_paused() getter function - Add comprehensive test suite for pause functionality - Emit PauseToggled events for transparency This provides the 'Big Red Button' emergency pause functionality to halt all withdrawals in case of discovered vulnerabilities. --- contracts/vesting_contracts/src/lib.rs | 141 +++++++++- contracts/vesting_contracts/src/test.rs | 353 +++++++++++++++++++++++- 2 files changed, 474 insertions(+), 20 deletions(-) diff --git a/contracts/vesting_contracts/src/lib.rs b/contracts/vesting_contracts/src/lib.rs index 71a09e3..6d701d3 100644 --- a/contracts/vesting_contracts/src/lib.rs +++ b/contracts/vesting_contracts/src/lib.rs @@ -7,6 +7,24 @@ pub enum WhitelistDataKey { WhitelistedTokens, } +// DataKey for contract storage +#[contracttype] +pub enum DataKey { + AdminAddress, + AdminBalance, + InitialSupply, + ProposedAdmin, + VaultCount, + VaultData(u64), + UserVaults(Address), + VaultMilestones(u64), + IsPaused, +} + +// Storage keys for auto-claim functionality +const VAULT_DATA: DataKey = DataKey::VaultData(0); // Placeholder, will be indexed dynamically +const KEEPER_FEES: Symbol = Symbol::short(&"keeper_fees"); + mod factory; pub use factory::{VestingFactory, VestingFactoryClient}; @@ -104,6 +122,9 @@ impl VestingContract { // Initialize vault count env.storage().instance().set(&DataKey::VaultCount, &0u64); + // Initialize pause state to false (unpaused) + env.storage().instance().set(&DataKey::IsPaused, &false); + // Initialize whitelisted tokens map let whitelist: Map = Map::new(&env); env.storage().instance().set(&WhitelistDataKey::WhitelistedTokens, &whitelist); @@ -197,7 +218,45 @@ impl VestingContract { env.storage().instance().get(&DataKey::ProposedAdmin) } + // Toggle pause state (Admin only) - "Big Red Button" for emergency pause + pub fn toggle_pause(env: Env) { + Self::require_admin(&env); + + let current_pause_state: bool = env.storage() + .instance() + .get(&DataKey::IsPaused) + .unwrap_or(false); + + let new_pause_state = !current_pause_state; + env.storage().instance().set(&DataKey::IsPaused, &new_pause_state); + + // Emit event for pause state change + env.events().publish( + Symbol::new(&env, "PauseToggled"), + (new_pause_state, env.ledger().timestamp()), + ); + } + + // Get current pause state + pub fn is_paused(env: Env) -> bool { + env.storage() + .instance() + .get(&DataKey::IsPaused) + .unwrap_or(false) + } + // Full initialization - writes all metadata immediately + pub fn create_vault_full( + env: Env, + owner: Address, + amount: i128, + start_time: u64, + end_time: u64, + keeper_fee: i128, + is_revocable: bool, + is_transferable: bool, + step_duration: u64, + ) -> u64 { Self::require_admin(&env); @@ -281,6 +340,17 @@ impl VestingContract { } // Lazy initialization - writes minimal data initially + pub fn create_vault_lazy( + env: Env, + owner: Address, + amount: i128, + start_time: u64, + end_time: u64, + keeper_fee: i128, + is_revocable: bool, + is_transferable: bool, + step_duration: u64, + ) -> u64 { // Get next vault ID let mut vault_count: u64 = env @@ -353,6 +423,14 @@ impl VestingContract { } // Initialize vault metadata when needed (on-demand) + fn initialize_vault_metadata(env: &Env, vault_id: u64) -> bool { + let vault: Vault = env + .storage() + .instance() + .get(&DataKey::VaultData(vault_id)) + .unwrap_or_else(|| { + panic!("Vault not found"); + }); // Only initialize if not already initialized if !vault.is_initialized { @@ -395,18 +473,34 @@ impl VestingContract { return vault.total_amount; } let elapsed = now - vault.start_time; - let effective_elapsed = if vault.step_duration > 0 { - (elapsed / vault.step_duration) * vault.step_duration - } else { - elapsed - }; - // Use i128 math - (vault.total_amount * effective_elapsed as i128) / duration as i128 + if vault.step_duration > 0 { + // Periodic vesting: calculate vested = (elapsed / step_duration) * rate * step_duration + // Rate is total_amount / duration, so: vested = (elapsed / step_duration) * (total_amount / duration) * step_duration + // This simplifies to: vested = (elapsed / step_duration) * total_amount * step_duration / duration + let completed_steps = elapsed / vault.step_duration; + let rate_per_second = vault.total_amount / duration as i128; + let vested = completed_steps as i128 * rate_per_second * vault.step_duration as i128; + + // Ensure we don't exceed total amount + if vested > vault.total_amount { + vault.total_amount + } else { + vested + } + } else { + // Linear vesting + (vault.total_amount * elapsed as i128) / duration as i128 + } } // Claim tokens from vault pub fn claim_tokens(env: Env, vault_id: u64, claim_amount: i128) -> i128 { + // Check if contract is paused + if Self::is_paused(env.clone()) { + panic!("Contract is paused - all withdrawals are disabled"); + } + let mut vault: Vault = env .storage() .instance() @@ -562,6 +656,11 @@ impl VestingContract { // Claim tokens as delegate (tokens still go to owner) pub fn claim_as_delegate(env: Env, vault_id: u64, claim_amount: i128) -> i128 { + // Check if contract is paused + if Self::is_paused(env.clone()) { + panic!("Contract is paused - all withdrawals are disabled"); + } + let vault: Vault = env .storage() .instance() @@ -862,6 +961,13 @@ impl VestingContract { // Get vault info (initializes if needed) pub fn get_vault(env: Env, vault_id: u64) -> Vault { + let vault: Vault = env + .storage() + .instance() + .get(&DataKey::VaultData(vault_id)) + .unwrap_or_else(|| { + panic!("Vault not found"); + }); // Auto-initialize if lazy if !vault.is_initialized { @@ -886,7 +992,12 @@ impl VestingContract { // Initialize all lazy vaults for this user for vault_id in vault_ids.iter() { - + let vault: Vault = env + .storage() + .instance() + .get(&DataKey::VaultData(vault_id)) + .unwrap_or_else(|| { + panic!("Vault not found"); }); if !vault.is_initialized { @@ -1328,7 +1439,7 @@ impl VestingContract { // Calculate currently claimable tokens based on linear vesting pub fn get_claimable_amount(env: Env, vault_id: u64) -> i128 { let vault: Vault = env.storage().instance() - .get(&VAULT_DATA, &vault_id) + .get(&DataKey::VaultData(vault_id)) .unwrap_or_else(|| panic!("Vault not found")); let vested = Self::calculate_time_vested_amount(&env, &vault); @@ -1344,21 +1455,25 @@ impl VestingContract { // Tokens go to beneficiary, but keeper can get a tip. pub fn auto_claim(env: Env, vault_id: u64, keeper: Address) { let mut vault: Vault = env.storage().instance() - .get(&VAULT_DATA, &vault_id) + .get(&DataKey::VaultData(vault_id)) .unwrap_or_else(|| panic!("Vault not found")); - require!(vault.is_initialized, "Vault not initialized"); + if !vault.is_initialized { + panic!("Vault not initialized"); + } let claimable = Self::get_claimable_amount(env.clone(), vault_id); // Ensure there's enough to cover the fee and something left for beneficiary - require!(claimable > vault.keeper_fee, "Insufficient claimable tokens to cover fee"); + if claimable <= vault.keeper_fee { + panic!("Insufficient claimable tokens to cover fee"); + } let beneficiary_amount = claimable - vault.keeper_fee; // Update vault vault.released_amount += claimable; - env.storage().instance().set(&VAULT_DATA, &vault_id, &vault); + env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); // Update keeper fees let mut fees: Map = env.storage().instance().get(&KEEPER_FEES).unwrap_or(Map::new(&env)); diff --git a/contracts/vesting_contracts/src/test.rs b/contracts/vesting_contracts/src/test.rs index b7a64d5..523bece 100644 --- a/contracts/vesting_contracts/src/test.rs +++ b/contracts/vesting_contracts/src/test.rs @@ -83,6 +83,237 @@ fn test_admin_ownership_transfer() { assert_eq!(client.get_proposed_admin(), Some(another_admin)); } +#[test] +fn test_periodic_vesting_monthly_steps() { + 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 beneficiary = Address::generate(&env); + + // Initialize contract + let initial_supply = 1000000i128; + client.initialize(&admin, &initial_supply); + + // Set admin as caller + env.as_contract(&contract_id, || { + env.current_contract_address().set(&admin); + }); + + // Create vault with monthly vesting (30 days = 2,592,000 seconds) + let amount = 1200000i128; // 1,200,000 tokens over 12 months = 100,000 per month + let start_time = 1000000u64; + let end_time = start_time + (365 * 24 * 60 * 60); // 1 year + let step_duration = 30 * 24 * 60 * 60; // 30 days in seconds + let keeper_fee = 1000i128; + + let vault_id = client.create_vault_full( + &beneficiary, + &amount, + &start_time, + &end_time, + &keeper_fee, + &false, // revocable + &true, // transferable + &step_duration, + ); + + // Test 1: Before start time - no vesting + env.ledger().set_timestamp(start_time - 1000); + let claimable = client.get_claimable_amount(&vault_id); + assert_eq!(claimable, 0, "Should have no claimable tokens before start time"); + + // Test 2: After 15 days (less than one step) - still no vesting (rounds down) + env.ledger().set_timestamp(start_time + (15 * 24 * 60 * 60)); + let claimable = client.get_claimable_amount(&vault_id); + assert_eq!(claimable, 0, "Should have no claimable tokens before first step completes"); + + // Test 3: After exactly 30 days - one step completed + env.ledger().set_timestamp(start_time + step_duration); + let claimable = client.get_claimable_amount(&vault_id); + let expected_monthly = amount / 12; // 100,000 tokens per month + assert_eq!(claimable, expected_monthly, "Should have exactly one month of tokens after 30 days"); + + // Test 4: After 45 days - still only one step (rounds down) + env.ledger().set_timestamp(start_time + (45 * 24 * 60 * 60)); + let claimable = client.get_claimable_amount(&vault_id); + assert_eq!(claimable, expected_monthly, "Should still have only one month of tokens after 45 days"); + + // Test 5: After 60 days - two steps completed + env.ledger().set_timestamp(start_time + (2 * step_duration)); + let claimable = client.get_claimable_amount(&vault_id); + assert_eq!(claimable, 2 * expected_monthly, "Should have two months of tokens after 60 days"); + + // Test 6: After 6 months - 6 steps completed + env.ledger().set_timestamp(start_time + (6 * step_duration)); + let claimable = client.get_claimable_amount(&vault_id); + assert_eq!(claimable, 6 * expected_monthly, "Should have six months of tokens after 6 months"); + + // Test 7: After end time - all tokens vested + env.ledger().set_timestamp(end_time + 1000); + let claimable = client.get_claimable_amount(&vault_id); + assert_eq!(claimable, amount, "Should have all tokens vested after end time"); +} + +#[test] +fn test_periodic_vesting_weekly_steps() { + 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 beneficiary = Address::generate(&env); + + // Initialize contract + let initial_supply = 1000000i128; + client.initialize(&admin, &initial_supply); + + // Set admin as caller + env.as_contract(&contract_id, || { + env.current_contract_address().set(&admin); + }); + + // Create vault with weekly vesting (7 days = 604,800 seconds) + let amount = 520000i128; // 520,000 tokens over 52 weeks = 10,000 per week + let start_time = 1000000u64; + let end_time = start_time + (365 * 24 * 60 * 60); // 1 year + let step_duration = 7 * 24 * 60 * 60; // 7 days in seconds + let keeper_fee = 100i128; + + let vault_id = client.create_vault_full( + &beneficiary, + &amount, + &start_time, + &end_time, + &keeper_fee, + &false, // revocable + &true, // transferable + &step_duration, + ); + + // Test: After 3 weeks - 3 steps completed + env.ledger().set_timestamp(start_time + (3 * step_duration)); + let claimable = client.get_claimable_amount(&vault_id); + let expected_weekly = 10000i128; // 10,000 tokens per week + assert_eq!(claimable, 3 * expected_weekly, "Should have three weeks of tokens after 3 weeks"); + + // Test: After 10 weeks - 10 steps completed + env.ledger().set_timestamp(start_time + (10 * step_duration)); + let claimable = client.get_claimable_amount(&vault_id); + assert_eq!(claimable, 10 * expected_weekly, "Should have ten weeks of tokens after 10 weeks"); +} + +#[test] +fn test_linear_vesting_step_duration_zero() { + 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 beneficiary = Address::generate(&env); + + // Initialize contract + let initial_supply = 1000000i128; + client.initialize(&admin, &initial_supply); + + // Set admin as caller + env.as_contract(&contract_id, || { + env.current_contract_address().set(&admin); + }); + + // Create vault with linear vesting (step_duration = 0) + let amount = 1200000i128; + let start_time = 1000000u64; + let end_time = start_time + (365 * 24 * 60 * 60); // 1 year + let step_duration = 0u64; // Linear vesting + let keeper_fee = 1000i128; + + let vault_id = client.create_vault_full( + &beneficiary, + &amount, + &start_time, + &end_time, + &keeper_fee, + &false, // revocable + &true, // transferable + &step_duration, + ); + + // Test: After 6 months (half the duration) - should have 50% vested + env.ledger().set_timestamp(start_time + (182 * 24 * 60 * 60)); // ~6 months + let claimable = client.get_claimable_amount(&vault_id); + let expected_half = amount / 2; // 50% of tokens + assert_eq!(claimable, expected_half, "Should have 50% of tokens after half the time for linear vesting"); + + // Test: After 3 months (quarter of the duration) - should have 25% vested + env.ledger().set_timestamp(start_time + (91 * 24 * 60 * 60)); // ~3 months + let claimable = client.get_claimable_amount(&vault_id); + let expected_quarter = amount / 4; // 25% of tokens + assert_eq!(claimable, expected_quarter, "Should have 25% of tokens after quarter of the time for linear vesting"); +} + +#[test] +fn test_periodic_vesting_claim_partial() { + 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 beneficiary = Address::generate(&env); + + // Initialize contract + let initial_supply = 1000000i128; + client.initialize(&admin, &initial_supply); + + // Set beneficiary as caller for claiming + env.as_contract(&contract_id, || { + env.current_contract_address().set(&beneficiary); + }); + + // Create vault with monthly vesting + let amount = 120000i128; // 120,000 tokens over 12 months = 10,000 per month + let start_time = 1000000u64; + let end_time = start_time + (365 * 24 * 60 * 60); // 1 year + let step_duration = 30 * 24 * 60 * 60; // 30 days + let keeper_fee = 100i128; + + let vault_id = client.create_vault_full( + &beneficiary, + &amount, + &start_time, + &end_time, + &keeper_fee, + &false, // revocable + &true, // transferable + &step_duration, + ); + + // Move time to 3 months + env.ledger().set_timestamp(start_time + (3 * step_duration)); + + // Claim partial amount + let claim_amount = 15000i128; // Less than the 30,000 available + let claimed = client.claim_tokens(&vault_id, &claim_amount); + assert_eq!(claimed, claim_amount, "Should claim the requested amount"); + + // Check remaining claimable + let remaining_claimable = client.get_claimable_amount(&vault_id); + assert_eq!(remaining_claimable, 15000i128, "Should have 15,000 tokens remaining claimable"); + + // Claim the rest + let final_claim = client.claim_tokens(&vault_id, &remaining_claimable); + assert_eq!(final_claim, remaining_claimable, "Should claim remaining tokens"); + + // Check no more tokens available + let no_more_claimable = client.get_claimable_amount(&vault_id); + assert_eq!(no_more_claimable, 0, "Should have no more claimable tokens"); +} + #[test] fn test_admin_access_control() { let env = Env::default(); @@ -104,12 +335,16 @@ fn test_admin_access_control() { }); let result = std::panic::catch_unwind(|| { - - }); - assert!(result.is_err()); - - let result = std::panic::catch_unwind(|| { - + client.create_vault_full( + &vault_owner, + &1000i128, + &100u64, + &200u64, + &0i128, + &false, + &true, + &0u64, + ); }); assert!(result.is_err()); @@ -118,7 +353,16 @@ fn test_admin_access_control() { env.current_contract_address().set(&admin); }); - + let vault_id2 = client.create_vault_full( + &vault_owner, + &1000i128, + &100u64, + &200u64, + &0i128, + &false, + &true, + &0u64, + ); assert_eq!(vault_id2, 2); } @@ -180,6 +424,7 @@ fn test_milestone_unlocking_and_claim_limits() { let contract_id = env.register(VestingContract, ()); let client = VestingContractClient::new(&env, &contract_id); + let admin = Address::generate(&env); let initial_supply = 1000000i128; client.initialize(&admin, &initial_supply); @@ -457,3 +702,97 @@ fn test_vault_start_time_immutable() { assert_eq!(updated_vault.start_time, original_start_time); assert_eq!(updated_vault.cliff_duration, original_cliff_duration); } + +#[test] +fn test_global_pause_functionality() { + 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 beneficiary = Address::generate(&env); + let unauthorized_user = Address::generate(&env); + + // Initialize contract with admin + let initial_supply = 1000000i128; + client.initialize(&admin, &initial_supply); + + // Verify initial state is unpaused + assert_eq!(client.is_paused(), false); + + // Test: Unauthorized user cannot toggle pause + env.as_contract(&contract_id, || { + env.current_contract_address().set(&unauthorized_user); + }); + + let result = std::panic::catch_unwind(|| { + client.toggle_pause(); + }); + assert!(result.is_err()); + assert_eq!(client.is_paused(), false); // Should still be unpaused + + // Test: Admin can pause the contract + env.as_contract(&contract_id, || { + env.current_contract_address().set(&admin); + }); + + client.toggle_pause(); + assert_eq!(client.is_paused(), true); // Should now be paused + + // Create a vault for testing claims + let now = env.ledger().timestamp(); + let vault_id = client.create_vault_full( + &beneficiary, + &1000i128, + &now, + &(now + 1000), + &0i128, + &false, + &true, + &0u64, + ); + + // Move time to make tokens claimable + env.ledger().set_timestamp(now + 1001); + + // Set beneficiary as caller + env.as_contract(&contract_id, || { + env.current_contract_address().set(&beneficiary); + }); + + // Test: Claims should fail when paused + let result = std::panic::catch_unwind(|| { + client.claim_tokens(&vault_id, &100i128); + }); + assert!(result.is_err()); + + // Test: Delegate claims should also fail when paused + let delegate = Address::generate(&env); + client.set_delegate(&vault_id, &Some(delegate.clone())); + + env.as_contract(&contract_id, || { + env.current_contract_address().set(&delegate); + }); + + let result = std::panic::catch_unwind(|| { + client.claim_as_delegate(&vault_id, &100i128); + }); + assert!(result.is_err()); + + // Test: Admin can unpause the contract + env.as_contract(&contract_id, || { + env.current_contract_address().set(&admin); + }); + + client.toggle_pause(); + assert_eq!(client.is_paused(), false); // Should be unpaused + + // Test: Claims should work after unpausing + env.as_contract(&contract_id, || { + env.current_contract_address().set(&beneficiary); + }); + + let claimed = client.claim_tokens(&vault_id, &100i128); + assert_eq!(claimed, 100i128); // Should succeed +}