fix(token-swap): mint LP pro-rata, key first deposit on LP supply, enforce mint ordering

tokens/token-swap/anchor/Anchor.toml

@@ -6,6 +6,9 @@ solana_version = "3.1.8"
 resolution = true
 skip-lint = false
 
+[programs.localnet]
+swap_example = "UPxp2moQFWsGqfFd3ynqG2W9mj9CTjH68NN2bYUAqV1"
+
 [programs.devnet]
 swap_example = "AsGVFxWqEn8icRBFQApxJe68x3r9zvfSbmiEzYFATGYn"
 

tokens/token-swap/anchor/programs/token-swap/src/errors.rs

@@ -11,6 +11,9 @@ pub enum TutorialError {
     #[msg("Depositing too little liquidity")]
     DepositTooSmall,
 
+    #[msg("Pool reserves are empty")]
+    EmptyPoolReserves,
+
     #[msg("Output is below the minimum expected")]
     OutputTooSmall,
 

tokens/token-swap/anchor/programs/token-swap/src/instructions/create_pool.rs

@@ -6,6 +6,7 @@ use anchor_spl::{
 
 use crate::{
     constants::{AUTHORITY_SEED, LIQUIDITY_SEED},
+    errors::TutorialError,
     state::{Amm, Pool},
 };
 
@@ -38,6 +39,7 @@ pub struct CreatePool<'info> {
             mint_b.key().as_ref(),
         ],
         bump,
+        constraint = mint_a.key() < mint_b.key() @ TutorialError::InvalidMint,
     )]
     pub pool: Box<Account<'info, Pool>>,
 

tokens/token-swap/anchor/programs/token-swap/src/instructions/deposit_liquidity.rs

@@ -10,11 +10,7 @@ use crate::{
     state::Pool,
 };
 
