Contracts: Add safe math guards and overflow assertions for fee calcu…

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

contracts/vault/src/lib.rs

@@ -192,10 +192,8 @@ pub enum VaultError {
     NoPendingWithdrawal = 8,
     /// Strategy allocation would leave idle liquidity below the configured buffer.
     LiquidityBufferNotMet = 9,
-    /// Strategy allocation exceeds configured cap.
-    ExceedsStrategyCap = 10,
-    /// Strategy allocation exceeds configured risk threshold.
-    ExceedsRiskThreshold = 11,
+    /// Arithmetic overflow or underflow occurred while calculating fees or accounting values.
+    MathOverflow = 10,
 }
 
 #[contractclient(name = "KoreanDebtStrategyClient")]
@@ -1040,20 +1038,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
@@ -1063,7 +1067,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)?,
             );
         }
 
@@ -1074,12 +1080,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<i128, VaultError> {
+        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 ──────────────────────────────────────────────────

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() {