program: add token 2022 support

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Features
 
 - program: add switchboard on demand integration ([#1154](https://github.com/drift-labs/protocol-v2/pull/1154))
+- program: add support for token 2022 ([#1125](https://github.com/drift-labs/protocol-v2/pull/1125))
 
 ### Fixes
 

programs/drift/src/controller/insurance.rs

@@ -1,5 +1,5 @@
 use anchor_lang::prelude::*;
-use anchor_spl::token::{Token, TokenAccount};
+use anchor_spl::token_interface::{Mint, TokenAccount, TokenInterface};
 use solana_program::msg;
 
 use crate::controller::spot_balance::{
@@ -633,13 +633,14 @@ pub fn transfer_protocol_insurance_fund_stake(
 }
 
 pub fn attempt_settle_revenue_to_insurance_fund<'info>(
-    spot_market_vault: &Account<'info, TokenAccount>,
-    insurance_fund_vault: &Account<'info, TokenAccount>,
+    spot_market_vault: &InterfaceAccount<'info, TokenAccount>,
+    insurance_fund_vault: &InterfaceAccount<'info, TokenAccount>,
     spot_market: &mut SpotMarket,
     now: i64,
-    token_program: &Program<'info, Token>,
+    token_program: &Interface<'info, TokenInterface>,
     drift_signer: &AccountInfo<'info>,
     state: &State,
+    mint: &Option<InterfaceAccount<'info, Mint>>,
 ) -> Result<()> {
     let valid_revenue_settle_time = if spot_market.insurance_fund.revenue_settle_period > 0 {
         let time_until_next_update = on_the_hour_update(
@@ -680,6 +681,7 @@ pub fn attempt_settle_revenue_to_insurance_fund<'info>(
                 drift_signer,
                 state.signer_nonce,
                 token_amount.cast()?,
+                mint,
             )?;
         }
 

programs/drift/src/controller/token.rs

@@ -1,47 +1,95 @@
+use crate::error::ErrorCode;
 use crate::signer::get_signer_seeds;
+use crate::validate;
 use anchor_lang::prelude::*;
-use anchor_spl::token::{self, CloseAccount, Token, TokenAccount, Transfer};
+use anchor_spl::token_2022::spl_token_2022::extension::transfer_fee::TransferFeeConfig;
+use anchor_spl::token_2022::spl_token_2022::extension::{
+    BaseStateWithExtensions, StateWithExtensions,
+};
+use anchor_spl::token_2022::spl_token_2022::state::Mint as MintInner;
+use anchor_spl::token_interface::{
+    self, CloseAccount, Mint, TokenAccount, TokenInterface, Transfer, TransferChecked,
+};
 
 pub fn send_from_program_vault<'info>(
-    token_program: &Program<'info, Token>,
-    from: &Account<'info, TokenAccount>,
-    to: &Account<'info, TokenAccount>,
+    token_program: &Interface<'info, TokenInterface>,
+    from: &InterfaceAccount<'info, TokenAccount>,
+    to: &InterfaceAccount<'info, TokenAccount>,
     authority: &AccountInfo<'info>,
     nonce: u8,
     amount: u64,
+    mint: &Option<InterfaceAccount<'info, Mint>>,
 ) -> Result<()> {
     let signature_seeds = get_signer_seeds(&nonce);
     let signers = &[&signature_seeds[..]];
-    let cpi_accounts = Transfer {
-        from: from.to_account_info().clone(),
-        to: to.to_account_info().clone(),
-        authority: authority.to_account_info().clone(),
-    };
-    let cpi_program = token_program.to_account_info();
-    let cpi_context = CpiContext::new_with_signer(cpi_program, cpi_accounts, signers);
-    token::transfer(cpi_context, amount)
+
+    if let Some(mint) = mint {
+        let mint_account_info = mint.to_account_info().clone();
+
+        validate_mint_fee(&mint_account_info)?;
+
+        let cpi_accounts = TransferChecked {
+            from: from.to_account_info().clone(),
+            mint: mint_account_info,
+            to: to.to_account_info().clone(),
+            authority: authority.to_account_info().clone(),
+        };
+
+        let cpi_program = token_program.to_account_info();
+        let cpi_context = CpiContext::new_with_signer(cpi_program, cpi_accounts, signers);
+        token_interface::transfer_checked(cpi_context, amount, mint.decimals)
+    } else {
+        let cpi_accounts = Transfer {
+            from: from.to_account_info().clone(),
+            to: to.to_account_info().clone(),
+            authority: authority.to_account_info().clone(),
+        };
+
+        let cpi_program = token_program.to_account_info();
+        let cpi_context = CpiContext::new_with_signer(cpi_program, cpi_accounts, signers);
+        #[allow(deprecated)]
+        token_interface::transfer(cpi_context, amount)
+    }
 }
 
 pub fn receive<'info>(
-    token_program: &Program<'info, Token>,
-    from: &Account<'info, TokenAccount>,
-    to: &Account<'info, TokenAccount>,
+    token_program: &Interface<'info, TokenInterface>,
+    from: &InterfaceAccount<'info, TokenAccount>,
+    to: &InterfaceAccount<'info, TokenAccount>,
     authority: &AccountInfo<'info>,
     amount: u64,
+    mint: &Option<InterfaceAccount<'info, Mint>>,
 ) -> Result<()> {
-    let cpi_accounts = Transfer {
-        from: from.to_account_info().clone(),
-        to: to.to_account_info().clone(),
-        authority: authority.to_account_info().clone(),
-    };
-    let cpi_program = token_program.to_account_info();
-    let cpi_context = CpiContext::new(cpi_program, cpi_accounts);
-    token::transfer(cpi_context, amount)
+    if let Some(mint) = mint {
+        let mint_account_info = mint.to_account_info().clone();
+
+        validate_mint_fee(&mint_account_info)?;
+
+        let cpi_accounts = TransferChecked {
+            from: from.to_account_info().clone(),
+            to: to.to_account_info().clone(),
+            mint: mint_account_info,
+            authority: authority.to_account_info().clone(),
+        };
+        let cpi_program = token_program.to_account_info();
+        let cpi_context = CpiContext::new(cpi_program, cpi_accounts);
+        token_interface::transfer_checked(cpi_context, amount, mint.decimals)
+    } else {
+        let cpi_accounts = Transfer {
+            from: from.to_account_info().clone(),
+            to: to.to_account_info().clone(),
+            authority: authority.to_account_info().clone(),
+        };
+        let cpi_program = token_program.to_account_info();
+        let cpi_context = CpiContext::new(cpi_program, cpi_accounts);
+        #[allow(deprecated)]
+        token_interface::transfer(cpi_context, amount)
+    }
 }
 
 pub fn close_vault<'info>(
-    token_program: &Program<'info, Token>,
-    account: &Account<'info, TokenAccount>,
+    token_program: &Interface<'info, TokenInterface>,
+    account: &InterfaceAccount<'info, TokenAccount>,
     destination: &AccountInfo<'info>,
     authority: &AccountInfo<'info>,
     nonce: u8,
@@ -55,5 +103,20 @@ pub fn close_vault<'info>(
     };
     let cpi_program = token_program.to_account_info();
     let cpi_context = CpiContext::new_with_signer(cpi_program, cpi_accounts, signers);
-    token::close_account(cpi_context)
+    token_interface::close_account(cpi_context)
+}
+
+pub fn validate_mint_fee(account_info: &AccountInfo) -> Result<()> {
+    let mint_data = account_info.try_borrow_data()?;
+    let mint_with_extension = StateWithExtensions::<MintInner>::unpack(&mint_data)?;
+    if let Ok(fee_config) = mint_with_extension.get_extension::<TransferFeeConfig>() {
+        let fee = u16::from(
+            fee_config
+                .get_epoch_fee(Clock::get()?.epoch)
+                .transfer_fee_basis_points,
+        );
+        validate!(fee == 0, ErrorCode::NonZeroTransferFee)?
+    }
+
+    Ok(())
 }

programs/drift/src/error.rs

@@ -569,6 +569,8 @@ pub enum ErrorCode {
     InvalidOpenbookV2Program,
     #[msg("InvalidOpenbookV2Market")]
     InvalidOpenbookV2Market,
+    #[msg("Non zero transfer fee")]
+    NonZeroTransferFee,
 }
 
 #[macro_export]

programs/drift/src/instructions/admin.rs

@@ -2,7 +2,9 @@ use std::convert::identity;
 use std::mem::size_of;
 
 use anchor_lang::prelude::*;
-use anchor_spl::token::{Mint, Token, TokenAccount};
+use anchor_spl::token::Token;
+use anchor_spl::token_2022::Token2022;
+use anchor_spl::token_interface::{Mint, TokenAccount, TokenInterface};
 use phoenix::quantities::WrapperU64;
 use pyth_solana_receiver_sdk::cpi::accounts::InitPriceUpdate;
 use pyth_solana_receiver_sdk::program::PythSolanaReceiver;
@@ -29,6 +31,7 @@ use crate::math::safe_math::SafeMath;
 use crate::math::spot_balance::get_token_amount;
 use crate::math::{amm, bn};
 use crate::math_error;
+use crate::optional_accounts::get_token_mint;
 use crate::state::events::CurveRecord;
 use crate::state::fulfillment_params::openbook_v2::{
     OpenbookV2Context, OpenbookV2FulfillmentConfig,
@@ -227,6 +230,15 @@ pub fn handle_initialize_spot_market(
 
     let decimals = ctx.accounts.spot_market_mint.decimals.cast::<u32>()?;
 
+    let token_program = if ctx.accounts.token_program.key() == Token2022::id() {
+        1_u8
+    } else if ctx.accounts.token_program.key() == Token::id() {
+        0_u8
+    } else {
+        msg!("unexpected program {:?}", ctx.accounts.token_program.key());
+        return Err(ErrorCode::DefaultError.into());
+    };
+
     **spot_market = SpotMarket {
         market_index: spot_market_index,
         pubkey: spot_market_pubkey,
@@ -296,7 +308,8 @@ pub fn handle_initialize_spot_market(
         fuel_boost_taker: 0,
         fuel_boost_maker: 0,
         fuel_boost_insurance: 0,
-        padding: [0; 42],
+        token_program,
+        padding: [0; 41],
         insurance_fund: InsuranceFund {
             vault: *ctx.accounts.insurance_fund_vault.to_account_info().key,
             unstaking_period: THIRTEEN_DAY,
@@ -1608,12 +1621,16 @@ pub fn handle_settle_expired_market_pools_to_revenue_pool(
 #[access_control(
     perp_market_valid(&ctx.accounts.perp_market)
 )]
-pub fn handle_deposit_into_perp_market_fee_pool(
-    ctx: Context<DepositIntoMarketFeePool>,
+pub fn handle_deposit_into_perp_market_fee_pool<'c: 'info, 'info>(
+    ctx: Context<'_, '_, 'c, 'info, DepositIntoMarketFeePool<'info>>,
     amount: u64,
 ) -> Result<()> {
     let perp_market = &mut load_mut!(ctx.accounts.perp_market)?;
 
+    let remaining_accounts_iter = &mut ctx.remaining_accounts.iter().peekable();
+
+    let mint = get_token_mint(remaining_accounts_iter)?;
+
     msg!(
         "depositing {} into perp market {} fee pool",
         amount,
@@ -1650,6 +1667,7 @@ pub fn handle_deposit_into_perp_market_fee_pool(
         &ctx.accounts.spot_market_vault,
         &ctx.accounts.admin.to_account_info(),
         amount,
+        &mint,
     )?;
 
     Ok(())
@@ -3932,12 +3950,12 @@ pub struct Initialize<'info> {
         payer = admin
     )]
     pub state: Box<Account<'info, State>>,
-    pub quote_asset_mint: Box<Account<'info, Mint>>,
+    pub quote_asset_mint: Box<InterfaceAccount<'info, Mint>>,
     /// CHECK: checked in `initialize`
     pub drift_signer: AccountInfo<'info>,
     pub rent: Sysvar<'info, Rent>,
     pub system_program: Program<'info, System>,
-    pub token_program: Program<'info, Token>,
+    pub token_program: Interface<'info, TokenInterface>,
 }
 
 #[derive(Accounts)]
@@ -3950,7 +3968,7 @@ pub struct InitializeSpotMarket<'info> {
         payer = admin
     )]
     pub spot_market: AccountLoader<'info, SpotMarket>,
-    pub spot_market_mint: Box<Account<'info, Mint>>,
+    pub spot_market_mint: Box<InterfaceAccount<'info, Mint>>,
     #[account(
         init,
         seeds = [b"spot_market_vault".as_ref(), state.number_of_spot_markets.to_le_bytes().as_ref()],
@@ -3959,7 +3977,7 @@ pub struct InitializeSpotMarket<'info> {
         token::mint = spot_market_mint,
         token::authority = drift_signer
     )]
