diff --git a/packages/token-staking/src/idl.ts b/packages/token-staking/src/idl.ts index ed6c05c..1ef508c 100644 --- a/packages/token-staking/src/idl.ts +++ b/packages/token-staking/src/idl.ts @@ -356,9 +356,11 @@ export const _SplTokenStakingIDL = { }, { name: "from", - isMut: false, + isMut: true, isSigner: false, - docs: ["[dead] left in for backwards compatibility."], + docs: [ + "Token Account holding weighted stake representation token to burn", + ], }, { name: "destination", diff --git a/programs/spl-token-staking/src/instructions/withdraw.rs b/programs/spl-token-staking/src/instructions/withdraw.rs index 2f232bb..5a48619 100644 --- a/programs/spl-token-staking/src/instructions/withdraw.rs +++ b/programs/spl-token-staking/src/instructions/withdraw.rs @@ -1,7 +1,7 @@ use anchor_lang::prelude::*; -use anchor_spl::token::{self, Mint, TokenAccount, Transfer}; +use anchor_spl::token::{self, Burn, Mint, TokenAccount, Transfer}; -use crate::{errors::ErrorCode, stake_pool_signer_seeds}; +use crate::{errors::ErrorCode, stake_pool_signer_seeds, state::StakeDepositReceipt}; use super::claim_base::*; use crate::state::u128; @@ -18,9 +18,9 @@ pub struct Withdraw<'info> { #[account(mut)] pub stake_mint: Account<'info, Mint>, - /// [dead] left in for backwards compatibility. - /// CHECK: unused account, no check needed - pub from: UncheckedAccount<'info>, + /// Token Account holding weighted stake representation token to burn + #[account(mut)] + pub from: Account<'info, TokenAccount>, /// Token account to transfer the previously staked token to #[account(mut)] @@ -39,6 +39,10 @@ impl<'info> Withdraw<'info> { stake_pool.stake_mint.key() == self.stake_mint.key(), ErrorCode::InvalidStakeMint ); + require!( + self.from.owner.key() == self.claim_base.owner.key(), + ErrorCode::InvalidAuthority + ); Ok(()) } /// Transfer the owner's previously staked tokens back. @@ -60,6 +64,23 @@ impl<'info> Withdraw<'info> { ) } + pub fn burn_stake_weight_tokens_from_owner(&self) -> Result<()> { + let stake_pool = self.claim_base.stake_pool.load()?; + let cpi_ctx = CpiContext::new( + self.claim_base.token_program.to_account_info(), + Burn { + mint: self.stake_mint.to_account_info(), + from: self.from.to_account_info(), + authority: self.claim_base.owner.to_account_info(), + }, + ); + let effective_stake_token_amount = StakeDepositReceipt::get_token_amount_from_stake( + self.claim_base.stake_deposit_receipt.effective_stake_u128(), + stake_pool.max_weight, + ); + token::burn(cpi_ctx, effective_stake_token_amount) + } + pub fn close_stake_deposit_receipt(&self) -> Result<()> { self.claim_base .stake_deposit_receipt @@ -90,6 +111,7 @@ pub fn handler<'info>(ctx: Context<'_, '_, '_, 'info, Withdraw<'info>>) -> Resul stake_pool.total_weighted_stake = u128(total_staked.to_le_bytes()); } ctx.accounts.transfer_staked_tokens_to_owner()?; + ctx.accounts.burn_stake_weight_tokens_from_owner()?; // claim all unclaimed rewards let claimed_amounts = ctx .accounts diff --git a/tests/withdraw.ts b/tests/withdraw.ts index 4eea135..9499809 100644 --- a/tests/withdraw.ts +++ b/tests/withdraw.ts @@ -77,7 +77,6 @@ describe("withdraw", () => { it("withdraw unlocked tokens", async () => { const receiptNonce = 0; - const depositAmount = new anchor.BN(1_000_000_000); const [stakeReceiptKey] = anchor.web3.PublicKey.findProgramAddressSync( [ depositor1.publicKey.toBuffer(), @@ -101,7 +100,7 @@ describe("withdraw", () => { depositor1, mintToBeStakedAccountKey, stakeMintAccountKey, - depositAmount, + new anchor.BN(1_000_000_000), new anchor.BN(0), receiptNonce ); @@ -117,7 +116,7 @@ describe("withdraw", () => { }, vault: vaultKey, stakeMint, - from: anchor.web3.PublicKey.default, + from: stakeMintAccountKey, destination: mintToBeStakedAccountKey, }) .remainingAccounts([ @@ -140,12 +139,14 @@ describe("withdraw", () => { depositerMintAccount, sTokenAccountAfter, vaultAfter, + stakeMintAfter, stakeDepositReceipt, ] = await Promise.all([ program.account.stakePool.fetch(stakePoolKey), tokenProgramInstance.account.account.fetch(mintToBeStakedAccountKey), tokenProgramInstance.account.account.fetch(stakeMintAccountKey), tokenProgramInstance.account.account.fetch(vaultKey), + tokenProgramInstance.account.mint.fetch(stakeMint), program.provider.connection.getAccountInfo(stakeReceiptKey), ]); assertBNEqual( @@ -156,8 +157,9 @@ describe("withdraw", () => { depositerMintAccount.amount, depositerMintAccountBefore.amount ); - assertBNEqual(sTokenAccountAfter.amount, sTokenAccountBefore.amount.add(depositAmount)); + assertBNEqual(sTokenAccountAfter.amount, sTokenAccountBefore.amount); assertBNEqual(vaultAfter.amount, 0); + assertBNEqual(stakeMintAfter.supply, 0); assert.isNull( stakeDepositReceipt, "StakeDepositReceipt account not closed" @@ -166,7 +168,6 @@ describe("withdraw", () => { it("withdraw claims unclaimed rewards", async () => { const receiptNonce = 1; - const depositAmount = new anchor.BN(1_000_000_000); const [stakeReceiptKey] = anchor.web3.PublicKey.findProgramAddressSync( [ depositor1.publicKey.toBuffer(), @@ -191,7 +192,7 @@ describe("withdraw", () => { depositor1, mintToBeStakedAccountKey, stakeMintAccountKey, - depositAmount, + new anchor.BN(1_000_000_000), new anchor.BN(0), receiptNonce, [rewardVaultKey] @@ -229,7 +230,7 @@ describe("withdraw", () => { }, vault: vaultKey, stakeMint, - from: anchor.web3.PublicKey.default, + from: stakeMintAccountKey, destination: mintToBeStakedAccountKey, }) .remainingAccounts([ @@ -252,12 +253,14 @@ describe("withdraw", () => { depositerMintAccount, sTokenAccountAfter, vaultAfter, + stakeMintAfter, depositorReward1AccountAfter, ] = await Promise.all([ program.account.stakePool.fetch(stakePoolKey), tokenProgramInstance.account.account.fetch(mintToBeStakedAccountKey), tokenProgramInstance.account.account.fetch(stakeMintAccountKey), tokenProgramInstance.account.account.fetch(vaultKey), + tokenProgramInstance.account.mint.fetch(stakeMint), tokenProgramInstance.account.account.fetch(depositorReward1AccountKey), ]); assertBNEqual( @@ -268,8 +271,9 @@ describe("withdraw", () => { depositerMintAccount.amount, depositerMintAccountBefore.amount ); - assertBNEqual(sTokenAccountAfter.amount, sTokenAccountBefore.amount.add(depositAmount)); + assertBNEqual(sTokenAccountAfter.amount, sTokenAccountBefore.amount); assertBNEqual(vaultAfter.amount, 0); + assertBNEqual(stakeMintAfter.supply, 0); assertBNEqual(depositorReward1AccountAfter.amount, totalReward1); }); @@ -310,7 +314,7 @@ describe("withdraw", () => { }, vault: vaultKey, stakeMint, - from: anchor.web3.PublicKey.default, + from: stakeMintAccountKey, destination: mintToBeStakedAccountKey, }) .remainingAccounts([ @@ -333,121 +337,4 @@ describe("withdraw", () => { } assert.isTrue(false, "TX should have failed"); }); - - describe("After burning stake_mint tokens", () => { - const receiptNonce = 3; - before(async () => { - // deposit 1 token - await deposit( - program, - stakePoolNonce, - mintToBeStaked, - depositor1, - mintToBeStakedAccountKey, - stakeMintAccountKey, - new anchor.BN(1_000_000_000), - new anchor.BN(0), - receiptNonce, - [rewardVaultKey] - ); - const stakeMintTokenAccount = - await tokenProgramInstance.account.account.fetch(stakeMintAccountKey); - - // Burn the staking tokens - await tokenProgramInstance.methods - .burn(stakeMintTokenAccount.amount) - .accounts({ - account: stakeMintAccountKey, - mint: stakeMint, - authority: depositor1.publicKey, - }) - .signers([depositor1]) - .rpc(); - }); - - it("should still withdraw", async () => { - const [stakeReceiptKey] = anchor.web3.PublicKey.findProgramAddressSync( - [ - depositor1.publicKey.toBuffer(), - stakePoolKey.toBuffer(), - new anchor.BN(receiptNonce).toArrayLike(Buffer, "le", 4), - Buffer.from("stakeDepositReceipt", "utf-8"), - ], - program.programId - ); - const [stakePoolBefore, stakeReceipt, depositerMintAccountBefore, sTokenAccountBefore, vaultBefore] = - await Promise.all([ - program.account.stakePool.fetch(stakePoolKey), - program.account.stakeDepositReceipt.fetch(stakeReceiptKey), - tokenProgramInstance.account.account.fetch(mintToBeStakedAccountKey), - tokenProgramInstance.account.account.fetch(stakeMintAccountKey, 'processed'), - tokenProgramInstance.account.account.fetch(vaultKey), - ]); - assert.equal(sTokenAccountBefore.amount.toString(), "0"); - - // Withdraw - try { - await program.methods - .withdraw() - .accounts({ - claimBase: { - owner: depositor1.publicKey, - stakePool: stakePoolKey, - stakeDepositReceipt: stakeReceiptKey, - tokenProgram: TOKEN_PROGRAM_ID, - }, - vault: vaultKey, - stakeMint, - from: anchor.web3.PublicKey.default, - destination: mintToBeStakedAccountKey, - }) - .remainingAccounts([ - { - pubkey: rewardVaultKey, - isWritable: true, - isSigner: false, - }, - { - pubkey: depositorReward1AccountKey, - isWritable: true, - isSigner: false, - }, - ]) - .signers([depositor1]) - .rpc({ skipPreflight: true }); - } catch (err) { - console.error(err); - assert.ok(false); - } - - const [ - stakePoolAfter, - depositerMintAccount, - sTokenAccountAfter, - vaultAfter, - stakeDepositReceipt, - ] = await Promise.all([ - program.account.stakePool.fetch(stakePoolKey), - tokenProgramInstance.account.account.fetch(mintToBeStakedAccountKey), - tokenProgramInstance.account.account.fetch(stakeMintAccountKey, 'processed'), - tokenProgramInstance.account.account.fetch(vaultKey), - program.provider.connection.getAccountInfo(stakeReceiptKey), - ]); - assertBNEqual( - stakePoolBefore.totalWeightedStake.sub(stakeReceipt.effectiveStake), - stakePoolAfter.totalWeightedStake - ); - assertBNEqual( - depositerMintAccount.amount, - depositerMintAccountBefore.amount.add(stakeReceipt.depositAmount) - ); - // No change to the stake token because it's not burning - assertBNEqual(sTokenAccountAfter.amount, sTokenAccountBefore.amount); - assertBNEqual(vaultAfter.amount, vaultBefore.amount.sub(stakeReceipt.depositAmount)); - assert.isNull( - stakeDepositReceipt, - "StakeDepositReceipt account not closed" - ); - }); - }); });