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
This commit is contained in:
Jon Cinque 2022-03-22 22:58:10 +01:00 committed by GitHub
parent b7a3fc6243
commit 182c0532e4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 899 additions and 51 deletions

View File

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

View File

@ -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],
)
}

View File

@ -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 =

View File

@ -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

View File

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

View File

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

View File

@ -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;
}

View File

@ -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();