-    pub spot_market_vault: Box<Account<'info, TokenAccount>>,
+    pub spot_market_vault: Box<InterfaceAccount<'info, TokenAccount>>,
     #[account(
         init,
         seeds = [b"insurance_fund_vault".as_ref(), state.number_of_spot_markets.to_le_bytes().as_ref()],
@@ -3968,7 +3986,7 @@ pub struct InitializeSpotMarket<'info> {
         token::mint = spot_market_mint,
         token::authority = drift_signer
     )]
-    pub insurance_fund_vault: Box<Account<'info, TokenAccount>>,
+    pub insurance_fund_vault: Box<InterfaceAccount<'info, TokenAccount>>,
     #[account(
         constraint = state.signer.eq(&drift_signer.key())
     )]
@@ -3985,7 +4003,7 @@ pub struct InitializeSpotMarket<'info> {
     pub admin: Signer<'info>,
     pub rent: Sysvar<'info, Rent>,
     pub system_program: Program<'info, System>,
-    pub token_program: Program<'info, Token>,
+    pub token_program: Interface<'info, TokenInterface>,
 }
 
 #[derive(Accounts)]
@@ -4005,16 +4023,16 @@ pub struct DeleteInitializedSpotMarket<'info> {
         seeds = [b"spot_market_vault".as_ref(), market_index.to_le_bytes().as_ref()],
         bump,
     )]
