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

Commit 25556ce

Browse files
joncinquejoncinque
committed
Step 4: Use appropriate token program
1 parent e77efb3 commit 25556ce

File tree

1 file changed

+87
-39
lines changed

1 file changed

+87
-39
lines changed

‎token-swap/program/src/processor.rs

Lines changed: 87 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ use solana_program::{
2727
pubkey::Pubkey,
2828
};
2929
use spl_token_2022::{
30+
check_spl_token_program_account,
3031
error::TokenError,
3132
extension::StateWithExtensions,
3233
state::{Account, Mint},
@@ -41,7 +42,9 @@ impl Processor {
4142
account_info: &AccountInfo,
4243
token_program_id: &Pubkey,
4344
) -> Result<Account, SwapError> {
44-
if account_info.owner != token_program_id {
45+
if account_info.owner != token_program_id
46+
&& check_spl_token_program_account(account_info.owner).is_err()
47+
{
4548
Err(SwapError::IncorrectTokenProgramId)
4649
} else {
4750
StateWithExtensions::<Account>::unpack(&account_info.data.borrow())
@@ -55,7 +58,9 @@ impl Processor {
5558
account_info: &AccountInfo,
5659
token_program_id: &Pubkey,
5760
) -> Result<Mint, SwapError> {
58-
if account_info.owner != token_program_id {
61+
if account_info.owner != token_program_id
62+
&& check_spl_token_program_account(account_info.owner).is_err()
63+
{
5964
Err(SwapError::IncorrectTokenProgramId)
6065
} else {
6166
StateWithExtensions::<Mint>::unpack(&account_info.data.borrow())
@@ -171,7 +176,7 @@ impl Processor {
171176
token_a_info: &AccountInfo,
172177
token_b_info: &AccountInfo,
173178
pool_mint_info: &AccountInfo,
174-
token_program_info: &AccountInfo,
179+
pool_token_program_info: &AccountInfo,
175180
user_token_a_info: Option<&AccountInfo>,
176181
user_token_b_info: Option<&AccountInfo>,
177182
pool_fee_account_info: Option<&AccountInfo>,
@@ -193,7 +198,7 @@ impl Processor {
193198
if *pool_mint_info.key != *token_swap.pool_mint() {
194199
return Err(SwapError::IncorrectPoolMint.into());
195200
}
196-
if *token_program_info.key != *token_swap.token_program_id() {
201+
if *pool_token_program_info.key != *token_swap.token_program_id() {
197202
return Err(SwapError::IncorrectTokenProgramId.into());
198203
}
199204
if let Some(user_token_a_info) = user_token_a_info {
@@ -433,7 +438,7 @@ impl Processor {
433438

434439
Self::token_transfer(
435440
swap_info.key,
436-
pool_token_program_info.clone(),
441+
source_token_program_info.clone(),
437442
source_info.clone(),
438443
swap_source_info.clone(),
439444
user_transfer_authority_info.clone(),
@@ -495,7 +500,7 @@ impl Processor {
495500

496501
Self::token_transfer(
497502
swap_info.key,
498-
pool_token_program_info.clone(),
503+
destination_token_program_info.clone(),
499504
swap_destination_info.clone(),
500505
destination_info.clone(),
501506
authority_info.clone(),
@@ -585,7 +590,7 @@ impl Processor {
585590

586591
Self::token_transfer(
587592
swap_info.key,
588-
pool_token_program_info.clone(),
593+
token_a_program_info.clone(),
589594
source_a_info.clone(),
590595
token_a_info.clone(),
591596
user_transfer_authority_info.clone(),
@@ -594,7 +599,7 @@ impl Processor {
594599
)?;
595600
Self::token_transfer(
596601
swap_info.key,
597-
pool_token_program_info.clone(),
602+
token_b_program_info.clone(),
598603
source_b_info.clone(),
599604
token_b_info.clone(),
600605
user_transfer_authority_info.clone(),
@@ -721,7 +726,7 @@ impl Processor {
721726
if token_a_amount > 0 {
722727
Self::token_transfer(
723728
swap_info.key,
724-
pool_token_program_info.clone(),
729+
token_a_program_info.clone(),
725730
token_a_info.clone(),
726731
dest_token_a_info.clone(),
727732
authority_info.clone(),
@@ -732,7 +737,7 @@ impl Processor {
732737
if token_b_amount > 0 {
733738
Self::token_transfer(
734739
swap_info.key,
735-
pool_token_program_info.clone(),
740+
token_b_program_info.clone(),
736741
token_b_info.clone(),
737742
dest_token_b_info.clone(),
738743
authority_info.clone(),
@@ -831,7 +836,7 @@ impl Processor {
831836
TradeDirection::AtoB => {
832837
Self::token_transfer(
833838
swap_info.key,
834-
pool_token_program_info.clone(),
839+
source_token_program_info.clone(),
835840
source_info.clone(),
836841
swap_token_a_info.clone(),
837842
user_transfer_authority_info.clone(),
@@ -842,7 +847,7 @@ impl Processor {
842847
TradeDirection::BtoA => {
843848
Self::token_transfer(
844849
swap_info.key,
845-
pool_token_program_info.clone(),
850+
source_token_program_info.clone(),
846851
source_info.clone(),
847852
swap_token_b_info.clone(),
848853
user_transfer_authority_info.clone(),
@@ -980,7 +985,7 @@ impl Processor {
980985
TradeDirection::AtoB => {
981986
Self::token_transfer(
982987
swap_info.key,
983-
pool_token_program_info.clone(),
988+
destination_token_program_info.clone(),
984989
swap_token_a_info.clone(),
985990
destination_info.clone(),
986991
authority_info.clone(),
@@ -991,7 +996,7 @@ impl Processor {
991996
TradeDirection::BtoA => {
992997
Self::token_transfer(
993998
swap_info.key,
994-
pool_token_program_info.clone(),
999+
destination_token_program_info.clone(),
9951000
swap_token_b_info.clone(),
9961001
destination_info.clone(),
9971002
authority_info.clone(),
@@ -1521,9 +1526,10 @@ mod tests {
15211526
maximum_token_b_amount: u64,
15221527
) -> ProgramResult {
15231528
let user_transfer_authority = Pubkey::new_unique();
1529+
let token_a_program_id = depositor_token_a_account.owner;
15241530
do_process_instruction(
15251531
approve(
1526-
&self.token_a_program_id,
1532+
&token_a_program_id,
15271533
depositor_token_a_key,
15281534
&user_transfer_authority,
15291535
depositor_key,
@@ -1539,9 +1545,10 @@ mod tests {
15391545
)
15401546
.unwrap();
15411547

1548+
let token_b_program_id = depositor_token_b_account.owner;
15421549
do_process_instruction(
15431550
approve(
1544-
&self.token_b_program_id,
1551+
&token_b_program_id,
15451552
depositor_token_b_key,
15461553
&user_transfer_authority,
15471554
depositor_key,
@@ -1557,12 +1564,13 @@ mod tests {
15571564
)
15581565
.unwrap();
15591566

1567+
let pool_token_program_id = depositor_pool_account.owner;
15601568
do_process_instruction(
15611569
deposit_all_token_types(
15621570
&SWAP_PROGRAM_ID,
1563-
&self.token_a_program_id,
1564-
&self.token_b_program_id,
1565-
&self.pool_token_program_id,
1571+
&token_a_program_id,
1572+
&token_b_program_id,
1573+
&pool_token_program_id,
15661574
&self.swap_key,
15671575
&self.authority_key,
15681576
&user_transfer_authority,
@@ -1611,10 +1619,11 @@ mod tests {
16111619
minimum_token_b_amount: u64,
16121620
) -> ProgramResult {
16131621
let user_transfer_authority_key = Pubkey::new_unique();
1622+
let pool_token_program_id = pool_account.owner;
16141623
// approve user transfer authority to take out pool tokens
16151624
do_process_instruction(
16161625
approve(
1617-
&self.pool_token_program_id,
1626+
&pool_token_program_id,
16181627
pool_key,
16191628
&user_transfer_authority_key,
16201629
user_key,
@@ -1631,12 +1640,14 @@ mod tests {
16311640
.unwrap();
16321641

16331642
// withdraw token a and b correctly
1643+
let token_a_program_id = token_a_account.owner;
1644+
let token_b_program_id = token_b_account.owner;
16341645
do_process_instruction(
16351646
withdraw_all_token_types(
16361647
&SWAP_PROGRAM_ID,
1637-
&self.pool_token_program_id,
1638-
&self.token_a_program_id,
1639-
&self.token_b_program_id,
1648+
&pool_token_program_id,
1649+
&token_a_program_id,
1650+
&token_b_program_id,
16401651
&self.swap_key,
16411652
&self.authority_key,
16421653
&user_transfer_authority_key,
@@ -1703,11 +1714,12 @@ mod tests {
17031714
)
17041715
.unwrap();
17051716

1717+
let pool_token_program_id = deposit_pool_account.owner;
17061718
do_process_instruction(
17071719
deposit_single_token_type_exact_amount_in(
17081720
&SWAP_PROGRAM_ID,
17091721
&source_token_program_id,
1710-
&self.pool_token_program_id,
1722+
&pool_token_program_id,
17111723
&self.swap_key,
17121724
&self.authority_key,
17131725
&user_transfer_authority_key,
@@ -1749,10 +1761,11 @@ mod tests {
17491761
maximum_pool_token_amount: u64,
17501762
) -> ProgramResult {
17511763
let user_transfer_authority_key = Pubkey::new_unique();
1764+
let pool_token_program_id = pool_account.owner;
17521765
// approve user transfer authority to take out pool tokens
17531766
do_process_instruction(
17541767
approve(
1755-
&self.pool_token_program_id,
1768+
&pool_token_program_id,
17561769
pool_key,
17571770
&user_transfer_authority_key,
17581771
user_key,
@@ -1772,7 +1785,7 @@ mod tests {
17721785
do_process_instruction(
17731786
withdraw_single_token_type_exact_amount_out(
17741787
&SWAP_PROGRAM_ID,
1775-
&self.pool_token_program_id,
1788+
&pool_token_program_id,
17761789
&destination_token_program_id,
17771790
&self.swap_key,
17781791
&self.authority_key,
@@ -2412,7 +2425,7 @@ mod tests {
24122425
// pool fee account has wrong mint
24132426
{
24142427
let (_pool_fee_key, pool_fee_account) = mint_token(
2415-
&pool_token_program_id,
2428+
&token_a_program_id,
24162429
&accounts.token_a_mint_key,
24172430
&mut accounts.token_a_mint_account,
24182431
&user_key,
@@ -2576,7 +2589,7 @@ mod tests {
25762589
{
25772590
let wrong_program_id = Pubkey::new_unique();
25782591
assert_eq!(
2579-
Err(SwapError::IncorrectTokenProgramId.into()),
2592+
Err(ProgramError::IncorrectProgramId),
25802593
do_process_instruction(
25812594
initialize(
25822595
&SWAP_PROGRAM_ID,
@@ -3231,8 +3244,13 @@ mod tests {
32313244
pool_key,
32323245
mut pool_account,
32333246
) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3247+
let expected_error: ProgramError = if token_a_account.owner == token_b_account.owner {
3248+
TokenError::MintMismatch.into()
3249+
} else {
3250+
ProgramError::InvalidAccountData
3251+
};
32343252
assert_eq!(
3235-
Err(TokenError::MintMismatch.into()),
3253+
Err(expected_error),
32363254
accounts.deposit_all_token_types(
32373255
&depositor_key,
32383256
&token_b_key,
@@ -3264,10 +3282,15 @@ mod tests {
32643282
_token_b_key,
32653283
mut _token_b_account,
32663284
_pool_key,
3267-
mut _pool_account,
3285+
pool_account,
32683286
) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
3287+
let expected_error: ProgramError = if token_a_account.owner == pool_account.owner {
3288+
TokenError::MintMismatch.into()
3289+
} else {
3290+
SwapError::IncorrectTokenProgramId.into()
3291+
};
32693292
assert_eq!(
3270-
Err(TokenError::MintMismatch.into()),
3293+
Err(expected_error),
32713294
accounts.deposit_all_token_types(
32723295
&depositor_key,
32733296
&token_a_key,
@@ -3841,8 +3864,13 @@ mod tests {
38413864
initial_b,
38423865
withdraw_amount.try_into().unwrap(),
38433866
);
3867+
let expected_error: ProgramError = if token_a_account.owner == token_b_account.owner {
3868+
TokenError::MintMismatch.into()
3869+
} else {
3870+
ProgramError::InvalidAccountData
3871+
};
38443872
assert_eq!(
3845-
Err(TokenError::MintMismatch.into()),
3873+
Err(expected_error),
38463874
accounts.withdraw_all_token_types(
38473875
&withdrawer_key,
38483876
&pool_key,
@@ -3880,16 +3908,21 @@ mod tests {
38803908
_token_b_key,
38813909
_token_b_account,
38823910
_pool_key,
3883-
_pool_account,
3911+
pool_account,
38843912
) = accounts.setup_token_accounts(
38853913
&user_key,
38863914
&withdrawer_key,
38873915
withdraw_amount.try_into().unwrap(),
38883916
initial_b,
38893917
withdraw_amount.try_into().unwrap(),
38903918
);
3919+
let expected_error: ProgramError = if token_a_account.owner == pool_account.owner {
3920+
TokenError::MintMismatch.into()
3921+
} else {
3922+
SwapError::IncorrectTokenProgramId.into()
3923+
};
38913924
assert_eq!(
3892-
Err(TokenError::MintMismatch.into()),
3925+
Err(expected_error),
38933926
accounts.withdraw_all_token_types(
38943927
&withdrawer_key,
38953928
&wrong_token_a_key,
@@ -4661,10 +4694,15 @@ mod tests {
46614694
token_b_key,
46624695
mut token_b_account,
46634696
_pool_key,
4664-
mut _pool_account,
4697+
pool_account,
46654698
) = accounts.setup_token_accounts(&user_key, &depositor_key, deposit_a, deposit_b, 0);
4699+
let expected_error: ProgramError = if token_b_account.owner == pool_account.owner {
4700+
TokenError::MintMismatch.into()
4701+
} else {
4702+
SwapError::IncorrectTokenProgramId.into()
4703+
};
46664704
assert_eq!(
4667-
Err(TokenError::MintMismatch.into()),
4705+
Err(expected_error),
46684706
accounts.deposit_single_token_type_exact_amount_in(
46694707
&depositor_key,
46704708
&token_a_key,
@@ -5184,16 +5222,21 @@ mod tests {
51845222
token_b_key,
51855223
mut token_b_account,
51865224
_pool_key,
5187-
_pool_account,
5225+
pool_account,
51885226
) = accounts.setup_token_accounts(
51895227
&user_key,
51905228
&withdrawer_key,
51915229
maximum_pool_token_amount,
51925230
initial_b,
51935231
maximum_pool_token_amount,
51945232
);
5233+
let expected_error: ProgramError = if token_a_account.owner == pool_account.owner {
5234+
TokenError::MintMismatch.into()
5235+
} else {
5236+
SwapError::IncorrectTokenProgramId.into()
5237+
};
51955238
assert_eq!(
5196-
Err(TokenError::MintMismatch.into()),
5239+
Err(expected_error),
51975240
accounts.withdraw_single_token_type_exact_amount_out(
51985241
&withdrawer_key,
51995242
&token_a_key,
@@ -6448,8 +6491,13 @@ mod tests {
64486491
_pool_key,
64496492
_pool_account,
64506493
) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
6494+
let expected_error: ProgramError = if token_a_account.owner == token_b_account.owner {
6495+
TokenError::MintMismatch.into()
6496+
} else {
6497+
ProgramError::IncorrectProgramId
6498+
};
64516499
assert_eq!(
6452-
Err(TokenError::MintMismatch.into()),
6500+
Err(expected_error),
64536501
accounts.swap(
64546502
&swapper_key,
64556503
&token_b_key,

0 commit comments

Comments
 (0)