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

Author: joncinque

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]);