fix(contracts): harden treasury withdrawals with enforceable guardrails

contracts/src/lib.rs

@@ -1017,6 +1017,31 @@ impl NesteraContract {
         treasury::get_reserve_balance(&env)
     }
 
+    /// Returns treasury withdrawal security limits.
+    pub fn get_treasury_limits(env: Env) -> treasury::types::TreasurySecurityConfig {
+        treasury::get_treasury_limits(&env)
+    }
+
+    /// Updates treasury withdrawal limits (admin only).
+    pub fn set_treasury_limits(
+        env: Env,
+        admin: Address,
+        max_withdrawal_per_tx: i128,
+        daily_withdrawal_cap: i128,
+    ) -> Result<treasury::types::TreasurySecurityConfig, SavingsError> {
+        treasury::set_treasury_limits(&env, &admin, max_withdrawal_per_tx, daily_withdrawal_cap)
+    }
+
+    /// Withdraws from a treasury pool with per-tx and daily caps (admin only).
+    pub fn withdraw_treasury(
+        env: Env,
+        admin: Address,
+        pool: treasury::types::TreasuryPool,
+        amount: i128,
+    ) -> Result<treasury::types::Treasury, SavingsError> {
+        treasury::withdraw_treasury(&env, &admin, pool, amount)
+    }
+
     /// Allocates the unallocated treasury balance into reserves, rewards, and operations.
     /// Percentages are in basis points and must sum to 10_000.
     pub fn allocate_treasury(

contracts/src/storage_types.rs

@@ -149,6 +149,10 @@ pub enum DataKey {
     ConfigInitialized,
     /// Treasury allocation config (reserve/rewards/operations percentages)
     AllocationConfig,
+    /// Treasury security limits for admin withdrawals
+    TreasurySecurityConfig,
+    /// Daily treasury withdrawal tracker (timestamp + amount)
+    TreasuryDailyWithdrawal,
     /// Early break fee (basis points) for goal saves
     EarlyBreakFeeBps,
     /// Fee recipient for protocol/treasury fees

contracts/src/treasury/mod.rs

@@ -1,12 +1,18 @@
 pub mod types;
 
+#[cfg(test)]
+mod security_tests;
 #[cfg(test)]
 mod views_tests;
 
 use crate::errors::SavingsError;
 use crate::storage_types::DataKey;
 use soroban_sdk::{symbol_short, Address, Env};
-use types::{AllocationConfig, Treasury};
+use types::{
+    AllocationConfig, Treasury, TreasuryDailyWithdrawal, TreasuryPool, TreasurySecurityConfig,
+};
+
+const TREASURY_DAILY_WINDOW_SECS: u64 = 24 * 60 * 60;
 
 // ========== Treasury Storage Helpers ==========
 
@@ -23,13 +29,68 @@ fn set_treasury(env: &Env, treasury: &Treasury) {
     env.storage().persistent().set(&DataKey::Treasury, treasury);
 }
 
+fn require_admin(env: &Env, admin: &Address) -> Result<(), SavingsError> {
+    let stored_admin: Address = env
+        .storage()
+        .instance()
+        .get(&DataKey::Admin)
+        .ok_or(SavingsError::Unauthorized)?;
+    if stored_admin != *admin {
+        return Err(SavingsError::Unauthorized);
+    }
+    admin.require_auth();
+    Ok(())
+}
+
+fn validate_treasury_state(treasury: &Treasury) -> Result<(), SavingsError> {
+    if treasury.total_fees_collected < 0
+        || treasury.total_yield_earned < 0
+        || treasury.reserve_balance < 0
+        || treasury.treasury_balance < 0
+        || treasury.rewards_balance < 0
+        || treasury.operations_balance < 0
+    {
+        return Err(SavingsError::InvariantViolation);
+    }
+    Ok(())
+}
+
+fn get_treasury_limits_internal(env: &Env) -> TreasurySecurityConfig {
+    env.storage()
+        .persistent()
+        .get(&DataKey::TreasurySecurityConfig)
+        .unwrap_or(TreasurySecurityConfig::default_limits())
+}
+
+fn set_treasury_limits_internal(env: &Env, limits: &TreasurySecurityConfig) {
+    env.storage()
+        .persistent()
+        .set(&DataKey::TreasurySecurityConfig, limits);
+}
+
+fn get_daily_withdrawal_tracker(env: &Env) -> TreasuryDailyWithdrawal {
+    let now = env.ledger().timestamp();
+    env.storage()
+        .persistent()
+        .get(&DataKey::TreasuryDailyWithdrawal)
+        .unwrap_or(TreasuryDailyWithdrawal::new(now))
+}
+
+fn set_daily_withdrawal_tracker(env: &Env, tracker: &TreasuryDailyWithdrawal) {
+    env.storage()
+        .persistent()
+        .set(&DataKey::TreasuryDailyWithdrawal, tracker);
+}
+
 // ========== Treasury Initialization ==========
 
 /// Initializes the treasury with default zero values.
 /// Called during `initialize_config`.
 pub fn initialize_treasury(env: &Env) {
     let treasury = Treasury::new();
     set_treasury(env, &treasury);
+    set_treasury_limits_internal(env, &TreasurySecurityConfig::default_limits());
+    set_daily_withdrawal_tracker(env, &TreasuryDailyWithdrawal::new(env.ledger().timestamp()));
 }
 
 // ========== Fee Recording ==========
@@ -45,14 +106,19 @@ pub fn record_fee(env: &Env, amount: i128, fee_type: soroban_sdk::Symbol) {
         return;
     }
     let mut treasury = get_treasury(env);
-    treasury.total_fees_collected = treasury
-        .total_fees_collected
-        .checked_add(amount)
-        .unwrap_or(treasury.total_fees_collected);
-    treasury.treasury_balance = treasury
-        .treasury_balance
-        .checked_add(amount)
-        .unwrap_or(treasury.treasury_balance);
+    if let Some(updated_total_fees) = treasury.total_fees_collected.checked_add(amount) {
+        treasury.total_fees_collected = updated_total_fees;
+    } else {
+        return;
+    }
+    if let Some(updated_treasury_balance) = treasury.treasury_balance.checked_add(amount) {
+        treasury.treasury_balance = updated_treasury_balance;
+    } else {
+        return;
+    }
+    if validate_treasury_state(&treasury).is_err() {
+        return;
+    }
     set_treasury(env, &treasury);
 
     env.events()
@@ -65,10 +131,14 @@ pub fn record_yield(env: &Env, amount: i128) {
         return;
     }
     let mut treasury = get_treasury(env);
-    treasury.total_yield_earned = treasury
-        .total_yield_earned
-        .checked_add(amount)
-        .unwrap_or(treasury.total_yield_earned);
+    if let Some(updated_total_yield) = treasury.total_yield_earned.checked_add(amount) {
+        treasury.total_yield_earned = updated_total_yield;
+    } else {
+        return;
+    }
+    if validate_treasury_state(&treasury).is_err() {
+        return;
+    }
     set_treasury(env, &treasury);
 }
 
@@ -94,6 +164,112 @@ pub fn get_reserve_balance(env: &Env) -> i128 {
     get_treasury(env).reserve_balance
 }
 
+/// Returns current treasury withdrawal safety limits.
+pub fn get_treasury_limits(env: &Env) -> TreasurySecurityConfig {
+    get_treasury_limits_internal(env)
+}
+
+/// Updates treasury withdrawal limits (admin only).
+pub fn set_treasury_limits(
+    env: &Env,
+    admin: &Address,
+    max_withdrawal_per_tx: i128,
+    daily_withdrawal_cap: i128,
+) -> Result<TreasurySecurityConfig, SavingsError> {
+    require_admin(env, admin)?;
+
+    if max_withdrawal_per_tx <= 0 || daily_withdrawal_cap <= 0 {
+        return Err(SavingsError::InvalidAmount);
+    }
+    if max_withdrawal_per_tx > daily_withdrawal_cap {
+        return Err(SavingsError::AmountExceedsLimit);
+    }
+
+    let limits = TreasurySecurityConfig {
+        max_withdrawal_per_tx,
+        daily_withdrawal_cap,
+    };
+    set_treasury_limits_internal(env, &limits);
+    env.events().publish(
+        (symbol_short!("trs_lim"),),
+        (max_withdrawal_per_tx, daily_withdrawal_cap),
+    );
+
+    Ok(limits)
+}
+
+/// Withdraws from a treasury sub-balance with per-tx and daily safety caps (admin only).
+pub fn withdraw_treasury(
+    env: &Env,
+    admin: &Address,
+    pool: TreasuryPool,
+    amount: i128,
+) -> Result<Treasury, SavingsError> {
+    require_admin(env, admin)?;
+    if amount <= 0 {
+        return Err(SavingsError::InvalidAmount);
+    }
+
+    let limits = get_treasury_limits_internal(env);
+    if amount > limits.max_withdrawal_per_tx {
+        return Err(SavingsError::AmountExceedsLimit);
+    }
+
+    let now = env.ledger().timestamp();
+    let mut daily = get_daily_withdrawal_tracker(env);
+    if now.saturating_sub(daily.window_start_ts) >= TREASURY_DAILY_WINDOW_SECS {
+        daily = TreasuryDailyWithdrawal::new(now);
+    }
+
+    let new_daily_total = daily
+        .withdrawn_amount
+        .checked_add(amount)
+        .ok_or(SavingsError::Overflow)?;
+    if new_daily_total > limits.daily_withdrawal_cap {
+        return Err(SavingsError::AmountExceedsLimit);
+    }
+
+    let mut treasury = get_treasury(env);
+    match pool.clone() {
+        TreasuryPool::Reserve => {
+            if amount > treasury.reserve_balance {
+                return Err(SavingsError::InsufficientBalance);
+            }
+            treasury.reserve_balance = treasury
+                .reserve_balance
+                .checked_sub(amount)
+                .ok_or(SavingsError::InsufficientBalance)?;
+        }
+        TreasuryPool::Rewards => {
+            if amount > treasury.rewards_balance {
+                return Err(SavingsError::InsufficientBalance);
+            }
+            treasury.rewards_balance = treasury
+                .rewards_balance
+                .checked_sub(amount)
+                .ok_or(SavingsError::InsufficientBalance)?;
+        }
+        TreasuryPool::Operations => {
+            if amount > treasury.operations_balance {
+                return Err(SavingsError::InsufficientBalance);
+            }
+            treasury.operations_balance = treasury
+                .operations_balance
+                .checked_sub(amount)
+                .ok_or(SavingsError::InsufficientBalance)?;
+        }
+    }
+
+    validate_treasury_state(&treasury)?;
+    daily.withdrawn_amount = new_daily_total;
+    set_treasury(env, &treasury);
+    set_daily_withdrawal_tracker(env, &daily);
+    env.events()
+        .publish((symbol_short!("trs_wth"),), (pool, amount, new_daily_total));
+
+    Ok(treasury)
+}
+
 // ========== Allocation Logic ==========
 
 /// Allocates the unallocated treasury balance into reserves, rewards, and operations.
@@ -118,16 +294,7 @@ pub fn allocate_treasury(
     rewards_percent: u32,
     operations_percent: u32,
 ) -> Result<Treasury, SavingsError> {
-    // Verify admin
-    let stored_admin: Address = env
-        .storage()
-        .instance()
-        .get(&DataKey::Admin)
-        .ok_or(SavingsError::Unauthorized)?;
-    if stored_admin != *admin {
-        return Err(SavingsError::Unauthorized);
-    }
-    admin.require_auth();
+    require_admin(env, admin)?;
 
     // Validate percentages sum to 100%
     let total = reserve_percent
@@ -177,6 +344,7 @@ pub fn allocate_treasury(
 
     // Zero out the unallocated balance
     treasury.treasury_balance = 0;
+    validate_treasury_state(&treasury)?;
 
     // Store the allocation config for reference
     let alloc_config = AllocationConfig {