fix: improve token2022 support (#188)

Co-authored-by: BretasArthur1 <[email protected]>

Author: lgalabru

.cursor/mcp.json

@@ -1,5 +1,5 @@
 [workspace.package]
-version = "0.8.0"
+version = "0.8.1"
 edition = "2021"
 description = "Surfpool is the best place to train before surfing Solana."
 license = "Apache-2.0"

Cargo.lock

@@ -70,8 +70,8 @@ solana-version = { workspace = true }
 spl-associated-token-account = { workspace = true }
 surfpool-subgraph = { workspace = true }
 surfpool-types = { workspace = true }
-# surfpool-types = { version = "0.2.0", default-features = false }
-# surfpool-subgraph = { version = "0.2.0", default-features = false }
+# surfpool-types = { version = "0.8.1", default-features = false }
+# surfpool-subgraph = { version = "0.8.1", default-features = false }
 symlink = "0.1.0"
 tokio = { version = "1.45.0", features = ["full"] }
 tokio-util = "0.7.13"

Cargo.toml

@@ -20,7 +20,7 @@ use super::{not_implemented_err, RunloopContext, SurfnetRpcContext};
 use crate::{
     error::{SurfpoolError, SurfpoolResult},
     rpc::{utils::verify_pubkey, State},
-    surfnet::locker::SvmAccessContext,
+    surfnet::{locker::SvmAccessContext, GetAccountResult},
 };
 
 #[rpc]
@@ -377,12 +377,20 @@ impl AccountsData for SurfpoolAccountsDataRpc {
                 inner: account_update,
                 ..
             } = svm_locker.get_account(&remote_ctx, &pubkey, None).await?;
-
             svm_locker.write_account_update(account_update.clone());
 
+            let SvmAccessContext {
+                inner: associated_data,
+                ..
+            } = svm_locker.get_local_account_associated_data(&account_update);
+
             Ok(RpcResponse {
                 context: RpcResponseContext::new(slot),
-                value: account_update.try_into_ui_account(config.encoding, config.data_slice),
+                value: account_update.try_into_ui_account(
+                    config.encoding,
+                    config.data_slice,
+                    associated_data,
+                ),
             })
         })
     }