-    pub spot_market_vault: Box<Account<'info, TokenAccount>>,
+    pub spot_market_vault: Box<InterfaceAccount<'info, TokenAccount>>,
     #[account(
         mut,
         seeds = [b"insurance_fund_vault".as_ref(), market_index.to_le_bytes().as_ref()],
         bump,
     )]
-    pub insurance_fund_vault: Box<Account<'info, TokenAccount>>,
+    pub insurance_fund_vault: Box<InterfaceAccount<'info, TokenAccount>>,
     /// CHECK: program signer
     pub drift_signer: AccountInfo<'info>,
-    pub token_program: Program<'info, Token>,
+    pub token_program: Interface<'info, TokenInterface>,
 }
 
 #[derive(Accounts)]
@@ -4139,7 +4157,7 @@ pub struct UpdateSerumVault<'info> {
     pub state: Box<Account<'info, State>>,
     #[account(mut)]
     pub admin: Signer<'info>,
-    pub srm_vault: Box<Account<'info, TokenAccount>>,
+    pub srm_vault: Box<InterfaceAccount<'info, TokenAccount>>,
 }
 
 #[derive(Accounts)]
@@ -4238,7 +4256,7 @@ pub struct DepositIntoMarketFeePool<'info> {
         mut,
         token::authority = admin
     )]
-    pub source_vault: Box<Account<'info, TokenAccount>>,
+    pub source_vault: Box<InterfaceAccount<'info, TokenAccount>>,
     #[account(
         constraint = state.signer.eq(&drift_signer.key())
     )]
@@ -4255,8 +4273,8 @@ pub struct DepositIntoMarketFeePool<'info> {
         seeds = [b"spot_market_vault".as_ref(), 0_u16.to_le_bytes().as_ref()],
         bump,
     )]
-    pub spot_market_vault: Box<Account<'info, TokenAccount>>,
-    pub token_program: Program<'info, Token>,
+    pub spot_market_vault: Box<InterfaceAccount<'info, TokenAccount>>,
+    pub token_program: Interface<'info, TokenInterface>,
 }
 
 #[derive(Accounts)]