diff --git a/associated-token-account/program/src/instruction.rs b/associated-token-account/program/src/instruction.rs index f3f1f67f..cdbe2190 100644 --- a/associated-token-account/program/src/instruction.rs +++ b/associated-token-account/program/src/instruction.rs @@ -34,6 +34,25 @@ pub enum AssociatedTokenAccountInstruction { /// 4. `[]` System program /// 5. `[]` SPL Token program CreateIdempotent, + /// Transfers from and closes a nested associated token account: an + /// associated token account owned by an associated token account. + /// + /// The tokens are moved from the nested associated token account to the + /// wallet's associated token account, and the nested account lamports are + /// moved to the wallet. + /// + /// Note: Nested token accounts are an anti-pattern, and almost always + /// created unintentionally, so this instruction should only be used to + /// recover from errors. + /// + /// 0. `[writeable]` Nested associated token account, must be owned by `3` + /// 1. `[]` Token mint for the nested associated token account + /// 2. `[writeable]` Wallet's associated token account + /// 3. `[]` Owner associated token account address, must be owned by `5` + /// 4. `[]` Token mint for the owner associated token account + /// 5. `[writeable, signer]` Wallet address for the owner associated token account + /// 6. `[]` SPL Token program + RecoverNested, } fn build_associated_token_account_instruction( @@ -99,3 +118,43 @@ pub fn create_associated_token_account_idempotent( AssociatedTokenAccountInstruction::CreateIdempotent, ) } + +/// Creates a `RecoverNested` instruction +pub fn recover_nested( + wallet_address: &Pubkey, + owner_token_mint_address: &Pubkey, + nested_token_mint_address: &Pubkey, + token_program_id: &Pubkey, +) -> Instruction { + let owner_associated_account_address = get_associated_token_address_with_program_id( + wallet_address, + owner_token_mint_address, + token_program_id, + ); + let destination_associated_account_address = get_associated_token_address_with_program_id( + wallet_address, + nested_token_mint_address, + token_program_id, + ); + let nested_associated_account_address = get_associated_token_address_with_program_id( + &owner_associated_account_address, // ATA is wrongly used as a wallet_address + nested_token_mint_address, + token_program_id, + ); + + let instruction_data = AssociatedTokenAccountInstruction::RecoverNested; + + Instruction { + program_id: id(), + accounts: vec![ + AccountMeta::new(nested_associated_account_address, false), + AccountMeta::new_readonly(*nested_token_mint_address, false), + AccountMeta::new(destination_associated_account_address, false), + AccountMeta::new_readonly(owner_associated_account_address, false), + AccountMeta::new_readonly(*owner_token_mint_address, false), + AccountMeta::new(*wallet_address, true), + AccountMeta::new_readonly(*token_program_id, false), + ], + data: instruction_data.try_to_vec().unwrap(), + } +} diff --git a/associated-token-account/program/src/processor.rs b/associated-token-account/program/src/processor.rs index 58f1e022..20767a24 100644 --- a/associated-token-account/program/src/processor.rs +++ b/associated-token-account/program/src/processor.rs @@ -12,7 +12,7 @@ use { account_info::{next_account_info, AccountInfo}, entrypoint::ProgramResult, msg, - program::invoke, + program::{invoke, invoke_signed}, program_error::ProgramError, pubkey::Pubkey, rent::Rent, @@ -21,7 +21,7 @@ use { }, spl_token_2022::{ extension::{ExtensionType, StateWithExtensions}, - state::Account, + state::{Account, Mint}, }, }; @@ -56,6 +56,9 @@ pub fn process_instruction( AssociatedTokenAccountInstruction::CreateIdempotent => { process_create_associated_token_account(program_id, accounts, CreateMode::Idempotent) } + AssociatedTokenAccountInstruction::RecoverNested => { + process_recover_nested(program_id, accounts) + } } } @@ -157,3 +160,150 @@ fn process_create_associated_token_account( ], ) } + +/// Processes `RecoverNested` instruction +pub fn process_recover_nested(program_id: &Pubkey, accounts: &[AccountInfo]) -> ProgramResult { + let account_info_iter = &mut accounts.iter(); + + let nested_associated_token_account_info = next_account_info(account_info_iter)?; + let nested_token_mint_info = next_account_info(account_info_iter)?; + let destination_associated_token_account_info = next_account_info(account_info_iter)?; + let owner_associated_token_account_info = next_account_info(account_info_iter)?; + let owner_token_mint_info = next_account_info(account_info_iter)?; + let wallet_account_info = next_account_info(account_info_iter)?; + let spl_token_program_info = next_account_info(account_info_iter)?; + let spl_token_program_id = spl_token_program_info.key; + + // Check owner address derivation + let (owner_associated_token_address, bump_seed) = + get_associated_token_address_and_bump_seed_internal( + wallet_account_info.key, + owner_token_mint_info.key, + program_id, + spl_token_program_id, + ); + if owner_associated_token_address != *owner_associated_token_account_info.key { + msg!("Error: Owner associated address does not match seed derivation"); + return Err(ProgramError::InvalidSeeds); + } + + // Check nested address derivation + let (nested_associated_token_address, _) = get_associated_token_address_and_bump_seed_internal( + owner_associated_token_account_info.key, + nested_token_mint_info.key, + program_id, + spl_token_program_id, + ); + if nested_associated_token_address != *nested_associated_token_account_info.key { + msg!("Error: Nested associated address does not match seed derivation"); + return Err(ProgramError::InvalidSeeds); + } + + // Check destination address derivation + let (destination_associated_token_address, _) = + get_associated_token_address_and_bump_seed_internal( + wallet_account_info.key, + nested_token_mint_info.key, + program_id, + spl_token_program_id, + ); + if destination_associated_token_address != *destination_associated_token_account_info.key { + msg!("Error: Destination associated address does not match seed derivation"); + return Err(ProgramError::InvalidSeeds); + } + + if !wallet_account_info.is_signer { + msg!("Wallet of the owner associated token account must sign"); + return Err(ProgramError::MissingRequiredSignature); + } + + if owner_token_mint_info.owner != spl_token_program_id { + msg!("Owner mint not owned by provided token program"); + return Err(ProgramError::IllegalOwner); + } + + // Account data is dropped at the end of this, so the CPI can succeed + // without a double-borrow + let (amount, decimals) = { + // Check owner associated token account data + if owner_associated_token_account_info.owner != spl_token_program_id { + msg!("Owner associated token account not owned by provided token program, recreate the owner associated token account first"); + return Err(ProgramError::IllegalOwner); + } + let owner_account_data = owner_associated_token_account_info.data.borrow(); + let owner_account = StateWithExtensions::::unpack(&owner_account_data)?; + if owner_account.base.owner != *wallet_account_info.key { + msg!("Owner associated token account not owned by provided wallet"); + return Err(AssociatedTokenAccountError::InvalidOwner.into()); + } + + // Check nested associated token account data + if nested_associated_token_account_info.owner != spl_token_program_id { + msg!("Nested associated token account not owned by provided token program"); + return Err(ProgramError::IllegalOwner); + } + let nested_account_data = nested_associated_token_account_info.data.borrow(); + let nested_account = StateWithExtensions::::unpack(&nested_account_data)?; + if nested_account.base.owner != *owner_associated_token_account_info.key { + msg!("Nested associated token account not owned by provided associated token account"); + return Err(AssociatedTokenAccountError::InvalidOwner.into()); + } + let amount = nested_account.base.amount; + + // Check nested token mint data + if nested_token_mint_info.owner != spl_token_program_id { + msg!("Nested mint account not owned by provided token program"); + return Err(ProgramError::IllegalOwner); + } + let nested_mint_data = nested_token_mint_info.data.borrow(); + let nested_mint = StateWithExtensions::::unpack(&nested_mint_data)?; + let decimals = nested_mint.base.decimals; + (amount, decimals) + }; + + // Transfer everything out + let owner_associated_token_account_signer_seeds: &[&[_]] = &[ + &wallet_account_info.key.to_bytes(), + &spl_token_program_id.to_bytes(), + &owner_token_mint_info.key.to_bytes(), + &[bump_seed], + ]; + invoke_signed( + &spl_token_2022::instruction::transfer_checked( + spl_token_program_id, + nested_associated_token_account_info.key, + nested_token_mint_info.key, + destination_associated_token_account_info.key, + owner_associated_token_account_info.key, + &[], + amount, + decimals, + )?, + &[ + nested_associated_token_account_info.clone(), + nested_token_mint_info.clone(), + destination_associated_token_account_info.clone(), + owner_associated_token_account_info.clone(), + spl_token_program_info.clone(), + ], + &[owner_associated_token_account_signer_seeds], + )?; + + // Close the nested account so it's never used again + invoke_signed( + &spl_token_2022::instruction::close_account( + spl_token_program_id, + nested_associated_token_account_info.key, + wallet_account_info.key, + owner_associated_token_account_info.key, + &[], + )?, + &[ + nested_associated_token_account_info.clone(), + wallet_account_info.clone(), + owner_associated_token_account_info.clone(), + spl_token_program_info.clone(), + ], + &[owner_associated_token_account_signer_seeds], + ) +} diff --git a/associated-token-account/program/tests/create_idempotent.rs b/associated-token-account/program/tests/create_idempotent.rs index e5e7a187..9d2c29f9 100644 --- a/associated-token-account/program/tests/create_idempotent.rs +++ b/associated-token-account/program/tests/create_idempotent.rs @@ -3,7 +3,7 @@ mod program_test; use { - program_test::program_test, + program_test::program_test_2022, solana_program::{instruction::*, pubkey::Pubkey}, solana_program_test::*, solana_sdk::{ @@ -40,7 +40,7 @@ async fn success_account_exists() { ); let (mut banks_client, payer, recent_blockhash) = - program_test(token_mint_address, true).start().await; + program_test_2022(token_mint_address, true).start().await; let rent = banks_client.get_rent().await.unwrap(); let expected_token_account_len = ExtensionType::get_account_len::(&[ExtensionType::ImmutableOwner]); @@ -150,7 +150,7 @@ async fn fail_account_exists_with_wrong_owner() { close_authority: COption::None, }; Account::pack(token_account, &mut associated_token_account.data).unwrap(); - let mut pt = program_test(token_mint_address, true); + let mut pt = program_test_2022(token_mint_address, true); pt.add_account(associated_token_address, associated_token_account); let (mut banks_client, payer, recent_blockhash) = pt.start().await; @@ -185,7 +185,7 @@ async fn fail_account_exists_with_wrong_owner() { async fn fail_non_ata() { let token_mint_address = Pubkey::new_unique(); let (mut banks_client, payer, recent_blockhash) = - program_test(token_mint_address, true).start().await; + program_test_2022(token_mint_address, true).start().await; let rent = banks_client.get_rent().await.unwrap(); let token_account_len = diff --git a/associated-token-account/program/tests/extended_mint.rs b/associated-token-account/program/tests/extended_mint.rs index 9c427e19..cc131c0d 100644 --- a/associated-token-account/program/tests/extended_mint.rs +++ b/associated-token-account/program/tests/extended_mint.rs @@ -4,7 +4,7 @@ mod program_test; use { - program_test::program_test, + program_test::program_test_2022, solana_program::{instruction::*, pubkey::Pubkey, system_instruction}, solana_program_test::*, solana_sdk::{ @@ -29,7 +29,7 @@ async fn test_associated_token_account_with_transfer_fees() { let wallet_receiver = Keypair::new(); let wallet_address_receiver = wallet_receiver.pubkey(); let (mut banks_client, payer, recent_blockhash) = - program_test(Pubkey::new_unique(), true).start().await; + program_test_2022(Pubkey::new_unique(), true).start().await; let rent = banks_client.get_rent().await.unwrap(); // create extended mint diff --git a/associated-token-account/program/tests/process_create_associated_token_account.rs b/associated-token-account/program/tests/process_create_associated_token_account.rs index a46eddb6..c988e6e2 100644 --- a/associated-token-account/program/tests/process_create_associated_token_account.rs +++ b/associated-token-account/program/tests/process_create_associated_token_account.rs @@ -4,7 +4,7 @@ mod program_test; use { - program_test::program_test, + program_test::program_test_2022, solana_program::{instruction::*, pubkey::Pubkey, system_instruction, sysvar}, solana_program_test::*, solana_sdk::{ @@ -28,7 +28,7 @@ async fn test_associated_token_address() { ); let (mut banks_client, payer, recent_blockhash) = - program_test(token_mint_address, true).start().await; + program_test_2022(token_mint_address, true).start().await; let rent = banks_client.get_rent().await.unwrap(); let expected_token_account_len = @@ -78,7 +78,7 @@ async fn test_create_with_fewer_lamports() { ); let (mut banks_client, payer, recent_blockhash) = - program_test(token_mint_address, true).start().await; + program_test_2022(token_mint_address, true).start().await; let rent = banks_client.get_rent().await.unwrap(); let expected_token_account_len = ExtensionType::get_account_len::(&[ExtensionType::ImmutableOwner]); @@ -138,7 +138,7 @@ async fn test_create_with_excess_lamports() { ); let (mut banks_client, payer, recent_blockhash) = - program_test(token_mint_address, true).start().await; + program_test_2022(token_mint_address, true).start().await; let rent = banks_client.get_rent().await.unwrap(); let expected_token_account_len = @@ -198,7 +198,7 @@ async fn test_create_account_mismatch() { ); let (mut banks_client, payer, recent_blockhash) = - program_test(token_mint_address, true).start().await; + program_test_2022(token_mint_address, true).start().await; let mut instruction = create_associated_token_account( &payer.pubkey(), @@ -269,7 +269,7 @@ async fn test_create_associated_token_account_using_legacy_implicit_instruction( ); let (mut banks_client, payer, recent_blockhash) = - program_test(token_mint_address, true).start().await; + program_test_2022(token_mint_address, true).start().await; let rent = banks_client.get_rent().await.unwrap(); let expected_token_account_len = ExtensionType::get_account_len::(&[ExtensionType::ImmutableOwner]); diff --git a/associated-token-account/program/tests/program_test.rs b/associated-token-account/program/tests/program_test.rs index 134ccfde..f72f378c 100644 --- a/associated-token-account/program/tests/program_test.rs +++ b/associated-token-account/program/tests/program_test.rs @@ -4,6 +4,7 @@ use { spl_associated_token_account::{id, processor::process_instruction}, }; +#[allow(dead_code)] pub fn program_test(token_mint_address: Pubkey, use_latest_spl_token: bool) -> ProgramTest { let mut pc = ProgramTest::new( "spl_associated_token_account", @@ -12,6 +13,45 @@ pub fn program_test(token_mint_address: Pubkey, use_latest_spl_token: bool) -> P ); if use_latest_spl_token { + // TODO: Remove when spl-token is available by default in program-test + pc.add_program( + "spl_token", + spl_token::id(), + processor!(spl_token::processor::Processor::process), + ); + } + + // Add a token mint account + // + // The account data was generated by running: + // $ solana account EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v \ + // --output-file tests/fixtures/token-mint-data.bin + // + pc.add_account_with_file_data( + token_mint_address, + 1461600, + spl_token::id(), + "token-mint-data.bin", + ); + + // Dial down the BPF compute budget to detect if the program gets bloated in the future + pc.set_compute_max_units(50_000); + + pc +} + +#[allow(dead_code)] +pub fn program_test_2022( + token_mint_address: Pubkey, + use_latest_spl_token_2022: bool, +) -> ProgramTest { + let mut pc = ProgramTest::new( + "spl_associated_token_account", + id(), + processor!(process_instruction), + ); + + if use_latest_spl_token_2022 { // TODO: Remove when spl-token-2022 is available by default in program-test pc.add_program( "spl_token_2022", diff --git a/associated-token-account/program/tests/recover_nested.rs b/associated-token-account/program/tests/recover_nested.rs new file mode 100644 index 00000000..d86c59eb --- /dev/null +++ b/associated-token-account/program/tests/recover_nested.rs @@ -0,0 +1,632 @@ +// Mark this test as BPF-only due to current `ProgramTest` limitations when CPIing into the system program +#![cfg(feature = "test-bpf")] + +mod program_test; + +use { + program_test::{program_test, program_test_2022}, + solana_program::{pubkey::Pubkey, system_instruction}, + solana_program_test::*, + solana_sdk::{ + instruction::{AccountMeta, InstructionError}, + signature::Signer, + signer::keypair::Keypair, + transaction::{Transaction, TransactionError}, + }, + spl_associated_token_account::{get_associated_token_address_with_program_id, instruction}, + spl_token_2022::{ + extension::{ExtensionType, StateWithExtensionsOwned}, + state::{Account, Mint}, + }, +}; + +async fn create_mint(context: &mut ProgramTestContext, program_id: &Pubkey) -> (Pubkey, Keypair) { + let mint_account = Keypair::new(); + let token_mint_address = mint_account.pubkey(); + let mint_authority = Keypair::new(); + let space = ExtensionType::get_account_len::(&[]); + let rent = context.banks_client.get_rent().await.unwrap(); + let transaction = Transaction::new_signed_with_payer( + &[ + system_instruction::create_account( + &context.payer.pubkey(), + &mint_account.pubkey(), + rent.minimum_balance(space), + space as u64, + program_id, + ), + spl_token_2022::instruction::initialize_mint( + program_id, + &token_mint_address, + &mint_authority.pubkey(), + Some(&mint_authority.pubkey()), + 0, + ) + .unwrap(), + ], + Some(&context.payer.pubkey()), + &[&context.payer, &mint_account], + context.last_blockhash, + ); + context + .banks_client + .process_transaction(transaction) + .await + .unwrap(); + (token_mint_address, mint_authority) +} + +async fn create_associated_token_account( + context: &mut ProgramTestContext, + owner: &Pubkey, + mint: &Pubkey, + program_id: &Pubkey, +) -> Pubkey { + let transaction = Transaction::new_signed_with_payer( + &[instruction::create_associated_token_account( + &context.payer.pubkey(), + owner, + mint, + program_id, + )], + Some(&context.payer.pubkey()), + &[&context.payer], + context.last_blockhash, + ); + context + .banks_client + .process_transaction(transaction) + .await + .unwrap(); + + get_associated_token_address_with_program_id(owner, mint, program_id) +} + +#[allow(clippy::too_many_arguments)] +async fn try_recover_nested( + context: &mut ProgramTestContext, + program_id: &Pubkey, + nested_mint: Pubkey, + nested_mint_authority: Keypair, + nested_associated_token_address: Pubkey, + destination_token_address: Pubkey, + wallet: Keypair, + recover_transaction: Transaction, + expected_error: Option, +) { + let nested_account = context + .banks_client + .get_account(nested_associated_token_address) + .await + .unwrap() + .unwrap(); + let lamports = nested_account.lamports; + + // mint to nested account + let amount = 100; + let transaction = Transaction::new_signed_with_payer( + &[spl_token_2022::instruction::mint_to( + program_id, + &nested_mint, + &nested_associated_token_address, + &nested_mint_authority.pubkey(), + &[], + amount, + ) + .unwrap()], + Some(&context.payer.pubkey()), + &[&context.payer, &nested_mint_authority], + context.last_blockhash, + ); + context + .banks_client + .process_transaction(transaction) + .await + .unwrap(); + + // transfer / close nested account + let result = context + .banks_client + .process_transaction(recover_transaction) + .await; + + if let Some(expected_error) = expected_error { + let error = result.unwrap_err().unwrap(); + assert_eq!(error, TransactionError::InstructionError(0, expected_error)); + } else { + result.unwrap(); + // nested account is gone + assert!(context + .banks_client + .get_account(nested_associated_token_address) + .await + .unwrap() + .is_none()); + let destination_account = context + .banks_client + .get_account(destination_token_address) + .await + .unwrap() + .unwrap(); + let destination_state = + StateWithExtensionsOwned::::unpack(destination_account.data).unwrap(); + assert_eq!(destination_state.base.amount, amount); + let wallet_account = context + .banks_client + .get_account(wallet.pubkey()) + .await + .unwrap() + .unwrap(); + assert_eq!(wallet_account.lamports, lamports); + } +} + +async fn check_same_mint(context: &mut ProgramTestContext, program_id: &Pubkey) { + let wallet = Keypair::new(); + let (mint, mint_authority) = create_mint(context, program_id).await; + + let owner_associated_token_address = + create_associated_token_account(context, &wallet.pubkey(), &mint, program_id).await; + let nested_associated_token_address = create_associated_token_account( + context, + &owner_associated_token_address, + &mint, + program_id, + ) + .await; + + let transaction = Transaction::new_signed_with_payer( + &[instruction::recover_nested( + &wallet.pubkey(), + &mint, + &mint, + program_id, + )], + Some(&context.payer.pubkey()), + &[&context.payer, &wallet], + context.last_blockhash, + ); + try_recover_nested( + context, + program_id, + mint, + mint_authority, + nested_associated_token_address, + owner_associated_token_address, + wallet, + transaction, + None, + ) + .await; +} + +#[tokio::test] +async fn success_same_mint_2022() { + let dummy_mint = Pubkey::new_unique(); + let pt = program_test_2022(dummy_mint, true); + let mut context = pt.start_with_context().await; + check_same_mint(&mut context, &spl_token_2022::id()).await; +} + +#[tokio::test] +async fn success_same_mint() { + let dummy_mint = Pubkey::new_unique(); + let pt = program_test(dummy_mint, true); + let mut context = pt.start_with_context().await; + check_same_mint(&mut context, &spl_token::id()).await; +} + +async fn check_different_mints(context: &mut ProgramTestContext, program_id: &Pubkey) { + let wallet = Keypair::new(); + let (owner_mint, _owner_mint_authority) = create_mint(context, program_id).await; + let (nested_mint, nested_mint_authority) = create_mint(context, program_id).await; + + let owner_associated_token_address = + create_associated_token_account(context, &wallet.pubkey(), &owner_mint, program_id).await; + let nested_associated_token_address = create_associated_token_account( + context, + &owner_associated_token_address, + &nested_mint, + program_id, + ) + .await; + let destination_token_address = + create_associated_token_account(context, &wallet.pubkey(), &nested_mint, program_id).await; + + let transaction = Transaction::new_signed_with_payer( + &[instruction::recover_nested( + &wallet.pubkey(), + &owner_mint, + &nested_mint, + program_id, + )], + Some(&context.payer.pubkey()), + &[&context.payer, &wallet], + context.last_blockhash, + ); + try_recover_nested( + context, + program_id, + nested_mint, + nested_mint_authority, + nested_associated_token_address, + destination_token_address, + wallet, + transaction, + None, + ) + .await; +} + +#[tokio::test] +async fn success_different_mints() { + let dummy_mint = Pubkey::new_unique(); + let pt = program_test(dummy_mint, true); + let mut context = pt.start_with_context().await; + check_different_mints(&mut context, &spl_token::id()).await; +} + +#[tokio::test] +async fn success_different_mints_2022() { + let dummy_mint = Pubkey::new_unique(); + let pt = program_test_2022(dummy_mint, true); + let mut context = pt.start_with_context().await; + check_different_mints(&mut context, &spl_token_2022::id()).await; +} + +async fn check_missing_wallet_signature(context: &mut ProgramTestContext, program_id: &Pubkey) { + let wallet = Keypair::new(); + let (mint, mint_authority) = create_mint(context, program_id).await; + + let owner_associated_token_address = + create_associated_token_account(context, &wallet.pubkey(), &mint, program_id).await; + + let nested_associated_token_address = create_associated_token_account( + context, + &owner_associated_token_address, + &mint, + program_id, + ) + .await; + + let mut recover = instruction::recover_nested(&wallet.pubkey(), &mint, &mint, program_id); + recover.accounts[5] = AccountMeta::new(wallet.pubkey(), false); + let transaction = Transaction::new_signed_with_payer( + &[recover], + Some(&context.payer.pubkey()), + &[&context.payer], + context.last_blockhash, + ); + try_recover_nested( + context, + program_id, + mint, + mint_authority, + nested_associated_token_address, + owner_associated_token_address, + wallet, + transaction, + Some(InstructionError::MissingRequiredSignature), + ) + .await; +} + +#[tokio::test] +async fn fail_missing_wallet_signature_2022() { + let dummy_mint = Pubkey::new_unique(); + let pt = program_test_2022(dummy_mint, true); + let mut context = pt.start_with_context().await; + check_missing_wallet_signature(&mut context, &spl_token_2022::id()).await; +} + +#[tokio::test] +async fn fail_missing_wallet_signature() { + let dummy_mint = Pubkey::new_unique(); + let pt = program_test(dummy_mint, true); + let mut context = pt.start_with_context().await; + check_missing_wallet_signature(&mut context, &spl_token::id()).await; +} + +async fn check_wrong_signer(context: &mut ProgramTestContext, program_id: &Pubkey) { + let wallet = Keypair::new(); + let wrong_wallet = Keypair::new(); + let (mint, mint_authority) = create_mint(context, program_id).await; + + let owner_associated_token_address = + create_associated_token_account(context, &wallet.pubkey(), &mint, program_id).await; + let nested_associated_token_address = create_associated_token_account( + context, + &owner_associated_token_address, + &mint, + program_id, + ) + .await; + + let transaction = Transaction::new_signed_with_payer( + &[instruction::recover_nested( + &wrong_wallet.pubkey(), + &mint, + &mint, + program_id, + )], + Some(&context.payer.pubkey()), + &[&context.payer, &wrong_wallet], + context.last_blockhash, + ); + try_recover_nested( + context, + program_id, + mint, + mint_authority, + nested_associated_token_address, + owner_associated_token_address, + wrong_wallet, + transaction, + Some(InstructionError::IllegalOwner), + ) + .await; +} + +#[tokio::test] +async fn fail_wrong_signer_2022() { + let dummy_mint = Pubkey::new_unique(); + let pt = program_test_2022(dummy_mint, true); + let mut context = pt.start_with_context().await; + check_wrong_signer(&mut context, &spl_token_2022::id()).await; +} + +#[tokio::test] +async fn fail_wrong_signer() { + let dummy_mint = Pubkey::new_unique(); + let pt = program_test(dummy_mint, true); + let mut context = pt.start_with_context().await; + check_wrong_signer(&mut context, &spl_token::id()).await; +} + +async fn check_not_nested(context: &mut ProgramTestContext, program_id: &Pubkey) { + let wallet = Keypair::new(); + let wrong_wallet = Keypair::new(); + let (mint, mint_authority) = create_mint(context, program_id).await; + + let owner_associated_token_address = + create_associated_token_account(context, &wallet.pubkey(), &mint, program_id).await; + let nested_associated_token_address = + create_associated_token_account(context, &wrong_wallet.pubkey(), &mint, program_id).await; + + let transaction = Transaction::new_signed_with_payer( + &[instruction::recover_nested( + &wallet.pubkey(), + &mint, + &mint, + program_id, + )], + Some(&context.payer.pubkey()), + &[&context.payer, &wallet], + context.last_blockhash, + ); + try_recover_nested( + context, + program_id, + mint, + mint_authority, + nested_associated_token_address, + owner_associated_token_address, + wallet, + transaction, + Some(InstructionError::IllegalOwner), + ) + .await; +} + +#[tokio::test] +async fn fail_not_nested_2022() { + let dummy_mint = Pubkey::new_unique(); + let pt = program_test_2022(dummy_mint, true); + let mut context = pt.start_with_context().await; + check_not_nested(&mut context, &spl_token_2022::id()).await; +} + +#[tokio::test] +async fn fail_not_nested() { + let dummy_mint = Pubkey::new_unique(); + let pt = program_test(dummy_mint, true); + let mut context = pt.start_with_context().await; + check_not_nested(&mut context, &spl_token::id()).await; +} + +async fn check_wrong_address_derivation_owner( + context: &mut ProgramTestContext, + program_id: &Pubkey, +) { + let wallet = Keypair::new(); + let wrong_wallet = Keypair::new(); + let (mint, mint_authority) = create_mint(context, program_id).await; + + let owner_associated_token_address = + create_associated_token_account(context, &wallet.pubkey(), &mint, program_id).await; + let nested_associated_token_address = create_associated_token_account( + context, + &owner_associated_token_address, + &mint, + program_id, + ) + .await; + + let wrong_owner_associated_token_address = + get_associated_token_address_with_program_id(&mint, &wrong_wallet.pubkey(), program_id); + let mut recover = instruction::recover_nested(&wallet.pubkey(), &mint, &mint, program_id); + recover.accounts[3] = AccountMeta::new(wrong_owner_associated_token_address, false); + let transaction = Transaction::new_signed_with_payer( + &[recover], + Some(&context.payer.pubkey()), + &[&context.payer, &wallet], + context.last_blockhash, + ); + try_recover_nested( + context, + program_id, + mint, + mint_authority, + nested_associated_token_address, + wrong_owner_associated_token_address, + wallet, + transaction, + Some(InstructionError::InvalidSeeds), + ) + .await; +} + +#[tokio::test] +async fn fail_wrong_address_derivation_owner_2022() { + let dummy_mint = Pubkey::new_unique(); + let pt = program_test_2022(dummy_mint, true); + let mut context = pt.start_with_context().await; + check_wrong_address_derivation_owner(&mut context, &spl_token_2022::id()).await; +} + +#[tokio::test] +async fn fail_wrong_address_derivation_owner() { + let dummy_mint = Pubkey::new_unique(); + let pt = program_test(dummy_mint, true); + let mut context = pt.start_with_context().await; + check_wrong_address_derivation_owner(&mut context, &spl_token::id()).await; +} + +async fn check_owner_account_does_not_exist(context: &mut ProgramTestContext, program_id: &Pubkey) { + let wallet = Keypair::new(); + let (mint, mint_authority) = create_mint(context, program_id).await; + + let owner_associated_token_address = + get_associated_token_address_with_program_id(&wallet.pubkey(), &mint, program_id); + let nested_associated_token_address = create_associated_token_account( + context, + &owner_associated_token_address, + &mint, + program_id, + ) + .await; + + let transaction = Transaction::new_signed_with_payer( + &[instruction::recover_nested( + &wallet.pubkey(), + &mint, + &mint, + program_id, + )], + Some(&context.payer.pubkey()), + &[&context.payer, &wallet], + context.last_blockhash, + ); + try_recover_nested( + context, + program_id, + mint, + mint_authority, + nested_associated_token_address, + owner_associated_token_address, + wallet, + transaction, + Some(InstructionError::IllegalOwner), + ) + .await; +} + +#[tokio::test] +async fn fail_owner_account_does_not_exist() { + let dummy_mint = Pubkey::new_unique(); + let pt = program_test_2022(dummy_mint, true); + let mut context = pt.start_with_context().await; + check_owner_account_does_not_exist(&mut context, &spl_token_2022::id()).await; +} + +#[tokio::test] +async fn fail_wrong_spl_token_program() { + let wallet = Keypair::new(); + let dummy_mint = Pubkey::new_unique(); + let pt = program_test_2022(dummy_mint, true); + let mut context = pt.start_with_context().await; + let program_id = spl_token_2022::id(); + let wrong_program_id = spl_token::id(); + let (mint, mint_authority) = create_mint(&mut context, &program_id).await; + + let owner_associated_token_address = + create_associated_token_account(&mut context, &wallet.pubkey(), &mint, &program_id).await; + let nested_associated_token_address = create_associated_token_account( + &mut context, + &owner_associated_token_address, + &mint, + &program_id, + ) + .await; + + let transaction = Transaction::new_signed_with_payer( + &[instruction::recover_nested( + &wallet.pubkey(), + &mint, + &mint, + &wrong_program_id, + )], + Some(&context.payer.pubkey()), + &[&context.payer, &wallet], + context.last_blockhash, + ); + try_recover_nested( + &mut context, + &program_id, + mint, + mint_authority, + nested_associated_token_address, + owner_associated_token_address, + wallet, + transaction, + Some(InstructionError::IllegalOwner), + ) + .await; +} + +#[tokio::test] +async fn fail_destination_not_wallet_ata() { + let wallet = Keypair::new(); + let wrong_wallet = Keypair::new(); + let dummy_mint = Pubkey::new_unique(); + let pt = program_test_2022(dummy_mint, true); + let program_id = spl_token_2022::id(); + let mut context = pt.start_with_context().await; + let (mint, mint_authority) = create_mint(&mut context, &program_id).await; + + let owner_associated_token_address = + create_associated_token_account(&mut context, &wallet.pubkey(), &mint, &program_id).await; + let nested_associated_token_address = create_associated_token_account( + &mut context, + &owner_associated_token_address, + &mint, + &program_id, + ) + .await; + let wrong_destination_associated_token_account_address = + create_associated_token_account(&mut context, &wrong_wallet.pubkey(), &mint, &program_id) + .await; + + let mut recover = instruction::recover_nested(&wallet.pubkey(), &mint, &mint, &program_id); + recover.accounts[2] = + AccountMeta::new(wrong_destination_associated_token_account_address, false); + + let transaction = Transaction::new_signed_with_payer( + &[recover], + Some(&context.payer.pubkey()), + &[&context.payer, &wallet], + context.last_blockhash, + ); + try_recover_nested( + &mut context, + &program_id, + mint, + mint_authority, + nested_associated_token_address, + owner_associated_token_address, + wallet, + transaction, + Some(InstructionError::InvalidSeeds), + ) + .await; +} diff --git a/associated-token-account/program/tests/spl_token_create.rs b/associated-token-account/program/tests/spl_token_create.rs index a4840d94..a1e465ce 100644 --- a/associated-token-account/program/tests/spl_token_create.rs +++ b/associated-token-account/program/tests/spl_token_create.rs @@ -1,13 +1,15 @@ // Mark this test as BPF-only due to current `ProgramTest` limitations when CPIing into the system program #![cfg(feature = "test-bpf")] +mod program_test; + use { + program_test::program_test, solana_program::pubkey::Pubkey, solana_program_test::*, solana_sdk::{program_pack::Pack, signature::Signer, transaction::Transaction}, spl_associated_token_account::{ - get_associated_token_address, id, instruction::create_associated_token_account, - processor::process_instruction, + get_associated_token_address, instruction::create_associated_token_account, }, spl_token::state::Account, }; @@ -15,41 +17,6 @@ use { #[allow(deprecated)] use spl_associated_token_account::create_associated_token_account as deprecated_create_associated_token_account; -fn program_test(token_mint_address: Pubkey, use_latest_spl_token: bool) -> ProgramTest { - let mut pc = ProgramTest::new( - "spl_associated_token_account", - id(), - processor!(process_instruction), - ); - - if use_latest_spl_token { - // TODO: Remove after Token >3.3.0 is available by default in program-test - pc.add_program( - "spl_token", - spl_token::id(), - processor!(spl_token::processor::Processor::process), - ); - } - - // Add a token mint account - // - // The account data was generated by running: - // $ solana account EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v \ - // --output-file tests/fixtures/token-mint-data.bin - // - pc.add_account_with_file_data( - token_mint_address, - 1461600, - spl_token::id(), - "token-mint-data.bin", - ); - - // Dial down the BPF compute budget to detect if the program gets bloated in the future - pc.set_compute_max_units(50_000); - - pc -} - #[tokio::test] async fn success_create() { let wallet_address = Pubkey::new_unique();