From 9e18fc9121d06ba675d93dba1eb7c97784712f3e Mon Sep 17 00:00:00 2001 From: OluRemiFour Date: Sat, 30 May 2026 22:08:44 +0100 Subject: [PATCH 1/2] Contracts: Add safe math guards and overflow assertions for fee calculations --- .../vault/proptest-regressions/fuzz_math.txt | 1 + contracts/vault/src/lib.rs | 46 +++++++++++++++---- contracts/vault/src/test.rs | 41 +++++++++++++++++ 3 files changed, 78 insertions(+), 10 deletions(-) diff --git a/contracts/vault/proptest-regressions/fuzz_math.txt b/contracts/vault/proptest-regressions/fuzz_math.txt index 57c519eb..293ada94 100644 --- a/contracts/vault/proptest-regressions/fuzz_math.txt +++ b/contracts/vault/proptest-regressions/fuzz_math.txt @@ -5,3 +5,4 @@ # It is recommended to check this file in to source control so that # everyone who runs the test benefits from these saved cases. cc 5fb6e9c997750435f990105f21d40fa1e16d428d3130a6b68acf02703a302c11 # shrinks to deposit_amount = 1 +cc 76bcf8ca3dcb1b07962c6d984fe8ba32d6fe23c613492f646a7b084a726e9656 # shrinks to deposit_amount = 6384692 diff --git a/contracts/vault/src/lib.rs b/contracts/vault/src/lib.rs index 490fbbc9..a7d5b87f 100644 --- a/contracts/vault/src/lib.rs +++ b/contracts/vault/src/lib.rs @@ -190,6 +190,8 @@ pub enum VaultError { NoPendingWithdrawal = 8, /// Strategy allocation would leave idle liquidity below the configured buffer. LiquidityBufferNotMet = 9, + /// Arithmetic overflow or underflow occurred while calculating fees or accounting values. + MathOverflow = 10, } #[contractclient(name = "KoreanDebtStrategyClient")] @@ -1019,20 +1021,26 @@ impl YieldVault { /// Admin function to artificially accrue yield, deducting the protocol fee. /// The fee portion is credited to the treasury balance. - pub fn accrue_yield(env: Env, amount: i128) { + pub fn accrue_yield(env: Env, amount: i128) -> Result<(), VaultError> { let admin: Address = get_admin(&env).expect("Admin not set"); admin.require_auth(); + if amount <= 0 { + return Err(VaultError::InvalidAmount); + } + + // Goal 1: deduct protocol fee before distributing to depositors. + // Calculate before the token transfer so overflow/underflow reverts cleanly. + let fee_bps: i128 = env.storage().instance().get(&DataKey::FeeBps).unwrap_or(0); + let fee_amount = Self::calculate_fee_amount(amount, fee_bps)?; + let net_yield = amount + .checked_sub(fee_amount) + .ok_or(VaultError::MathOverflow)?; + let token_addr = Self::token(env.clone()); let token_client = token::Client::new(&env, &token_addr); - token_client.transfer(&admin, &env.current_contract_address(), &amount); - // Goal 1: deduct protocol fee before distributing to depositors - let fee_bps: i128 = env.storage().instance().get(&DataKey::FeeBps).unwrap_or(0); - let fee_amount = amount.checked_mul(fee_bps).expect("overflow") / 10_000; - let net_yield = amount.checked_sub(fee_amount).expect("underflow"); - // Accumulate fee in treasury balance if fee_amount > 0 { let treasury_bal: i128 = env @@ -1042,7 +1050,9 @@ impl YieldVault { .unwrap_or(0); env.storage().instance().set( &DataKey::TreasuryBalance, - &treasury_bal.checked_add(fee_amount).expect("overflow"), + &treasury_bal + .checked_add(fee_amount) + .ok_or(VaultError::MathOverflow)?, ); } @@ -1053,12 +1063,28 @@ impl YieldVault { .unwrap_or(0); env.storage().instance().set( &DataKey::TotalAssets, - &ta.checked_add(net_yield).expect("overflow"), + &ta.checked_add(net_yield).ok_or(VaultError::MathOverflow)?, ); let mut state = Self::get_state(&env); - state.total_assets = state.total_assets.checked_add(net_yield).expect("overflow"); + state.total_assets = state + .total_assets + .checked_add(net_yield) + .ok_or(VaultError::MathOverflow)?; env.storage().instance().set(&DataKey::State, &state); + Ok(()) + } + + fn calculate_fee_amount(amount: i128, fee_bps: i128) -> Result { + if amount <= 0 || !(0..=10_000).contains(&fee_bps) { + return Err(VaultError::InvalidAmount); + } + + amount + .checked_mul(fee_bps) + .ok_or(VaultError::MathOverflow)? + .checked_div(10_000) + .ok_or(VaultError::MathOverflow) } // ── Goal 1: Protocol fee ────────────────────────────────────────────────── diff --git a/contracts/vault/src/test.rs b/contracts/vault/src/test.rs index 1e8fcd05..736acba7 100644 --- a/contracts/vault/src/test.rs +++ b/contracts/vault/src/test.rs @@ -408,6 +408,47 @@ fn test_accrue_yield_increases_total_assets() { // ─── 5. report_benji_yield ─────────────────────────────────────────────────── +#[test] +fn test_accrue_yield_rejects_zero_amount() { + let env = Env::default(); + env.mock_all_auths(); + + let (vault, _, _, _) = setup_vault(&env); + + let result = vault.try_accrue_yield(&0); + assert!(matches!(result, Err(Ok(VaultError::InvalidAmount)))); +} + +#[test] +fn test_accrue_yield_fee_math_overflow_reverts_before_transfer() { + let env = Env::default(); + env.mock_all_auths(); + + let (vault, usdc, _, admin) = setup_vault(&env); + vault.set_fee_bps(&10_000); + + let result = vault.try_accrue_yield(&i128::MAX); + assert!(matches!(result, Err(Ok(VaultError::MathOverflow)))); + assert_eq!(usdc.balance(&admin), 0); + assert_eq!(vault.total_assets(), 0); + assert_eq!(vault.treasury_balance(), 0); +} + +#[test] +fn test_accrue_yield_full_fee_accumulates_to_treasury_only() { + let env = Env::default(); + env.mock_all_auths(); + + let (vault, _, usdc_sa, admin) = setup_vault(&env); + usdc_sa.mint(&admin, &250); + + vault.set_fee_bps(&10_000); + vault.accrue_yield(&250); + + assert_eq!(vault.total_assets(), 0); + assert_eq!(vault.treasury_balance(), 250); +} + #[test] #[should_panic] fn test_report_benji_yield_wrong_strategy_panics() { From 4a2434b73895367020c69f6a8782474fe453118a Mon Sep 17 00:00:00 2001 From: OluRemiFour Date: Sat, 30 May 2026 22:11:25 +0100 Subject: [PATCH 2/2] quick fix [ci skip]