program: token 22 scaled ui support (#1894)

* init

* program: scaled ui working

* cargo fmt --

* ts lint

Author: crispheaney

programs/drift/src/controller/token.rs

@@ -199,3 +199,21 @@ pub fn transfer_checked_with_transfer_hook<'info>(
 
     solana_program::program::invoke_signed(&ix, &account_infos, signer_seeds).map_err(Into::into)
 }
+
+pub fn initialize_token_account<'info>(
+    token_program: &Interface<'info, TokenInterface>,
+    account: &AccountInfo<'info>,
+    owner: &AccountInfo<'info>,
+    mint: &InterfaceAccount<'info, Mint>,
+) -> Result<()> {
+    let cpi_program = token_program.to_account_info();
+    let accounts = ::anchor_spl::token_interface::InitializeAccount3 {
+        account: account.to_account_info(),
+        mint: mint.to_account_info(),
+        authority: owner.to_account_info(),
+    };
+    let cpi_ctx = anchor_lang::context::CpiContext::new(cpi_program, accounts);
+    ::anchor_spl::token_interface::initialize_account3(cpi_ctx)?;
+
+    Ok(())
+}

programs/drift/src/instructions/admin.rs

@@ -10,7 +10,7 @@ use pyth_solana_receiver_sdk::cpi::accounts::InitPriceUpdate;
 use pyth_solana_receiver_sdk::program::PythSolanaReceiver;
 use serum_dex::state::ToAlignedBytes;
 
-use crate::controller::token::close_vault;
+use crate::controller::token::{close_vault, initialize_token_account};
 use crate::error::ErrorCode;
 use crate::ids::{admin_hot_wallet, amm_spread_adjust_wallet, mm_oracle_crank_wallet};
 use crate::instructions::constraints::*;
@@ -146,15 +146,19 @@ pub fn handle_initialize_spot_market(
     let state = &mut ctx.accounts.state;
     let spot_market_pubkey = ctx.accounts.spot_market.key();
 
-    // protocol must be authority of collateral vault
-    if ctx.accounts.spot_market_vault.owner != state.signer {
-        return Err(ErrorCode::InvalidSpotMarketAuthority.into());
-    }
+    initialize_token_account(
+        &ctx.accounts.token_program,
+        &ctx.accounts.spot_market_vault,
+        &ctx.accounts.drift_signer,
+        &ctx.accounts.spot_market_mint,
+    )?;
 
-    // protocol must be authority of collateral vault
-    if ctx.accounts.insurance_fund_vault.owner != state.signer {
-        return Err(ErrorCode::InvalidInsuranceFundAuthority.into());
-    }
+    initialize_token_account(
+        &ctx.accounts.token_program,
+        &ctx.accounts.insurance_fund_vault,
+        &ctx.accounts.drift_signer,
+        &ctx.accounts.spot_market_mint,
+    )?;
 
     validate_borrow_rate(optimal_utilization, optimal_borrow_rate, max_borrow_rate, 0)?;
 
@@ -280,7 +284,7 @@ pub fn handle_initialize_spot_market(
         historical_oracle_data: historical_oracle_data_default,
         historical_index_data: historical_index_data_default,
         mint: ctx.accounts.spot_market_mint.key(),
-        vault: *ctx.accounts.spot_market_vault.to_account_info().key,
+        vault: ctx.accounts.spot_market_vault.key(),
         revenue_pool: PoolBalance {
             scaled_balance: 0,
             market_index: spot_market_index,
@@ -337,7 +341,7 @@ pub fn handle_initialize_spot_market(
         pool_id: 0,
         padding: [0; 40],
         insurance_fund: InsuranceFund {
-            vault: *ctx.accounts.insurance_fund_vault.to_account_info().key,
+            vault: ctx.accounts.insurance_fund_vault.key(),
             unstaking_period: THIRTEEN_DAY,
             total_factor: if_total_factor,
             user_factor: if_total_factor / 2,
@@ -4945,25 +4949,28 @@ pub struct InitializeSpotMarket<'info> {
         payer = admin
     )]
     pub spot_market: AccountLoader<'info, SpotMarket>,
+    #[account(
+        mint::token_program = token_program,
+    )]
     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()],
         bump,
         payer = admin,
-        token::mint = spot_market_mint,
-        token::authority = drift_signer
+        space = get_vault_len(&spot_market_mint)?,
+        owner = token_program.key()
     )]
-    pub spot_market_vault: Box<InterfaceAccount<'info, TokenAccount>>,
+    pub spot_market_vault: AccountInfo<'info>,
     #[account(
         init,
         seeds = [b"insurance_fund_vault".as_ref(), state.number_of_spot_markets.to_le_bytes().as_ref()],
         bump,
         payer = admin,
-        token::mint = spot_market_mint,
-        token::authority = drift_signer
+        space = get_vault_len(&spot_market_mint)?,
+        owner = token_program.key()
     )]
-    pub insurance_fund_vault: Box<InterfaceAccount<'info, TokenAccount>>,
+    pub insurance_fund_vault: AccountInfo<'info>,
     #[account(
         constraint = state.signer.eq(&drift_signer.key())
     )]

programs/drift/src/instructions/constraints.rs

@@ -1,7 +1,9 @@
 use anchor_lang::accounts::account::Account;
 use anchor_lang::accounts::account_loader::AccountLoader;
 use anchor_lang::accounts::signer::Signer;
+use anchor_lang::prelude::*;
 use anchor_lang::prelude::{AccountInfo, Pubkey};
+use anchor_spl::token_interface::Mint;
 
 use crate::error::ErrorCode;
 use crate::msg;
@@ -145,3 +147,28 @@ pub fn exchange_not_paused(state: &Account<State>) -> anchor_lang::Result<()> {
     }
     Ok(())
 }
+
+pub fn get_vault_len(mint: &InterfaceAccount<Mint>) -> anchor_lang::Result<usize> {
+    let mint_info = mint.to_account_info();
+    let len = if *mint_info.owner == ::anchor_spl::token_2022::Token2022::id() {
+        use ::anchor_spl::token_2022::spl_token_2022::extension::{
+            BaseStateWithExtensions, ExtensionType, StateWithExtensions,
+        };
+        use ::anchor_spl::token_2022::spl_token_2022::state::{Account, Mint};
+        let mint_data = mint_info.try_borrow_data()?;
+        let mint_state = StateWithExtensions::<Mint>::unpack(&mint_data)?;
+        let mint_extensions = match mint_state.get_extension_types() {
+            Ok(extensions) => extensions,
+            // If we cant deserialize the mint, we use the default token account length
+            // Init token will fail if this size doesnt work, so worst case init account just fails
+            Err(_) => return Ok(::anchor_spl::token::TokenAccount::LEN),
+        };
+        let required_extensions =
+            ExtensionType::get_required_init_account_extensions(&mint_extensions);
+        ExtensionType::try_calculate_account_len::<Account>(&required_extensions)?
+    } else {
+        ::anchor_spl::token::TokenAccount::LEN
+    };
+
+    Ok(len)
+}

test-scripts/single-anchor-test.sh

@@ -6,7 +6,9 @@ fi
 
 export ANCHOR_WALLET=~/.config/solana/id.json
 
-test_files=(admin.ts)
+test_files=(
+  spotDepositWithdraw22ScaledUI.ts
+)
 
 for test_file in ${test_files[@]}; do
   ts-mocha -t 300000 ./tests/${test_file}