From 296b1772aa7650c47ad827f475850335f3c26e93 Mon Sep 17 00:00:00 2001 From: Christian Krueger Date: Thu, 23 Jan 2025 21:35:30 -0700 Subject: [PATCH] fix --- .../tests/vault/burn_withdrawal_ticket.rs | 316 ++++++++++++++++++ vault_core/src/vault.rs | 118 +++++-- vault_program/src/burn_withdrawal_ticket.rs | 10 +- 3 files changed, 424 insertions(+), 20 deletions(-) diff --git a/integration_tests/tests/vault/burn_withdrawal_ticket.rs b/integration_tests/tests/vault/burn_withdrawal_ticket.rs index fddab05e..29693cd9 100644 --- a/integration_tests/tests/vault/burn_withdrawal_ticket.rs +++ b/integration_tests/tests/vault/burn_withdrawal_ticket.rs @@ -478,6 +478,167 @@ mod tests { assert_eq!(vault.vrt_cooling_down_amount(), 0); } + /// Tests that the program fee is non deducted when the staker is the program fee wallet + #[tokio::test] + async fn test_burn_withdrawal_ticket_with_staker_as_program_fee_wallet() { + const MINT_AMOUNT: u64 = 100_000; + const PROGRAM_FEE_BPS: u16 = 10; // 0.1% + + let deposit_fee_bps = 0; + let withdraw_fee_bps = 0; + let reward_fee_bps = 0; + let num_operators = 1; + let slasher_amounts = vec![]; + + let mut fixture = TestBuilder::new().await; + let ConfiguredVault { + mut vault_program_client, + restaking_program_client: _, + vault_config_admin, + vault_root, + restaking_config_admin: _, + operator_roots, + } = fixture + .setup_vault_with_ncn_and_operators( + deposit_fee_bps, + withdraw_fee_bps, + reward_fee_bps, + num_operators, + &slasher_amounts, + ) + .await + .unwrap(); + + // Set program fee + vault_program_client + .set_program_fee(&vault_config_admin, PROGRAM_FEE_BPS) + .await + .unwrap(); + + // Depositor is the program fee wallet + let depositor = vault_config_admin; + vault_program_client + .configure_depositor(&vault_root, &depositor.pubkey(), MINT_AMOUNT) + .await + .unwrap(); + vault_program_client + .do_mint_to(&vault_root, &depositor, MINT_AMOUNT, MINT_AMOUNT) + .await + .unwrap(); + + let config = vault_program_client + .get_config(&Config::find_program_address(&jito_vault_program::id()).0) + .await + .unwrap(); + + // Delegate all funds to the operator + vault_program_client + .do_add_delegation(&vault_root, &operator_roots[0].operator_pubkey, MINT_AMOUNT) + .await + .unwrap(); + + let VaultStakerWithdrawalTicketRoot { base } = vault_program_client + .do_enqueue_withdrawal(&vault_root, &depositor, MINT_AMOUNT) + .await + .unwrap(); + + // If this breaks: MINT_AMOUNT * 9990 / 10000 - 1, + vault_program_client + .do_cooldown_delegation(&vault_root, &operator_roots[0].operator_pubkey, MINT_AMOUNT) + .await + .unwrap(); + + // Warp forward two epochs + fixture + .warp_slot_incremental(config.epoch_length()) + .await + .unwrap(); + vault_program_client + .do_full_vault_update( + &vault_root.vault_pubkey, + &[operator_roots[0].operator_pubkey], + ) + .await + .unwrap(); + fixture + .warp_slot_incremental(config.epoch_length()) + .await + .unwrap(); + vault_program_client + .do_full_vault_update( + &vault_root.vault_pubkey, + &[operator_roots[0].operator_pubkey], + ) + .await + .unwrap(); + + let vault = vault_program_client + .get_vault(&vault_root.vault_pubkey) + .await + .unwrap(); + + let initial_depositor_balance = fixture + .get_token_account(&get_associated_token_address( + &depositor.pubkey(), + &vault.supported_mint, + )) + .await + .unwrap() + .amount; + + vault_program_client + .do_burn_withdrawal_ticket(&vault_root, &depositor, &base, &config.program_fee_wallet) + .await + .unwrap(); + + // Calculate expected fee + let expected_fee = 0; + let expected_withdrawal = MINT_AMOUNT - expected_fee; + + // Check final balances + let final_depositor_balance = fixture + .get_token_account(&get_associated_token_address( + &depositor.pubkey(), + &vault.supported_mint, + )) + .await + .unwrap() + .amount; + let program_fee_balance = fixture + .get_token_account(&get_associated_token_address( + &config.program_fee_wallet, + &vault.vrt_mint, + )) + .await + .unwrap() + .amount; + + // Assert correct amounts were transferred + assert_eq!( + final_depositor_balance - initial_depositor_balance, + expected_withdrawal + ); + assert_eq!(program_fee_balance, expected_fee); + + // Check that the vault state is correct + let vault = vault_program_client + .get_vault(&vault_root.vault_pubkey) + .await + .unwrap(); + assert_eq!( + vault.tokens_deposited() - Vault::DEFAULT_INITIALIZATION_TOKEN_AMOUNT, + expected_fee + ); + assert_eq!( + vault.vrt_supply() - Vault::DEFAULT_INITIALIZATION_TOKEN_AMOUNT, + expected_fee + ); + assert_eq!(vault.delegation_state, DelegationState::default()); + assert_eq!(vault.vrt_enqueued_for_cooldown_amount(), 0); + assert_eq!(vault.vrt_ready_to_claim_amount(), 0); + assert_eq!(vault.vrt_cooling_down_amount(), 0); + } + /// Tests that the program fee is correctly deducted and transferred during burn_withdrawal_ticket #[tokio::test] async fn test_burn_withdrawal_ticket_withdrawal_fee() { @@ -633,6 +794,161 @@ mod tests { assert_eq!(vault.vrt_cooling_down_amount(), 0); } + /// Tests that the program fee is non deducted when the staker is the vault fee wallet + #[tokio::test] + async fn test_burn_withdrawal_ticket_with_staker_as_vault_fee_wallet() { + const MINT_AMOUNT: u64 = 100_000; + const WITHDRAWAL_FEE_BPS: u16 = 1000; // 10% + + let deposit_fee_bps = 0; + let withdraw_fee_bps = WITHDRAWAL_FEE_BPS; + let reward_fee_bps = 0; + let num_operators = 1; + let slasher_amounts = vec![]; + + let mut fixture = TestBuilder::new().await; + let ConfiguredVault { + mut vault_program_client, + restaking_program_client: _, + vault_config_admin: _, + vault_root, + restaking_config_admin: _, + operator_roots, + } = fixture + .setup_vault_with_ncn_and_operators( + deposit_fee_bps, + withdraw_fee_bps, + reward_fee_bps, + num_operators, + &slasher_amounts, + ) + .await + .unwrap(); + + // Initial deposit + mint + let depositor = vault_root.vault_admin.insecure_clone(); + vault_program_client + .configure_depositor(&vault_root, &depositor.pubkey(), MINT_AMOUNT) + .await + .unwrap(); + vault_program_client + .do_mint_to(&vault_root, &depositor, MINT_AMOUNT, MINT_AMOUNT) + .await + .unwrap(); + + let config = vault_program_client + .get_config(&Config::find_program_address(&jito_vault_program::id()).0) + .await + .unwrap(); + + // Delegate all funds to the operator + vault_program_client + .do_add_delegation(&vault_root, &operator_roots[0].operator_pubkey, MINT_AMOUNT) + .await + .unwrap(); + + let VaultStakerWithdrawalTicketRoot { base } = vault_program_client + .do_enqueue_withdrawal(&vault_root, &depositor, MINT_AMOUNT) + .await + .unwrap(); + + // If this breaks: MINT_AMOUNT * 9990 / 10000 - 1, + vault_program_client + .do_cooldown_delegation(&vault_root, &operator_roots[0].operator_pubkey, MINT_AMOUNT) + .await + .unwrap(); + + // Warp forward two epochs + fixture + .warp_slot_incremental(config.epoch_length()) + .await + .unwrap(); + vault_program_client + .do_full_vault_update( + &vault_root.vault_pubkey, + &[operator_roots[0].operator_pubkey], + ) + .await + .unwrap(); + fixture + .warp_slot_incremental(config.epoch_length()) + .await + .unwrap(); + vault_program_client + .do_full_vault_update( + &vault_root.vault_pubkey, + &[operator_roots[0].operator_pubkey], + ) + .await + .unwrap(); + + let vault = vault_program_client + .get_vault(&vault_root.vault_pubkey) + .await + .unwrap(); + + let initial_depositor_balance = fixture + .get_token_account(&get_associated_token_address( + &depositor.pubkey(), + &vault.supported_mint, + )) + .await + .unwrap() + .amount; + + vault_program_client + .do_burn_withdrawal_ticket(&vault_root, &depositor, &base, &config.program_fee_wallet) + .await + .unwrap(); + + // Calculate expected fee + let expected_fee = 0; + let expected_withdrawal = MINT_AMOUNT - expected_fee; + + // Check final balances + let final_depositor_balance = fixture + .get_token_account(&get_associated_token_address( + &depositor.pubkey(), + &vault.supported_mint, + )) + .await + .unwrap() + .amount; + let program_fee_balance = fixture + .get_token_account(&get_associated_token_address( + &vault.fee_wallet, + &vault.vrt_mint, + )) + .await + .unwrap() + .amount; + + // Assert correct amounts were transferred + assert_eq!( + final_depositor_balance - initial_depositor_balance, + expected_withdrawal + ); + assert_eq!(program_fee_balance, expected_fee); + + // Check that the vault state is correct + let vault = vault_program_client + .get_vault(&vault_root.vault_pubkey) + .await + .unwrap(); + assert_eq!( + vault.tokens_deposited() - Vault::DEFAULT_INITIALIZATION_TOKEN_AMOUNT, + expected_fee + ); + assert_eq!( + vault.vrt_supply() - Vault::DEFAULT_INITIALIZATION_TOKEN_AMOUNT, + expected_fee + ); + assert_eq!(vault.delegation_state, DelegationState::default()); + assert_eq!(vault.vrt_enqueued_for_cooldown_amount(), 0); + assert_eq!(vault.vrt_ready_to_claim_amount(), 0); + assert_eq!(vault.vrt_cooling_down_amount(), 0); + } + #[tokio::test] async fn test_burn_withdrawal_ticket_with_extra_vrt_sent_to_ticket() { const MINT_AMOUNT: u64 = 100_000; diff --git a/vault_core/src/vault.rs b/vault_core/src/vault.rs index df28dcee..fc015390 100644 --- a/vault_core/src/vault.rs +++ b/vault_core/src/vault.rs @@ -1007,8 +1007,14 @@ impl Vault { }) } - pub fn calculate_burn_summary(&self, amount_in: u64) -> Result { - let program_fee_amount = Config::calculate_program_fee(self.program_fee_bps(), amount_in)?; + pub fn calculate_burn_summary( + &self, + is_staker_program_fee_wallet: bool, + is_staker_vault_fee_wallet: bool, + amount_in: u64, + ) -> Result { + let mut program_fee_amount = + Config::calculate_program_fee(self.program_fee_bps(), amount_in)?; let mut vault_fee_amount = self.calculate_withdrawal_fee(amount_in)?; // Prioritize program fee over vault fee if together they exceed the amount in @@ -1022,6 +1028,14 @@ impl Vault { .ok_or(VaultError::VaultUnderflow)?; } + if is_staker_program_fee_wallet { + program_fee_amount = 0; + } + + if is_staker_vault_fee_wallet { + vault_fee_amount = 0; + } + let amount_to_burn = amount_in .checked_sub(program_fee_amount) .and_then(|x| x.checked_sub(vault_fee_amount)) @@ -1041,7 +1055,12 @@ impl Vault { }) } - pub fn burn_with_fee(&mut self, amount_in: u64) -> Result { + pub fn burn_with_fee( + &mut self, + is_staker_program_fee_wallet: bool, + is_staker_vault_fee_wallet: bool, + amount_in: u64, + ) -> Result { if amount_in == 0 { msg!("Amount in is zero"); return Err(VaultError::VaultBurnZero); @@ -1054,7 +1073,11 @@ impl Vault { vault_fee_amount, burn_amount, out_amount, - } = self.calculate_burn_summary(amount_in)?; + } = self.calculate_burn_summary( + is_staker_program_fee_wallet, + is_staker_vault_fee_wallet, + amount_in, + )?; let max_withdrawable = self .tokens_deposited() @@ -1102,7 +1125,7 @@ impl Vault { let BurnSummary { out_amount: amount_to_reserve_for_vrts, .. - } = self.calculate_burn_summary(vrt_reserve)?; + } = self.calculate_burn_summary(false, false, vrt_reserve)?; Ok(amount_to_reserve_for_vrts) } @@ -1594,12 +1617,60 @@ mod tests { program_fee_amount: _, burn_amount, out_amount, - } = vault.burn_with_fee(100).unwrap(); + } = vault.burn_with_fee(false, false, 100).unwrap(); assert_eq!(fee_amount, 1); assert_eq!(burn_amount, 99); assert_eq!(out_amount, 99); } + #[test] + fn test_burn_with_staker_as_program_fee_wallet() { + let mut vault = make_test_vault(0, 100, 100, 100, 100, DelegationState::default()); + + let BurnSummary { + vault_fee_amount: fee_amount, + program_fee_amount, + burn_amount, + out_amount, + } = vault.burn_with_fee(true, false, 100).unwrap(); + assert_eq!(fee_amount, 1); + assert_eq!(program_fee_amount, 0); + assert_eq!(burn_amount, 99); + assert_eq!(out_amount, 99); + } + + #[test] + fn test_burn_with_staker_as_vault_fee_wallet() { + let mut vault = make_test_vault(0, 100, 100, 100, 100, DelegationState::default()); + + let BurnSummary { + vault_fee_amount, + program_fee_amount, + burn_amount, + out_amount, + } = vault.burn_with_fee(false, true, 100).unwrap(); + assert_eq!(vault_fee_amount, 0); + assert_eq!(program_fee_amount, 1); + assert_eq!(burn_amount, 99); + assert_eq!(out_amount, 99); + } + + #[test] + fn test_burn_with_staker_as_both_program_and_vault_fee_wallet() { + let mut vault = make_test_vault(0, 100, 100, 100, 100, DelegationState::default()); + + let BurnSummary { + vault_fee_amount, + program_fee_amount, + burn_amount, + out_amount, + } = vault.burn_with_fee(true, true, 100).unwrap(); + assert_eq!(vault_fee_amount, 0); + assert_eq!(program_fee_amount, 0); + assert_eq!(burn_amount, 100); + assert_eq!(out_amount, 100); + } + #[test] fn test_burn_with_program_fee_ok() { let mut vault = make_test_vault(0, 100, 200, 100, 100, DelegationState::default()); @@ -1609,7 +1680,7 @@ mod tests { program_fee_amount, burn_amount, out_amount, - } = vault.burn_with_fee(100).unwrap(); + } = vault.burn_with_fee(false, false, 100).unwrap(); assert_eq!(vault_fee_amount, 1); assert_eq!(program_fee_amount, 2); assert_eq!(burn_amount, 97); @@ -1625,7 +1696,7 @@ mod tests { program_fee_amount, burn_amount, out_amount, - } = vault.burn_with_fee(100).unwrap(); + } = vault.burn_with_fee(false, false, 100).unwrap(); assert_eq!(program_fee_amount, 90); assert_eq!(vault_fee_amount, 10); assert_eq!(burn_amount, 0); @@ -1641,7 +1712,7 @@ mod tests { program_fee_amount, burn_amount, out_amount, - } = vault.burn_with_fee(100).unwrap(); + } = vault.burn_with_fee(false, false, 100).unwrap(); assert_eq!(vault_fee_amount, 0); assert_eq!(program_fee_amount, 100); assert_eq!(burn_amount, 0); @@ -1653,7 +1724,7 @@ mod tests { let mut vault = make_test_vault(0, 100, 0, 100, 100, DelegationState::default()); assert_eq!( - vault.burn_with_fee(101), + vault.burn_with_fee(false, false, 101), Err(VaultError::VaultInsufficientFunds) ); } @@ -1661,7 +1732,10 @@ mod tests { #[test] fn test_burn_zero_fails() { let mut vault = make_test_vault(0, 100, 0, 100, 100, DelegationState::default()); - assert_eq!(vault.burn_with_fee(0), Err(VaultError::VaultBurnZero)); + assert_eq!( + vault.burn_with_fee(false, false, 0), + Err(VaultError::VaultBurnZero) + ); } #[test] @@ -1673,7 +1747,7 @@ mod tests { program_fee_amount: _, burn_amount, out_amount, - } = vault.burn_with_fee(50).unwrap(); + } = vault.burn_with_fee(false, false, 50).unwrap(); assert_eq!(fee_amount, 0); assert_eq!(burn_amount, 50); assert_eq!(out_amount, 50); @@ -1685,14 +1759,17 @@ mod tests { fn test_burn_more_than_withdrawable_fails() { let mut vault = make_test_vault(0, 0, 0, 100, 100, DelegationState::new(50, 0, 0)); - assert_eq!(vault.burn_with_fee(51), Err(VaultError::VaultUnderflow)); + assert_eq!( + vault.burn_with_fee(false, false, 51), + Err(VaultError::VaultUnderflow) + ); } #[test] fn test_burn_all_delegated() { let mut vault = make_test_vault(0, 0, 0, 100, 100, DelegationState::new(100, 0, 0)); - let result = vault.burn_with_fee(1); + let result = vault.burn_with_fee(false, false, 1); assert_eq!(result, Err(VaultError::VaultUnderflow)); } @@ -1700,7 +1777,7 @@ mod tests { fn test_burn_rounding_issues() { let mut vault = make_test_vault(0, 0, 0, 1_000_000, 1_000_000, DelegationState::default()); - let result = vault.burn_with_fee(1).unwrap(); + let result = vault.burn_with_fee(false, false, 1).unwrap(); assert_eq!(result.out_amount, 1); assert_eq!(vault.tokens_deposited(), 999_999); assert_eq!(vault.vrt_supply(), 999_999); @@ -1709,7 +1786,7 @@ mod tests { #[test] fn test_burn_max_values() { let mut vault = make_test_vault(0, 100, 0, u64::MAX, u64::MAX, DelegationState::default()); - let result = vault.burn_with_fee(u64::MAX).unwrap(); + let result = vault.burn_with_fee(false, false, u64::MAX).unwrap(); let fee_amount = (((u64::MAX as u128) * 100).div_ceil(10000)) as u64; assert_eq!(result.vault_fee_amount, fee_amount); } @@ -1718,7 +1795,7 @@ mod tests { fn test_burn_different_fees() { let mut vault = make_test_vault(0, 500, 0, 10000, 10000, DelegationState::default()); - let result = vault.burn_with_fee(1000).unwrap(); + let result = vault.burn_with_fee(false, false, 1000).unwrap(); assert_eq!(result.vault_fee_amount, 50); assert_eq!(result.burn_amount, 950); assert_eq!(result.out_amount, 950); @@ -1804,7 +1881,7 @@ mod tests { program_fee_amount: _, burn_amount, out_amount, - } = vault.burn_with_fee(1).unwrap(); + } = vault.burn_with_fee(false, false, 1).unwrap(); assert_eq!(fee_amount, 1); assert_eq!(burn_amount, 0); assert_eq!(out_amount, 0); @@ -2339,7 +2416,10 @@ mod tests { #[test] fn test_burn_with_fee_zero_amount() { let mut vault = make_test_vault(0, 0, 0, 1000, 1000, DelegationState::default()); - assert_eq!(vault.burn_with_fee(0), Err(VaultError::VaultBurnZero)); + assert_eq!( + vault.burn_with_fee(false, false, 0), + Err(VaultError::VaultBurnZero) + ); } // ---------- REWARD FEE HELPERS ------------ diff --git a/vault_program/src/burn_withdrawal_ticket.rs b/vault_program/src/burn_withdrawal_ticket.rs index 6f21f193..e1bc748d 100644 --- a/vault_program/src/burn_withdrawal_ticket.rs +++ b/vault_program/src/burn_withdrawal_ticket.rs @@ -83,12 +83,20 @@ pub fn process_burn_withdrawal_ticket( return Err(VaultError::VaultStakerWithdrawalTicketNotWithdrawable.into()); } + let is_staker_program_fee_wallet = config.program_fee_wallet.eq(staker.key); + let is_staker_vault_fee_wallet = vault.fee_wallet.eq(staker.key); + let amount_in = vault_staker_withdrawal_ticket.vrt_amount(); + let BurnSummary { vault_fee_amount, program_fee_amount, burn_amount, out_amount, - } = vault.burn_with_fee(vault_staker_withdrawal_ticket.vrt_amount())?; + } = vault.burn_with_fee( + is_staker_program_fee_wallet, + is_staker_vault_fee_wallet, + amount_in, + )?; // To close the token account, the balance needs to be 0. // The only way for vault_staker_withdrawal_ticket.vrt_amount() != ticket_vrt_amount