diff --git a/contracts/vesting_contracts/src/lib.rs b/contracts/vesting_contracts/src/lib.rs index 4a721d8..e06e97e 100644 --- a/contracts/vesting_contracts/src/lib.rs +++ b/contracts/vesting_contracts/src/lib.rs @@ -22,6 +22,10 @@ pub struct Vault { pub keeper_fee: i128, // Fee paid to anyone who triggers auto_claim pub is_initialized: bool, // Lazy initialization flag pub is_irrevocable: bool, // Security flag to prevent admin withdrawal + pub creation_time: u64, // Timestamp of creation for clawback grace period + pub is_transferable: bool, // Can the beneficiary transfer this vault? + pub step_duration: u64, // Duration of each vesting step in seconds (0 = linear) + pub staked_amount: i128, // Amount currently staked in external contract } #[contracttype] @@ -39,6 +43,7 @@ pub struct BatchCreateData { pub start_times: Vec, pub end_times: Vec, pub keeper_fees: Vec, + pub step_durations: Vec, } #[contracttype] @@ -195,6 +200,8 @@ impl VestingContract { .instance() .set(&DataKey::AdminBalance, &admin_balance); + let now = env.ledger().timestamp(); + // Create vault with full initialization let vault = Vault { owner: owner.clone(), @@ -203,7 +210,13 @@ impl VestingContract { released_amount: 0, start_time, end_time, -< + keeper_fee, + is_initialized: true, + is_irrevocable: !is_revocable, + creation_time: now, + is_transferable, + step_duration, + staked_amount: 0, }; // Store vault data immediately (expensive gas usage) @@ -228,7 +241,6 @@ impl VestingContract { .set(&DataKey::VaultCount, &vault_count); // Emit VaultCreated event with strictly typed fields - let now = env.ledger().timestamp(); let cliff_duration = start_time.saturating_sub(now); let vault_created = VaultCreated { vault_id: vault_count, @@ -269,6 +281,8 @@ impl VestingContract { .instance() .set(&DataKey::AdminBalance, &admin_balance); + let now = env.ledger().timestamp(); + // Create vault with lazy initialization (minimal storage) let vault = Vault { owner: owner.clone(), @@ -280,6 +294,10 @@ impl VestingContract { keeper_fee, is_initialized: false, // Mark as lazy initialized is_irrevocable: !is_revocable, // Convert from is_revocable parameter + creation_time: now, + is_transferable, + step_duration, + staked_amount: 0, }; // Store only essential data initially (cheaper gas) @@ -295,7 +313,6 @@ impl VestingContract { // Don't update user vaults list yet (lazy) // Emit VaultCreated event with strictly typed fields - let now = env.ledger().timestamp(); let cliff_duration = start_time.saturating_sub(now); let vault_created = VaultCreated { vault_id: vault_count, @@ -341,6 +358,30 @@ impl VestingContract { } } + // Helper to calculate vested amount based on time (linear or step) + fn calculate_time_vested_amount(env: &Env, vault: &Vault) -> i128 { + let now = env.ledger().timestamp(); + if now < vault.start_time { + return 0; + } + if now >= vault.end_time { + return vault.total_amount; + } + let duration = vault.end_time - vault.start_time; + if duration == 0 { + 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 + } + // Claim tokens from vault pub fn claim_tokens(env: Env, vault_id: u64, claim_amount: i128) -> i128 { let mut vault: Vault = env @@ -358,9 +399,35 @@ impl VestingContract { panic!("Claim amount must be positive"); } - let milestones = Self::require_milestones_configured(&env, vault_id); - let unlocked_pct = Self::unlocked_percentage(&milestones); - let unlocked_amount = Self::unlocked_amount(vault.total_amount, unlocked_pct); + // Check if milestones are configured + let unlocked_amount = if env.storage().instance().has(&DataKey::VaultMilestones(vault_id)) { + let milestones = Self::require_milestones_configured(&env, vault_id); + let unlocked_pct = Self::unlocked_percentage(&milestones); + Self::unlocked_amount(vault.total_amount, unlocked_pct) + } else { + // Fallback to time-based vesting + Self::calculate_time_vested_amount(&env, &vault) + }; + + // Auto-unstake logic if needed + let liquid_balance = vault.total_amount - vault.released_amount - vault.staked_amount; + if claim_amount > liquid_balance { + let deficit = claim_amount - liquid_balance; + + // Get staking contract + let staking_contract: Address = env.storage().instance() + .get(&Symbol::new(&env, "StakingContract")) + .expect("Staking contract not set"); + + // Call unstake on external contract + let args = vec![&env, vault_id.into_val(&env), deficit.into_val(&env)]; + env.invoke_contract::<()>(&staking_contract, &Symbol::new(&env, "unstake"), args); + + // Update local state + vault.staked_amount -= deficit; + // Note: We don't save vault here yet, it's saved at the end of function + } + let available_to_claim = unlocked_amount - vault.released_amount; if available_to_claim <= 0 { panic!("No tokens available to claim"); @@ -428,9 +495,9 @@ impl VestingContract { .instance() .set(&DataKey::VaultData(vault_id), &vault); - // Emit BeneficiaryChanged event + // Emit BeneficiaryUpdated event env.events().publish( - (Symbol::new(&env, "BeneficiaryChanged"), vault_id), + (Symbol::new(&env, "BeneficiaryUpdated"), vault_id), (old_owner.clone(), new_address), ); } @@ -455,11 +522,19 @@ impl VestingContract { panic!("Only vault owner can set delegate"); } + let old_delegate = vault.delegate.clone(); + // Update delegate - vault.delegate = delegate; + vault.delegate = delegate.clone(); env.storage() .instance() .set(&DataKey::VaultData(vault_id), &vault); + + // Emit DelegateUpdated event + env.events().publish( + (Symbol::new(&env, "DelegateUpdated"), vault_id), + (old_delegate, delegate), + ); } // Claim tokens as delegate (tokens still go to owner) @@ -626,6 +701,7 @@ impl VestingContract { .instance() .set(&DataKey::AdminBalance, &admin_balance); + let now = env.ledger().timestamp(); for i in 0..batch_data.recipients.len() { let vault_id = initial_count + i as u64 + 1; @@ -640,6 +716,10 @@ impl VestingContract { keeper_fee: batch_data.keeper_fees.get(i).unwrap(), is_initialized: false, // Lazy initialization is_irrevocable: false, // Default to revocable for batch operations + creation_time: now, + is_transferable: false, // Default to non-transferable for batch + step_duration: batch_data.step_durations.get(i).unwrap_or(0), + staked_amount: 0, }; // Store vault data (minimal writes) @@ -648,7 +728,6 @@ impl VestingContract { .set(&DataKey::VaultData(vault_id), &vault); vault_ids.push_back(vault_id); // Emit VaultCreated event for each created vault - let now = env.ledger().timestamp(); let start_time = batch_data.start_times.get(i).unwrap(); let cliff_duration = start_time.saturating_sub(now); let vault_created = VaultCreated { @@ -697,6 +776,7 @@ impl VestingContract { .instance() .set(&DataKey::AdminBalance, &admin_balance); + let now = env.ledger().timestamp(); for i in 0..batch_data.recipients.len() { let vault_id = initial_count + i as u64 + 1; @@ -708,7 +788,14 @@ impl VestingContract { released_amount: 0, start_time: batch_data.start_times.get(i).unwrap(), end_time: batch_data.end_times.get(i).unwrap(), - }; + keeper_fee: batch_data.keeper_fees.get(i).unwrap(), + is_initialized: true, + is_irrevocable: false, // Default to revocable for batch operations + creation_time: now, + is_transferable: false, // Default to non-transferable for batch + step_duration: batch_data.step_durations.get(i).unwrap_or(0), + staked_amount: 0, + }; // Store vault data (expensive writes) env.storage() @@ -728,7 +815,6 @@ impl VestingContract { vault_ids.push_back(vault_id); // Emit VaultCreated event for each created vault - let now = env.ledger().timestamp(); let start_time = batch_data.start_times.get(i).unwrap(); let cliff_duration = start_time.saturating_sub(now); let vault_created = VaultCreated { @@ -895,6 +981,174 @@ impl VestingContract { amount } + // Clawback a vault within the grace period (1 hour) + pub fn clawback_vault(env: Env, vault_id: u64) -> i128 { + Self::require_admin(&env); + + let mut vault: Vault = env + .storage() + .instance() + .get(&DataKey::VaultData(vault_id)) + .unwrap_or_else(|| { + panic!("Vault not found"); + }); + + let now = env.ledger().timestamp(); + let grace_period = 3600; // 1 hour in seconds + + if now > vault.creation_time + grace_period { + panic!("Grace period expired"); + } + + if vault.released_amount > 0 { + panic!("Tokens already claimed"); + } + + // Refund admin + let mut admin_balance: i128 = env + .storage() + .instance() + .get(&DataKey::AdminBalance) + .unwrap_or(0); + admin_balance += vault.total_amount; + env.storage() + .instance() + .set(&DataKey::AdminBalance, &admin_balance); + + // Mark as released/revoked so it can't be claimed + vault.released_amount = vault.total_amount; + env.storage() + .instance() + .set(&DataKey::VaultData(vault_id), &vault); + + // Emit event + env.events().publish( + (Symbol::new(&env, "VaultClawedBack"), vault_id), + vault.total_amount, + ); + + vault.total_amount + } + + // Transfer vault ownership to another beneficiary (if transferable) + pub fn transfer_vault(env: Env, vault_id: u64, new_beneficiary: Address) { + let mut vault: Vault = env + .storage() + .instance() + .get(&DataKey::VaultData(vault_id)) + .unwrap_or_else(|| { + panic!("Vault not found"); + }); + + if !vault.is_initialized { + panic!("Vault not initialized"); + } + + if !vault.is_transferable { + panic!("Vault is non-transferable"); + } + + // Check if caller is the vault owner + let caller = env.current_contract_address(); + if caller != vault.owner { + panic!("Only vault owner can transfer"); + } + + let old_owner = vault.owner.clone(); + + // Update UserVaults + // Remove from old owner + let mut old_user_vaults: Vec = env + .storage() + .instance() + .get(&DataKey::UserVaults(old_owner.clone())) + .unwrap_or(Vec::new(&env)); + + let mut new_old_user_vaults = Vec::new(&env); + for id in old_user_vaults.iter() { + if id != vault_id { + new_old_user_vaults.push_back(id); + } + } + env.storage() + .instance() + .set(&DataKey::UserVaults(old_owner.clone()), &new_old_user_vaults); + + // Add to new owner + let mut new_user_vaults: Vec = env + .storage() + .instance() + .get(&DataKey::UserVaults(new_beneficiary.clone())) + .unwrap_or(Vec::new(&env)); + new_user_vaults.push_back(vault_id); + env.storage() + .instance() + .set(&DataKey::UserVaults(new_beneficiary.clone()), &new_user_vaults); + + // Update vault + vault.owner = new_beneficiary.clone(); + vault.delegate = None; // Reset delegate on transfer + + env.storage() + .instance() + .set(&DataKey::VaultData(vault_id), &vault); + + // Emit event + env.events().publish( + (Symbol::new(&env, "BeneficiaryUpdated"), vault_id), + (old_owner, new_beneficiary), + ); + } + + // Set the whitelisted staking contract address + pub fn set_staking_contract(env: Env, contract: Address) { + Self::require_admin(&env); + env.storage().instance().set(&Symbol::new(&env, "StakingContract"), &contract); + } + + // Stake unvested tokens to the whitelisted staking contract + pub fn stake_tokens(env: Env, vault_id: u64, amount: i128, validator: Address) { + let mut vault: Vault = env + .storage() + .instance() + .get(&DataKey::VaultData(vault_id)) + .unwrap_or_else(|| { + panic!("Vault not found"); + }); + + if !vault.is_initialized { + panic!("Vault not initialized"); + } + + // Check auth (owner or delegate?) - usually owner + let caller = env.current_contract_address(); + if caller != vault.owner { + panic!("Only vault owner can stake"); + } + + // Check available balance (total - released - staked) + let available = vault.total_amount - vault.released_amount - vault.staked_amount; + if amount <= 0 { + panic!("Amount must be positive"); + } + if amount > available { + panic!("Insufficient funds to stake"); + } + + // Get staking contract + let staking_contract: Address = env.storage().instance() + .get(&Symbol::new(&env, "StakingContract")) + .expect("Staking contract not set"); + + // Call stake on external contract + let args = vec![&env, vault_id.into_val(&env), amount.into_val(&env), validator.into_val(&env)]; + env.invoke_contract::<()>(&staking_contract, &Symbol::new(&env, "stake"), args); + + // Update vault state + vault.staked_amount += amount; + env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); + } + // Mark a vault as irrevocable to prevent admin withdrawal pub fn mark_irrevocable(env: Env, vault_id: u64) { Self::require_admin(&env); @@ -991,26 +1245,7 @@ impl VestingContract { .get(&VAULT_DATA, &vault_id) .unwrap_or_else(|| panic!("Vault not found")); - let now = env.ledger().timestamp(); - - if now <= vault.start_time { - return 0; - } - - let elapsed = if now >= vault.end_time { - vault.end_time - vault.start_time - } else { - now - vault.start_time - }; - - let total_duration = vault.end_time - vault.start_time; - - let vested = if total_duration > 0 { - // Use i128 for calculation to prevent overflow then back to i128 - (vault.total_amount * elapsed as i128) / total_duration as i128 - } else { - vault.total_amount - }; + let vested = Self::calculate_time_vested_amount(&env, &vault); if vested > vault.released_amount { vested - vault.released_amount diff --git a/contracts/vesting_contracts/src/test.rs b/contracts/vesting_contracts/src/test.rs index f89fea5..d62ea07 100644 --- a/contracts/vesting_contracts/src/test.rs +++ b/contracts/vesting_contracts/src/test.rs @@ -1,7 +1,7 @@ #![cfg(test)] use super::*; -use soroban_sdk::{vec, Env, Address}; +use soroban_sdk::{vec, Env, Address, contract, contractimpl}; #[test] fn test_admin_ownership_transfer() { @@ -145,6 +145,7 @@ fn test_batch_operations_admin_control() { start_times: vec![&env, 100u64, 150u64], end_times: vec![&env, 200u64, 250u64], keeper_fees: vec![&env, 0i128, 0i128], + step_durations: vec![&env, 0u64, 0u64], }; // Test: Unauthorized user cannot create batch vaults @@ -187,3 +188,123 @@ fn test_milestone_unlocking_and_claim_limits() { }); } + +#[test] +fn test_step_vesting_fuzz() { + let env = Env::default(); + let contract_id = env.register(VestingContract, ()); + let client = VestingContractClient::new(&env, &contract_id); + + let admin = Address::generate(&env); + let beneficiary = Address::generate(&env); + + let initial_supply = 1_000_000_000_000i128; + client.initialize(&admin, &initial_supply); + + env.as_contract(&contract_id, || { + env.current_contract_address().set(&admin); + }); + + // Fuzz testing with prime numbers to check for truncation errors + // Primes: 1009 (amount), 17 (step), 101 (duration) + let total_amount = 1009i128; + let start_time = 1000u64; + let duration = 101u64; // Prime duration + let end_time = start_time + duration; + let step_duration = 17u64; // Prime step + + let vault_id = client.create_vault_full( + &beneficiary, + &total_amount, + &start_time, + &end_time, + &0i128, + &true, + &true, + &step_duration, + ); + + // Advance time to end + env.ledger().with_mut(|li| { + li.timestamp = end_time + 1; + }); + + // Claim all + let claimed = client.claim_tokens(&vault_id, &total_amount); + + // Assert full amount is claimed + assert_eq!(claimed, total_amount); + + // Verify vault state + let vault = client.get_vault(&vault_id); + assert_eq!(vault.released_amount, total_amount); +} + +// Mock Staking Contract for testing cross-contract calls +#[contract] +pub struct MockStakingContract; + +#[contractimpl] +impl MockStakingContract { + pub fn stake(env: Env, vault_id: u64, amount: i128, _validator: Address) { + env.events().publish((Symbol::new(&env, "stake"), vault_id), amount); + } + pub fn unstake(env: Env, vault_id: u64, amount: i128) { + env.events().publish((Symbol::new(&env, "unstake"), vault_id), amount); + } +} + +#[test] +fn test_staking_integration() { + let env = Env::default(); + let contract_id = env.register(VestingContract, ()); + let client = VestingContractClient::new(&env, &contract_id); + + // Register mock staking contract + let staking_contract_id = env.register(MockStakingContract, ()); + let staking_client = MockStakingContractClient::new(&env, &staking_contract_id); + + let admin = Address::generate(&env); + let beneficiary = Address::generate(&env); + let validator = Address::generate(&env); + + let initial_supply = 1_000_000i128; + client.initialize(&admin, &initial_supply); + + // Set staking contract + env.as_contract(&contract_id, || { + env.current_contract_address().set(&admin); + }); + client.set_staking_contract(&staking_contract_id); + + // Create vault + let total_amount = 1000i128; + let now = env.ledger().timestamp(); + let vault_id = client.create_vault_full( + &beneficiary, &total_amount, &now, &(now + 1000), &0i128, &true, &true, &0u64 + ); + + // Stake tokens as beneficiary + env.as_contract(&contract_id, || { + env.current_contract_address().set(&beneficiary); + }); + + let stake_amount = 500i128; + client.stake_tokens(&vault_id, &stake_amount, &validator); + + // Verify vault state + let vault = client.get_vault(&vault_id); + assert_eq!(vault.staked_amount, stake_amount); + + // Fast forward to end of vesting + env.ledger().with_mut(|li| { + li.timestamp = now + 1001; + }); + + // Claim ALL tokens (should trigger auto-unstake) + client.claim_tokens(&vault_id, &total_amount); + + let vault_final = client.get_vault(&vault_id); + assert_eq!(vault_final.staked_amount, 0); + assert_eq!(vault_final.released_amount, total_amount); +}