feat(contract): add bounded withdrawal limits and anti-drain policy (#77)

CHANGELOG.md

@@ -7,6 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Bounded withdrawal / anti-drain controls: per-operation and rolling-period
+  asset caps (`set_withdraw_limits`), auditable period reset
+  (`reset_withdraw_period`), emergency override (`set_withdraw_limits_override`),
+  and `withdraw_batch` that cannot bypass the period aggregate. Defaults remain
+  unlimited (`0`) until configured. On-chain `version()` bumped to 3.
+
 ### Changed
 
 - Aggregate totals (`total_shares`, `total_assets`, user balances) now use

src/error.rs

@@ -31,4 +31,14 @@ pub enum Error {
     Paused = 9,
     /// The provided Wasm hash does not match the admin-approved expected hash.
     WasmHashMismatch = 10,
+    /// A single withdraw (or batch leg) exceeds the configured per-operation
+    /// asset limit.
+    WithdrawLimitExceeded = 11,
+    /// The rolling-period aggregate withdrawal limit would be exceeded.
+    WithdrawPeriodLimitExceeded = 12,
+    /// A batch withdraw was invoked with an empty share list.
+    EmptyBatch = 13,
+    /// Withdrawal-limit configuration is invalid (e.g. a positive period cap
+    /// with a zero-length period).
+    InvalidWithdrawLimit = 14,
 }

src/events.rs

@@ -56,3 +56,27 @@ pub fn upgrade(env: &Env, admin: &Address, new_wasm_hash: &BytesN<32>) {
     let topics = (Symbol::new(env, "upgrade"), admin.clone());
     env.events().publish(topics, new_wasm_hash.clone());
 }
+
+/// Publishes a `wd_limits` event recording the new per-operation and rolling-
+/// period withdrawal caps (units: underlying assets) and the period length
+/// in seconds.
+pub fn withdraw_limits(env: &Env, max_per_op: u128, max_per_period: u128, period_secs: u64) {
+    let topics = (Symbol::new(env, "wd_limits"),);
+    env.events()
+        .publish(topics, (max_per_op, max_per_period, period_secs));
+}
+
+/// Publishes a `wd_reset` event after an authorized rolling-period usage reset.
+/// `cleared` is the withdrawn amount that was zeroed; `at` is the new period
+/// start timestamp.
+pub fn withdraw_period_reset(env: &Env, cleared: u128, at: u64) {
+    let topics = (Symbol::new(env, "wd_reset"),);
+    env.events().publish(topics, (cleared, at));
+}
+
+/// Publishes a `wd_override` event recording whether the emergency withdrawal-
+/// limit override is enabled.
+pub fn withdraw_override(env: &Env, enabled: bool) {
+    let topics = (Symbol::new(env, "wd_override"),);
+    env.events().publish(topics, enabled);
+}

src/lib.rs

@@ -20,7 +20,7 @@ mod test;
 
 pub use error::Error;
 
-use soroban_sdk::{contract, contractimpl, contractmeta, token, Address, BytesN, Env};
+use soroban_sdk::{contract, contractimpl, contractmeta, token, Address, BytesN, Env, Vec};
 
 contractmeta!(
     key = "Description",
@@ -48,6 +48,14 @@ impl YieldVault {
         // from the outset rather than relying solely on read-time fallbacks.
         storage::set_min_deposit(&env, types::DEFAULT_MIN_DEPOSIT);
         storage::set_paused(&env, false);
+        // Withdrawal limits default to unlimited (`0`) with a 1-day rolling
+        // window ready for operators to tighten without a storage migration.
+        storage::set_max_withdraw_per_op(&env, 0);
+        storage::set_max_withdraw_per_period(&env, 0);
+        storage::set_withdraw_period_secs(&env, types::DEFAULT_WITHDRAW_PERIOD_SECS);
+        storage::set_period_withdrawn(&env, 0);
+        storage::set_period_started_at(&env, env.ledger().timestamp());
+        storage::set_withdraw_limits_override(&env, false);
         storage::extend_instance(&env);
         events::initialize(&env, &admin, &token);
         Ok(())
@@ -265,7 +273,9 @@ impl YieldVault {
     /// underlying assets, transferring them back to `from`.
     ///
     /// Requires authorization from `from`. Returns [`Error::InsufficientShares`]
-    /// if `from` does not hold enough shares.
+    /// if `from` does not hold enough shares. Subject to the configured
+    /// per-operation and rolling-period withdrawal limits (units: underlying
+    /// assets) unless the admin override is enabled.
     pub fn withdraw(env: Env, from: Address, shares: u128) -> Result<u128, Error> {
         storage::require_initialized(&env)?;
         from.require_auth();
@@ -286,21 +296,187 @@ impl YieldVault {
             return Err(Error::ZeroAmount);
         }
 
-        let new_total_shares = total_shares.saturating_sub(shares);
-        let new_total_assets = total_assets.saturating_sub(assets);
-        let new_user_balance = user_balance.saturating_sub(shares);
+        // Enforce limits before any state mutation so over-limit attempts fail
+        // atomically (host rolls back the whole invocation).
+        Self::enforce_withdraw_limits(&env, assets, assets)?;
+
+        Self::apply_withdraw(
+            &env,
+            &from,
+            shares,
+            assets,
+            user_balance,
+            total_shares,
+            total_assets,
+        );
+        Ok(assets)
+    }
 
-        storage::set_total_shares(&env, new_total_shares);
-        storage::set_total_assets(&env, new_total_assets);
-        storage::set_balance(&env, &from, new_user_balance);
+    /// Redeems multiple share amounts for `from` in one invocation.
+    ///
+    /// Each leg is checked against the per-operation asset cap; the *sum* of
+    /// redeemed assets is checked against the rolling-period aggregate so a
+    /// batch cannot bypass the period limit by splitting. On any failure the
+    /// whole batch reverts with no partial burns or transfers.
+    ///
+    /// Asset amounts are computed at the pre-batch exchange rate (same as a
+    /// sequence of single withdraws before any share-price change).
+    pub fn withdraw_batch(env: Env, from: Address, shares_list: Vec<u128>) -> Result<u128, Error> {
+        storage::require_initialized(&env)?;
+        from.require_auth();
 
-        let token_address = storage::get_token(&env);
-        let client = token::Client::new(&env, &token_address);
-        client.transfer(&env.current_contract_address(), &from, &(assets as i128));
+        if shares_list.is_empty() {
+            return Err(Error::EmptyBatch);
+        }
+
+        let user_balance = storage::get_balance(&env, &from);
+        let total_shares = storage::get_total_shares(&env);
+        let total_assets_vault = storage::get_total_assets(&env);
+
+        let mut legs: Vec<(u128, u128)> = Vec::new(&env);
+        let mut batch_assets: u128 = 0;
+        let mut shares_needed: u128 = 0;
+        let mut i = 0u32;
+        while i < shares_list.len() {
+            let shares = shares_list.get(i).unwrap();
+            if shares == 0 {
+                return Err(Error::ZeroShares);
+            }
+            shares_needed = shares_needed.saturating_add(shares);
+            let assets = math::convert_to_assets(shares, total_shares, total_assets_vault)?;
+            if assets == 0 {
+                return Err(Error::ZeroAmount);
+            }
+            // Per-operation limit applies to each leg individually.
+            let max_per_op = storage::get_max_withdraw_per_op(&env);
+            if !storage::is_withdraw_limits_override(&env) && max_per_op > 0 && assets > max_per_op
+            {
+                return Err(Error::WithdrawLimitExceeded);
+            }
+            batch_assets = batch_assets.saturating_add(assets);
+            legs.push_back((shares, assets));
+            i += 1;
+        }
+
+        if user_balance < shares_needed {
+            return Err(Error::InsufficientShares);
+        }
+
+        // Period aggregate sees the full batch sum — splitting cannot bypass it.
+        // Pass `op_assets = 0` here because per-op was already checked per leg.
+        Self::enforce_withdraw_limits(&env, 0, batch_assets)?;
+
+        let mut remaining_user = user_balance;
+        let mut rem_shares = total_shares;
+        let mut rem_assets = total_assets_vault;
+        let mut m = 0u32;
+        while m < legs.len() {
+            let (shares, assets) = legs.get(m).unwrap();
+            Self::apply_withdraw(
+                &env,
+                &from,
+                shares,
+                assets,
+                remaining_user,
+                rem_shares,
+                rem_assets,
+            );
+            remaining_user = remaining_user.saturating_sub(shares);
+            rem_shares = rem_shares.saturating_sub(shares);
+            rem_assets = rem_assets.saturating_sub(assets);
+            m += 1;
+        }
+
+        Ok(batch_assets)
+    }
 
+    /// Configures per-operation and rolling-period withdrawal limits.
+    ///
+    /// Units are underlying assets. `0` for either cap means unlimited.
+    /// `period_secs` must be non-zero when `max_per_period > 0`. Admin-only.
+    /// Emits a `wd_limits` event. Does not reset period usage; call
+    /// [`Self::reset_withdraw_period`] for an auditable usage clear.
+    pub fn set_withdraw_limits(
+        env: Env,
+        max_per_op: u128,
+        max_per_period: u128,
+        period_secs: u64,
+    ) -> Result<(), Error> {
+        storage::require_initialized(&env)?;
+        let admin = storage::get_admin(&env);
+        admin.require_auth();
+
+        if max_per_period > 0 && period_secs == 0 {
+            return Err(Error::InvalidWithdrawLimit);
+        }
+
+        storage::set_max_withdraw_per_op(&env, max_per_op);
+        storage::set_max_withdraw_per_period(&env, max_per_period);
+        storage::set_withdraw_period_secs(&env, period_secs);
         storage::extend_instance(&env);
-        events::withdraw(&env, &from, shares, assets);
-        Ok(assets)
+        events::withdraw_limits(&env, max_per_op, max_per_period, period_secs);
+        Ok(())
+    }
+
+    /// Clears rolling-period withdrawal usage and starts a fresh window at the
+    /// current ledger timestamp. Admin-only. Emits `wd_reset`.
+    pub fn reset_withdraw_period(env: Env) -> Result<(), Error> {
+        storage::require_initialized(&env)?;
+        let admin = storage::get_admin(&env);
+        admin.require_auth();
+
+        let cleared = storage::get_period_withdrawn(&env);
+        let at = env.ledger().timestamp();
+        storage::set_period_withdrawn(&env, 0);
+        storage::set_period_started_at(&env, at);
+        storage::extend_instance(&env);
+        events::withdraw_period_reset(&env, cleared, at);
+        Ok(())
+    }
+
+    /// Enables or disables the emergency withdrawal-limit override.
+    ///
+    /// When enabled, per-operation and rolling-period checks are skipped so
+    /// operators can unblock legitimate exits. Admin-only. Emits `wd_override`.
+    pub fn set_withdraw_limits_override(env: Env, enabled: bool) -> Result<(), Error> {
+        storage::require_initialized(&env)?;
+        let admin = storage::get_admin(&env);
+        admin.require_auth();
+
+        storage::set_withdraw_limits_override(&env, enabled);
+        storage::extend_instance(&env);
+        events::withdraw_override(&env, enabled);
+        Ok(())
+    }
+
+    /// Returns the per-operation withdrawal asset cap (`0` = unlimited).
+    pub fn get_max_withdraw_per_op(env: Env) -> u128 {
+        storage::get_max_withdraw_per_op(&env)
+    }
+
+    /// Returns the rolling-period withdrawal asset cap (`0` = unlimited).
+    pub fn get_max_withdraw_per_period(env: Env) -> u128 {
+        storage::get_max_withdraw_per_period(&env)
+    }
+
+    /// Returns the rolling withdrawal window length in seconds.
+    pub fn get_withdraw_period_secs(env: Env) -> u64 {
+        storage::get_withdraw_period_secs(&env)
+    }
+
+    /// Returns assets withdrawn so far in the current rolling period.
+    pub fn get_period_withdrawn(env: Env) -> u128 {
+        storage::get_period_withdrawn(&env)
+    }
+
+    /// Returns the ledger timestamp when the current rolling period started.
+    pub fn get_period_started_at(env: Env) -> u64 {
+        storage::get_period_started_at(&env)
+    }
+
+    /// Returns `true` when the emergency withdrawal-limit override is on.
+    pub fn is_withdraw_limits_override(env: Env) -> bool {
+        storage::is_withdraw_limits_override(&env)
     }
 
     /// Mocks yield accrual by increasing the vault's total assets by `amount`
@@ -389,3 +565,78 @@ impl YieldVault {
         Ok(())
     }
 }
+
+impl YieldVault {
+    /// Checks per-operation and rolling-period limits for `op_assets` (the
+    /// amount counting against the per-op cap) and `period_assets` (the amount
+    /// added to the rolling aggregate). Records period usage on success.
+    ///
+    /// When the emergency override is enabled, checks are skipped and usage is
+    /// not recorded so operators can drain/exit without polluting the window.
+    fn enforce_withdraw_limits(
+        env: &Env,
+        op_assets: u128,
+        period_assets: u128,
+    ) -> Result<(), Error> {
+        if storage::is_withdraw_limits_override(env) {
+            return Ok(());
+        }
+
+        let max_per_op = storage::get_max_withdraw_per_op(env);
+        if max_per_op > 0 && op_assets > max_per_op {
+            return Err(Error::WithdrawLimitExceeded);
+        }
+
+        let max_per_period = storage::get_max_withdraw_per_period(env);
+        if max_per_period == 0 {
+            return Ok(());
+        }
+
+        let period_secs = storage::get_withdraw_period_secs(env);
+        let now = env.ledger().timestamp();
+        let started = storage::get_period_started_at(env);
+        let mut withdrawn = storage::get_period_withdrawn(env);
+
+        // Roll the window forward when the configured period has elapsed.
+        if period_secs > 0 && now.saturating_sub(started) >= period_secs {
+            withdrawn = 0;
+            storage::set_period_started_at(env, now);
+            storage::set_period_withdrawn(env, 0);
+        }
+
+        let new_withdrawn = withdrawn.saturating_add(period_assets);
+        if new_withdrawn > max_per_period {
+            return Err(Error::WithdrawPeriodLimitExceeded);
+        }
+        storage::set_period_withdrawn(env, new_withdrawn);
+        Ok(())
+    }
+
+    /// Applies share burn, aggregate updates, token push, and withdraw event
+    /// for a single redemption leg. Caller must have already enforced limits
+    /// and validated balances.
+    fn apply_withdraw(
+        env: &Env,
+        from: &Address,
+        shares: u128,
+        assets: u128,
+        user_balance: u128,
+        total_shares: u128,
+        total_assets: u128,
+    ) {
+        let new_total_shares = total_shares.saturating_sub(shares);
+        let new_total_assets = total_assets.saturating_sub(assets);
+        let new_user_balance = user_balance.saturating_sub(shares);
+
+        storage::set_total_shares(env, new_total_shares);
+        storage::set_total_assets(env, new_total_assets);
+        storage::set_balance(env, from, new_user_balance);
+
+        let token_address = storage::get_token(env);
+        let client = token::Client::new(env, &token_address);
+        client.transfer(&env.current_contract_address(), from, &(assets as i128));
+
+        storage::extend_instance(env);
+        events::withdraw(env, from, shares, assets);
+    }
+}