Skip to content
This repository was archived by the owner on Mar 11, 2025. It is now read-only.

Commit eaaed0d

Browse files
CriesofCarrotsCriesofCarrots
authored
Handle extended token Accounts in ATA program (#2738)
* Use token-2022 in ProgramTest * Add get_account_len() helper with CPI to spl_token::get_account_data_size() * Add test for extended mints/accounts using transfer_fee
1 parent 9d759b7 commit eaaed0d

File tree

5 files changed

+236
-11
lines changed

5 files changed

+236
-11
lines changed

‎associated-token-account/program/src/lib.rs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ pub mod tools;
1111
pub use solana_program;
1212
use solana_program::{
1313
instruction::{AccountMeta, Instruction},
14-
program_pack::Pack,
1514
pubkey::Pubkey,
1615
sysvar,
1716
};

‎associated-token-account/program/src/processor.rs

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
//! Program state processor
22
33
use crate::*;
4-
use crate::{instruction::AssociatedTokenAccountInstruction, tools::account::create_pda_account};
4+
use crate::{
5+
instruction::AssociatedTokenAccountInstruction,
6+
tools::account::{create_pda_account, get_account_len},
7+
};
58
use borsh::BorshDeserialize;
69
use solana_program::{
710
account_info::{next_account_info, AccountInfo},
@@ -71,10 +74,12 @@ pub fn process_create_associated_token_account(
7174
&[bump_seed],
7275
];
7376

77+
let account_len = get_account_len(spl_token_mint_info, spl_token_program_info)?;
78+
7479
create_pda_account(
7580
funder_info,
7681
&rent,
77-
spl_token::state::Account::LEN,
82+
account_len,
7883
spl_token_program_id,
7984
system_program_info,
8085
associated_token_account_info,

‎associated-token-account/program/src/tools/account.rs

Lines changed: 32 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,17 @@
11
//! Account utility functions
22
3-
use solana_program::{
4-
account_info::AccountInfo,
5-
entrypoint::ProgramResult,
6-
program::{invoke, invoke_signed},
7-
pubkey::Pubkey,
8-
rent::Rent,
9-
system_instruction,
3+
use {
4+
solana_program::{
5+
account_info::AccountInfo,
6+
entrypoint::ProgramResult,
7+
program::{get_return_data, invoke, invoke_signed},
8+
program_error::ProgramError,
9+
pubkey::Pubkey,
10+
rent::Rent,
11+
system_instruction,
12+
},
13+
spl_token::check_program_account,
14+
std::convert::TryInto,
1015
};
1116

1217
/// Creates associated token account using Program Derived Address for the given seeds
@@ -65,3 +70,23 @@ pub fn create_pda_account<'a>(
6570
)
6671
}
6772
}
73+
74+
/// Determines the required initial data length for a new token account based on the extensions
75+
/// initialized on the Mint
76+
pub fn get_account_len<'a>(
77+
mint: &AccountInfo<'a>,
78+
spl_token_program: &AccountInfo<'a>,
79+
) -> Result<usize, ProgramError> {
80+
invoke(
81+
&spl_token::instruction::get_account_data_size(spl_token_program.key, mint.key)?,
82+
&[mint.clone(), spl_token_program.clone()],
83+
)?;
84+
get_return_data()
85+
.ok_or(ProgramError::InvalidInstructionData)
86+
.and_then(|(key, data)| {
87+
check_program_account(&key)?;
88+
data.try_into()
89+
.map(usize::from_le_bytes)
90+
.map_err(|_| ProgramError::InvalidInstructionData)
91+
})
92+
}