-pub fn deposit_liquidity(
-    ctx: Context<DepositLiquidity>,
-    amount_a: u64,
-    amount_b: u64,
-) -> Result<()> {
+pub fn deposit_liquidity(ctx: Context<DepositLiquidity>, amount_a: u64, amount_b: u64) -> Result<()> {
     // Prevent depositing assets the depositor does not own
     let mut amount_a = if amount_a > ctx.accounts.depositor_account_a.amount {
         ctx.accounts.depositor_account_a.amount
@@ -30,12 +26,18 @@ pub fn deposit_liquidity(
     // Making sure they are provided in the same proportion as existing liquidity
     let pool_a = &ctx.accounts.pool_account_a;
     let pool_b = &ctx.accounts.pool_account_b;
-    // Defining pool creation like this allows attackers to frontrun pool creation with bad ratios
-    let pool_creation = pool_a.amount == 0 && pool_b.amount == 0;
+    // Keyed on LP supply rather than reserves so tokens sent directly to the
+    // pool accounts cannot force the ratio path (and a division by zero).
+    let lp_supply = ctx.accounts.mint_liquidity.supply;
+    let pool_creation = lp_supply == 0;
     (amount_a, amount_b) = if pool_creation {
         // Add as is if there is no liquidity
         (amount_a, amount_b)
     } else {
+        if pool_a.amount == 0 || pool_b.amount == 0 {
+            return err!(TutorialError::EmptyPoolReserves);
+        }
+
         // u128 is enough precision here
         let amount_a_u128 = amount_a as u128;
         let amount_b_u128 = amount_b as u128;
@@ -45,9 +47,9 @@ pub fn deposit_liquidity(
         // Calculate the amount of B required if we deposit all of A provided
         let amount_b_required = amount_a_u128
             .checked_mul(pool_b_u128)
-            .unwrap()
+            .ok_or(TutorialError::MathOverflow)?
             .checked_div(pool_a_u128)
-            .unwrap();
+            .ok_or(TutorialError::MathOverflow)?;
 
         if amount_b_required <= amount_b_u128 {
             // We have enough B to match the A provided
@@ -56,28 +58,49 @@ pub fn deposit_liquidity(
             // We don't have enough B, so we must limit by B and calculate A required
             let amount_a_required = amount_b_u128
                 .checked_mul(pool_a_u128)
-                .unwrap()
+                .ok_or(TutorialError::MathOverflow)?
                 .checked_div(pool_b_u128)
-                .unwrap();
+                .ok_or(TutorialError::MathOverflow)?;
             (amount_a_required as u64, amount_b)
         }
     };
 
     // Computing the amount of liquidity about to be deposited.
-    // Multiply in u128 so the product of two u64 amounts cannot overflow.
-    let mut liquidity = (amount_a as u128)
-        .checked_mul(amount_b as u128)
-        .unwrap()
-        .isqrt() as u64;
-
-    // Lock some minimum liquidity on the first deposit
-    if pool_creation {
+    let liquidity = if pool_creation {
+        // Multiply in u128 so the product of two u64 amounts cannot overflow.
+        let liquidity =
+            (amount_a as u128).checked_mul(amount_b as u128).ok_or(TutorialError::MathOverflow)?.isqrt() as u64;
+
+        // Lock some minimum liquidity on the first deposit
         if liquidity < MINIMUM_LIQUIDITY {
             return err!(TutorialError::DepositTooSmall);
         }
 
-        liquidity -= MINIMUM_LIQUIDITY;
-    }
+        liquidity - MINIMUM_LIQUIDITY
+    } else {
+        // Pro-rata share of the existing supply, so fees accrued to the
+        // reserves stay with the LPs who earned them. The locked minimum
+        // liquidity is part of the supply, matching withdraw_liquidity.
+        let total_liquidity =
+            (lp_supply as u128).checked_add(MINIMUM_LIQUIDITY as u128).ok_or(TutorialError::MathOverflow)?;
+        let liquidity_a = (amount_a as u128)
+            .checked_mul(total_liquidity)
+            .ok_or(TutorialError::MathOverflow)?
+            .checked_div(pool_a.amount as u128)
+            .ok_or(TutorialError::MathOverflow)?;
+        let liquidity_b = (amount_b as u128)
+            .checked_mul(total_liquidity)
+            .ok_or(TutorialError::MathOverflow)?
+            .checked_div(pool_b.amount as u128)
+            .ok_or(TutorialError::MathOverflow)?;
+        let liquidity = u64::try_from(liquidity_a.min(liquidity_b)).map_err(|_| TutorialError::MathOverflow)?;
+
+        if liquidity == 0 {
+            return err!(TutorialError::DepositTooSmall);
+        }
+
+        liquidity
+    };
 
     // Transfer tokens to the pool
     token::transfer(

tokens/token-swap/anchor/tests/create-pool.ts

@@ -1,8 +1,9 @@
 import type { Program } from '@anchor-lang/core';
 import * as anchor from '@anchor-lang/core';
 import { PublicKey } from '@solana/web3.js';
+import { expect } from 'chai';
 import type { SwapExample } from '../target/types/swap_example';
-import { createValues, expectRevert, mintingTokens, type TestValues } from './utils';
+import { createValues, expectAnchorError, expectRevert, mintingTokens, type TestValues } from './utils';
 
 describe('Create pool', () => {
     const provider = anchor.AnchorProvider.env();
@@ -83,4 +84,30 @@ describe('Create pool', () => {
                 .rpc(),
         );
     });
+
+    it('Rejects mints out of order', async () => {
+        const swapped = createValues({
+            id: values.id,
+            mintAKeypair: values.mintBKeypair,
+            mintBKeypair: values.mintAKeypair,
+        });
+
+        await expectAnchorError(
+            program.methods
+                .createPool()
+                .accountsPartial({
+                    amm: swapped.ammKey,
+                    pool: swapped.poolKey,
+                    poolAuthority: swapped.poolAuthority,
+                    mintLiquidity: swapped.mintLiquidity,
+                    mintA: swapped.mintAKeypair.publicKey,
+                    mintB: swapped.mintBKeypair.publicKey,
+                    poolAccountA: swapped.poolAccountA,
+                    poolAccountB: swapped.poolAccountB,
+                })
+                .rpc(),
+            'InvalidMint',
+        );
+        expect(await connection.getAccountInfo(swapped.poolKey)).to.be.null;
+    });
 });

tokens/token-swap/anchor/tests/deposit-liquidity.ts

@@ -1,8 +1,10 @@
 import type { Program } from '@anchor-lang/core';
 import * as anchor from '@anchor-lang/core';
+import { getAssociatedTokenAddressSync, transfer } from '@solana/spl-token';
+import { Keypair } from '@solana/web3.js';
 import { expect } from 'chai';
 import type { SwapExample } from '../target/types/swap_example';
-import { createValues, mintingTokens, type TestValues } from './utils';
+import { createValues, mintToHolder, mintingTokens, type TestValues } from './utils';
 
 describe('Deposit liquidity', () => {
     const provider = anchor.AnchorProvider.env();
@@ -181,4 +183,194 @@ describe('Deposit liquidity', () => {
         // Total B: 5,000,000 + 500,000 = 5,500,000
         expect(poolAccountB.value.amount).to.equal(initialAmountB.add(secondDepositBInput).toString());
     });
+
+    it('Second depositor cannot capture fees accrued by earlier depositors', async () => {
+        const depositor = Keypair.generate();
+        await connection.confirmTransaction(await connection.requestAirdrop(depositor.publicKey, 10 ** 10));
+        await mintToHolder({
+            connection,
+            creator: values.admin,
+            holder: depositor,
+            mintAKeypair: values.mintAKeypair,
+            mintBKeypair: values.mintBKeypair,
+        });
+        const depositorAccountA = getAssociatedTokenAddressSync(values.mintAKeypair.publicKey, depositor.publicKey);
+        const depositorAccountB = getAssociatedTokenAddressSync(values.mintBKeypair.publicKey, depositor.publicKey);
+        const depositorAccountLiquidity = getAssociatedTokenAddressSync(values.mintLiquidity, depositor.publicKey);
+
+        // 1. Admin seeds the pool
+        const initialAmount = new anchor.BN(10_000_000);
+        await program.methods
+            .depositLiquidity(initialAmount, initialAmount)
+            .accountsPartial({
+                pool: values.poolKey,
+                poolAuthority: values.poolAuthority,
+                depositor: values.admin.publicKey,
+                mintLiquidity: values.mintLiquidity,
+                mintA: values.mintAKeypair.publicKey,
+                mintB: values.mintBKeypair.publicKey,
+                poolAccountA: values.poolAccountA,
+                poolAccountB: values.poolAccountB,
+                depositorAccountLiquidity: values.liquidityAccount,
+                depositorAccountA: values.holderAccountA,
+                depositorAccountB: values.holderAccountB,
+            })
+            .signers([values.admin])
+            .rpc({ skipPreflight: true });
+
+        // 2. Swaps accrue fees to the pool while LP supply stays fixed
+        for (const swapA of [true, false, true, false]) {
+            await program.methods
+                .swapExactTokensForTokens(swapA, new anchor.BN(1_000_000), new anchor.BN(1))
+                .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 reserveABefore = new anchor.BN(
+            (await connection.getTokenAccountBalance(values.poolAccountA)).value.amount,
+        );
+        const reserveBBefore = new anchor.BN(
+            (await connection.getTokenAccountBalance(values.poolAccountB)).value.amount,
+        );
+        const supplyBefore = new anchor.BN((await connection.getTokenSupply(values.mintLiquidity)).value.amount);
+        const totalBefore = supplyBefore.add(values.minimumLiquidity);
+
+        // 3. Second depositor joins
+        await program.methods
+            .depositLiquidity(initialAmount, initialAmount)
+            .accountsPartial({
+                pool: values.poolKey,
+                poolAuthority: values.poolAuthority,
+                depositor: depositor.publicKey,
+                mintLiquidity: values.mintLiquidity,
+                mintA: values.mintAKeypair.publicKey,
+                mintB: values.mintBKeypair.publicKey,
+                poolAccountA: values.poolAccountA,
+                poolAccountB: values.poolAccountB,
+                depositorAccountLiquidity,
+                depositorAccountA,
+                depositorAccountB,
+            })
+            .signers([depositor])
+            .rpc({ skipPreflight: true });
+
+        const depositedA = new anchor.BN(
+            (await connection.getTokenAccountBalance(values.poolAccountA)).value.amount,
+        ).sub(reserveABefore);
+        const depositedB = new anchor.BN(
+            (await connection.getTokenAccountBalance(values.poolAccountB)).value.amount,
+        ).sub(reserveBBefore);
+        const liquidity = new anchor.BN(
+            (await connection.getTokenAccountBalance(depositorAccountLiquidity)).value.amount,
+        );
+
+        // LP minted must be pro-rata to the existing supply, not sqrt(a * b)
+        const expectedLiquidity = anchor.BN.min(
+            depositedA.mul(totalBefore).div(reserveABefore),
+            depositedB.mul(totalBefore).div(reserveBBefore),
+        );
+        expect(liquidity.toString()).to.equal(expectedLiquidity.toString());
+
+        // 4. Withdrawing everything must not return more than was deposited
+        await program.methods
+            .withdrawLiquidity(liquidity)
+            .accountsPartial({
+                amm: values.ammKey,
+                pool: values.poolKey,
+                poolAuthority: values.poolAuthority,
+                depositor: depositor.publicKey,
+                mintLiquidity: values.mintLiquidity,
+                mintA: values.mintAKeypair.publicKey,
+                mintB: values.mintBKeypair.publicKey,
+                poolAccountA: values.poolAccountA,
+                poolAccountB: values.poolAccountB,
+                depositorAccountLiquidity,
+                depositorAccountA,
+                depositorAccountB,
+            })
+            .signers([depositor])
+            .rpc({ skipPreflight: true });
+
+        const receivedA = new anchor.BN((await connection.getTokenAccountBalance(depositorAccountA)).value.amount)
+            .sub(values.defaultSupply)
+            .add(depositedA);
+        const receivedB = new anchor.BN((await connection.getTokenAccountBalance(depositorAccountB)).value.amount)
+            .sub(values.defaultSupply)
+            .add(depositedB);
+        expect(receivedA.lte(depositedA), `received ${receivedA} A for a ${depositedA} deposit`).to.be.true;
+        expect(receivedB.lte(depositedB), `received ${receivedB} B for a ${depositedB} deposit`).to.be.true;
+        // Sanity: rounding only costs dust
+        expect(receivedA.gt(depositedA.muln(999).divn(1000))).to.be.true;
+        expect(receivedB.gt(depositedB.muln(999).divn(1000))).to.be.true;
+    });
+
+    it('First deposit succeeds after token B is donated to the pool', async () => {
+        await transfer(connection, values.admin, values.holderAccountB, values.poolAccountB, values.admin, 1);
+
+        await program.methods
+            .depositLiquidity(values.depositAmountA, values.depositAmountA)
+            .accountsPartial({
+                pool: values.poolKey,
+                poolAuthority: values.poolAuthority,
+                depositor: values.admin.publicKey,
+                mintLiquidity: values.mintLiquidity,
+                mintA: values.mintAKeypair.publicKey,
+                mintB: values.mintBKeypair.publicKey,
+                poolAccountA: values.poolAccountA,
+                poolAccountB: values.poolAccountB,
+                depositorAccountLiquidity: values.liquidityAccount,
+                depositorAccountA: values.holderAccountA,
+                depositorAccountB: values.holderAccountB,
+            })
+            .signers([values.admin])
+            .rpc();
+
+        const liquidity = await connection.getTokenAccountBalance(values.liquidityAccount);
+        expect(liquidity.value.amount).to.equal(values.depositAmountA.sub(values.minimumLiquidity).toString());
+        const poolAccountA = await connection.getTokenAccountBalance(values.poolAccountA);
+        expect(poolAccountA.value.amount).to.equal(values.depositAmountA.toString());
+        const poolAccountB = await connection.getTokenAccountBalance(values.poolAccountB);
+        expect(poolAccountB.value.amount).to.equal(values.depositAmountA.addn(1).toString());
+    });
+
+    it('First deposit mints liquidity after token A is donated to the pool', async () => {
+        await transfer(connection, values.admin, values.holderAccountA, values.poolAccountA, values.admin, 1);
+
+        await program.methods
+            .depositLiquidity(values.depositAmountA, values.depositAmountA)
+            .accountsPartial({
+                pool: values.poolKey,
+                poolAuthority: values.poolAuthority,
+                depositor: values.admin.publicKey,
+                mintLiquidity: values.mintLiquidity,
+                mintA: values.mintAKeypair.publicKey,
+                mintB: values.mintBKeypair.publicKey,
+                poolAccountA: values.poolAccountA,
+                poolAccountB: values.poolAccountB,
+                depositorAccountLiquidity: values.liquidityAccount,
+                depositorAccountA: values.holderAccountA,
+                depositorAccountB: values.holderAccountB,
+            })
+            .signers([values.admin])
+            .rpc();
+
+        const liquidity = await connection.getTokenAccountBalance(values.liquidityAccount);
+        expect(liquidity.value.amount).to.equal(values.depositAmountA.sub(values.minimumLiquidity).toString());
+        const poolAccountA = await connection.getTokenAccountBalance(values.poolAccountA);
+        expect(poolAccountA.value.amount).to.equal(values.depositAmountA.addn(1).toString());
+        const poolAccountB = await connection.getTokenAccountBalance(values.poolAccountB);
+        expect(poolAccountB.value.amount).to.equal(values.depositAmountA.toString());
+    });
 });