Add Token Extensions Program (Token22) support

core/src/loader.rs

@@ -2,7 +2,7 @@
 use solana_program::{
     account_info::AccountInfo, msg, program_error::ProgramError, pubkey::Pubkey, system_program,
 };
-use spl_associated_token_account::get_associated_token_address;
+use spl_associated_token_account::get_associated_token_address_with_program_id;
 use spl_token_2022::extension::StateWithExtensions;
 
 /// Loads the account as a signer, returning an error if it is not or if it is not writable while
@@ -59,15 +59,16 @@ pub fn load_associated_token_account_program(info: &AccountInfo) -> Result<(), P
     Ok(())
 }
 
-/// Loads the account as the `spl_token` program, returning an error if it is not.
+/// Loads the account as the `spl_token` or `spl_token_2022` program, returning an error if it is
+/// neither.
 ///
 /// # Arguments
 /// * `info` - The account to load the token program from
 ///
 /// # Returns
 /// * `Result<(), ProgramError>` - The result of the operation
 pub fn load_token_program(info: &AccountInfo) -> Result<(), ProgramError> {
-    if info.key.ne(&spl_token::id()) {
+    if info.key.ne(&spl_token::id()) && info.key.ne(&spl_token_2022::id()) {
         msg!("Account is not the spl token program");
         return Err(ProgramError::IncorrectProgramId);
     }
@@ -134,7 +135,7 @@ pub fn load_associated_token_account(
     owner: &Pubkey,
     mint: &Pubkey,
 ) -> Result<(), ProgramError> {
-    if token_account.owner.ne(&spl_token::id()) {
+    if token_account.owner.ne(&spl_token::id()) && token_account.owner.ne(&spl_token_2022::id()) {
         msg!("Account is not owned by the spl token program");
         return Err(ProgramError::InvalidAccountOwner);
     }
@@ -144,7 +145,8 @@ pub fn load_associated_token_account(
         return Err(ProgramError::InvalidAccountData);
     }
 
-    let associated_token_account = get_associated_token_address(owner, mint);
+    let associated_token_account =
+        get_associated_token_address_with_program_id(owner, mint, token_account.owner);
     if token_account.key.ne(&associated_token_account) {
         msg!("Account is not the associated token account");
         return Err(ProgramError::InvalidAccountData);
@@ -178,12 +180,12 @@ pub fn load_token_account(
     mint: &Pubkey,
     token_program: &AccountInfo,
 ) -> Result<(), ProgramError> {
-    if token_program.key.ne(&spl_token::id()) {
-        msg!("Account is not owned by the spl token program");
+    if token_program.key.ne(&spl_token::id()) && token_program.key.ne(&spl_token_2022::id()) {
+        msg!("Account is not the spl token program");
         return Err(ProgramError::IncorrectProgramId);
     }
 
-    if token_account.owner.ne(&spl_token::id()) {
+    if token_account.owner.ne(token_program.key) {
         msg!("Account is not owned by the token program");
         return Err(ProgramError::InvalidAccountOwner);
     }
@@ -224,7 +226,7 @@ pub fn load_token_account(
 /// # Returns
 /// * `Result<(), ProgramError>` - The result of the operation
 pub fn load_token_mint(info: &AccountInfo) -> Result<(), ProgramError> {
-    if info.owner.ne(&spl_token::id()) {
+    if info.owner.ne(&spl_token::id()) && info.owner.ne(&spl_token_2022::id()) {
         msg!("Account is not owned by the spl token program");
         return Err(ProgramError::InvalidAccountOwner);
     }

integration_tests/tests/restaking/ncn_delegate_token_account.rs

@@ -68,7 +68,7 @@ mod tests {
     }
 
     #[test_case(spl_token::id(); "token")]
-    // #[test_case(spl_token_2022::id(); "token-2022")]
+    #[test_case(spl_token_2022::id(); "token-2022")]
     #[tokio::test]
     async fn test_ncn_delegate_token_account_ok(token_program_id: Pubkey) {
         let (mut fixture, ncn_root, random_mint, operator_token_account) =
@@ -118,7 +118,7 @@ mod tests {
     }
 
     #[test_case(spl_token::id(); "token")]
-    // #[test_case(spl_token_2022::id(); "token-2022")]
+    #[test_case(spl_token_2022::id(); "token-2022")]
     #[tokio::test]
     async fn test_ncn_delegate_token_account_wrong_delegate_admin_fails(token_program_id: Pubkey) {
         let (fixture, ncn_root, random_mint, operator_token_account) =

integration_tests/tests/restaking/operator_delegate_token_account.rs

@@ -72,7 +72,7 @@ mod tests {
     }
 
     #[test_case(spl_token::id(); "token")]
-    // #[test_case(spl_token_2022::id(); "token-2022")]
+    #[test_case(spl_token_2022::id(); "token-2022")]
     #[tokio::test]
     async fn test_operator_delegate_token_account_ok(token_program_id: Pubkey) {
         let (mut fixture, operator_root, random_mint, operator_token_account) =
@@ -126,7 +126,7 @@ mod tests {
     }
 
     #[test_case(spl_token::id(); "token")]
-    // #[test_case(spl_token_2022::id(); "token-2022")]
+    #[test_case(spl_token_2022::id(); "token-2022")]
     #[tokio::test]
     async fn test_operator_delegate_token_account_wrong_delegate_admin_fails(
         token_program_id: Pubkey,

integration_tests/tests/vault/delegate_token_account.rs

@@ -80,7 +80,7 @@ mod tests {
     }
 
     #[test_case(spl_token::id(); "token")]
-    // #[test_case(spl_token_2022::id(); "token-2022")]
+    #[test_case(spl_token_2022::id(); "token-2022")]
     #[tokio::test]
     async fn test_delegate_token_account_ok(token_program_id: Pubkey) {
         let (mut fixture, vault_pubkey, vault_admin, random_mint, vault_token_account) =
@@ -133,7 +133,7 @@ mod tests {
     }
 
     #[test_case(spl_token::id(); "token")]
-    // #[test_case(spl_token_2022::id(); "token-2022")]
+    #[test_case(spl_token_2022::id(); "token-2022")]
     #[tokio::test]
     async fn test_delegate_vault_wrong_delegate_asset_admin_fails(token_program_id: Pubkey) {
         let (fixture, vault_pubkey, _vault_admin, random_mint, vault_token_account) =
@@ -176,7 +176,7 @@ mod tests {
     }
 
     #[test_case(spl_token::id(); "token")]
-    // #[test_case(spl_token_2022::id(); "token-2022")]
+    #[test_case(spl_token_2022::id(); "token-2022")]
     #[tokio::test]
     async fn test_delegate_vault_account_supported_token_account_fails(token_program_id: Pubkey) {
         let (fixture, vault_pubkey, vault_admin, random_mint, vault_token_account) =
@@ -220,7 +220,7 @@ mod tests {
     }
 
     #[test_case(spl_token::id(); "token")]
-    // #[test_case(spl_token_2022::id(); "token-2022")]
+    #[test_case(spl_token_2022::id(); "token-2022")]
     #[tokio::test]
     async fn test_delegate_vault_token_account_does_not_match_token_mint_fails(
         token_program_id: Pubkey,

integration_tests/tests/vault/revoke_delegate_token_account.rs

@@ -79,7 +79,7 @@ mod tests {
     }
 
     #[test_case(spl_token::id(); "token")]
-    // #[test_case(spl_token_2022::id(); "token-2022")]
+    #[test_case(spl_token_2022::id(); "token-2022")]
     #[tokio::test]
     async fn test_delegate_token_account_ok(token_program_id: Pubkey) {
         let (mut fixture, vault_pubkey, vault_admin, random_mint, vault_token_account) =
@@ -171,7 +171,7 @@ mod tests {
     }
 
     #[test_case(spl_token::id(); "token")]
-    // #[test_case(spl_token_2022::id(); "token-2022")]
+    #[test_case(spl_token_2022::id(); "token-2022")]
     #[tokio::test]
     async fn test_delegate_vault_wrong_delegate_asset_admin_fails(token_program_id: Pubkey) {
         let (mut fixture, vault_pubkey, vault_admin, random_mint, vault_token_account) =

restaking_program/src/ncn_delegate_token_account.rs

@@ -36,7 +36,6 @@ pub fn process_ncn_delegate_token_account(
         token_mint.key,
         token_program_info,
     )?;
-    // Only the original spl token program is allowed
     load_token_program(token_program_info)?;
 
     // The owner of token mint and token account must match

restaking_program/src/operator_delegate_token_account.rs

@@ -37,7 +37,6 @@ pub fn process_operator_delegate_token_account(
         token_mint.key,
         token_program_info,
     )?;
-    // Only the original spl token program is allowed
     load_token_program(token_program_info)?;
 
     // The owner of token mint and token account must match

vault_program/src/burn_withdrawal_ticket.rs

@@ -13,11 +13,14 @@ use jito_vault_core::{
 use jito_vault_sdk::error::VaultError;
 use solana_program::{
     account_info::AccountInfo, clock::Clock, entrypoint::ProgramResult, msg,
-    program::invoke_signed, program_error::ProgramError, program_pack::Pack, pubkey::Pubkey,
-    sysvar::Sysvar,
+    program::invoke_signed, program_error::ProgramError, pubkey::Pubkey, sysvar::Sysvar,
+};
+use spl_token::instruction::transfer;
+use spl_token_2022::{
+    extension::StateWithExtensions,
+    instruction::{burn, close_account},
+    state::Account,
 };
-use spl_token::instruction::{burn, close_account, transfer};
-use spl_token_2022::state::Account;
 
 /// Burns the withdrawal ticket, transferring the assets to the staker and closing the withdrawal ticket.
 ///
@@ -61,17 +64,16 @@ pub fn process_burn_withdrawal_ticket(
         &vault.vrt_mint,
     )?;
 
-    let ticket_vrt_account =
-        Account::unpack(&vault_staker_withdrawal_ticket_token_account.data.borrow())?;
-    let ticket_vrt_amount = ticket_vrt_account.amount;
+    let ticket_vrt_data = vault_staker_withdrawal_ticket_token_account.data.borrow();
+    let ticket_vrt_account = StateWithExtensions::<Account>::unpack(&ticket_vrt_data)?;
+    let ticket_vrt_amount = ticket_vrt_account.base.amount;
 
     load_associated_token_account(vault_fee_token_account, &vault.fee_wallet, &vault.vrt_mint)?;
     load_associated_token_account(
         program_fee_token_account,
         &config.program_fee_wallet,
         &vault.vrt_mint,
     )?;
-    // Only the original spl token program is allowed
     load_token_program(token_program)?;
 
     load_system_program(system_program)?;
@@ -135,44 +137,52 @@ pub fn process_burn_withdrawal_ticket(
     // );
 
     // transfer fee to fee wallet
-    invoke_signed(
-        &transfer(
+    {
+        let mut ix = transfer(
             &spl_token::id(),
             vault_staker_withdrawal_ticket_token_account.key,
             vault_fee_token_account.key,
             vault_staker_withdrawal_ticket_info.key,
             &[],
             vault_fee_amount,
-        )?,
-        &[
-            vault_staker_withdrawal_ticket_token_account.clone(),
-            vault_fee_token_account.clone(),
-            vault_staker_withdrawal_ticket_info.clone(),
-        ],
-        &[&seed_slices],
-    )?;
+        )?;
+        ix.program_id = *token_program.key;
+        invoke_signed(
+            &ix,
+            &[
+                vault_staker_withdrawal_ticket_token_account.clone(),
+                vault_fee_token_account.clone(),
+                vault_staker_withdrawal_ticket_info.clone(),
+            ],
+            &[&seed_slices],
+        )?;
+    }
     // Transfer program fee to program fee wallet
-    invoke_signed(
-        &transfer(
+    {
+        let mut ix = transfer(
             &spl_token::id(),
             vault_staker_withdrawal_ticket_token_account.key,
             program_fee_token_account.key,
             vault_staker_withdrawal_ticket_info.key,
             &[],
             program_fee_amount,
-        )?,
-        &[
-            vault_staker_withdrawal_ticket_token_account.clone(),
-            program_fee_token_account.clone(),
-            vault_staker_withdrawal_ticket_info.clone(),
-        ],
-        &[&seed_slices],
-    )?;
+        )?;
+        ix.program_id = *token_program.key;
+        invoke_signed(
+            &ix,
+            &[
+                vault_staker_withdrawal_ticket_token_account.clone(),
+                program_fee_token_account.clone(),
+                vault_staker_withdrawal_ticket_info.clone(),
+            ],
+            &[&seed_slices],
+        )?;
+    }
 
     // burn the VRT tokens
     invoke_signed(
         &burn(
-            &spl_token::id(),
+            token_program.key,
             vault_staker_withdrawal_ticket_token_account.key,
             vrt_mint.key,
             vault_staker_withdrawal_ticket_info.key,
@@ -190,7 +200,7 @@ pub fn process_burn_withdrawal_ticket(
     // close token account
     invoke_signed(
         &close_account(
-            &spl_token::id(),
+            token_program.key,
             vault_staker_withdrawal_ticket_token_account.key,
             staker.key,
             vault_staker_withdrawal_ticket_info.key,
@@ -214,15 +224,17 @@ pub fn process_burn_withdrawal_ticket(
 
     drop(vault_data); // avoid double borrow
 
+    let mut ix = transfer(
+        &spl_token::id(),
+        vault_token_account.key,
+        staker_token_account.key,
+        vault_info.key,
+        &[],
+        out_amount,
+    )?;
+    ix.program_id = *token_program.key;
     invoke_signed(
-        &transfer(
-            &spl_token::id(),
-            vault_token_account.key,
-            staker_token_account.key,
-            vault_info.key,
-            &[],
-            out_amount,
-        )?,
+        &ix,
         &[
             vault_token_account.clone(),
             staker_token_account.clone(),

vault_program/src/delegate_token_account.rs

@@ -38,7 +38,6 @@ pub fn process_delegate_token_account(
         token_mint.key,
         token_program_info,
     )?;
-    // Only the original spl token program is allowed
     load_token_program(token_program_info)?;
 
     // The owner of token mint and token account must match

vault_program/src/enqueue_withdrawal.rs

@@ -124,15 +124,17 @@ pub fn process_enqueue_withdrawal(
 
     // Withdraw funds from the staker's VRT account, transferring them to an ATA owned
     // by the VaultStakerWithdrawalTicket
+    let mut ix = transfer(
+        &spl_token::id(),
+        staker_vrt_token_account.key,
+        vault_staker_withdrawal_ticket_token_account.key,
+        staker.key,
+        &[],
+        vrt_amount,
+    )?;
+    ix.program_id = *token_program.key;
     invoke(
-        &transfer(
-            &spl_token::id(),
-            staker_vrt_token_account.key,
-            vault_staker_withdrawal_ticket_token_account.key,
-            staker.key,
-            &[],
-            vrt_amount,
-        )?,
+        &ix,
         &[
             staker_vrt_token_account.clone(),
             vault_staker_withdrawal_ticket_token_account.clone(),