Security: Fix critical vulnerabilities in AMM program

programs/ammv2/src/error.rs

@@ -10,4 +10,6 @@ pub enum ErrorCode {
     BurnTooMuch,
     #[msg("Not enough out")]
     NotEnoughOut,
+    #[msg("Invalid fee parameters: denominator must be > 0, numerator must be <= denominator")]
+    InvalidFees,
 }

programs/ammv2/src/instructions/init_pool.rs

@@ -4,17 +4,22 @@ use anchor_spl::{
     token::{Mint, Token, TokenAccount},
 };
 use crate::state::PoolState;
+use crate::error::ErrorCode;
 
 pub fn handler(
     ctx: Context<InitializePool>, 
     fee_numerator: u64,
     fee_denominator: u64,
 ) -> Result<()> {
 
+    // Validate fee parameters to prevent DoS and abuse
+    require!(fee_denominator > 0, ErrorCode::InvalidFees);
+    require!(fee_numerator <= fee_denominator, ErrorCode::InvalidFees);
+
     let pool_state = &mut ctx.accounts.pool_state;
     pool_state.fee_numerator = fee_numerator;
     pool_state.fee_denominator = fee_denominator;
-    pool_state.total_amount_minted = 0; 
+    pool_state.total_amount_minted = 0;
 
     Ok(())
 }

programs/ammv2/src/instructions/liquidity.rs

@@ -37,14 +37,16 @@ pub fn add_liquidity(
     msg!("init deposits: {} {}", amount_liq0, amount_liq1);
 
     if vault_balance0 == 0 && vault_balance1 == 0 {
-        // bit shift (a + b)/2
-        amount_to_mint = (amount_liq0 + amount_liq1) >> 1; 
+        // Use checked arithmetic to prevent overflow in release mode
+        amount_to_mint = amount_liq0.checked_add(amount_liq1).unwrap() >> 1;
         deposit1 = amount_liq1;
     } else { 
         // require equal amount deposit based on pool exchange rate 
-        let exchange10 = vault_balance1.checked_div(vault_balance0).unwrap();
-        let amount_deposit_1 = amount_liq0.checked_mul(exchange10).unwrap();
-        msg!("new deposits: {} {} {}", exchange10, amount_liq0, amount_deposit_1);
+        // Multiply before dividing to avoid integer truncation precision loss
+        let amount_deposit_1 = (amount_liq0 as u128)
+            .checked_mul(vault_balance1 as u128).unwrap()
+            .checked_div(vault_balance0 as u128).unwrap() as u64;
+        msg!("new deposits: {} {}", amount_liq0, amount_deposit_1);
 
         // enough funds + user is ok with it in single check 
         require!(amount_deposit_1 <= amount_liq1, ErrorCode::NotEnoughBalance);
@@ -66,7 +68,7 @@ pub fn add_liquidity(
     require!(amount_to_mint > 0, ErrorCode::NoPoolMintOutput);
 
     // give pool_mints 
-    pool_state.total_amount_minted += amount_to_mint;
+    pool_state.total_amount_minted = pool_state.total_amount_minted.checked_add(amount_to_mint).unwrap();
     let mint_ctx = CpiContext::new(
         ctx.accounts.token_program.to_account_info(), 
         MintTo {
@@ -162,7 +164,7 @@ pub fn remove_liquidity(
         }
     ).with_signer(&[pda_sign]), burn_amount)?;
 
-    state.total_amount_minted -= burn_amount; 
+    state.total_amount_minted = state.total_amount_minted.checked_sub(burn_amount).unwrap();
 
     Ok(())
 }

programs/ammv2/src/instructions/swap.rs

@@ -78,9 +78,9 @@ pub struct Swap<'info> {
         constraint=vault_src.mint == user_src.mint,
     )]
     pub vault_src: Box<Account<'info, TokenAccount>>, 
-    #[account(mut, 
+    #[account(mut,
         constraint=vault_dst.owner == pool_authority.key(),
-        constraint=vault_src.mint == user_src.mint,
+        constraint=vault_dst.mint == user_dst.mint,
     )]
     pub vault_dst: Box<Account<'info, TokenAccount>>,