Token 2022 upgrade

Anchor.toml

@@ -12,9 +12,16 @@ cluster = "Localnet"
 wallet = "~/.config/solana/id.json"
 
 [scripts]
-# test = "yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/**/*.ts -r tests/hooks.ts"
-#test = "yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/**/initialize-stake-pool.ts -r tests/hooks.ts"
- test = "yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/token-2022/*.ts -r tests/token-2022/hooks22.ts"
+# Classic Token program tests
+# test = "yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/*.ts -r tests/hooks.ts"
+# test = "yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/initialize-stake-pool.ts -r tests/hooks.ts"
+# test = "yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/claim-all.ts -r tests/hooks.ts"
+
+# Token 2022 specific tests (currently do not play nice with classic tests)
+test = "yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/token-2022/*.ts -r tests/token-2022/hooks22.ts"
+# test = "yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/token-2022/01*.ts -r tests/token-2022/hooks22.ts"
+# test = "yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/token-2022/03*.ts -r tests/token-2022/hooks22.ts"
+# test = "yarn run ts-mocha -p ./tsconfig.json -t 1000000 tests/token-2022/04*.ts -r tests/token-2022/hooks22.ts"
 
 [test.validator]
 [[test.genesis]]

packages/token-staking/src/idl.ts

@@ -3,7 +3,7 @@ type Mutable<T> = {
 };
 
 export const _SplTokenStakingIDL = {
-  version: "2.0.0",
+  version: "2.0.1",
   name: "spl_token_staking",
   instructions: [
     {
@@ -209,6 +209,11 @@ export const _SplTokenStakingIDL = {
             "from the account staking.",
           ],
         },
+        {
+          name: "mint",
+          isMut: false,
+          isSigner: false,
+        },
         {
           name: "from",
           isMut: true,
@@ -386,6 +391,12 @@ export const _SplTokenStakingIDL = {
           isSigner: false,
           docs: ["Token account to transfer the previously staked token to"],
         },
+        {
+          name: "mint",
+          isMut: false,
+          isSigner: false,
+          docs: ["Staking pool's mint"],
+        },
       ],
       args: [],
     },
@@ -736,10 +747,14 @@ export const _SplTokenStakingIDL = {
             docs: ["latest amount of tokens in the vault"],
             type: "u64",
           },
+          {
+            name: "decimals",
+            type: "u8",
+          },
           {
             name: "padding0",
             type: {
-              array: ["u8", 8],
+              array: ["u8", 7],
             },
           },
         ],

packages/token-staking/src/instructions.ts

@@ -15,6 +15,7 @@ import { SplTokenStakingV0 } from "./idl_v0";
  * @param maxDuration
  * @param authority - defaults to `program.provider.publicKey`
  * @param registar - Registrar key if this StakePool uses SPL-Governance
