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

[NikkiAung]

Summary

swap_exact_tokens_for_tokens's swap_a == false branch (B→A swaps) has two independent bugs in the same ~20-line block. Neither was ever exercised by the existing tests — the only swap test always calls swap_a = true.

Bug 1 (the severe one): transfer amounts are swapped in the B→A branch

} else {
    token::transfer(/* pool_account_a -> trader_account_a */ .., input)?;   // should be `output`
    token::transfer(/* trader_account_b -> pool_account_b */ .., output)?;  // should be `input`
}

Compare to the correct swap_a == true branch just above it, which moves input trader→pool and output pool→trader. In the else branch the two amounts are swapped: the pool pays out input (the raw, caller-chosen B amount, reused as if it were the A payout) instead of the formula-computed output, and the trader only pays output instead of the full input they claimed to be swapping.

Hand-verified this drains real value: on a pool with pool_a=1,000,000 / pool_b=4,000,000 (5% fee), input=500,000 lets a trader pull out half the pool's entire A reserve while only paying ~106,145 B — the invariant drops from 4e12 to ~2.05e12 in one call.

Bug 2: the post-trade invariant check compares the wrong accounts

if invariant > ctx.accounts.pool_account_a.amount * ctx.accounts.pool_account_a.amount {
    return err!(TutorialError::InvariantViolated);
}

pool_account_a.amount is multiplied by itself instead of by pool_account_b.amount — comparing pre-trade A0*B0 against post-trade A1*A1, not A1*B1. This check is meant to be the AMM's final safety net (per its own comment) but as written never actually verifies the real constant-product invariant.

Why both had to be fixed together

These two bugs currently mask each other by coincidence: severe drains via Bug 1 always shrink pool_account_a a lot, and shrinking A a lot also shrinks A² a lot — so Bug 2's check, wrong as it is, still happens to reject the most extreme drain attempts. A more moderate trade can still be profitably exploitable while slipping under that coincidental threshold. Fixing either bug alone — without the other — would remove the accidental protection and make the remaining bug fully exploitable. Both are fixed in this PR.

Secondary fix: overflow-safe arithmetic

Also switched the pre/post-trade invariant computation and the fee calculation to u128 arithmetic, matching the checked-math style already used for the output-amount calculation in the same function. The raw u64 multiplications could otherwise overflow and panic (overflow-checks = true in this workspace's release profile) for large-supply pools.

Test changes

Added 'Swap from B to A' to tests/swap.ts — the first test in this example to exercise that branch at all. Verified it fails against the unpatched code (reverts with InvariantViolated, since the trade genuinely destroys value under the old code) and passes after the fix, with all other existing tests unaffected.

Verification

• cargo check clean.
• anchor build — no account-struct or instruction-signature changes, IDL unaffected beyond recompilation.
• Full anchor test suite (create-amm.ts, create-pool.ts, deposit-liquidity.ts, withdraw-liquidity.ts, swap.ts) — all 10 tests pass post-fix.
• Ran the pre-fix/post-fix matrix: captured the fix as a patch, reverted only the program file, rebuilt, confirmed the new test fails (and only that test — nothing else regresses), reapplied the patch, confirmed all 10 pass.
• pnpm exec tsc --noEmit and prettier --check clean on the modified test file.

Explicitly out of scope

• create_amm.rs's admin field is never actually used to gate any instruction (dead authority field) — a minor design smell, not exploitable since nothing checks it, not touched here.
• Pool creation is fully permissionless (anyone can create a pool for any AMM+mint pair) — standard, intentional AMM-factory design, not a bug.
• deposit_liquidity.rs's MINIMUM_LIQUIDITY handling (withheld from the first LP mint rather than minted-then-burned) — investigated and confirmed self-consistent with withdraw_liquidity.rs's redemption formula, which compensates for it in the supply denominator. Not a bug.

[greptile-apps]

Greptile Summary

This PR corrects B-to-A token transfer amounts, fixes the constant-product invariant comparison, and widens fee and invariant arithmetic to avoid overflow.

• Transfers the calculated output from pool A while collecting the full input in token B.
• Compares the pre-trade invariant against the actual post-trade A×B reserves using u128.
• Adds an isolated B-to-A regression test with exact balance assertions.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
tokens/token-swap/anchor/programs/token-swap/src/instructions/swap_exact_tokens_for_tokens.rs	Corrects B-to-A transfer directions and amounts, uses overflow-safe arithmetic, and validates the true post-swap constant-product invariant.
tokens/token-swap/anchor/tests/swap.ts	Adds an isolated B-to-A swap test whose exact expected output and final balances cover the corrected behavior and resolve the prior review finding.

_{Reviews (3): Last reviewed commit: "fix(#690): use checked_ math for the fee..." | Re-trigger Greptile}

[greptile-apps]

Output assertion accepts wrong amounts

The broad balance range accepts any positive A output below input, so regressions in the output formula, fee application, reserve selection, or rounding can pass without verifying the exact corrected result.

Suggested change

	expect(Number(traderTokenAccountA.value.amount)).to.be.greaterThan(
	values.defaultSupply.sub(values.depositAmountA).toNumber(),
	);
	expect(Number(traderTokenAccountA.value.amount)).to.be.lessThan(
	values.defaultSupply.sub(values.depositAmountA).add(input).toNumber(),
	);
	expect(traderTokenAccountA.value.amount).to.equal(
	values.defaultSupply
	.sub(values.depositAmountA)
	.add(new anchor.BN(2_961_038))
	.toString(),
	);

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[NikkiAung]

Good catch, fixed in `2ee337a`: the new test now computes the expected output by mirroring the on-chain constant-product formula (fee, taxed input, then the same mul/div) in BN arithmetic, and asserts exact equality instead of a broad range. Re-ran the full suite — all 10 pass.

[NikkiAung]

Hey @dev-jodee — this one's ready whenever you have a chance to take a look. CI is green and the Greptile feedback from earlier has been addressed.

[dev-jodee]

should use checked_ math for better handling of overflow