Revert withdraw changes (#8)

* Revert "add from back for client side backwards compatibility (#5)"

This reverts commit 9350447.

* Revert "No token withdraw (#4)"

This reverts commit e199e26.

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",

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

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"
-      );
-    });
-  });
 });