diff --git a/contracts/vesting_contracts/src/lib.rs b/contracts/vesting_contracts/src/lib.rs index d168be5..8b09448 100644 --- a/contracts/vesting_contracts/src/lib.rs +++ b/contracts/vesting_contracts/src/lib.rs @@ -1,6 +1,5 @@ -#![no_std] + #![no_std] use soroban_sdk::{contract, contractimpl, contracttype, token, vec, Address, Env, IntoVal, Map, Symbol, Vec}; -use soroban_sdk::{contract, contractimpl, contracttype, Address, Env, Map, Symbol, Vec, String}; // DataKey for whitelisted tokens #[contracttype] @@ -8,24 +7,6 @@ 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}; @@ -43,26 +24,27 @@ pub enum DataKey { VaultMilestones(u64), UserVaults(Address), KeeperFees, + Token, // yield-bearing token + TotalShares, // remaining initial_deposit_shares + TotalStaked, } -// Vault structure with lazy initialization #[contracttype] #[derive(Clone)] pub struct Vault { pub owner: Address, - pub delegate: Option
, // Optional delegate address for claiming - pub total_amount: i128, + pub delegate: Option
, + pub total_amount: i128, // = initial_deposit_shares pub released_amount: i128, pub start_time: u64, pub end_time: u64, - pub keeper_fee: i128, // Fee paid to anyone who triggers auto_claim - pub title: String, // Short human-readable title (max 32 chars) - 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 + pub keeper_fee: i128, + pub is_initialized: bool, + pub is_irrevocable: bool, + pub creation_time: u64, + pub is_transferable: bool, + pub step_duration: u64, + pub staked_amount: i128, } #[contracttype] @@ -99,13 +81,11 @@ pub struct VaultCreated { pub total_amount: i128, pub cliff_duration: u64, pub start_time: u64, - pub title: String, } #[contractimpl] #[allow(deprecated)] impl VestingContract { - // Admin-only: Add token to whitelist pub fn add_to_whitelist(env: Env, token: Address) { Self::require_admin(&env); let mut whitelist: Map = env @@ -119,7 +99,6 @@ impl VestingContract { .set(&WhitelistDataKey::WhitelistedTokens, &whitelist); } - // Check if token is whitelisted fn is_token_whitelisted(env: &Env, token: &Address) -> bool { let whitelist: Map = env .storage() @@ -129,31 +108,36 @@ impl VestingContract { whitelist.get(token.clone()).unwrap_or(false) } - // Initialize contract with initial supply pub fn initialize(env: Env, admin: Address, initial_supply: i128) { - env.storage() - .instance() - .set(&DataKey::InitialSupply, &initial_supply); - - env.storage() - .instance() - .set(&DataKey::AdminBalance, &initial_supply); - + env.storage().instance().set(&DataKey::InitialSupply, &initial_supply); + env.storage().instance().set(&DataKey::AdminBalance, &initial_supply); env.storage().instance().set(&DataKey::AdminAddress, &admin); - 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() + env.storage().instance().set(&WhitelistDataKey::WhitelistedTokens, &whitelist); + + env.storage().instance().set(&DataKey::TotalShares, &0i128); + env.storage().instance().set(&DataKey::TotalStaked, &0i128); + } + + pub fn set_token(env: Env, token: Address) { + Self::require_admin(&env); + if env.storage().instance().has(&DataKey::Token) { + panic!("Token already set"); + } + env.storage().instance().set(&DataKey::Token, &token); + } + + fn get_token_client(env: &Env) -> token::Client { + let token: Address = env + .storage() .instance() - .set(&WhitelistDataKey::WhitelistedTokens, &whitelist); + .get(&DataKey::Token) + .unwrap_or_else(|| panic!("Token not set - call set_token first")); + token::Client::new(env, &token) } - // Helper function to check if caller is admin fn require_admin(env: &Env) { let admin: Address = env .storage() @@ -182,83 +166,37 @@ impl VestingContract { pct = pct.saturating_add(m.percentage); } } - if pct > 100 { - 100 - } else { - pct - } + if pct > 100 { 100 } else { pct } } fn unlocked_amount(total_amount: i128, unlocked_percentage: u32) -> i128 { (total_amount * unlocked_percentage as i128) / 100i128 } - // Propose a new admin (first step of two-step process) pub fn propose_new_admin(env: Env, new_admin: Address) { Self::require_admin(&env); - env.storage() - .instance() - .set(&DataKey::ProposedAdmin, &new_admin); + env.storage().instance().set(&DataKey::ProposedAdmin, &new_admin); } - // Accept admin ownership (second step of two-step process) pub fn accept_ownership(env: Env) { let proposed_admin: Address = env .storage() .instance() .get(&DataKey::ProposedAdmin) .unwrap_or_else(|| panic!("No proposed admin found")); - proposed_admin.require_auth(); - - env.storage() - .instance() - .set(&DataKey::AdminAddress, &proposed_admin); - + env.storage().instance().set(&DataKey::AdminAddress, &proposed_admin); env.storage().instance().remove(&DataKey::ProposedAdmin); } - // Get current admin address pub fn get_admin(env: Env) -> Address { - env.storage() - .instance() - .get(&DataKey::AdminAddress) - .unwrap_or_else(|| panic!("Admin not set")) + env.storage().instance().get(&DataKey::AdminAddress).unwrap_or_else(|| panic!("Admin not set")) } - // Get proposed admin address (if any) pub fn get_proposed_admin(env: Env) -> Option
{ 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, @@ -270,28 +208,15 @@ impl VestingContract { is_transferable: bool, step_duration: u64, ) -> u64 { - Self::require_admin(&env); - let mut vault_count: u64 = env - .storage() - .instance() - .get(&DataKey::VaultCount) - .unwrap_or(0); + let mut vault_count: u64 = env.storage().instance().get(&DataKey::VaultCount).unwrap_or(0); vault_count += 1; - let mut admin_balance: i128 = env - .storage() - .instance() - .get(&DataKey::AdminBalance) - .unwrap_or(0); - if admin_balance < amount { - panic!("Insufficient admin balance"); - } + let mut admin_balance: i128 = env.storage().instance().get(&DataKey::AdminBalance).unwrap_or(0); + if admin_balance < amount { panic!("Insufficient admin balance"); } admin_balance -= amount; - env.storage() - .instance() - .set(&DataKey::AdminBalance, &admin_balance); + env.storage().instance().set(&DataKey::AdminBalance, &admin_balance); let now = env.ledger().timestamp(); @@ -303,7 +228,6 @@ impl VestingContract { start_time, end_time, keeper_fee, - title: String::from_slice(&env, ""), is_initialized: true, is_irrevocable: !is_revocable, creation_time: now, @@ -312,42 +236,25 @@ impl VestingContract { staked_amount: 0, }; - env.storage() - .instance() - .set(&DataKey::VaultData(vault_count), &vault); + env.storage().instance().set(&DataKey::VaultData(vault_count), &vault); - let mut user_vaults: Vec = env - .storage() - .instance() - .get(&DataKey::UserVaults(owner.clone())) - .unwrap_or(Vec::new(&env)); + let mut user_vaults: Vec = env.storage().instance().get(&DataKey::UserVaults(owner.clone())).unwrap_or(Vec::new(&env)); user_vaults.push_back(vault_count); - env.storage() - .instance() - .set(&DataKey::UserVaults(owner.clone()), &user_vaults); + env.storage().instance().set(&DataKey::UserVaults(owner.clone()), &user_vaults); - env.storage() - .instance() - .set(&DataKey::VaultCount, &vault_count); + env.storage().instance().set(&DataKey::VaultCount, &vault_count); + + let mut total_shares: i128 = env.storage().instance().get(&DataKey::TotalShares).unwrap_or(0); + total_shares += amount; + env.storage().instance().set(&DataKey::TotalShares, &total_shares); let cliff_duration = start_time.saturating_sub(now); - let vault_created = VaultCreated { - vault_id: vault_count, - beneficiary: owner, - total_amount: amount, - cliff_duration, - start_time, - title: String::from_slice(&env, ""), - }; - env.events().publish( - (Symbol::new(&env, "VaultCreated"), vault_count), - vault_created, - ); + let vault_created = VaultCreated { vault_id: vault_count, beneficiary: owner, total_amount: amount, cliff_duration, start_time }; + env.events().publish((Symbol::new(&env, "VaultCreated"), vault_count), vault_created); vault_count } - // Lazy initialization - writes minimal data initially pub fn create_vault_lazy( env: Env, owner: Address, @@ -361,25 +268,13 @@ impl VestingContract { ) -> u64 { Self::require_admin(&env); - let mut vault_count: u64 = env - .storage() - .instance() - .get(&DataKey::VaultCount) - .unwrap_or(0); + let mut vault_count: u64 = env.storage().instance().get(&DataKey::VaultCount).unwrap_or(0); vault_count += 1; - let mut admin_balance: i128 = env - .storage() - .instance() - .get(&DataKey::AdminBalance) - .unwrap_or(0); - if admin_balance < amount { - panic!("Insufficient admin balance"); - } + let mut admin_balance: i128 = env.storage().instance().get(&DataKey::AdminBalance).unwrap_or(0); + if admin_balance < amount { panic!("Insufficient admin balance"); } admin_balance -= amount; - env.storage() - .instance() - .set(&DataKey::AdminBalance, &admin_balance); + env.storage().instance().set(&DataKey::AdminBalance, &admin_balance); let now = env.ledger().timestamp(); @@ -391,8 +286,7 @@ impl VestingContract { start_time, end_time, keeper_fee, - title: String::from_slice(&env, ""), - is_initialized: false, // Mark as lazy initialized + is_initialized: false, is_irrevocable: !is_revocable, creation_time: now, is_transferable, @@ -400,137 +294,73 @@ impl VestingContract { staked_amount: 0, }; - env.storage() - .instance() - .set(&DataKey::VaultData(vault_count), &vault); + env.storage().instance().set(&DataKey::VaultData(vault_count), &vault); + env.storage().instance().set(&DataKey::VaultCount, &vault_count); - // Don't update user vaults list yet (lazy) - env.storage() - .instance() - .set(&DataKey::VaultCount, &vault_count); + let mut total_shares: i128 = env.storage().instance().get(&DataKey::TotalShares).unwrap_or(0); + total_shares += amount; + env.storage().instance().set(&DataKey::TotalShares, &total_shares); let cliff_duration = start_time.saturating_sub(now); - let vault_created = VaultCreated { - vault_id: vault_count, - beneficiary: owner.clone(), - total_amount: amount, - cliff_duration, - start_time, - title: String::from_slice(&env, ""), - }; - env.events().publish( - (Symbol::new(&env, "VaultCreated"), vault_count), - vault_created, - ); + let vault_created = VaultCreated { vault_id: vault_count, beneficiary: owner.clone(), total_amount: amount, cliff_duration, start_time }; + env.events().publish((Symbol::new(&env, "VaultCreated"), vault_count), vault_created); vault_count } - // 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")); - .unwrap_or_else(|| { - panic!("Vault not found"); - }); + let vault: Vault = env.storage().instance().get(&DataKey::VaultData(vault_id)).unwrap_or_else(|| panic!("Vault not found")); if !vault.is_initialized { let mut updated_vault = vault.clone(); updated_vault.is_initialized = true; - env.storage() - .instance() - .set(&DataKey::VaultData(vault_id), &updated_vault); + env.storage().instance().set(&DataKey::VaultData(vault_id), &updated_vault); - let mut user_vaults: Vec = env - .storage() - .instance() - .get(&DataKey::UserVaults(updated_vault.owner.clone())) - .unwrap_or(Vec::new(env)); + let mut user_vaults: Vec = env.storage().instance().get(&DataKey::UserVaults(updated_vault.owner.clone())).unwrap_or(Vec::new(env)); user_vaults.push_back(vault_id); - env.storage() - .instance() - .set(&DataKey::UserVaults(updated_vault.owner), &user_vaults); + env.storage().instance().set(&DataKey::UserVaults(updated_vault.owner), &user_vaults); true } else { - false // Already initialized + false } } - // 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; - } + 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; - } + 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 }; - (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() .get(&DataKey::VaultData(vault_id)) .unwrap_or_else(|| panic!("Vault not found")); - if !vault.is_initialized { - panic!("Vault not initialized"); - } - if claim_amount <= 0 { - panic!("Claim amount must be positive"); - } + if !vault.is_initialized { panic!("Vault not initialized"); } + if claim_amount <= 0 { panic!("Claim amount must be positive"); } - 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 { - Self::calculate_time_vested_amount(&env, &vault) - }; + vault.owner.require_auth(); + + 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 { + Self::calculate_time_vested_amount(&env, &vault) + }; let liquid_balance = vault.total_amount - vault.released_amount - vault.staked_amount; if claim_amount > liquid_balance { @@ -546,122 +376,91 @@ impl VestingContract { env.invoke_contract::<()>(&staking_contract, &Symbol::new(&env, "unstake"), args); vault.staked_amount -= deficit; + + let mut total_staked: i128 = env.storage().instance().get(&DataKey::TotalStaked).unwrap_or(0); + total_staked -= deficit; + env.storage().instance().set(&DataKey::TotalStaked, &total_staked); } let available_to_claim = unlocked_amount - vault.released_amount; - if available_to_claim <= 0 { - panic!("No tokens available to claim"); - } - if claim_amount > available_to_claim { - panic!("Insufficient unlocked tokens to claim"); - } + if available_to_claim <= 0 { panic!("No tokens available to claim"); } + if claim_amount > available_to_claim { panic!("Insufficient unlocked tokens to claim"); } + + // YIELD DISTRIBUTION - only vault-owned portion + let token_client = Self::get_token_client(&env); + let current_balance = token_client.balance(&env.current_contract_address()); + let admin_balance: i128 = env.storage().instance().get(&DataKey::AdminBalance).unwrap_or(0); + + let total_shares: i128 = env.storage().instance().get(&DataKey::TotalShares).unwrap_or(0); + let total_staked: i128 = env.storage().instance().get(&DataKey::TotalStaked).unwrap_or(0); + let liquid_shares = total_shares - total_staked; + + let vault_portion = (current_balance - admin_balance).max(0); + let transfer_amount = if liquid_shares > 0 { + (claim_amount * vault_portion) / liquid_shares + } else { + claim_amount + }; vault.released_amount += claim_amount; - env.storage() - .instance() - .set(&DataKey::VaultData(vault_id), &vault); + let mut updated_total_shares = total_shares; + updated_total_shares -= claim_amount; + env.storage().instance().set(&DataKey::TotalShares, &updated_total_shares); + env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); - claim_amount + token_client.transfer(&env.current_contract_address(), &vault.owner, &transfer_amount); + + transfer_amount } - /// Transfers the beneficiary role of a vault to a new address. - /// Only the admin can perform this action (e.g., in case of lost keys). pub fn transfer_beneficiary(env: Env, vault_id: u64, new_address: Address) { Self::require_admin(&env); - let mut vault: Vault = env - .storage() - .instance() - .get(&DataKey::VaultData(vault_id)) - .unwrap_or_else(|| panic!("Vault not found")); + let mut vault: Vault = env.storage().instance().get(&DataKey::VaultData(vault_id)).unwrap_or_else(|| panic!("Vault not found")); let old_owner = vault.owner.clone(); if vault.is_initialized { - let old_vaults: Vec = env - .storage() - .instance() - .get(&DataKey::UserVaults(old_owner.clone())) - .unwrap_or(Vec::new(&env)); - + let old_vaults: Vec = env.storage().instance().get(&DataKey::UserVaults(old_owner.clone())).unwrap_or(Vec::new(&env)); let mut updated_old_vaults = Vec::new(&env); for id in old_vaults.iter() { if id != vault_id { updated_old_vaults.push_back(id); } } - env.storage() - .instance() - .set(&DataKey::UserVaults(old_owner.clone()), &updated_old_vaults); + env.storage().instance().set(&DataKey::UserVaults(old_owner.clone()), &updated_old_vaults); - let mut new_vaults: Vec = env - .storage() - .instance() - .get(&DataKey::UserVaults(new_address.clone())) - .unwrap_or(Vec::new(&env)); + let mut new_vaults: Vec = env.storage().instance().get(&DataKey::UserVaults(new_address.clone())).unwrap_or(Vec::new(&env)); new_vaults.push_back(vault_id); - env.storage() - .instance() - .set(&DataKey::UserVaults(new_address.clone()), &new_vaults); + env.storage().instance().set(&DataKey::UserVaults(new_address.clone()), &new_vaults); } vault.owner = new_address.clone(); - env.storage() - .instance() - .set(&DataKey::VaultData(vault_id), &vault); + env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); - env.events().publish( - (Symbol::new(&env, "BeneficiaryUpdated"), vault_id), - (old_owner.clone(), new_address), - ); + env.events().publish((Symbol::new(&env, "BeneficiaryUpdated"), vault_id), (old_owner.clone(), new_address)); } - // Set delegate address for a vault (only owner can call) pub fn set_delegate(env: Env, vault_id: u64, delegate: Option
) { - let mut vault: Vault = env - .storage() - .instance() - .get(&DataKey::VaultData(vault_id)) - .unwrap_or_else(|| panic!("Vault not found")); + 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_initialized { panic!("Vault not initialized"); } vault.owner.require_auth(); let old_delegate = vault.delegate.clone(); vault.delegate = delegate.clone(); - env.storage() - .instance() - .set(&DataKey::VaultData(vault_id), &vault); + env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); - env.events().publish( - (Symbol::new(&env, "DelegateUpdated"), vault_id), - (old_delegate, delegate), - ); + env.events().publish((Symbol::new(&env, "DelegateUpdated"), vault_id), (old_delegate, delegate)); } - // 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 mut vault: Vault = env.storage().instance().get(&DataKey::VaultData(vault_id)).unwrap_or_else(|| panic!("Vault not found")); - let 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 claim_amount <= 0 { - panic!("Claim amount must be positive"); - } + if !vault.is_initialized { panic!("Vault not initialized"); } + if claim_amount <= 0 { panic!("Claim amount must be positive"); } let delegate = vault.delegate.clone().unwrap_or_else(|| panic!("No delegate set for this vault")); delegate.require_auth(); @@ -670,81 +469,69 @@ impl VestingContract { let unlocked_pct = Self::unlocked_percentage(&milestones); let unlocked_amount = Self::unlocked_amount(vault.total_amount, unlocked_pct); let available_to_claim = unlocked_amount - vault.released_amount; - if available_to_claim <= 0 { - panic!("No tokens available to claim"); - } - if claim_amount > available_to_claim { - panic!("Insufficient unlocked tokens to claim"); - } + if available_to_claim <= 0 { panic!("No tokens available to claim"); } + if claim_amount > available_to_claim { panic!("Insufficient unlocked tokens to claim"); } + + // YIELD DISTRIBUTION - only vault-owned portion + let token_client = Self::get_token_client(&env); + let current_balance = token_client.balance(&env.current_contract_address()); + let admin_balance: i128 = env.storage().instance().get(&DataKey::AdminBalance).unwrap_or(0); + + let total_shares: i128 = env.storage().instance().get(&DataKey::TotalShares).unwrap_or(0); + let total_staked: i128 = env.storage().instance().get(&DataKey::TotalStaked).unwrap_or(0); + let liquid_shares = total_shares - total_staked; + + let vault_portion = (current_balance - admin_balance).max(0); + let transfer_amount = if liquid_shares > 0 { + (claim_amount * vault_portion) / liquid_shares + } else { + claim_amount + }; let mut updated_vault = vault.clone(); updated_vault.released_amount += claim_amount; - env.storage() - .instance() - .set(&DataKey::VaultData(vault_id), &updated_vault); - claim_amount + let mut updated_total_shares = total_shares; + updated_total_shares -= claim_amount; + env.storage().instance().set(&DataKey::TotalShares, &updated_total_shares); + env.storage().instance().set(&DataKey::VaultData(vault_id), &updated_vault); + + token_client.transfer(&env.current_contract_address(), &updated_vault.owner, &transfer_amount); + + transfer_amount } pub fn set_milestones(env: Env, vault_id: u64, milestones: Vec) { Self::require_admin(&env); - let 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"); - } + let 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 milestones.is_empty() { - panic!("No milestones provided"); - } + if milestones.is_empty() { panic!("No milestones provided"); } let mut total_pct: u32 = 0; let mut seen: Map = Map::new(&env); for m in milestones.iter() { - if m.percentage == 0 { - panic!("Milestone percentage must be positive"); - } - if m.percentage > 100 { - panic!("Milestone percentage too large"); - } - if seen.contains_key(m.id) { - panic!("Duplicate milestone id"); - } + if m.percentage == 0 { panic!("Milestone percentage must be positive"); } + if m.percentage > 100 { panic!("Milestone percentage too large"); } + if seen.contains_key(m.id) { panic!("Duplicate milestone id"); } seen.set(m.id, true); total_pct = total_pct.saturating_add(m.percentage); } - if total_pct > 100 { - panic!("Total milestone percentage exceeds 100"); - } + if total_pct > 100 { panic!("Total milestone percentage exceeds 100"); } - env.storage() - .instance() - .set(&DataKey::VaultMilestones(vault_id), &milestones); - env.events().publish( - (Symbol::new(&env, "MilestonesSet"), vault_id), - (milestones.len(), total_pct), - ); + env.storage().instance().set(&DataKey::VaultMilestones(vault_id), &milestones); + env.events().publish((Symbol::new(&env, "MilestonesSet"), vault_id), (milestones.len(), total_pct)); } pub fn get_milestones(env: Env, vault_id: u64) -> Vec { - env.storage() - .instance() - .get(&DataKey::VaultMilestones(vault_id)) - .unwrap_or(Vec::new(&env)) + env.storage().instance().get(&DataKey::VaultMilestones(vault_id)).unwrap_or(Vec::new(&env)) } pub fn unlock_milestone(env: Env, vault_id: u64, milestone_id: u64) { Self::require_admin(&env); - let _vault: Vault = env - .storage() - .instance() - .get(&DataKey::VaultData(vault_id)) - .unwrap_or_else(|| panic!("Vault not found")); + let _vault: Vault = env.storage().instance().get(&DataKey::VaultData(vault_id)).unwrap_or_else(|| panic!("Vault not found")); let milestones = Self::require_milestones_configured(&env, vault_id); @@ -753,75 +540,30 @@ impl VestingContract { for m in milestones.iter() { if m.id == milestone_id { found = true; - if m.is_unlocked { - panic!("Milestone already unlocked"); - } - updated.push_back(Milestone { - id: m.id, - percentage: m.percentage, - is_unlocked: true, - }); + if m.is_unlocked { panic!("Milestone already unlocked"); } + updated.push_back(Milestone { id: m.id, percentage: m.percentage, is_unlocked: true }); } else { updated.push_back(m); } } - if !found { - panic!("Milestone not found"); - } + if !found { panic!("Milestone not found"); } - env.storage() - .instance() - .set(&DataKey::VaultMilestones(vault_id), &updated); + env.storage().instance().set(&DataKey::VaultMilestones(vault_id), &updated); let timestamp = env.ledger().timestamp(); - env.events().publish( - (Symbol::new(&env, "MilestoneUnlocked"), vault_id), - (milestone_id, timestamp), - ); + env.events().publish((Symbol::new(&env, "MilestoneUnlocked"), vault_id), (milestone_id, timestamp)); } - // Admin-only: set a short title for a vault (max 32 bytes) - pub fn set_vault_title(env: Env, vault_id: u64, title: String) { - Self::require_admin(&env); - - // Enforce max length (32 bytes) - if title.len() > 32 { - panic!("Title too long"); - } - - let mut vault: Vault = env - .storage() - .instance() - .get(&DataKey::VaultData(vault_id)) - .unwrap_or_else(|| panic!("Vault not found")); - - vault.title = title; - env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); - } - - // Batch create vaults with lazy initialization pub fn batch_create_vaults_lazy(env: Env, batch_data: BatchCreateData) -> Vec { Self::require_admin(&env); let mut vault_ids = Vec::new(&env); - let initial_count: u64 = env - .storage() - .instance() - .get(&DataKey::VaultCount) - .unwrap_or(0); + let initial_count: u64 = env.storage().instance().get(&DataKey::VaultCount).unwrap_or(0); let total_amount: i128 = batch_data.amounts.iter().sum(); - let mut admin_balance: i128 = env - .storage() - .instance() - .get(&DataKey::AdminBalance) - .unwrap_or(0); - if admin_balance < total_amount { - panic!("Insufficient admin balance for batch"); - } + let mut admin_balance: i128 = env.storage().instance().get(&DataKey::AdminBalance).unwrap_or(0); + if admin_balance < total_amount { panic!("Insufficient admin balance for batch"); } admin_balance -= total_amount; - env.storage() - .instance() - .set(&DataKey::AdminBalance, &admin_balance); + env.storage().instance().set(&DataKey::AdminBalance, &admin_balance); let now = env.ledger().timestamp(); for i in 0..batch_data.recipients.len() { @@ -837,66 +579,42 @@ impl VestingContract { keeper_fee: batch_data.keeper_fees.get(i).unwrap(), is_initialized: false, is_irrevocable: false, - title: String::from_slice(&env, ""), - is_initialized: false, // Lazy initialization - is_irrevocable: false, // Default to revocable for batch operations creation_time: now, is_transferable: false, step_duration: batch_data.step_durations.get(i).unwrap_or(0), staked_amount: 0, }; - env.storage() - .instance() - .set(&DataKey::VaultData(vault_id), &vault); + env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); vault_ids.push_back(vault_id); let start_time = batch_data.start_times.get(i).unwrap(); let cliff_duration = start_time.saturating_sub(now); - let vault_created = VaultCreated { - vault_id, - beneficiary: vault.owner.clone(), - total_amount: vault.total_amount, - cliff_duration, - start_time, - title: String::from_slice(&env, ""), - }; - env.events() - .publish((Symbol::new(&env, "VaultCreated"), vault_id), vault_created); + let vault_created = VaultCreated { vault_id, beneficiary: vault.owner.clone(), total_amount: vault.total_amount, cliff_duration, start_time }; + env.events().publish((Symbol::new(&env, "VaultCreated"), vault_id), vault_created); } + let mut total_shares: i128 = env.storage().instance().get(&DataKey::TotalShares).unwrap_or(0); + total_shares += total_amount; + env.storage().instance().set(&DataKey::TotalShares, &total_shares); + let final_count = initial_count + batch_data.recipients.len() as u64; - env.storage() - .instance() - .set(&DataKey::VaultCount, &final_count); + env.storage().instance().set(&DataKey::VaultCount, &final_count); vault_ids } - // Batch create vaults with full initialization pub fn batch_create_vaults_full(env: Env, batch_data: BatchCreateData) -> Vec { Self::require_admin(&env); let mut vault_ids = Vec::new(&env); - let initial_count: u64 = env - .storage() - .instance() - .get(&DataKey::VaultCount) - .unwrap_or(0); + let initial_count: u64 = env.storage().instance().get(&DataKey::VaultCount).unwrap_or(0); let total_amount: i128 = batch_data.amounts.iter().sum(); - let mut admin_balance: i128 = env - .storage() - .instance() - .get(&DataKey::AdminBalance) - .unwrap_or(0); - if admin_balance < total_amount { - panic!("Insufficient admin balance for batch"); - } + let mut admin_balance: i128 = env.storage().instance().get(&DataKey::AdminBalance).unwrap_or(0); + if admin_balance < total_amount { panic!("Insufficient admin balance for batch"); } admin_balance -= total_amount; - env.storage() - .instance() - .set(&DataKey::AdminBalance, &admin_balance); + env.storage().instance().set(&DataKey::AdminBalance, &admin_balance); let now = env.ledger().timestamp(); for i in 0..batch_data.recipients.len() { @@ -910,7 +628,6 @@ impl VestingContract { 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(), - title: String::from_slice(&env, ""), is_initialized: true, is_irrevocable: false, creation_time: now, @@ -919,83 +636,46 @@ impl VestingContract { staked_amount: 0, }; - env.storage() - .instance() - .set(&DataKey::VaultData(vault_id), &vault); + env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); - let mut user_vaults: Vec = env - .storage() - .instance() - .get(&DataKey::UserVaults(vault.owner.clone())) - .unwrap_or(Vec::new(&env)); + let mut user_vaults: Vec = env.storage().instance().get(&DataKey::UserVaults(vault.owner.clone())).unwrap_or(Vec::new(&env)); user_vaults.push_back(vault_id); - env.storage() - .instance() - .set(&DataKey::UserVaults(vault.owner.clone()), &user_vaults); + env.storage().instance().set(&DataKey::UserVaults(vault.owner.clone()), &user_vaults); vault_ids.push_back(vault_id); let start_time = batch_data.start_times.get(i).unwrap(); let cliff_duration = start_time.saturating_sub(now); - let vault_created = VaultCreated { - vault_id, - beneficiary: vault.owner.clone(), - total_amount: vault.total_amount, - cliff_duration, - start_time, - title: String::from_slice(&env, ""), - }; - env.events() - .publish((Symbol::new(&env, "VaultCreated"), vault_id), vault_created); + let vault_created = VaultCreated { vault_id, beneficiary: vault.owner.clone(), total_amount: vault.total_amount, cliff_duration, start_time }; + env.events().publish((Symbol::new(&env, "VaultCreated"), vault_id), vault_created); } + let mut total_shares: i128 = env.storage().instance().get(&DataKey::TotalShares).unwrap_or(0); + total_shares += total_amount; + env.storage().instance().set(&DataKey::TotalShares, &total_shares); + let final_count = initial_count + batch_data.recipients.len() as u64; - env.storage() - .instance() - .set(&DataKey::VaultCount, &final_count); + env.storage().instance().set(&DataKey::VaultCount, &final_count); vault_ids } - // 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")); - .unwrap_or_else(|| { - panic!("Vault not found"); - }); + let vault: Vault = env.storage().instance().get(&DataKey::VaultData(vault_id)).unwrap_or_else(|| panic!("Vault not found")); if !vault.is_initialized { Self::initialize_vault_metadata(&env, vault_id); - env.storage() - .instance() - .get(&DataKey::VaultData(vault_id)) - .unwrap() + env.storage().instance().get(&DataKey::VaultData(vault_id)).unwrap() } else { vault } } - // Get user vaults (initializes all if needed) pub fn get_user_vaults(env: Env, user: Address) -> Vec { - let vault_ids: Vec = env - .storage() - .instance() - .get(&DataKey::UserVaults(user.clone())) - .unwrap_or(Vec::new(&env)); + let vault_ids: Vec = env.storage().instance().get(&DataKey::UserVaults(user.clone())).unwrap_or(Vec::new(&env)); 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")); - .unwrap_or_else(|| { - panic!("Vault not found"); - }); + let vault: Vault = env.storage().instance().get(&DataKey::VaultData(vault_id)).unwrap_or_else(|| panic!("Vault not found")); if !vault.is_initialized { Self::initialize_vault_metadata(&env, vault_id); @@ -1005,216 +685,99 @@ impl VestingContract { vault_ids } - // Revoke tokens from a vault and return them to admin - // Internal helper: revoke full unreleased amount from a vault and emit event. - // Does NOT update admin balance — caller is responsible for a single aggregated transfer. - fn internal_revoke_full(env: &Env, vault_id: u64) -> i128 { - let mut vault: Vault = env - .storage() - .instance() - .get(&DataKey::VaultData(vault_id)) - .unwrap_or_else(|| panic!("Vault not found")); + pub fn revoke_tokens(env: Env, vault_id: u64) -> i128 { + Self::require_admin(&env); - if vault.is_irrevocable { - panic!("Vault is irrevocable"); - } + let mut vault: Vault = env.storage().instance().get(&DataKey::VaultData(vault_id)).unwrap_or_else(|| panic!("Vault not found")); + + if vault.is_irrevocable { panic!("Vault is irrevocable"); } let unreleased_amount = vault.total_amount - vault.released_amount; - if unreleased_amount <= 0 { - panic!("No tokens available to revoke"); - } + if unreleased_amount <= 0 { panic!("No tokens available to revoke"); } vault.released_amount = vault.total_amount; env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); - let timestamp = env.ledger().timestamp(); - env.events().publish( - (Symbol::new(env, "TokensRevoked"), vault_id), - (unreleased_amount, timestamp), - ); + let mut admin_balance: i128 = env.storage().instance().get(&DataKey::AdminBalance).unwrap_or(0); + admin_balance += unreleased_amount; + env.storage().instance().set(&DataKey::AdminBalance, &admin_balance); - unreleased_amount - } - - // Admin-only: Revoke tokens from a vault and return them to admin - pub fn revoke_tokens(env: Env, vault_id: u64) -> i128 { - Self::require_admin(&env); - - let returned = Self::internal_revoke_full(&env, vault_id); - - // Single admin balance update for this call - let mut admin_balance: i128 = env - .storage() - .instance() - .get(&DataKey::AdminBalance) - .unwrap_or(0); - admin_balance += returned; - env.storage() - .instance() - .set(&DataKey::AdminBalance, &admin_balance); + let mut total_shares: i128 = env.storage().instance().get(&DataKey::TotalShares).unwrap_or(0); + total_shares -= unreleased_amount; + env.storage().instance().set(&DataKey::TotalShares, &total_shares); let timestamp = env.ledger().timestamp(); - env.events().publish( - (Symbol::new(&env, "TokensRevoked"), vault_id), - (unreleased_amount, timestamp), - ); + env.events().publish((Symbol::new(&env, "TokensRevoked"), vault_id), (unreleased_amount, timestamp)); unreleased_amount - returned } - // Revoke a specific amount of tokens from a vault and return them to admin pub fn revoke_partial(env: Env, vault_id: u64, amount: i128) -> i128 { Self::require_admin(&env); - let returned = Self::internal_revoke_partial(&env, vault_id, amount); - - // Single admin balance update for this call - let mut admin_balance: i128 = env - .storage() - .instance() - .get(&DataKey::AdminBalance) - .unwrap_or(0); - admin_balance += returned; - env.storage() - .instance() - .set(&DataKey::AdminBalance, &admin_balance); - - returned - } - - // Internal helper: revoke a specific amount from a vault and emit event. - // Does NOT update admin balance — caller is responsible for a single aggregated transfer. - fn internal_revoke_partial(env: &Env, vault_id: u64, amount: i128) -> i128 { - let mut vault: Vault = env - .storage() - .instance() - .get(&DataKey::VaultData(vault_id)) - .unwrap_or_else(|| panic!("Vault not found")); + let mut vault: Vault = env.storage().instance().get(&DataKey::VaultData(vault_id)).unwrap_or_else(|| panic!("Vault not found")); - if vault.is_irrevocable { - panic!("Vault is irrevocable"); - } + if vault.is_irrevocable { panic!("Vault is irrevocable"); } let unvested_balance = vault.total_amount - vault.released_amount; - if amount <= 0 { - panic!("Amount to revoke must be positive"); - } - if amount > unvested_balance { - panic!("Amount exceeds unvested balance"); - } + if amount <= 0 { panic!("Amount to revoke must be positive"); } + if amount > unvested_balance { panic!("Amount exceeds unvested balance"); } vault.released_amount += amount; env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); - let timestamp = env.ledger().timestamp(); - env.events().publish( - (Symbol::new(env, "TokensRevoked"), vault_id), - (amount, timestamp), - ); + let mut admin_balance: i128 = env.storage().instance().get(&DataKey::AdminBalance).unwrap_or(0); + admin_balance += amount; + env.storage().instance().set(&DataKey::AdminBalance, &admin_balance); - amount - } - - // Admin-only: Revoke many vaults in a single call and credit the admin once. - pub fn batch_revoke(env: Env, vault_ids: Vec) -> i128 { - Self::require_admin(&env); - - let mut total_returned: i128 = 0; - for id in vault_ids.iter() { - let returned = Self::internal_revoke_full(&env, id); - total_returned += returned; - } - - // Single admin balance update for the whole batch - let mut admin_balance: i128 = env - .storage() - .instance() - .get(&DataKey::AdminBalance) - .unwrap_or(0); - admin_balance += total_returned; - env.storage() - .instance() - .set(&DataKey::AdminBalance, &admin_balance); + let mut total_shares: i128 = env.storage().instance().get(&DataKey::TotalShares).unwrap_or(0); + total_shares -= amount; + env.storage().instance().set(&DataKey::TotalShares, &total_shares); let timestamp = env.ledger().timestamp(); - env.events().publish( - (Symbol::new(&env, "TokensRevoked"), vault_id), - (amount, timestamp), - ); + env.events().publish((Symbol::new(&env, "TokensRevoked"), vault_id), (amount, timestamp)); amount - total_returned } - // 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 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 = 3600u64; - if now > vault.creation_time + grace_period { - panic!("Grace period expired"); - } + if now > vault.creation_time + grace_period { panic!("Grace period expired"); } + if vault.released_amount > 0 { panic!("Tokens already claimed"); } - if vault.released_amount > 0 { - panic!("Tokens already claimed"); - } - - let mut admin_balance: i128 = env - .storage() - .instance() - .get(&DataKey::AdminBalance) - .unwrap_or(0); + 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); + env.storage().instance().set(&DataKey::AdminBalance, &admin_balance); vault.released_amount = vault.total_amount; - env.storage() - .instance() - .set(&DataKey::VaultData(vault_id), &vault); + env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); + + let mut total_shares: i128 = env.storage().instance().get(&DataKey::TotalShares).unwrap_or(0); + total_shares -= vault.total_amount; + env.storage().instance().set(&DataKey::TotalShares, &total_shares); - env.events().publish( - (Symbol::new(&env, "VaultClawedBack"), vault_id), - vault.total_amount, - ); + 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")); + 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"); - } + if !vault.is_initialized { panic!("Vault not initialized"); } + if !vault.is_transferable { panic!("Vault is non-transferable"); } vault.owner.require_auth(); let old_owner = vault.owner.clone(); - let old_user_vaults: Vec = env - .storage() - .instance() - .get(&DataKey::UserVaults(old_owner.clone())) - .unwrap_or(Vec::new(&env)); + let 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() { @@ -1222,53 +785,29 @@ impl VestingContract { new_old_user_vaults.push_back(id); } } - env.storage() - .instance() - .set(&DataKey::UserVaults(old_owner.clone()), &new_old_user_vaults); + env.storage().instance().set(&DataKey::UserVaults(old_owner.clone()), &new_old_user_vaults); - let mut new_user_vaults: Vec = env - .storage() - .instance() - .get(&DataKey::UserVaults(new_beneficiary.clone())) - .unwrap_or(Vec::new(&env)); + 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); + env.storage().instance().set(&DataKey::UserVaults(new_beneficiary.clone()), &new_user_vaults); vault.owner = new_beneficiary.clone(); vault.delegate = None; - env.storage() - .instance() - .set(&DataKey::VaultData(vault_id), &vault); + env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); - env.events().publish( - (Symbol::new(&env, "BeneficiaryUpdated"), vault_id), - (old_owner, new_beneficiary), - ); + env.events().publish((Symbol::new(&env, "BeneficiaryUpdated"), vault_id), (old_owner, new_beneficiary)); } - // Rotate beneficiary key (security feature, allows self-transfer even if non-transferable) pub fn rotate_beneficiary_key(env: Env, vault_id: u64, new_address: Address) { - let mut vault: Vault = env - .storage() - .instance() - .get(&DataKey::VaultData(vault_id)) - .unwrap_or_else(|| panic!("Vault not found")); + 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_initialized { panic!("Vault not initialized"); } vault.owner.require_auth(); let old_owner = vault.owner.clone(); - let old_user_vaults: Vec = env - .storage() - .instance() - .get(&DataKey::UserVaults(old_owner.clone())) - .unwrap_or(Vec::new(&env)); + let 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() { @@ -1276,141 +815,77 @@ impl VestingContract { new_old_user_vaults.push_back(id); } } - env.storage() - .instance() - .set(&DataKey::UserVaults(old_owner.clone()), &new_old_user_vaults); + env.storage().instance().set(&DataKey::UserVaults(old_owner.clone()), &new_old_user_vaults); - let mut new_user_vaults: Vec = env - .storage() - .instance() - .get(&DataKey::UserVaults(new_address.clone())) - .unwrap_or(Vec::new(&env)); + let mut new_user_vaults: Vec = env.storage().instance().get(&DataKey::UserVaults(new_address.clone())).unwrap_or(Vec::new(&env)); new_user_vaults.push_back(vault_id); - env.storage() - .instance() - .set(&DataKey::UserVaults(new_address.clone()), &new_user_vaults); + env.storage().instance().set(&DataKey::UserVaults(new_address.clone()), &new_user_vaults); vault.owner = new_address.clone(); vault.delegate = None; - env.storage() - .instance() - .set(&DataKey::VaultData(vault_id), &vault); + env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); - env.events().publish( - (Symbol::new(&env, "BeneficiaryRotated"), vault_id), - (old_owner, new_address), - ); + env.events().publish((Symbol::new(&env, "BeneficiaryRotated"), vault_id), (old_owner, new_address)); } - // 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); + 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")); + 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_initialized { panic!("Vault not initialized"); } vault.owner.require_auth(); 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"); - } + if amount <= 0 { panic!("Amount must be positive"); } + if amount > available { panic!("Insufficient funds to stake"); } - let staking_contract: Address = env - .storage() - .instance() - .get(&Symbol::new(&env, "StakingContract")) - .expect("Staking contract not set"); - - let args = vec![ - &env, - vault_id.into_val(&env), - amount.into_val(&env), - validator.into_val(&env), - ]; + let staking_contract: Address = env.storage().instance().get(&Symbol::new(&env, "StakingContract")).expect("Staking contract not set"); + + 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); vault.staked_amount += amount; - env.storage() - .instance() - .set(&DataKey::VaultData(vault_id), &vault); + + let mut total_staked: i128 = env.storage().instance().get(&DataKey::TotalStaked).unwrap_or(0); + total_staked += amount; + env.storage().instance().set(&DataKey::TotalStaked, &total_staked); + + 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); - let mut vault: Vault = env - .storage() - .instance() - .get(&DataKey::VaultData(vault_id)) - .unwrap_or_else(|| panic!("Vault not found")); + let mut vault: Vault = env.storage().instance().get(&DataKey::VaultData(vault_id)).unwrap_or_else(|| panic!("Vault not found")); - if vault.is_irrevocable { - panic!("Vault is already irrevocable"); - } + if vault.is_irrevocable { panic!("Vault is already irrevocable"); } vault.is_irrevocable = true; - env.storage() - .instance() - .set(&DataKey::VaultData(vault_id), &vault); + env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); let timestamp = env.ledger().timestamp(); - env.events().publish( - (Symbol::new(&env, "IrrevocableMarked"), vault_id), - 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(&DataKey::VaultData(vault_id)) - .unwrap_or_else(|| panic!("Vault not found")); - + let vault: Vault = env.storage().instance().get(&DataKey::VaultData(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 admin_balance: i128 = env - .storage() - .instance() - .get(&DataKey::AdminBalance) - .unwrap_or(0); + let admin_balance: i128 = env.storage().instance().get(&DataKey::AdminBalance).unwrap_or(0); - let vault_count: u64 = env - .storage() - .instance() - .get(&DataKey::VaultCount) - .unwrap_or(0); + let vault_count: u64 = env.storage().instance().get(&DataKey::VaultCount).unwrap_or(0); let mut total_locked = 0i128; let mut total_claimed = 0i128; for i in 1..=vault_count { - if let Some(vault) = env - .storage() - .instance() - .get::(&DataKey::VaultData(i)) - { + if let Some(vault) = env.storage().instance().get::(&DataKey::VaultData(i)) { total_locked += vault.total_amount - vault.released_amount; total_claimed += vault.released_amount; } @@ -1419,38 +894,16 @@ impl VestingContract { (total_locked, total_claimed, admin_balance) } - // Check invariant: Total Locked + Admin Balance + Tokens Paid Out = Initial Supply - // Tokens paid out = total_claimed minus any that were revoked (returned to admin_balance). - // Simplest correct form: total_locked + admin_balance <= initial_supply - // and total_locked + admin_balance + net_paid_out == initial_supply. - // We verify: admin_balance + total_locked == initial_supply - net_distributed - // The safe checkable invariant: total_locked + admin_balance must never exceed initial_supply, - // and (initial_supply - admin_balance - total_locked) must be non-negative (tokens claimed out). pub fn check_invariant(env: Env) -> bool { - let initial_supply: i128 = env - .storage() - .instance() - .get(&DataKey::InitialSupply) - .unwrap_or(0); + let initial_supply: i128 = env.storage().instance().get(&DataKey::InitialSupply).unwrap_or(0); let (total_locked, _total_claimed, admin_balance) = Self::get_contract_state(env); - // All tokens are either: locked in vaults, held by admin, or paid out to beneficiaries. - // locked + admin_balance must never exceed initial_supply (no tokens created from nothing). - // initial_supply - locked - admin_balance = net tokens paid out (must be >= 0). let net_paid_out = initial_supply - total_locked - admin_balance; net_paid_out >= 0 } - // --- Auto-Claim Logic --- - - // 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() - let vault: Vault = env.storage().instance() - .get(&DataKey::VaultData(vault_id)) - .unwrap_or_else(|| panic!("Vault not found")); + let vault: Vault = env.storage().instance().get(&DataKey::VaultData(vault_id)).unwrap_or_else(|| panic!("Vault not found")); let vested = Self::calculate_time_vested_amount(&env, &vault); @@ -1461,67 +914,62 @@ impl VestingContract { } } - // Auto-claim function that anyone can call. - // Tokens go to beneficiary, but keeper earns a fee. pub fn auto_claim(env: Env, vault_id: u64, keeper: Address) { - let mut vault: Vault = env - .storage() - .instance() - let mut vault: Vault = env.storage().instance() - .get(&DataKey::VaultData(vault_id)) - .unwrap_or_else(|| panic!("Vault not found")); + 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_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 - if claimable <= vault.keeper_fee { - panic!("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; let keeper_fee = vault.keeper_fee; + // YIELD DISTRIBUTION - only vault-owned portion + let token_client = Self::get_token_client(&env); + let current_balance = token_client.balance(&env.current_contract_address()); + let admin_balance: i128 = env.storage().instance().get(&DataKey::AdminBalance).unwrap_or(0); + + let total_shares: i128 = env.storage().instance().get(&DataKey::TotalShares).unwrap_or(0); + let total_staked: i128 = env.storage().instance().get(&DataKey::TotalStaked).unwrap_or(0); + let liquid_shares = total_shares - total_staked; + + let vault_portion = (current_balance - admin_balance).max(0); + + let beneficiary_tokens = if liquid_shares > 0 { + (beneficiary_amount * vault_portion) / liquid_shares + } else { + beneficiary_amount + }; + let keeper_tokens = if liquid_shares > 0 { + (keeper_fee * vault_portion) / liquid_shares + } else { + keeper_fee + }; + vault.released_amount += claimable; - env.storage() - .instance() - .set(&DataKey::VaultData(vault_id), &vault); + let mut updated_total_shares = total_shares; + updated_total_shares -= claimable; + env.storage().instance().set(&DataKey::TotalShares, &updated_total_shares); env.storage().instance().set(&DataKey::VaultData(vault_id), &vault); - let mut fees: Map = env - .storage() - .instance() - .get(&DataKey::KeeperFees) - .unwrap_or(Map::new(&env)); + token_client.transfer(&env.current_contract_address(), &vault.owner, &beneficiary_tokens); + token_client.transfer(&env.current_contract_address(), &keeper, &keeper_tokens); + + let mut fees: Map = env.storage().instance().get(&DataKey::KeeperFees).unwrap_or(Map::new(&env)); let current_fees = fees.get(keeper.clone()).unwrap_or(0); fees.set(keeper.clone(), current_fees + keeper_fee); - env.storage() - .instance() - .set(&DataKey::KeeperFees, &fees); + env.storage().instance().set(&DataKey::KeeperFees, &fees); - env.events().publish( - (Symbol::new(&env, "KeeperClaim"), vault_id), - (keeper, beneficiary_amount, keeper_fee), - ); + env.events().publish((Symbol::new(&env, "KeeperClaim"), vault_id), (keeper, beneficiary_amount, keeper_fee)); } - // Get accumulated fees for a keeper pub fn get_keeper_fee(env: Env, keeper: Address) -> i128 { - let fees: Map = env - .storage() - .instance() - .get(&DataKey::KeeperFees) - .unwrap_or(Map::new(&env)); + let fees: Map = env.storage().instance().get(&DataKey::KeeperFees).unwrap_or(Map::new(&env)); fees.get(keeper).unwrap_or(0) } - // Rescue tokens accidentally sent directly to the contract address. - // Calculates unallocated_balance = contract_token_balance - total_vault_liabilities - // and transfers it to the admin. pub fn rescue_unallocated_tokens(env: Env, token_address: Address) -> i128 { Self::require_admin(&env); @@ -1532,19 +980,17 @@ impl VestingContract { let token_client = token::Client::new(&env, &token_address); let contract_balance: i128 = token_client.balance(&env.current_contract_address()); - let vault_count: u64 = env - .storage() - .instance() - .get(&DataKey::VaultCount) - .unwrap_or(0); + if let Some(main_token) = env.storage().instance().get::<_, Address>(&DataKey::Token) { + if main_token == token_address { + panic!("Cannot rescue yield-bearing token. Yield is distributed to beneficiaries on claim."); + } + } + + let vault_count: u64 = env.storage().instance().get(&DataKey::VaultCount).unwrap_or(0); let mut total_liabilities: i128 = 0; for i in 1..=vault_count { - if let Some(vault) = env - .storage() - .instance() - .get::(&DataKey::VaultData(i)) - { + if let Some(vault) = env.storage().instance().get::(&DataKey::VaultData(i)) { let unreleased = vault.total_amount - vault.released_amount; if unreleased > 0 { total_liabilities += unreleased; @@ -1558,18 +1004,11 @@ impl VestingContract { panic!("No unallocated tokens to rescue"); } - let admin: Address = env - .storage() - .instance() - .get(&DataKey::AdminAddress) - .unwrap_or_else(|| panic!("Admin not set")); + let admin: Address = env.storage().instance().get(&DataKey::AdminAddress).unwrap_or_else(|| panic!("Admin not set")); token_client.transfer(&env.current_contract_address(), &admin, &unallocated_balance); - env.events().publish( - (Symbol::new(&env, "RescueExecuted"), token_address), - (unallocated_balance, admin), - ); + env.events().publish((Symbol::new(&env, "RescueExecuted"), token_address), (unallocated_balance, admin)); unallocated_balance } diff --git a/contracts/vesting_contracts/src/test.rs b/contracts/vesting_contracts/src/test.rs index e68528c..5d3f9ad 100644 --- a/contracts/vesting_contracts/src/test.rs +++ b/contracts/vesting_contracts/src/test.rs @@ -1,4 +1,4 @@ -#[cfg(test)] + #[cfg(test)] mod tests { use crate::{ BatchCreateData, Milestone, VestingContract, VestingContractClient, @@ -9,17 +9,28 @@ mod tests { }; // ------------------------------------------------------------------------- - // Helper: spin up a fresh contract + // Helper: fresh contract + yield-bearing token + tokens actually in contract // ------------------------------------------------------------------------- - fn setup() -> (Env, Address, VestingContractClient<'static>, Address) { + fn setup() -> (Env, Address, VestingContractClient<'static>, Address, Address) { let env = Env::default(); env.mock_all_auths(); + let contract_id = env.register(VestingContract, ()); let client = VestingContractClient::new(&env, &contract_id); + let admin = Address::generate(&env); client.initialize(&admin, &1_000_000i128); - (env, contract_id, client, admin) + + let token_addr = register_token(&env, &admin); + client.set_token(&token_addr); + client.add_to_whitelist(&token_addr); + + // Mint initial supply to contract + let stellar = token::StellarAssetClient::new(&env, &token_addr); + stellar.mint(&contract_id, &1_000_000i128); + + (env, contract_id, client, admin, token_addr) } fn register_token(env: &Env, admin: &Address) -> Address { @@ -32,12 +43,12 @@ mod tests { } // ------------------------------------------------------------------------- - // Admin ownership transfer + // Original tests // ------------------------------------------------------------------------- #[test] fn test_admin_ownership_transfer() { - let (env, _cid, client, admin) = setup(); + let (env, _cid, client, admin, _token) = setup(); let new_admin = Address::generate(&env); assert_eq!(client.get_admin(), admin); @@ -51,78 +62,59 @@ mod tests { assert_eq!(client.get_proposed_admin(), None); } - // ------------------------------------------------------------------------- - // Vault creation - // ------------------------------------------------------------------------- - #[test] fn test_create_vault_full_increments_count() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - let id1 = client.create_vault_full( - &beneficiary, &1_000i128, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); - let id2 = client.create_vault_full( - &beneficiary, &500i128, &(now + 10), &(now + 2_000), - &0i128, &true, &false, &0u64, - ); + let id1 = client.create_vault_full(&beneficiary, &1_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); + let id2 = client.create_vault_full(&beneficiary, &500i128, &(now + 10), &(now + 2_000), &0i128, &true, &false, &0u64); assert_eq!(id1, 1u64); assert_eq!(id2, 2u64); } #[test] fn test_create_vault_lazy_increments_count() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - let id = client.create_vault_lazy( - &beneficiary, &1_000i128, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); + let id = client.create_vault_lazy(&beneficiary, &1_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); assert_eq!(id, 1u64); } - // ------------------------------------------------------------------------- - // Batch vault creation - // ------------------------------------------------------------------------- - #[test] fn test_batch_create_vaults_lazy() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let r1 = Address::generate(&env); let r2 = Address::generate(&env); let batch = BatchCreateData { - recipients: vec![&env, r1.clone(), r2.clone()], - amounts: vec![&env, 1_000i128, 2_000i128], - start_times: vec![&env, 100u64, 150u64], - end_times: vec![&env, 200u64, 250u64], - keeper_fees: vec![&env, 0i128, 0i128], + recipients: vec![&env, r1.clone(), r2.clone()], + amounts: vec![&env, 1_000i128, 2_000i128], + start_times: vec![&env, 100u64, 150u64], + end_times: vec![&env, 200u64, 250u64], + keeper_fees: vec![&env, 0i128, 0i128], step_durations: vec![&env, 0u64, 0u64], }; let ids = client.batch_create_vaults_lazy(&batch); assert_eq!(ids.len(), 2); - assert_eq!(ids.get(0).unwrap(), 1u64); - assert_eq!(ids.get(1).unwrap(), 2u64); } #[test] fn test_batch_create_vaults_full() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let r1 = Address::generate(&env); let r2 = Address::generate(&env); let batch = BatchCreateData { - recipients: vec![&env, r1.clone(), r2.clone()], - amounts: vec![&env, 1_000i128, 2_000i128], - start_times: vec![&env, 100u64, 150u64], - end_times: vec![&env, 200u64, 250u64], - keeper_fees: vec![&env, 0i128, 0i128], + recipients: vec![&env, r1.clone(), r2.clone()], + amounts: vec![&env, 1_000i128, 2_000i128], + start_times: vec![&env, 100u64, 150u64], + end_times: vec![&env, 200u64, 250u64], + keeper_fees: vec![&env, 0i128, 0i128], step_durations: vec![&env, 0u64, 0u64], }; @@ -130,23 +122,16 @@ mod tests { assert_eq!(ids.len(), 2); } - // ------------------------------------------------------------------------- - // Step / lockup-only vesting - // ------------------------------------------------------------------------- - #[test] fn test_step_vesting_full_claim_at_end() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let beneficiary = Address::generate(&env); let start = 1_000u64; - let end = start + 101u64; - let step = 17u64; + let end = start + 101u64; + let step = 17u64; let total = 1_009i128; - let vault_id = client.create_vault_full( - &beneficiary, &total, &start, &end, - &0i128, &true, &true, &step, - ); + let vault_id = client.create_vault_full(&beneficiary, &total, &start, &end, &0i128, &true, &true, &step); env.ledger().with_mut(|l| l.timestamp = end + 1); let claimed = client.claim_tokens(&vault_id, &total); @@ -158,16 +143,13 @@ mod tests { #[test] fn test_lockup_only_claim_succeeds_at_end() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let beneficiary = Address::generate(&env); - let now = env.ledger().timestamp(); + let now = env.ledger().timestamp(); let duration = 1_000u64; - let total = 100_000i128; + let total = 100_000i128; - let vault_id = client.create_vault_full( - &beneficiary, &total, &now, &(now + duration), - &0i128, &true, &false, &duration, - ); + let vault_id = client.create_vault_full(&beneficiary, &total, &now, &(now + duration), &0i128, &true, &false, &duration); env.ledger().with_mut(|l| l.timestamp = now + duration); let claimed = client.claim_tokens(&vault_id, &total); @@ -177,15 +159,12 @@ mod tests { #[test] #[should_panic] fn test_lockup_only_claim_fails_before_end() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let beneficiary = Address::generate(&env); - let now = env.ledger().timestamp(); + let now = env.ledger().timestamp(); let duration = 1_000u64; - let vault_id = client.create_vault_full( - &beneficiary, &100_000i128, &now, &(now + duration), - &0i128, &true, &false, &duration, - ); + let vault_id = client.create_vault_full(&beneficiary, &100_000i128, &now, &(now + duration), &0i128, &true, &false, &duration); env.ledger().with_mut(|l| l.timestamp = now + duration - 1); client.claim_tokens(&vault_id, &1i128); @@ -635,20 +614,13 @@ impl MockStakingContract { env.events().publish((Symbol::new(&env, "stake"), vault_id), amount); } - // ------------------------------------------------------------------------- - // Irrevocable vault - // ------------------------------------------------------------------------- - #[test] fn test_mark_irrevocable_flag() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - let vault_id = client.create_vault_full( - &beneficiary, &1_000i128, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); + let vault_id = client.create_vault_full(&beneficiary, &1_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); assert!(!client.is_vault_irrevocable(&vault_id)); client.mark_irrevocable(&vault_id); @@ -658,33 +630,23 @@ impl MockStakingContract { #[test] #[should_panic] fn test_revoke_irrevocable_vault_panics() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - let vault_id = client.create_vault_full( - &beneficiary, &1_000i128, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); + let vault_id = client.create_vault_full(&beneficiary, &1_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); client.mark_irrevocable(&vault_id); client.revoke_tokens(&vault_id); } - // ------------------------------------------------------------------------- - // Clawback - // ------------------------------------------------------------------------- - #[test] fn test_clawback_within_grace_period() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - let vault_id = client.create_vault_full( - &beneficiary, &5_000i128, &(now + 100), &(now + 10_000), - &0i128, &true, &false, &0u64, - ); + let vault_id = client.create_vault_full(&beneficiary, &5_000i128, &(now + 100), &(now + 10_000), &0i128, &true, &false, &0u64); env.ledger().with_mut(|l| l.timestamp = now + 3_599); let returned = client.clawback_vault(&vault_id); @@ -694,39 +656,25 @@ impl MockStakingContract { #[test] #[should_panic] fn test_clawback_after_grace_period_panics() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - let vault_id = client.create_vault_full( - &beneficiary, &5_000i128, &(now + 100), &(now + 10_000), - &0i128, &true, &false, &0u64, - ); + let vault_id = client.create_vault_full(&beneficiary, &5_000i128, &(now + 100), &(now + 10_000), &0i128, &true, &false, &0u64); env.ledger().with_mut(|l| l.timestamp = now + 3_601); client.clawback_vault(&vault_id); } - // ------------------------------------------------------------------------- - // Milestones - // ------------------------------------------------------------------------- - #[test] fn test_milestone_unlock_and_claim() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - let vault_id = client.create_vault_full( - &beneficiary, &1_000i128, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); + let vault_id = client.create_vault_full(&beneficiary, &1_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); - let milestones = vec![ - &env, - Milestone { id: 1, percentage: 50, is_unlocked: false }, - Milestone { id: 2, percentage: 50, is_unlocked: false }, - ]; + let milestones = vec![&env, Milestone { id: 1, percentage: 50, is_unlocked: false }, Milestone { id: 2, percentage: 50, is_unlocked: false }]; client.set_milestones(&vault_id, &milestones); client.unlock_milestone(&vault_id, &1u64); @@ -741,92 +689,56 @@ impl MockStakingContract { #[test] #[should_panic] fn test_claim_before_any_milestone_unlocked_panics() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - let vault_id = client.create_vault_full( - &beneficiary, &1_000i128, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); + let vault_id = client.create_vault_full(&beneficiary, &1_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); - let milestones = vec![ - &env, - Milestone { id: 1, percentage: 100, is_unlocked: false }, - ]; + let milestones = vec![&env, Milestone { id: 1, percentage: 100, is_unlocked: false }]; client.set_milestones(&vault_id, &milestones); client.claim_tokens(&vault_id, &1i128); } - // ------------------------------------------------------------------------- - // Rotate beneficiary key - // ------------------------------------------------------------------------- - #[test] fn test_rotate_beneficiary_key() { - let (env, _cid, client, _admin) = setup(); - let beneficiary = Address::generate(&env); + let (env, _cid, client, _admin, _token) = setup(); + let beneficiary = Address::generate(&env); let new_beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - let vault_id = client.create_vault_full( - &beneficiary, &1_000i128, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); + let vault_id = client.create_vault_full(&beneficiary, &1_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); client.rotate_beneficiary_key(&vault_id, &new_beneficiary); let vault = client.get_vault(&vault_id); assert_eq!(vault.owner, new_beneficiary); - assert_eq!(vault.delegate, None); } - // ------------------------------------------------------------------------- - // Invariant - // ------------------------------------------------------------------------- - #[test] fn test_invariant_holds_after_operations() { - let (env, _cid, client, _admin) = setup(); + let (env, _cid, client, _admin, _token) = setup(); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - let initial_supply = 1_000_000i128; - let vault_amount = 10_000i128; - // After vault creation: admin_balance = 990_000, vault locked = 10_000 - let vault_id = client.create_vault_full( - &beneficiary, &vault_amount, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); - assert!(client.check_invariant(), "invariant failed after vault creation"); - let (locked, _claimed, admin_bal) = client.get_contract_state(); - assert_eq!(locked + admin_bal, initial_supply, "locked + admin should equal initial supply before any claims"); + let vault_id = client.create_vault_full(&beneficiary, &10_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); + assert!(client.check_invariant()); - // After partial claim: 5_000 paid out, 5_000 still locked env.ledger().with_mut(|l| l.timestamp = now + 500); client.claim_tokens(&vault_id, &5_000i128); - assert!(client.check_invariant(), "invariant failed after partial claim"); - let (locked2, _claimed2, admin_bal2) = client.get_contract_state(); - // 5_000 was paid out to beneficiary, so locked + admin = initial - 5_000 - assert_eq!(locked2 + admin_bal2, initial_supply - 5_000i128, "5000 should have left the pool"); + assert!(client.check_invariant()); - // After revoke: remaining 5_000 returned to admin client.revoke_tokens(&vault_id); - assert!(client.check_invariant(), "invariant failed after revoke"); - let (locked3, _claimed3, admin_bal3) = client.get_contract_state(); - // Still only 5_000 paid out total (the revoked portion came back to admin) - assert_eq!(locked3, 0i128, "no tokens should remain locked after full revoke"); - assert_eq!(admin_bal3, initial_supply - 5_000i128, "admin should hold everything except what was claimed"); - assert_eq!(locked3 + admin_bal3, initial_supply - 5_000i128, "invariant: only claimed tokens are gone"); + assert!(client.check_invariant()); } // ========================================================================= - // rescue_unallocated_tokens + // rescue tests // ========================================================================= #[test] fn test_rescue_basic_no_vaults() { - let (env, contract_id, client, admin) = setup(); + let (env, contract_id, client, admin, _main_token) = setup(); let token_addr = register_token(&env, &admin); client.add_to_whitelist(&token_addr); @@ -834,52 +746,36 @@ impl MockStakingContract { let rescued = client.rescue_unallocated_tokens(&token_addr); assert_eq!(rescued, 5_000i128); - - let tok = token::Client::new(&env, &token_addr); - assert_eq!(tok.balance(&admin), 5_000i128); - assert_eq!(tok.balance(&contract_id), 0i128); } #[test] fn test_rescue_only_surplus_above_vault_liability() { - let (env, contract_id, client, admin) = setup(); + let (env, contract_id, client, admin, _main_token) = setup(); let token_addr = register_token(&env, &admin); client.add_to_whitelist(&token_addr); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - client.create_vault_full( - &beneficiary, &3_000i128, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); + client.create_vault_full(&beneficiary, &3_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); - // 3000 liability + 2000 stray mint_to(&env, &token_addr, &contract_id, 5_000i128); let rescued = client.rescue_unallocated_tokens(&token_addr); assert_eq!(rescued, 2_000i128); - - let tok = token::Client::new(&env, &token_addr); - assert_eq!(tok.balance(&admin), 2_000i128); - assert_eq!(tok.balance(&contract_id), 3_000i128); } #[test] fn test_rescue_after_partial_claim_adjusts_liability() { - let (env, contract_id, client, admin) = setup(); + let (env, contract_id, client, admin, _main_token) = setup(); let token_addr = register_token(&env, &admin); client.add_to_whitelist(&token_addr); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - let vault_id = client.create_vault_full( - &beneficiary, &4_000i128, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); + let vault_id = client.create_vault_full(&beneficiary, &4_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); - // Claim 1000 → remaining liability 3000 env.ledger().with_mut(|l| l.timestamp = now + 1_001); client.claim_tokens(&vault_id, &1_000i128); @@ -887,84 +783,58 @@ impl MockStakingContract { let rescued = client.rescue_unallocated_tokens(&token_addr); assert_eq!(rescued, 2_000i128); - - let tok = token::Client::new(&env, &token_addr); - assert_eq!(tok.balance(&admin), 2_000i128); - assert_eq!(tok.balance(&contract_id), 3_000i128); } #[test] fn test_rescue_multiple_vaults_correct_liability_sum() { - let (env, contract_id, client, admin) = setup(); + let (env, contract_id, client, admin, _main_token) = setup(); let token_addr = register_token(&env, &admin); client.add_to_whitelist(&token_addr); let now = env.ledger().timestamp(); - // Three vaults × 2000 = 6000 total liability for _ in 0..3 { let b = Address::generate(&env); - client.create_vault_full( - &b, &2_000i128, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); + client.create_vault_full(&b, &2_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); } - // 6000 liability + 1000 stray mint_to(&env, &token_addr, &contract_id, 7_000i128); let rescued = client.rescue_unallocated_tokens(&token_addr); assert_eq!(rescued, 1_000i128); - - let tok = token::Client::new(&env, &token_addr); - assert_eq!(tok.balance(&admin), 1_000i128); - assert_eq!(tok.balance(&contract_id), 6_000i128); } #[test] fn test_rescue_after_full_claim_all_tokens_rescuable() { - let (env, contract_id, client, admin) = setup(); + let (env, contract_id, client, admin, _main_token) = setup(); let token_addr = register_token(&env, &admin); client.add_to_whitelist(&token_addr); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - let vault_id = client.create_vault_full( - &beneficiary, &2_000i128, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); + let vault_id = client.create_vault_full(&beneficiary, &2_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); - // Claim everything → liability = 0 env.ledger().with_mut(|l| l.timestamp = now + 1_001); client.claim_tokens(&vault_id, &2_000i128); - // Stray deposit after full claim mint_to(&env, &token_addr, &contract_id, 500i128); let rescued = client.rescue_unallocated_tokens(&token_addr); assert_eq!(rescued, 500i128); - - let tok = token::Client::new(&env, &token_addr); - assert_eq!(tok.balance(&admin), 500i128); - assert_eq!(tok.balance(&contract_id), 0i128); } #[test] fn test_rescue_after_revoke_liability_drops_to_zero() { - let (env, contract_id, client, admin) = setup(); + let (env, contract_id, client, admin, _main_token) = setup(); let token_addr = register_token(&env, &admin); client.add_to_whitelist(&token_addr); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - let vault_id = client.create_vault_full( - &beneficiary, &3_000i128, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); + let vault_id = client.create_vault_full(&beneficiary, &3_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); - // Revoke → vault liability drops to 0 client.revoke_tokens(&vault_id); mint_to(&env, &token_addr, &contract_id, 3_000i128); @@ -975,66 +845,133 @@ impl MockStakingContract { #[test] fn test_rescue_tokens_go_to_current_admin_after_transfer() { - let (env, contract_id, client, admin) = setup(); + let (env, contract_id, client, admin, _main_token) = setup(); let token_addr = register_token(&env, &admin); client.add_to_whitelist(&token_addr); - // Transfer admin to new_admin let new_admin = Address::generate(&env); client.propose_new_admin(&new_admin); client.accept_ownership(); - assert_eq!(client.get_admin(), new_admin); mint_to(&env, &token_addr, &contract_id, 1_000i128); let rescued = client.rescue_unallocated_tokens(&token_addr); assert_eq!(rescued, 1_000i128); - - let tok = token::Client::new(&env, &token_addr); - assert_eq!(tok.balance(&new_admin), 1_000i128); // new admin gets tokens - assert_eq!(tok.balance(&admin), 0i128); // old admin gets nothing } #[test] #[should_panic] fn test_rescue_panics_when_no_surplus() { - let (env, contract_id, client, admin) = setup(); + let (env, contract_id, client, admin, _main_token) = setup(); let token_addr = register_token(&env, &admin); client.add_to_whitelist(&token_addr); let beneficiary = Address::generate(&env); let now = env.ledger().timestamp(); - client.create_vault_full( - &beneficiary, &3_000i128, &now, &(now + 1_000), - &0i128, &true, &false, &0u64, - ); + client.create_vault_full(&beneficiary, &3_000i128, &now, &(now + 1_000), &0i128, &true, &false, &0u64); - // Mint exactly the liability — zero surplus mint_to(&env, &token_addr, &contract_id, 3_000i128); - client.rescue_unallocated_tokens(&token_addr); // must panic + client.rescue_unallocated_tokens(&token_addr); } #[test] #[should_panic] fn test_rescue_panics_when_contract_balance_zero() { - let (env, _cid, client, admin) = setup(); - let token_addr = register_token(&env, &admin); + let (env, _cid, client, _admin, _token) = setup(); + let token_addr = register_token(&env, &_admin); client.add_to_whitelist(&token_addr); - client.rescue_unallocated_tokens(&token_addr); // must panic + client.rescue_unallocated_tokens(&token_addr); } #[test] #[should_panic] fn test_rescue_panics_for_non_whitelisted_token() { - let (env, contract_id, client, admin) = setup(); - // Register but do NOT whitelist + let (env, contract_id, client, admin, _main_token) = setup(); let token_addr = register_token(&env, &admin); mint_to(&env, &token_addr, &contract_id, 1_000i128); - client.rescue_unallocated_tokens(&token_addr); // must panic + client.rescue_unallocated_tokens(&token_addr); + } + + // ========================================================================= + // Yield demonstration tests + // ========================================================================= + + #[test] + fn test_yield_is_distributed_on_claim() { + let (env, contract_id, client, _admin, token) = setup(); + let beneficiary = Address::generate(&env); + let now = env.ledger().timestamp(); + + let vault_id = client.create_vault_full( + &beneficiary, &10_000i128, &now, &(now + 1_000), + &0i128, &true, &false, &0u64, + ); + + // Simulate yield + let stellar = token::StellarAssetClient::new(&env, &token); + stellar.mint(&contract_id, &2_000i128); + + env.ledger().with_mut(|l| l.timestamp = now + 1_001); + let claimed = client.claim_tokens(&vault_id, &10_000i128); + + assert_eq!(claimed, 12_000i128); // principal + all yield + } + + #[test] + fn test_yield_on_partial_claim() { + let (env, contract_id, client, _admin, token) = setup(); + let beneficiary = Address::generate(&env); + let now = env.ledger().timestamp(); + + let vault_id = client.create_vault_full( + &beneficiary, &10_000i128, &now, &(now + 1_000), + &0i128, &true, &false, &0u64, + ); + + let stellar = token::StellarAssetClient::new(&env, &token); + stellar.mint(&contract_id, &2_000i128); + + env.ledger().with_mut(|l| l.timestamp = now + 500); + let claimed = client.claim_tokens(&vault_id, &5_000i128); + + assert_eq!(claimed, 6_000i128); // 5k principal + 1k yield + } + + #[test] + fn test_yield_proportional_with_multiple_vaults() { + let (env, contract_id, client, _admin, token) = setup(); + let now = env.ledger().timestamp(); + + let v1 = client.create_vault_full( + &Address::generate(&env), &10_000i128, &now, &(now + 1_000), + &0i128, &true, &false, &0u64, + ); + let v2 = client.create_vault_full( + &Address::generate(&env), &20_000i128, &now, &(now + 1_000), + &0i128, &true, &false, &0u64, + ); + + let stellar = token::StellarAssetClient::new(&env, &token); + stellar.mint(&contract_id, &6_000i128); + + env.ledger().with_mut(|l| l.timestamp = now + 1_001); + + let claimed1 = client.claim_tokens(&v1, &10_000i128); + let claimed2 = client.claim_tokens(&v2, &20_000i128); + + assert_eq!(claimed1, 12_000i128); + assert_eq!(claimed2, 24_000i128); + } + + #[test] + #[should_panic(expected = "Cannot rescue yield-bearing token")] + fn test_rescue_yield_token_panics() { + let (env, _cid, client, _admin, token) = setup(); + client.rescue_unallocated_tokens(&token); } } });