swap: Avoid fuzz 1 token withdrawal errors by using fee account (#1058)

* Fix fuzz 1 token withdrawal errors with fee account

* Reintroduce special logic

* Cargo fmt
This commit is contained in:
Jon Cinque 2021-01-12 13:57:58 +01:00 committed by GitHub
parent bfe0f54852
commit 7f17a140b1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 31 additions and 19 deletions

View File

@ -1,5 +1,6 @@
use spl_token_swap_fuzz::{
native_account_data::NativeAccountData, native_token::get_token_balance,
native_account_data::NativeAccountData,
native_token::{get_token_balance, transfer},
native_token_swap::NativeTokenSwap,
};
@ -231,27 +232,24 @@ fn run_fuzz_instructions(fuzz_instructions: Vec<FuzzInstruction>) {
+ get_token_balance(&token_swap.token_b_account);
assert_eq!(before_total_token_b, after_total_token_b);
// final check to make sure that withdrawing everything works
let mut withdrawn_token_a_account = token_swap.create_token_a_account(0);
let mut withdrawn_token_b_account = token_swap.create_token_b_account(0);
// Final check to make sure that withdrawing everything works
//
// First, transfer all pool tokens to the fee account to avoid withdrawal
// fees and a potential crash when withdrawing just 1 pool token.
let mut fee_account = token_swap.pool_fee_account.clone();
for mut pool_account in pool_accounts.values_mut() {
token_swap
.withdraw_all(
&mut pool_account,
&mut withdrawn_token_a_account,
&mut withdrawn_token_b_account,
)
.unwrap();
let pool_token_amount = get_token_balance(&pool_account);
if pool_token_amount > 0 {
transfer(&mut pool_account, &mut fee_account, pool_token_amount);
}
}
let mut pool_account = token_swap.pool_token_account.clone();
token_swap
.withdraw_all(
&mut pool_account,
&mut withdrawn_token_a_account,
&mut withdrawn_token_b_account,
)
.unwrap();
let mut fee_account = token_swap.pool_fee_account.clone();
let pool_token_amount = get_token_balance(&pool_account);
transfer(&mut pool_account, &mut fee_account, pool_token_amount);
// Withdraw everything once again
let mut withdrawn_token_a_account = token_swap.create_token_a_account(0);
let mut withdrawn_token_b_account = token_swap.create_token_b_account(0);
token_swap
.withdraw_all(
&mut fee_account,

View File

@ -39,3 +39,17 @@ pub fn get_token_balance(account_data: &NativeAccountData) -> u64 {
let account = TokenAccount::unpack(&account_data.data).unwrap();
account.amount
}
pub fn transfer(
from_account: &mut NativeAccountData,
to_account: &mut NativeAccountData,
amount: u64,
) {
let mut from = TokenAccount::unpack(&from_account.data).unwrap();
let mut to = TokenAccount::unpack(&to_account.data).unwrap();
assert_eq!(from.mint, to.mint);
from.amount -= amount;
to.amount += amount;
TokenAccount::pack(from, &mut from_account.data[..]).unwrap();
TokenAccount::pack(to, &mut from_account.data[..]).unwrap();
}