+ * @param tokenProgram
  */
 export const initStakePool = async (
   program: anchor.Program<SplTokenStaking | SplTokenStakingV0>,
@@ -25,6 +26,7 @@ export const initStakePool = async (
   maxDuration = new anchor.BN("18446744073709551615"),
   authority?: anchor.Address,
   registrar: anchor.web3.PublicKey | null = null,
+  tokenProgram: anchor.web3.PublicKey = SPL_TOKEN_PROGRAM_ID
 ) => {
   const _authority = authority
     ? new anchor.web3.PublicKey(authority)
@@ -50,7 +52,7 @@ export const initStakePool = async (
       stakePool: stakePoolKey,
       mint,
       vault: vaultKey,
-      tokenProgram: SPL_TOKEN_PROGRAM_ID,
+      tokenProgram: tokenProgram,
       rent: anchor.web3.SYSVAR_RENT_PUBKEY,
       systemProgram: anchor.web3.SystemProgram.programId,
     })
@@ -65,6 +67,7 @@ export const initStakePool = async (
  * @param rewardMint
  * @param rewardPoolIndex
  * @param authority
+ * @param tokenProgram
  * @returns
  */
 export const addRewardPool = async (
@@ -73,7 +76,8 @@ export const addRewardPool = async (
   stakePoolMint: anchor.Address,
   rewardMint: anchor.web3.PublicKey,
   rewardPoolIndex = 0,
-  authority?: anchor.Address
+  authority?: anchor.Address,
+  tokenProgram: anchor.web3.PublicKey = SPL_TOKEN_PROGRAM_ID,
 ) => {
   const _authority = authority
     ? new anchor.web3.PublicKey(authority)
@@ -103,7 +107,7 @@ export const addRewardPool = async (
       rewardMint,
       stakePool: stakePoolKey,
       rewardVault: rewardVaultKey,
-      tokenProgram: SPL_TOKEN_PROGRAM_ID,
+      tokenProgram: tokenProgram,
       rent: anchor.web3.SYSVAR_RENT_PUBKEY,
       systemProgram: anchor.web3.SystemProgram.programId,
     })
@@ -115,6 +119,7 @@ export const addRewardPool = async (
  * @param program
  * @param payer
  * @param owner
+ * @param mint
  * @param stakePoolKey
  * @param from
  * @param amount
@@ -127,6 +132,7 @@ export const createStakeBuilder = (
   program: anchor.Program<SplTokenStaking | SplTokenStakingV0>,
   payer: anchor.web3.PublicKey,
   owner: anchor.web3.PublicKey,
+  mint: anchor.web3.PublicKey,
   stakePoolKey: anchor.Address,
   from: anchor.Address,
   amount: anchor.BN,
@@ -157,6 +163,7 @@ export const createStakeBuilder = (
     .accounts({
       payer,
       owner,
+      mint: mint,
       from,
       stakePool: stakePoolKey,
       vault: vaultKey,
@@ -178,6 +185,7 @@ export const createStakeBuilder = (
  * @param program
  * @param payer
  * @param owner
+ * @param mint
  * @param stakePoolKey
  * @param from
  * @param amount
@@ -190,6 +198,7 @@ export const createStakeInstruction = async (
   program: anchor.Program<SplTokenStaking | SplTokenStakingV0>,
   payer: anchor.web3.PublicKey,
   owner: anchor.web3.PublicKey,
+  mint: anchor.web3.PublicKey,
   stakePoolkey: anchor.Address,
   from: anchor.Address,
   amount: anchor.BN,
@@ -201,6 +210,7 @@ export const createStakeInstruction = async (
     program,
     payer,
     owner,
+    mint,
     stakePoolkey,
     from,
     amount,
@@ -227,6 +237,7 @@ export const deposit = async (
   program: anchor.Program<SplTokenStaking | SplTokenStakingV0>,
   payer: anchor.web3.PublicKey,
   owner: anchor.web3.PublicKey,
+  mint: anchor.web3.PublicKey,
   stakePoolKey: anchor.Address,
   from: anchor.Address,
   amount: anchor.BN,
@@ -245,6 +256,7 @@ export const deposit = async (
     program,
     payer,
     owner,
+    mint,
     stakePoolKey,
     from,
     amount,

programs/spl-token-staking/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "spl-token-staking"
-version = "2.0.0"
+version = "2.0.1"
 description = "Created with Anchor"
 edition = "2021"
 

programs/spl-token-staking/src/instructions/add_reward_pool.rs

@@ -1,51 +1,56 @@
 use anchor_lang::prelude::*;
-use anchor_spl::token::{Mint, Token, TokenAccount};
+use anchor_spl::token_interface::{
+    Mint as MintInterface, TokenAccount as TokenAccountInterface, TokenInterface,
+};
 
-use crate::state::{RewardPool, StakePool};
 use crate::errors::ErrorCode;
+use crate::state::{RewardPool, StakePool};
 
 #[derive(Accounts)]
 #[instruction(index: u8)]
 pub struct AddRewardPool<'info> {
-  /// Payer of rent
-  #[account(mut)]
-  pub payer: Signer<'info>,
-
-  /// Authority of the StakePool
-  pub authority: Signer<'info>,
-
-  /// SPL Token Mint of the token that will be distributed as rewards
-  pub reward_mint: Account<'info, Mint>,
-
-  /// StakePool where the RewardPool will be added
-  #[account(
-    mut, 
-    has_one = authority @ ErrorCode::InvalidAuthority,
-    constraint = stake_pool.load()?.reward_pools[usize::from(index)].reward_vault == Pubkey::default() 
-      @ ErrorCode::RewardPoolIndexOccupied,
-  )]
-  pub stake_pool: AccountLoader<'info, StakePool>,
-
-  /// An SPL token Account for holding rewards to be claimed
-  #[account(
+    /// Payer of rent
+    #[account(mut)]
+    pub payer: Signer<'info>,
+
+    /// Authority of the StakePool
+    pub authority: Signer<'info>,
+
+    /// SPL Token Mint of the token that will be distributed as rewards
+    pub reward_mint: InterfaceAccount<'info, MintInterface>,
+
+    /// StakePool where the RewardPool will be added
+    #[account(
+      mut,
+      has_one = authority @ ErrorCode::InvalidAuthority,
+      constraint = stake_pool.load()?.reward_pools[usize::from(index)].reward_vault == Pubkey::default()
+        @ ErrorCode::RewardPoolIndexOccupied,
+    )]
+    pub stake_pool: AccountLoader<'info, StakePool>,
+
+    /// An SPL token Account for holding rewards to be claimed
+    #[account(
     init,
     seeds = [stake_pool.key().as_ref(), reward_mint.key().as_ref(), b"rewardVault"],
     bump,
     payer = payer,
     token::mint = reward_mint,
     token::authority = stake_pool,
   )]
-  pub reward_vault: Account<'info, TokenAccount>,
+    pub reward_vault: InterfaceAccount<'info, TokenAccountInterface>,
 
-  pub token_program: Program<'info, Token>,
-  pub rent: Sysvar<'info, Rent>,
-  pub system_program: Program<'info, System>,
+    pub token_program: Interface<'info, TokenInterface>,
+    pub rent: Sysvar<'info, Rent>,
+    pub system_program: Program<'info, System>,
 }
 
 pub fn handler(ctx: Context<AddRewardPool>, index: u8) -> Result<()> {
-  let mut stake_pool = ctx.accounts.stake_pool.load_mut()?;
-  let reward_pool = RewardPool::new(&ctx.accounts.reward_vault.key());
-  stake_pool.reward_pools[usize::from(index)] = reward_pool;
-
-  Ok(())
-}
+    let mut stake_pool = ctx.accounts.stake_pool.load_mut()?;
+    let reward_pool = RewardPool::new(
+        &ctx.accounts.reward_vault.key(),
+        ctx.accounts.reward_mint.decimals,
+    );
+    stake_pool.reward_pools[usize::from(index)] = reward_pool;
+
+    Ok(())
+}

programs/spl-token-staking/src/instructions/claim_all.rs

@@ -1,4 +1,5 @@
 use anchor_lang::prelude::*;
+use anchor_spl::token_2022::Token2022;
 
 use super::claim_base::*;
 
@@ -10,7 +11,12 @@ pub struct ClaimAll<'info> {
 pub fn handler<'info>(ctx: Context<'_, '_, 'info, 'info, ClaimAll<'info>>) -> Result<()> {
     {
         let mut stake_pool = ctx.accounts.claim_base.stake_pool.load_mut()?;
-        stake_pool.recalculate_rewards_per_effective_stake(&ctx.remaining_accounts, 2usize)?;
+        let step = if ctx.accounts.claim_base.token_program.key() == Token2022::id() {
+            3usize
+        } else {
+            2usize
+        };
+        stake_pool.recalculate_rewards_per_effective_stake(&ctx.remaining_accounts, step)?;
     }
 
     let claimed_amounts = ctx