associated-token-account: Add recover nested account ix (#2889)

* associated-token-account: Add transfer / close nested accounts

* Swap wallet and wrong_wallet in test

* Use new error

* Force destination to wallet and ATA

* Fix merge conflicts

* Add more vanilla spl-token tests

* Improve test, fix instruction comments

* Address feedback

* Rename CloseNested -> RecoverNested, add comment

* Fix typo in comment

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(),
+    }
+}

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::<Account>::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::<Account>::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::<Mint>::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],
+    )
+}

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::<Account>(&[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 =

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

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::<Account>(&[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::<Account>(&[ExtensionType::ImmutableOwner]);

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",

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::<Mint>(&[]);
+    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<InstructionError>,
+) {
+    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::<Account>::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;
+}

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,
+        get_associated_token_address, instruction::create_associated_token_account,
-        processor::process_instruction,
     },
     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();