‎associated-token-account/program/tests/extended_mint.rs

Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
// Mark this test as BPF-only due to current `ProgramTest` limitations when CPIing into the system program
2+
#![cfg(feature = "test-bpf")]
3+
4+
mod program_test;
5+
6+
use {
7+
program_test::program_test,
8+
solana_program::{instruction::*, pubkey::Pubkey, system_instruction},
9+
solana_program_test::*,
10+
solana_sdk::{
11+
signature::Signer,
12+
signer::keypair::Keypair,
13+
transaction::{Transaction, TransactionError},
14+
},
15+
spl_associated_token_account::{
16+
get_associated_token_address, instruction::create_associated_token_account,
17+
},
18+
spl_token::{
19+
error::TokenError,
20+
extension::{transfer_fee, ExtensionType, StateWithExtensionsOwned},
21+
state::{Account, Mint},
22+
},
23+
};
24+
25+
#[tokio::test]
26+
async fn test_associated_token_account_with_transfer_fees() {
27+
let wallet_sender = Keypair::new();
28+
let wallet_address_sender = wallet_sender.pubkey();
29+
let wallet_receiver = Keypair::new();
30+
let wallet_address_receiver = wallet_receiver.pubkey();
31+
let (mut banks_client, payer, recent_blockhash) =
32+
program_test(Pubkey::new_unique(), true).start().await;
33+
let rent = banks_client.get_rent().await.unwrap();
34+
35+
// create extended mint
36+
// ... in the future, a mint can be pre-loaded in program_test.rs like the regular mint
37+
let mint_account = Keypair::new();
38+
let token_mint_address = mint_account.pubkey();
39+
let mint_authority = Keypair::new();
40+
let space = ExtensionType::get_account_len::<Mint>(&[ExtensionType::TransferFeeConfig]);
41+
let maximum_fee = 100;
42+
let mut transaction = Transaction::new_with_payer(
43+
&[
44+
system_instruction::create_account(
45+
&payer.pubkey(),
46+
&mint_account.pubkey(),
47+
rent.minimum_balance(space),
48+
space as u64,
49+
&spl_token::id(),
50+
),
51+
transfer_fee::instruction::initialize_transfer_fee_config(
52+
&spl_token::id(),
53+
&token_mint_address,
54+
Some(&mint_authority.pubkey()),
55+
Some(&mint_authority.pubkey()),
56+
1_000,
57+
maximum_fee,
58+
)
59+
.unwrap(),
60+
spl_token::instruction::initialize_mint(
61+
&spl_token::id(),
62+
&token_mint_address,
63+
&mint_authority.pubkey(),
64+
Some(&mint_authority.pubkey()),
65+
0,
66+
)
67+
.unwrap(),
68+
],
69+
Some(&payer.pubkey()),
70+
);
71+
transaction.sign(&[&payer, &mint_account], recent_blockhash);
72+
banks_client.process_transaction(transaction).await.unwrap();
73+
74+
// create extended ATAs
75+
let mut transaction = Transaction::new_with_payer(
76+
&[create_associated_token_account(
77+
&payer.pubkey(),
78+
&wallet_address_sender,
79+
&token_mint_address,
80+
&spl_token::id(),
81+
)],
82+
Some(&payer.pubkey()),
83+
);
84+
transaction.sign(&[&payer], recent_blockhash);
85+
banks_client.process_transaction(transaction).await.unwrap();
86+
87+
let mut transaction = Transaction::new_with_payer(
88+
&[create_associated_token_account(
89+
&payer.pubkey(),
90+
&wallet_address_receiver,
91+
&token_mint_address,
92+
&spl_token::id(),
93+
)],
94+
Some(&payer.pubkey()),
95+
);
96+
transaction.sign(&[&payer], recent_blockhash);
97+
banks_client.process_transaction(transaction).await.unwrap();
98+
99+
let associated_token_address_sender =
100+
get_associated_token_address(&wallet_address_sender, &token_mint_address);
101+
let associated_token_address_receiver =
102+
get_associated_token_address(&wallet_address_receiver, &token_mint_address);
103+
104+
// mint tokens
105+
let sender_amount = 50 * maximum_fee;
106+
let mut transaction = Transaction::new_with_payer(
107+
&[spl_token::instruction::mint_to(
108+
&spl_token::id(),
109+
&token_mint_address,
110+
&associated_token_address_sender,
111+
&mint_authority.pubkey(),
112+
&[],
113+
sender_amount,
114+
)
115+
.unwrap()],
116+
Some(&payer.pubkey()),
117+
);
118+
transaction.sign(&[&payer, &mint_authority], recent_blockhash);
119+
banks_client.process_transaction(transaction).await.unwrap();
120+
121+
// not enough tokens
122+
let mut transaction = Transaction::new_with_payer(
123+
&[transfer_fee::instruction::transfer_checked_with_fee(
124+
&spl_token::id(),
125+
&associated_token_address_sender,
126+
&token_mint_address,
127+
&associated_token_address_receiver,
128+
&wallet_address_sender,
129+
&[],
130+
10_001,
131+
0,
132+
maximum_fee,
133+
)
134+
.unwrap()],
135+
Some(&payer.pubkey()),
136+
);
137+
transaction.sign(&[&payer, &wallet_sender], recent_blockhash);
138+
let err = banks_client
139+
.process_transaction(transaction)
140+
.await
141+
.unwrap_err()
142+
.unwrap();
143+
assert_eq!(
144+
err,
145+
TransactionError::InstructionError(
146+
0,
147+
InstructionError::Custom(TokenError::InsufficientFunds as u32)
148+
)
149+
);
150+
151+
// success
152+
let transfer_amount = 500;
153+
let fee = 50;
154+
let mut transaction = Transaction::new_with_payer(
155+
&[transfer_fee::instruction::transfer_checked_with_fee(
156+
&spl_token::id(),
157+
&associated_token_address_sender,
158+
&token_mint_address,
159+
&associated_token_address_receiver,
160+
&wallet_address_sender,
161+
&[],
162+
transfer_amount,
163+
0,
164+
fee,
165+
)
166+
.unwrap()],
167+
Some(&payer.pubkey()),
168+
);
169+
transaction.sign(&[&payer, &wallet_sender], recent_blockhash);
170+
banks_client.process_transaction(transaction).await.unwrap();
171+
172+
let sender_account = banks_client
173+
.get_account(associated_token_address_sender)
174+
.await
175+
.unwrap()
176+
.unwrap();
177+
let sender_state = StateWithExtensionsOwned::<Account>::unpack(sender_account.data).unwrap();
178+
assert_eq!(sender_state.base.amount, sender_amount - transfer_amount);
179+
let extension = sender_state
180+
.get_extension::<transfer_fee::TransferFeeAmount>()
181+
.unwrap();
182+
assert_eq!(extension.withheld_amount, 0.into());
183+
184+
let receiver_account = banks_client
185+
.get_account(associated_token_address_receiver)
186+
.await
187+
.unwrap()
188+
.unwrap();
189+
let receiver_state =
190+
StateWithExtensionsOwned::<Account>::unpack(receiver_account.data).unwrap();
191+
assert_eq!(receiver_state.base.amount, transfer_amount - fee);
192+
let extension = receiver_state
193+
.get_extension::<transfer_fee::TransferFeeAmount>()
194+
.unwrap();
195+
assert_eq!(extension.withheld_amount, fee.into());
196+
}

‎associated-token-account/program/tests/program_test.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ pub fn program_test(token_mint_address: Pubkey, use_latest_spl_token: bool) -> P
1414
if use_latest_spl_token {
1515
// TODO: Remove after Token >3.2.0 is available by default in program-test
1616
pc.add_program(
17-
"spl_token",
17+
"spl_token_2022",
1818
spl_token::id(),
1919
processor!(spl_token::processor::Processor::process),
2020
);

0 commit comments

Comments
 (0)