security: fix swapped transfer amounts + broken invariant check in token-swap

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);
     }
 

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(),
+        );
+    });
 });