@@ -425,9 +433,11 @@ impl AccountsData for SurfpoolAccountsDataRpc {
             let mut ui_accounts = vec![];
             {
                 for account_update in account_updates.into_iter() {
-                    ui_accounts.push(
-                        account_update.try_into_ui_account(config.encoding, config.data_slice),
-                    );
+                    ui_accounts.push(account_update.try_into_ui_account(
+                        config.encoding,
+                        config.data_slice,
+                        None,
+                    ));
                 }
             }
 
@@ -648,13 +658,30 @@ impl AccountsData for SurfpoolAccountsDataRpc {
 #[cfg(test)]
 mod tests {
     use solana_account::Account;
+    use solana_client::rpc_client::RpcClient;
+    use solana_keypair::Keypair;
     use solana_pubkey::Pubkey;
-    use solana_sdk::{program_option::COption, program_pack::Pack};
+    use solana_sdk::{
+        program_option::COption, program_pack::Pack, system_instruction::create_account,
+    };
+    use solana_signer::Signer;
+    use solana_transaction::Transaction;
+    use spl_associated_token_account::{
+        get_associated_token_address_with_program_id, instruction::create_associated_token_account,
+    };
     use spl_token::state::{Account as TokenAccount, AccountState, Mint};
+    use spl_token_2022::{
+        extension::StateWithExtensions,
+        instruction::{initialize_mint2, mint_to, transfer_checked},
+        state::Account as Token2022Account,
+    };
 
     use super::*;
     use crate::{
-        surfnet::{remote::SurfnetRemoteClient, GetAccountResult},
+        surfnet::{
+            remote::{SomeRemoteCtx, SurfnetRemoteClient},
+            GetAccountResult,
+        },
         tests::helpers::TestSetup,
     };
 
@@ -1019,4 +1046,210 @@ mod tests {
         let error_msg = res.unwrap_err().to_string();
         println!("✅ Remote RPC failure handled: {}", error_msg);
     }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_transfer_token() {
+        // Create connection to local validator
+        let client = TestSetup::new(SurfpoolAccountsDataRpc);
+        let recent_blockhash = client
+            .context
+            .svm_locker
+            .with_svm_reader(|svm_reader| svm_reader.latest_blockhash());
+
+        // Generate a new keypair for the fee payer
+        let fee_payer = Keypair::new();
+
+        // Generate a second keypair for the token recipient
+        let recipient = Keypair::new();
+
+        // Airdrop 1 SOL to fee payer
+        client
+            .context
+            .svm_locker
+            .airdrop(&fee_payer.pubkey(), 1_000_000_000)
+            .unwrap();
+
+        // Airdrop 1 SOL to recipient for rent exemption
+        client
+            .context
+            .svm_locker
+            .airdrop(&recipient.pubkey(), 1_000_000_000)
+            .unwrap();
+
+        // Generate keypair to use as address of mint
+        let mint = Keypair::new();
+
+        // Get default mint account size (in bytes), no extensions enabled
+        let mint_space = Mint::LEN;
+        let mint_rent = client.context.svm_locker.with_svm_reader(|svm_reader| {
+            svm_reader
+                .inner
+                .minimum_balance_for_rent_exemption(mint_space)
+        });
+
+        // Instruction to create new account for mint (token 2022 program)
+        let create_account_instruction = create_account(
+            &fee_payer.pubkey(),   // payer
+            &mint.pubkey(),        // new account (mint)
+            mint_rent,             // lamports
+            mint_space as u64,     // space
+            &spl_token_2022::id(), // program id
+        );
+
+        // Instruction to initialize mint account data
+        let initialize_mint_instruction = initialize_mint2(
+            &spl_token_2022::id(),
+            &mint.pubkey(),            // mint
+            &fee_payer.pubkey(),       // mint authority
+            Some(&fee_payer.pubkey()), // freeze authority
+            2,                         // decimals
+        )
+        .unwrap();
+
+        // Calculate the associated token account address for fee_payer
+        let source_token_address = get_associated_token_address_with_program_id(
+            &fee_payer.pubkey(),   // owner
+            &mint.pubkey(),        // mint
+            &spl_token_2022::id(), // program_id
+        );
+
+        // Instruction to create associated token account for fee_payer
+        let create_source_ata_instruction = create_associated_token_account(
+            &fee_payer.pubkey(),   // funding address
+            &fee_payer.pubkey(),   // wallet address
+            &mint.pubkey(),        // mint address
+            &spl_token_2022::id(), // program id
+        );
+
+        // Calculate the associated token account address for recipient
+        let destination_token_address = get_associated_token_address_with_program_id(
+            &recipient.pubkey(),   // owner
+            &mint.pubkey(),        // mint
+            &spl_token_2022::id(), // program_id
+        );
+
+        // Instruction to create associated token account for recipient
+        let create_destination_ata_instruction = create_associated_token_account(
+            &fee_payer.pubkey(),   // funding address
+            &recipient.pubkey(),   // wallet address
+            &mint.pubkey(),        // mint address
+            &spl_token_2022::id(), // program id
+        );
+
+        // Amount of tokens to mint (100 tokens with 2 decimal places)
+        let amount = 100_00;
+
+        // Create mint_to instruction to mint tokens to the source token account
+        let mint_to_instruction = mint_to(
+            &spl_token_2022::id(),
+            &mint.pubkey(),         // mint
+            &source_token_address,  // destination
+            &fee_payer.pubkey(),    // authority
+            &[&fee_payer.pubkey()], // signer
+            amount,                 // amount
+        )
+        .unwrap();
+
+        // Create transaction and add instructions
+        let transaction = Transaction::new_signed_with_payer(
+            &[
+                create_account_instruction,
+                initialize_mint_instruction,
+                create_source_ata_instruction,
+                create_destination_ata_instruction,
+                mint_to_instruction,
+            ],
+            Some(&fee_payer.pubkey()),
+            &[&fee_payer, &mint],
+            recent_blockhash,
+        );
+
+        // Send and confirm transaction
+        client.context.svm_locker.with_svm_writer(|svm_writer| {
+            svm_writer
+                .send_transaction(transaction.clone().into(), false)
+                .unwrap();
+        });
+
+        println!("Mint Address: {}", mint.pubkey());
+        println!("Recipient Address: {}", recipient.pubkey());
+        println!("Source Token Account Address: {}", source_token_address);
+        println!(
+            "Destination Token Account Address: {}",
+            destination_token_address
+        );
+        println!("Minted {} tokens to the source token account", amount);
+
+        // Get the latest blockhash for the transfer transaction
+        let recent_blockhash = client
+            .context
+            .svm_locker
+            .with_svm_reader(|svm_reader| svm_reader.latest_blockhash());
+
+        // Amount of tokens to transfer (0.50 tokens with 2 decimals)
+        let transfer_amount = 50;
+
+        // Create transfer_checked instruction to send tokens from source to destination
+        let transfer_instruction = transfer_checked(
+            &spl_token_2022::id(),      // program id
+            &source_token_address,      // source
+            &mint.pubkey(),             // mint
+            &destination_token_address, // destination
+            &fee_payer.pubkey(),        // owner of source
+            &[&fee_payer.pubkey()],     // signers
+            transfer_amount,            // amount
+            2,                          // decimals
+        )
+        .unwrap();
+
+        // Create transaction for transferring tokens
+        let transaction = Transaction::new_signed_with_payer(
+            &[transfer_instruction],
+            Some(&fee_payer.pubkey()),
+            &[&fee_payer],
+            recent_blockhash,
+        );
+
+        // Send and confirm transaction
+        client.context.svm_locker.with_svm_writer(|svm_writer| {
+            svm_writer
+                .send_transaction(transaction.clone().into(), false)
+                .unwrap();
+        });
+
+        println!("Successfully transferred 0.50 tokens from sender to recipient");
+
+        // Get token account balances to verify the transfer
+        let source_token_account = client
+            .context
+            .svm_locker
+            .get_account_local(&source_token_address)
+            .inner;
+        let destination_token_account = client
+            .context
+            .svm_locker
+            .get_account_local(&destination_token_address)
+            .inner;
+
+        if let GetAccountResult::FoundAccount(_, source_account, _) = source_token_account {
+            let unpacked =
+                StateWithExtensions::<Token2022Account>::unpack(&source_account.data).unwrap();
+            println!(
+                "Source Token Account Balance: {} tokens",
+                unpacked.base.amount
+            );
+            assert_eq!(unpacked.base.amount, 9950);
+        }
+
+        if let GetAccountResult::FoundAccount(_, destination_account, _) = destination_token_account
+        {
+            let unpacked =
+                StateWithExtensions::<Token2022Account>::unpack(&destination_account.data).unwrap();
+            println!(
+                "Destination Token Account Balance: {} tokens",
+                unpacked.base.amount
+            );
+            assert_eq!(unpacked.base.amount, 50);
+        }
+    }
 }