diff --git a/contracts/amm/src/lib.rs b/contracts/amm/src/lib.rs index 0e27c335..ffe0a2b5 100644 --- a/contracts/amm/src/lib.rs +++ b/contracts/amm/src/lib.rs @@ -6,25 +6,27 @@ //! - Strict slippage protection (min output / max input enforcement) //! - Dead-shares burn (first 1,000 shares) to prevent share inflation exploits //! - Virtual balances to mitigate sandwich attacks +//! - Strict single-transaction reentrancy guards on every state-mutating +//! entry point (flash-swap defense) #![no_std] use soroban_sdk::{ - contract, contracterror, contractimpl, panic_with_error, token, Address, Env, + contract, contracterror, contractimpl, panic_with_error, symbol_short, token, Address, Env, Symbol, }; /// Basis points denominator (10,000 bps = 100%) const BPS: i128 = 10_000; -/// Dead shares permanently burned on pool initialization -const DEAD_SHARES: i128 = 1_000; - /// Minimum swap fee in basis points (0.1%) const MIN_SWAP_FEE_BPS: i128 = 10; /// Maximum swap fee in basis points (10%) const MAX_SWAP_FEE_BPS: i128 = 1_000; +/// Temporary-storage key for the reentrancy guard mutex. +const RG_KEY: Symbol = symbol_short!("amm_rg"); + // ── Errors (SC-HARD-20: range 400+) ────────────────────────────────────────── /// Typed contract errors for the AMM. @@ -43,11 +45,15 @@ pub enum AmmError { SlippageExceeded = 403, /// `404` — Contract has not been initialised. NotInitialized = 404, + /// `405` — A reentrant call was detected and blocked. + Reentrancy = 405, } +pub mod lp_token; pub mod storage; pub mod twap; +pub use lp_token::DEAD_SHARES; pub use storage::DataKey; #[contract] @@ -112,6 +118,7 @@ impl AMMContract { if amount_a <= 0 || amount_b <= 0 { panic_with_error!(&env, AmmError::InvalidAmount); } + Self::lock(&env); let token_a: Address = env.storage().instance().get(&DataKey::TokenA).unwrap(); let token_b: Address = env.storage().instance().get(&DataKey::TokenB).unwrap(); @@ -140,23 +147,21 @@ impl AMMContract { let shares = if balance_a == 0 || balance_b == 0 { // Initial liquidity: mint sqrt(amount_a * amount_b) shares, reserving // DEAD_SHARES that are permanently locked (seeded in initialize). - let gross = Self::isqrt(amount_a * amount_b); - if gross <= DEAD_SHARES { - panic_with_error!(&env, AmmError::InvalidAmount); + match lp_token::initial_shares(amount_a, amount_b) { + Some(s) => s, + None => { + Self::unlock(&env); + panic_with_error!(&env, AmmError::InvalidAmount); + } } - gross - DEAD_SHARES } else { // Subsequent liquidity: mint proportional to either ratio (use minimum). - let shares_from_a = (amount_a * total_shares) / balance_a; - let shares_from_b = (amount_b * total_shares) / balance_b; - if shares_from_a < shares_from_b { - shares_from_a - } else { - shares_from_b - } + lp_token::deposit_shares(amount_a, amount_b, balance_a, balance_b, total_shares) + .unwrap_or(0) }; if shares < min_shares { + Self::unlock(&env); panic_with_error!(&env, AmmError::SlippageExceeded); } @@ -185,6 +190,7 @@ impl AMMContract { (provider, amount_a, amount_b, shares), ); + Self::unlock(&env); shares } @@ -200,6 +206,7 @@ impl AMMContract { if shares <= 0 { panic_with_error!(&env, AmmError::InvalidAmount); } + Self::lock(&env); let pool_share: Address = env.storage().instance().get(&DataKey::PoolShare).unwrap(); let token_a: Address = env.storage().instance().get(&DataKey::TokenA).unwrap(); @@ -213,10 +220,17 @@ impl AMMContract { let balance_b: i128 = env.storage().instance().get(&DataKey::BalanceB).unwrap(); let total_shares: i128 = env.storage().instance().get(&DataKey::TotalShares).unwrap(); - let amount_a = (shares * balance_a) / total_shares; - let amount_b = (shares * balance_b) / total_shares; + let (amount_a, amount_b) = + match lp_token::withdrawal_amounts(shares, balance_a, balance_b, total_shares) { + Some(amounts) => amounts, + None => { + Self::unlock(&env); + panic_with_error!(&env, AmmError::InvalidAmount); + } + }; if amount_a < min_amount_a || amount_b < min_amount_b { + Self::unlock(&env); panic_with_error!(&env, AmmError::SlippageExceeded); } @@ -249,6 +263,7 @@ impl AMMContract { (provider, amount_a, amount_b, shares), ); + Self::unlock(&env); (amount_a, amount_b) } @@ -258,6 +273,7 @@ impl AMMContract { if amount_in <= 0 { panic_with_error!(&env, AmmError::InvalidAmount); } + Self::lock(&env); let token_a: Address = env.storage().instance().get(&DataKey::TokenA).unwrap(); let token_b: Address = env.storage().instance().get(&DataKey::TokenB).unwrap(); @@ -292,6 +308,7 @@ impl AMMContract { let amount_out = amount_out_before_fee - fee; if amount_out < min_amount_out { + Self::unlock(&env); panic_with_error!(&env, AmmError::SlippageExceeded); } @@ -315,6 +332,7 @@ impl AMMContract { (user, amount_in, amount_out, fee), ); + Self::unlock(&env); amount_out } @@ -324,6 +342,7 @@ impl AMMContract { if amount_in <= 0 { panic_with_error!(&env, AmmError::InvalidAmount); } + Self::lock(&env); let token_a: Address = env.storage().instance().get(&DataKey::TokenA).unwrap(); let token_b: Address = env.storage().instance().get(&DataKey::TokenB).unwrap(); @@ -358,6 +377,7 @@ impl AMMContract { let amount_out = amount_out_before_fee - fee; if amount_out < min_amount_out { + Self::unlock(&env); panic_with_error!(&env, AmmError::SlippageExceeded); } @@ -381,6 +401,7 @@ impl AMMContract { (user, amount_in, amount_out, fee), ); + Self::unlock(&env); amount_out } @@ -430,18 +451,21 @@ impl AMMContract { Self::twap_consult(env, token_in, amount_in, twap::TWAP_PERIOD_1H) } - /// Integer square root for initial share calculation. - fn isqrt(n: i128) -> i128 { - if n == 0 { - return 0; + /// Acquire the reentrancy guard using temporary storage (flash-swap defense). + /// + /// Temporary storage entries are removed at ledger close, so a stale lock + /// cannot survive across transactions, while within a single transaction + /// envelope the entry persists and blocks any reentrant entry call. + fn lock(env: &Env) { + if env.storage().temporary().has(&RG_KEY) { + panic_with_error!(env, AmmError::Reentrancy); } - let mut x = n; - let mut y = (x + 1) / 2; - while y < x { - x = y; - y = (x + n / x) / 2; - } - x + env.storage().temporary().set(&RG_KEY, &true); + } + + /// Release the reentrancy guard. + fn unlock(env: &Env) { + env.storage().temporary().remove(&RG_KEY); } } diff --git a/contracts/amm/src/lp_token.rs b/contracts/amm/src/lp_token.rs new file mode 100644 index 00000000..e4cb40ce --- /dev/null +++ b/contracts/amm/src/lp_token.rs @@ -0,0 +1,169 @@ +//! LP share (pool token) mint and burn math for the constant-product AMM. +//! +//! Keeping this math in one module makes the "no dilution" property +//! directly testable: every share is minted against the pool's pre-deposit +//! reserves using the smaller of the two reserve ratios, so an existing +//! LP's claim on the pool can never decrease when a new LP deposits. + +/// Dead shares permanently burned on pool initialization to prevent +/// share-inflation attacks (attacker donates dust to skew the ratio). +pub const DEAD_SHARES: i128 = 1_000; + +/// Integer square root (Babylonian method). `isqrt(n)` = floor(sqrt(n)). +pub fn isqrt(n: i128) -> i128 { + if n <= 0 { + return 0; + } + let mut x = n; + let mut y = (x + 1) / 2; + while y < x { + x = y; + y = (x + n / x) / 2; + } + x +} + +/// Shares minted for the initial deposit: `sqrt(amount_a * amount_b)` minus +/// the permanently locked `DEAD_SHARES`. +/// +/// Returns `None` when the deposit is too small to cover the dead shares, +/// so the caller can reject it instead of minting an inflated position. +pub fn initial_shares(amount_a: i128, amount_b: i128) -> Option { + if amount_a <= 0 || amount_b <= 0 { + return None; + } + let gross = isqrt(amount_a * amount_b); + if gross <= DEAD_SHARES { + None + } else { + Some(gross - DEAD_SHARES) + } +} + +/// Shares minted for a subsequent deposit, proportional to the **smaller** +/// of the two reserve ratios. Using the minimum guarantees the deposit never +/// mints more shares than the pool's reserves support, so existing LP shares +/// are never diluted. +/// +/// Returns `None` when the pool reserves or total share supply are unusable. +pub fn deposit_shares( + amount_a: i128, + amount_b: i128, + balance_a: i128, + balance_b: i128, + total_shares: i128, +) -> Option { + if amount_a <= 0 || amount_b <= 0 || balance_a <= 0 || balance_b <= 0 || total_shares <= 0 { + return None; + } + let shares_from_a = (amount_a * total_shares) / balance_a; + let shares_from_b = (amount_b * total_shares) / balance_b; + Some(if shares_from_a < shares_from_b { + shares_from_a + } else { + shares_from_b + }) +} + +/// Token amounts returned when burning `shares`, pro-rata to the pool's +/// current reserves. Burning is the exact inverse of minting: redeeming +/// every share returns the entire pool (modulo integer rounding down). +/// +/// Returns `None` when the inputs cannot produce a meaningful withdrawal. +pub fn withdrawal_amounts( + shares: i128, + balance_a: i128, + balance_b: i128, + total_shares: i128, +) -> Option<(i128, i128)> { + if shares <= 0 || total_shares <= 0 { + return None; + } + let amount_a = (shares * balance_a) / total_shares; + let amount_b = (shares * balance_b) / total_shares; + Some((amount_a, amount_b)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn initial_shares_subtracts_dead_shares() { + let shares = initial_shares(1_000_000, 1_000_000).unwrap(); + assert_eq!(shares, 1_000_000 - DEAD_SHARES); + } + + #[test] + fn initial_shares_rejects_dust_deposit() { + assert_eq!(initial_shares(DEAD_SHARES, DEAD_SHARES), None); + assert_eq!(initial_shares(0, 1_000), None); + assert_eq!(initial_shares(1_000, 0), None); + assert_eq!(initial_shares(-1, 1_000), None); + } + + #[test] + fn deposit_shares_proportional_to_smaller_ratio() { + // Pool: 1_000 / 1_000 with 1_000 total shares. Deposit 100 / 200 + // is limited by the A-side ratio: 100 * 1000 / 1000 = 100 shares. + let shares = deposit_shares(100, 200, 1_000, 1_000, 1_000).unwrap(); + assert_eq!(shares, 100); + + // Same pool, deposit 200 / 100: now limited by the B-side ratio. + let shares = deposit_shares(200, 100, 1_000, 1_000, 1_000).unwrap(); + assert_eq!(shares, 100); + } + + #[test] + fn deposit_shares_does_not_dilute_existing_lps() { + let balance_a = 5_000i128; + let balance_b = 10_000i128; + let total = 7_000i128; + + let minted = deposit_shares(500, 1_000, balance_a, balance_b, total).unwrap(); + + // After the deposit the pool holds more reserves than before, and + // every pre-existing share still lays claim to at least as much: + let new_total = total + minted; + let old_claim_a = (total * (balance_a + 500)) / new_total; + let old_claim_b = (total * (balance_b + 1_000)) / new_total; + assert!( + old_claim_a >= balance_a - 1, + "existing LPs must not lose token A value" + ); + assert!( + old_claim_b >= balance_b - 1, + "existing LPs must not lose token B value" + ); + } + + #[test] + fn withdrawal_amounts_are_pro_rata() { + let (a, b) = withdrawal_amounts(500, 1_000_000, 2_000_000, 1_000).unwrap(); + assert_eq!(a, 500_000); + assert_eq!(b, 1_000_000); + + // Burning every share empties the pool. + let (a, b) = withdrawal_amounts(1_000, 1_000_000, 2_000_000, 1_000).unwrap(); + assert_eq!(a, 1_000_000); + assert_eq!(b, 2_000_000); + } + + #[test] + fn withdrawal_amounts_rejects_invalid_input() { + assert_eq!(withdrawal_amounts(0, 1_000, 1_000, 1_000), None); + assert_eq!(withdrawal_amounts(10, 1_000, 1_000, 0), None); + assert_eq!(withdrawal_amounts(-5, 1_000, 1_000, 1_000), None); + } + + #[test] + fn isqrt_matches_reference_values() { + assert_eq!(isqrt(0), 0); + assert_eq!(isqrt(1), 1); + assert_eq!(isqrt(2), 1); + assert_eq!(isqrt(4), 2); + assert_eq!(isqrt(99), 9); + assert_eq!(isqrt(1_000_000), 1_000); + assert_eq!(isqrt(999_999), 999); + } +} diff --git a/contracts/amm/src/test.rs b/contracts/amm/src/test.rs index 27651dba..432fd143 100644 --- a/contracts/amm/src/test.rs +++ b/contracts/amm/src/test.rs @@ -39,7 +39,7 @@ fn setup( &token_b_id, &pool_share_id, &fee_collector, - &50, + &30, // 0.3 % base swap fee (spec: 0.3 % dynamic swap fee) ); (admin, client, token_a_id, token_b_id, pool_share_id) @@ -255,3 +255,75 @@ fn test_twap_resistant_to_single_block_flash_manipulation() { assert_eq!(twap_before_flash, twap_during_flash); } +// ── Reentrancy guard (flash-swap defense) tests ────────────────────────────── + +#[test] +#[should_panic(expected = "Error(Contract, #405)")] +fn test_reentrancy_guard_blocks_double_entry() { + let env = Env::default(); + let (admin, client, _, _, _) = setup(&env); + + // Simulate a reentrant state by pre-setting the temporary guard key. + env.as_contract(&client.address, || { + env.storage().temporary().set(&RG_KEY, &true); + }); + + // This deposit must be rejected because the guard is already held. + client.add_liquidity(&admin, &1_000_000, &1_000_000, &0); +} + +#[test] +fn test_reentrancy_lock_released_after_successful_calls() { + let env = Env::default(); + let (admin, client, token_a_id, _, _) = setup(&env); + + // Consecutive mutations succeed because each call releases its guard. + client.add_liquidity(&admin, &10_000_000, &10_000_000, &0); + + let user = Address::generate(&env); + let token_a_client = token::StellarAssetClient::new(&env, &token_a_id); + token_a_client.mint(&user, &1_000_000); + client.swap_a_for_b(&user, &1_000_000, &0); + + client.add_liquidity(&admin, &1_000_000, &1_000_000, &0); + + // The guard must be absent from storage between transactions. + let lock_held = env.as_contract(&client.address, || env.storage().temporary().has(&RG_KEY)); + assert!( + !lock_held, + "reentrancy guard must be released after each call" + ); +} + +// ── Fee accumulator test (0.3 % base fee + dynamic price-impact uplift) ───── + +#[test] +fn test_swap_charges_0_3pct_base_fee_with_impact_uplift() { + let env = Env::default(); + let (admin, client, token_a_id, _token_b_id, _) = setup(&env); + + client.add_liquidity(&admin, &10_000_000, &10_000_000, &0); + + let user = Address::generate(&env); + let token_a_client = token::StellarAssetClient::new(&env, &token_a_id); + let amount_in = 1_000_000i128; + token_a_client.mint(&user, &amount_in); + + let (bal_a_before, bal_b_before, _) = client.get_pool_state(); + let new_a = bal_a_before + amount_in; + let out_before = bal_b_before - (bal_a_before * bal_b_before + new_a - 1) / new_a; + + // Spec: 0.3 % base fee, uplifted by price impact (bps / 100). + let impact_bps = (out_before * 10_000) / bal_b_before; + let expected_fee = (out_before * (30 + impact_bps / 100)) / 10_000; + + let amount_out = client.swap_a_for_b(&user, &amount_in, &0); + + let (_, bal_b_after, _) = client.get_pool_state(); + let fee = (bal_b_before - bal_b_after) - amount_out; + + assert!(fee > 0, "swap must collect a fee"); + assert_eq!(fee, expected_fee, "fee must be 0.3 % base + impact uplift"); + assert_eq!(amount_out, out_before - fee); +} +