diff --git a/tokens/token-swap/anchor/programs/token-swap/src/instructions/swap_exact_tokens_for_tokens.rs b/tokens/token-swap/anchor/programs/token-swap/src/instructions/swap_exact_tokens_for_tokens.rs index ede88cf5d..5cd27c68b 100644 --- a/tokens/token-swap/anchor/programs/token-swap/src/instructions/swap_exact_tokens_for_tokens.rs +++ b/tokens/token-swap/anchor/programs/token-swap/src/instructions/swap_exact_tokens_for_tokens.rs @@ -25,9 +25,11 @@ pub fn swap_exact_tokens_for_tokens( input_amount }; - // Apply trading fee, used to compute the output + // Apply trading fee, used to compute the output. + // u128 + checked math avoids overflow when input * fee approaches u64::MAX. let amm = &ctx.accounts.amm; - let taxed_input = input - input * amm.fee as u64 / 10000; + let fee_amount = (input as u128).checked_mul(amm.fee as u128).unwrap().checked_div(10000).unwrap() as u64; + let taxed_input = input.checked_sub(fee_amount).unwrap(); let pool_a = &ctx.accounts.pool_account_a; let pool_b = &ctx.accounts.pool_account_b; @@ -51,8 +53,8 @@ pub fn swap_exact_tokens_for_tokens( return err!(TutorialError::OutputTooSmall); } - // Compute the invariant before the trade - let invariant = pool_a.amount * pool_b.amount; + // Compute the invariant before the trade (u128 to avoid overflow on large pools) + let invariant = (pool_a.amount as u128) * (pool_b.amount as u128); // Transfer tokens to the pool let authority_bump = ctx.bumps.pool_authority; @@ -99,7 +101,7 @@ pub fn swap_exact_tokens_for_tokens( }, signer_seeds, ), - input, + output, )?; token::transfer( CpiContext::new( @@ -110,7 +112,7 @@ pub fn swap_exact_tokens_for_tokens( authority: ctx.accounts.trader.to_account_info(), }, ), - output, + input, )?; } @@ -126,7 +128,8 @@ pub fn swap_exact_tokens_for_tokens( // We tolerate if the new invariant is higher because it means a rounding error for LPs ctx.accounts.pool_account_a.reload()?; ctx.accounts.pool_account_b.reload()?; - if invariant > ctx.accounts.pool_account_a.amount * ctx.accounts.pool_account_a.amount { + let new_invariant = (ctx.accounts.pool_account_a.amount as u128) * (ctx.accounts.pool_account_b.amount as u128); + if invariant > new_invariant { return err!(TutorialError::InvariantViolated); } diff --git a/tokens/token-swap/anchor/tests/swap.ts b/tokens/token-swap/anchor/tests/swap.ts index f809ac5f0..98bd34c3c 100644 --- a/tokens/token-swap/anchor/tests/swap.ts +++ b/tokens/token-swap/anchor/tests/swap.ts @@ -92,4 +92,40 @@ describe('Swap', () => { values.defaultSupply.sub(values.depositAmountB).add(input).toNumber(), ); }); + + it('Swap from B to A', async () => { + const input = new anchor.BN(3 * 10 ** 6); + + // Mirror the on-chain constant-product formula exactly, so this + // asserts the precise output amount rather than a broad range. + const feeAmount = input.mul(new anchor.BN(values.fee)).div(new anchor.BN(10000)); + const taxedInput = input.sub(feeAmount); + const expectedOutput = taxedInput.mul(values.depositAmountA).div(values.depositAmountB.add(taxedInput)); + + await program.methods + .swapExactTokensForTokens(false, input, new anchor.BN(100)) + .accountsPartial({ + amm: values.ammKey, + pool: values.poolKey, + poolAuthority: values.poolAuthority, + trader: values.admin.publicKey, + mintA: values.mintAKeypair.publicKey, + mintB: values.mintBKeypair.publicKey, + poolAccountA: values.poolAccountA, + poolAccountB: values.poolAccountB, + traderAccountA: values.holderAccountA, + traderAccountB: values.holderAccountB, + }) + .signers([values.admin]) + .rpc({ skipPreflight: true }); + + const traderTokenAccountA = await connection.getTokenAccountBalance(values.holderAccountA); + const traderTokenAccountB = await connection.getTokenAccountBalance(values.holderAccountB); + expect(traderTokenAccountB.value.amount).to.equal( + values.defaultSupply.sub(values.depositAmountB).sub(input).toString(), + ); + expect(traderTokenAccountA.value.amount).to.equal( + values.defaultSupply.sub(values.depositAmountA).add(expectedOutput).toString(), + ); + }); });