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

Commit 20864c0

Browse files
joncinquejoncinque
authored
token: Support both token programs in associated-token-account (#2904)
* token: Support both token programs in associated-token-account * Rename spl-token -> spl-token-2022, spl-token-legacy -> spl-token * Use spl-token by default * Check mint in get_account_data_len * Add checks for either token program * Rebase fix * Update version comments in tests
1 parent 93c8086 commit 20864c0

File tree

17 files changed

+534
-195
lines changed

17 files changed

+534
-195
lines changed

‎Cargo.lock

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎associated-token-account/program/Cargo.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,8 @@ borsh = "0.9.1"
1717
num-derive = "0.3"
1818
num-traits = "0.2"
1919
solana-program = "1.9.5"
20-
spl-token = { version = "0.2", path = "../../token/program-2022", package = "spl-token-2022", features = ["no-entrypoint"] }
20+
spl-token = { version = "3.3", path = "../../token/program", features = ["no-entrypoint"] }
21+
spl-token-2022 = { version = "0.2", path = "../../token/program-2022", features = ["no-entrypoint"] }
2122
thiserror = "1.0"
2223

2324
[dev-dependencies]

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

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
//! Program instructions
22
33
use {
4+
crate::{get_associated_token_address_with_program_id, id},
45
assert_matches::assert_matches,
56
borsh::{BorshDeserialize, BorshSchema, BorshSerialize},
67
solana_program::{
@@ -9,8 +10,6 @@ use {
910
},
1011
};
1112

12-
use crate::{get_associated_token_address, id};
13-
1413
/// Instructions supported by the AssociatedTokenAccount program
1514
#[derive(Clone, Debug, PartialEq, BorshDeserialize, BorshSerialize, BorshSchema)]
1615
pub enum AssociatedTokenAccountInstruction {
@@ -44,8 +43,11 @@ fn build_associated_token_account_instruction(
4443
token_program_id: &Pubkey,
4544
instruction: AssociatedTokenAccountInstruction,
4645
) -> Instruction {
47-
let associated_account_address =
48-
get_associated_token_address(wallet_address, token_mint_address);
46+
let associated_account_address = get_associated_token_address_with_program_id(
47+
wallet_address,
48+
token_mint_address,
49+
token_program_id,
50+
);
4951
// safety check, assert if not a creation instruction
5052
assert_matches!(
5153
instruction,

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

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,10 @@ use {
1919
system_program,
2020
sysvar::Sysvar,
2121
},
22-
spl_token::{extension::StateWithExtensions, state::Account},
22+
spl_token_2022::{
23+
extension::{ExtensionType, StateWithExtensions},
24+
state::Account,
25+
},
2326
};
2427

2528
/// Specify when to create the associated token account
@@ -115,7 +118,7 @@ fn process_create_associated_token_account(
115118
let account_len = get_account_len(
116119
spl_token_mint_info,
117120
spl_token_program_info,
118-
&[spl_token::extension::ExtensionType::ImmutableOwner],
121+
&[ExtensionType::ImmutableOwner],
119122
)?;
120123

121124
create_pda_account(
@@ -130,7 +133,7 @@ fn process_create_associated_token_account(
130133

131134
msg!("Initialize the associated token account");
132135
invoke(
133-
&spl_token::instruction::initialize_immutable_owner(
136+
&spl_token_2022::instruction::initialize_immutable_owner(
134137
spl_token_program_id,
135138
associated_token_account_info.key,
136139
)?,
@@ -140,7 +143,7 @@ fn process_create_associated_token_account(
140143
],
141144
)?;
142145
invoke(
143-
&spl_token::instruction::initialize_account3(
146+
&spl_token_2022::instruction::initialize_account3(
144147
spl_token_program_id,
145148
associated_token_account_info.key,
146149
spl_token_mint_info.key,

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ use {
1010
rent::Rent,
1111
system_instruction,
1212
},
13-
spl_token::{check_program_account, extension::ExtensionType},
13+
spl_token_2022::extension::ExtensionType,
1414
std::convert::TryInto,
1515
};
1616

@@ -79,7 +79,7 @@ pub fn get_account_len<'a>(
7979
extension_types: &[ExtensionType],
8080
) -> Result<usize, ProgramError> {
8181
invoke(
82-
&spl_token::instruction::get_account_data_size(
82+
&spl_token_2022::instruction::get_account_data_size(
8383
spl_token_program.key,
8484
mint.key,
8585
extension_types,
@@ -89,7 +89,9 @@ pub fn get_account_len<'a>(
8989
get_return_data()
9090
.ok_or(ProgramError::InvalidInstructionData)
9191
.and_then(|(key, data)| {
92-
check_program_account(&key)?;
92+
if key != *spl_token_program.key {
93+
return Err(ProgramError::IncorrectProgramId);
94+
}
9395
data.try_into()
9496
.map(usize::from_le_bytes)
9597
.map_err(|_| ProgramError::InvalidInstructionData)

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

Lines changed: 25 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,12 @@ use {
1717
},
1818
spl_associated_token_account::{
1919
error::AssociatedTokenAccountError,
20-
get_associated_token_address,
20+
get_associated_token_address_with_program_id,
2121
instruction::{
2222
create_associated_token_account, create_associated_token_account_idempotent,
2323
},
2424
},
25-
spl_token::{
25+
spl_token_2022::{
2626
extension::ExtensionType,
2727
instruction::initialize_account,
2828
state::{Account, AccountState},
@@ -33,23 +33,24 @@ use {
3333
async fn success_account_exists() {
3434
let wallet_address = Pubkey::new_unique();
3535
let token_mint_address = Pubkey::new_unique();
36-
let associated_token_address =
37-
get_associated_token_address(&wallet_address, &token_mint_address);
36+
let associated_token_address = get_associated_token_address_with_program_id(
37+
&wallet_address,
38+
&token_mint_address,
39+
&spl_token_2022::id(),
40+
);
3841

3942
let (mut banks_client, payer, recent_blockhash) =
4043
program_test(token_mint_address, true).start().await;
4144
let rent = banks_client.get_rent().await.unwrap();
4245
let expected_token_account_len =
43-
ExtensionType::get_account_len::<spl_token::state::Account>(&[
44-
ExtensionType::ImmutableOwner,
45-
]);
46+
ExtensionType::get_account_len::<Account>(&[ExtensionType::ImmutableOwner]);
4647
let expected_token_account_balance = rent.minimum_balance(expected_token_account_len);
4748

4849
let instruction = create_associated_token_account_idempotent(
4950
&payer.pubkey(),
5051
&wallet_address,
5152
&token_mint_address,
52-
&spl_token::id(),
53+
&spl_token_2022::id(),
5354
);
5455

5556
let transaction = Transaction::new_signed_with_payer(
@@ -67,15 +68,15 @@ async fn success_account_exists() {
6768
.expect("get_account")
6869
.expect("associated_account not none");
6970
assert_eq!(associated_account.data.len(), expected_token_account_len);
70-
assert_eq!(associated_account.owner, spl_token::id());
71+
assert_eq!(associated_account.owner, spl_token_2022::id());
7172
assert_eq!(associated_account.lamports, expected_token_account_balance);
7273

7374
// Unchecked instruction fails
7475
let instruction = create_associated_token_account(
7576
&payer.pubkey(),
7677
&wallet_address,
7778
&token_mint_address,
78-
&spl_token::id(),
79+
&spl_token_2022::id(),
7980
);
8081

8182
let transaction = Transaction::new_signed_with_payer(
@@ -103,7 +104,7 @@ async fn success_account_exists() {
103104
&payer.pubkey(),
104105
&wallet_address,
105106
&token_mint_address,
106-
&spl_token::id(),
107+
&spl_token_2022::id(),
107108
);
108109

109110
let transaction = Transaction::new_signed_with_payer(
@@ -121,20 +122,23 @@ async fn success_account_exists() {
121122
.expect("get_account")
122123
.expect("associated_account not none");
123124
assert_eq!(associated_account.data.len(), expected_token_account_len);
124-
assert_eq!(associated_account.owner, spl_token::id());
125+
assert_eq!(associated_account.owner, spl_token_2022::id());
125126
assert_eq!(associated_account.lamports, expected_token_account_balance);
126127
}
127128

128129
#[tokio::test]
129130
async fn fail_account_exists_with_wrong_owner() {
130131
let wallet_address = Pubkey::new_unique();
131132
let token_mint_address = Pubkey::new_unique();
132-
let associated_token_address =
133-
get_associated_token_address(&wallet_address, &token_mint_address);
133+
let associated_token_address = get_associated_token_address_with_program_id(
134+
&wallet_address,
135+
&token_mint_address,
136+
&spl_token_2022::id(),
137+
);
134138

135139
let wrong_owner = Pubkey::new_unique();
136140
let mut associated_token_account =
137-
SolanaAccount::new(1_000_000_000, Account::LEN, &spl_token::id());
141+
SolanaAccount::new(1_000_000_000, Account::LEN, &spl_token_2022::id());
138142
let token_account = Account {
139143
mint: token_mint_address,
140144
owner: wrong_owner,
@@ -155,7 +159,7 @@ async fn fail_account_exists_with_wrong_owner() {
155159
&payer.pubkey(),
156160
&wallet_address,
157161
&token_mint_address,
158-
&spl_token::id(),
162+
&spl_token_2022::id(),
159163
);
160164
let transaction = Transaction::new_signed_with_payer(
161165
&[instruction],
@@ -184,9 +188,8 @@ async fn fail_non_ata() {
184188
program_test(token_mint_address, true).start().await;
185189

186190
let rent = banks_client.get_rent().await.unwrap();
187-
let token_account_len = ExtensionType::get_account_len::<spl_token::state::Account>(&[
188-
ExtensionType::ImmutableOwner,
189-
]);
191+
let token_account_len =
192+
ExtensionType::get_account_len::<Account>(&[ExtensionType::ImmutableOwner]);
190193
let token_account_balance = rent.minimum_balance(token_account_len);
191194

192195
let wallet_address = Pubkey::new_unique();
@@ -198,10 +201,10 @@ async fn fail_non_ata() {
198201
&account.pubkey(),
199202
token_account_balance,
200203
token_account_len as u64,
201-
&spl_token::id(),
204+
&spl_token_2022::id(),
202205
),
203206
initialize_account(
204-
&spl_token::id(),
207+
&spl_token_2022::id(),
205208
&account.pubkey(),
206209
&token_mint_address,
207210
&wallet_address,
@@ -218,7 +221,7 @@ async fn fail_non_ata() {
218221
&payer.pubkey(),
219222
&wallet_address,
220223
&token_mint_address,
221-
&spl_token::id(),
224+
&spl_token_2022::id(),
222225
);
223226
instruction.accounts[1] = AccountMeta::new(account.pubkey(), false); // <-- Invalid associated_account_address
224227

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

Lines changed: 22 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,9 @@ use {
1313
transaction::{Transaction, TransactionError},
1414
},
1515
spl_associated_token_account::{
16-
get_associated_token_address, instruction::create_associated_token_account,
16+
get_associated_token_address_with_program_id, instruction::create_associated_token_account,
1717
},
18-
spl_token::{
18+
spl_token_2022::{
1919
error::TokenError,
2020
extension::{transfer_fee, ExtensionType, StateWithExtensionsOwned},
2121
state::{Account, Mint},
@@ -46,19 +46,19 @@ async fn test_associated_token_account_with_transfer_fees() {
4646
&mint_account.pubkey(),
4747
rent.minimum_balance(space),
4848
space as u64,
49-
&spl_token::id(),
49+
&spl_token_2022::id(),
5050
),
5151
transfer_fee::instruction::initialize_transfer_fee_config(
52-
&spl_token::id(),
52+
&spl_token_2022::id(),
5353
&token_mint_address,
5454
Some(&mint_authority.pubkey()),
5555
Some(&mint_authority.pubkey()),
5656
1_000,
5757
maximum_fee,
5858
)
5959
.unwrap(),
60-
spl_token::instruction::initialize_mint(
61-
&spl_token::id(),
60+
spl_token_2022::instruction::initialize_mint(
61+
&spl_token_2022::id(),
6262
&token_mint_address,
6363
&mint_authority.pubkey(),
6464
Some(&mint_authority.pubkey()),
@@ -77,7 +77,7 @@ async fn test_associated_token_account_with_transfer_fees() {
7777
&payer.pubkey(),
7878
&wallet_address_sender,
7979
&token_mint_address,
80-
&spl_token::id(),
80+
&spl_token_2022::id(),
8181
)],
8282
Some(&payer.pubkey()),
8383
);
@@ -89,23 +89,29 @@ async fn test_associated_token_account_with_transfer_fees() {
8989
&payer.pubkey(),
9090
&wallet_address_receiver,
9191
&token_mint_address,
92-
&spl_token::id(),
92+
&spl_token_2022::id(),
9393
)],
9494
Some(&payer.pubkey()),
9595
);
9696
transaction.sign(&[&payer], recent_blockhash);
9797
banks_client.process_transaction(transaction).await.unwrap();
9898

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);
99+
let associated_token_address_sender = get_associated_token_address_with_program_id(
100+
&wallet_address_sender,
101+
&token_mint_address,
102+
&spl_token_2022::id(),
103+
);
104+
let associated_token_address_receiver = get_associated_token_address_with_program_id(
105+
&wallet_address_receiver,
106+
&token_mint_address,
107+
&spl_token_2022::id(),
108+
);
103109

104110
// mint tokens
105111
let sender_amount = 50 * maximum_fee;
106112
let mut transaction = Transaction::new_with_payer(
107-
&[spl_token::instruction::mint_to(
108-
&spl_token::id(),
113+
&[spl_token_2022::instruction::mint_to(
114+
&spl_token_2022::id(),
109115
&token_mint_address,
110116
&associated_token_address_sender,
111117
&mint_authority.pubkey(),
@@ -121,7 +127,7 @@ async fn test_associated_token_account_with_transfer_fees() {
121127
// not enough tokens
122128
let mut transaction = Transaction::new_with_payer(
123129
&[transfer_fee::instruction::transfer_checked_with_fee(
124-
&spl_token::id(),
130+
&spl_token_2022::id(),
125131
&associated_token_address_sender,
126132
&token_mint_address,
127133
&associated_token_address_receiver,
@@ -153,7 +159,7 @@ async fn test_associated_token_account_with_transfer_fees() {
153159
let fee = 50;
154160
let mut transaction = Transaction::new_with_payer(
155161
&[transfer_fee::instruction::transfer_checked_with_fee(
156-
&spl_token::id(),
162+
&spl_token_2022::id(),
157163
&associated_token_address_sender,
158164
&token_mint_address,
159165
&associated_token_address_receiver,

0 commit comments

Comments
 (0)