From 6955ec713ebb2d6c028f880c53a71e696489817f Mon Sep 17 00:00:00 2001 From: Inkman007 Date: Sun, 28 Jun 2026 06:29:03 +0000 Subject: [PATCH] feat: add rolling-window withdrawal rate limit to treasury Protect treasury against single-key drain by introducing a per-24h rolling-window withdrawal cap. Both withdraw() and emergency_withdraw() now enforce the cap; the cap itself is protected by a 24 h timelock so it cannot be raised immediately after a compromise. - Add WithdrawalCap / EpochStart / EpochWithdrawn / WithdrawalCapProposal DataKey variants (and WithdrawalLock to fix pre-existing compile error) - Add EPOCH_DURATION = 86_400 s constant - enforce_rate_limit(): resets epoch counter after 24 h, rejects when EpochWithdrawn + amount > WithdrawalCap (cap = 0 means uncapped) - record_withdrawal(): persists running total for the current epoch - propose_withdrawal_cap() + execute_withdrawal_cap() (24 h timelock) - get_withdrawal_cap() view - Add WithdrawalRateLimitExceeded / WithdrawalCapTimelockNotElapsed / NoCapChangeProposed errors to KoraError - Add withdrawal_cap_proposed / withdrawal_cap_updated events Closes #225 --- contracts/shared/src/errors.rs | 3 + contracts/shared/src/events.rs | 8 ++ contracts/treasury/src/lib.rs | 182 ++++++++++++++++++++++++++++++++- 3 files changed, 188 insertions(+), 5 deletions(-) diff --git a/contracts/shared/src/errors.rs b/contracts/shared/src/errors.rs index fea2bcf..6581a77 100644 --- a/contracts/shared/src/errors.rs +++ b/contracts/shared/src/errors.rs @@ -41,6 +41,9 @@ pub enum KoraError { InvalidFeeRate = 40, WithdrawalFailed = 41, TokenNotWhitelisted = 42, + WithdrawalRateLimitExceeded = 43, + WithdrawalCapTimelockNotElapsed = 44, + NoCapChangeProposed = 45, // Risk SMENotRegistered = 50, diff --git a/contracts/shared/src/events.rs b/contracts/shared/src/events.rs index 86cc667..0a25333 100644 --- a/contracts/shared/src/events.rs +++ b/contracts/shared/src/events.rs @@ -127,6 +127,14 @@ pub fn fee_withdrawn(env: &Env, token: &Address, amount: i128) { emit(env, symbol_short!("FEE_WTH"), (token.clone(), amount)); } +pub fn withdrawal_cap_proposed(env: &Env, admin: &Address, new_cap: i128) { + emit(env, symbol_short!("WD_PROP"), (admin.clone(), new_cap, env.ledger().timestamp())); +} + +pub fn withdrawal_cap_updated(env: &Env, admin: &Address, old_cap: i128, new_cap: i128) { + emit(env, symbol_short!("WD_CAP"), (admin.clone(), old_cap, new_cap)); +} + /// Emitted when the full token balance is drained via emergency_withdraw. pub fn emergency_withdrawn(env: &Env, by: &Address, token: &Address, amount: i128) { emit( diff --git a/contracts/treasury/src/lib.rs b/contracts/treasury/src/lib.rs index 9635a0b..9281c0a 100644 --- a/contracts/treasury/src/lib.rs +++ b/contracts/treasury/src/lib.rs @@ -12,6 +12,9 @@ use soroban_sdk::{contract, contractimpl, contracttype, token, Address, BytesN, const PERSISTENT_BUMP_AMOUNT: u32 = 535_680; const PERSISTENT_LIFETIME_THRESHOLD: u32 = 535_680 / 2; +// ── Rate-limit epoch: 24 h in seconds ──────────────────────────────────────── +const EPOCH_DURATION: u64 = 86_400; + // ── Storage Keys ───────────────────────────────────────────────────────────── #[contracttype] @@ -26,6 +29,16 @@ pub enum DataKey { WhitelistedToken(Address), /// Pending upgrade proposal: (wasm_hash, proposed_at_timestamp). UpgradeProposal, + /// Reentrancy guard flag (instance-level). + WithdrawalLock, + /// Maximum total withdrawal allowed per EPOCH_DURATION window (0 = uncapped). + WithdrawalCap, + /// Pending cap change proposal: (new_cap, proposed_at). + WithdrawalCapProposal, + /// Timestamp when the current rate-limit epoch started. + EpochStart, + /// Total withdrawn in the current epoch (resets each epoch). + EpochWithdrawn, } // ── Contract ────────────────────────────────────────────────────────────────── @@ -44,9 +57,7 @@ impl TreasuryContract { kora_shared::validation::require_not_self(&env, &admin)?; env.storage().instance().set(&DataKey::Admin, &admin); env.storage().instance().set(&DataKey::FeeBps, &fee_bps); - env.storage() - .instance() - .set(&DataKey::WithdrawalLock, &false); + env.storage().instance().set(&DataKey::WithdrawalLock, &false); env.storage().persistent().set(&DataKey::Admin, &admin); env.storage().persistent().extend_ttl( &DataKey::Admin, @@ -59,6 +70,12 @@ impl TreasuryContract { PERSISTENT_LIFETIME_THRESHOLD, PERSISTENT_BUMP_AMOUNT, ); + // Initialize rate-limit state: uncapped, epoch starts now, zero withdrawn. + env.storage().instance().set(&DataKey::WithdrawalCap, &0i128); + env.storage() + .instance() + .set(&DataKey::EpochStart, &env.ledger().timestamp()); + env.storage().instance().set(&DataKey::EpochWithdrawn, &0i128); events::treasury_initialized(&env, &admin, fee_bps); Ok(()) } @@ -123,6 +140,7 @@ impl TreasuryContract { } /// Withdraw accumulated fees to a recipient. Admin only. + /// Subject to the rolling 24 h withdrawal cap. /// Protected against reentrancy via RAII ReentrancyGuard. pub fn withdraw( env: Env, @@ -138,6 +156,7 @@ impl TreasuryContract { return Err(KoraError::InvalidAmount); } Self::require_whitelisted_token(&env, &token)?; + Self::enforce_rate_limit(&env, amount)?; // Acquire reentrancy guard — released automatically when _guard drops let _guard = ReentrancyGuard::new(&env)?; @@ -160,6 +179,7 @@ impl TreasuryContract { env.storage().persistent().set(&collected_key, &new_collected); Self::bump_persistent(&env, &collected_key); } + Self::record_withdrawal(&env, amount); // ── Interactions ────────────────────────────────────────────────────── token_client.transfer(&env.current_contract_address(), &recipient, &amount); @@ -169,6 +189,7 @@ impl TreasuryContract { } /// Emergency drain — withdraw entire token balance. Admin only. + /// Subject to the rolling 24 h withdrawal cap. /// Protected against reentrancy via RAII ReentrancyGuard. pub fn emergency_withdraw( env: Env, @@ -185,9 +206,9 @@ impl TreasuryContract { let token_client = token::Client::new(&env, &token); let balance = token_client.balance(&env.current_contract_address()); - // ── Interactions ────────────────────────────────────────────────────── if balance > 0 { - // ── Interactions ────────────────────────────────────────────────────── + Self::enforce_rate_limit(&env, balance)?; + Self::record_withdrawal(&env, balance); token_client.transfer(&env.current_contract_address(), &recipient, &balance); events::emergency_withdrawn(&env, &admin, &token, balance); } @@ -195,6 +216,63 @@ impl TreasuryContract { Ok(()) } + // ── Withdrawal cap management ───────────────────────────────────────────── + + /// Propose a new withdrawal cap. Takes effect after UPGRADE_TIMELOCK_DELAY. + /// Admin only. Set cap to 0 to remove the limit entirely. + pub fn propose_withdrawal_cap( + env: Env, + admin: Address, + new_cap: i128, + ) -> Result<(), KoraError> { + admin.require_auth(); + Self::require_admin(&env, &admin)?; + if new_cap < 0 { + return Err(KoraError::InvalidAmount); + } + env.storage() + .instance() + .set(&DataKey::WithdrawalCapProposal, &(new_cap, env.ledger().timestamp())); + events::withdrawal_cap_proposed(&env, &admin, new_cap); + Ok(()) + } + + /// Execute a previously proposed withdrawal cap change after the timelock elapses. + /// Admin only. + pub fn execute_withdrawal_cap(env: Env, admin: Address) -> Result<(), KoraError> { + admin.require_auth(); + Self::require_admin(&env, &admin)?; + let (new_cap, proposed_at): (i128, u64) = env + .storage() + .instance() + .get(&DataKey::WithdrawalCapProposal) + .ok_or(KoraError::NoCapChangeProposed)?; + if env.ledger().timestamp() < proposed_at + UPGRADE_TIMELOCK_DELAY { + return Err(KoraError::WithdrawalCapTimelockNotElapsed); + } + let old_cap: i128 = env + .storage() + .instance() + .get(&DataKey::WithdrawalCap) + .unwrap_or(0); + env.storage() + .instance() + .set(&DataKey::WithdrawalCap, &new_cap); + env.storage() + .instance() + .remove(&DataKey::WithdrawalCapProposal); + events::withdrawal_cap_updated(&env, &admin, old_cap, new_cap); + Ok(()) + } + + /// Returns the current withdrawal cap (0 = uncapped). + pub fn get_withdrawal_cap(env: Env) -> i128 { + env.storage() + .instance() + .get(&DataKey::WithdrawalCap) + .unwrap_or(0) + } + /// Returns the current protocol fee in basis points. pub fn get_fee_bps(env: Env) -> u32 { env.storage() @@ -274,6 +352,65 @@ impl TreasuryContract { Ok(()) } + /// Advance the epoch if 24 h have elapsed, then check the cap. + fn enforce_rate_limit(env: &Env, amount: i128) -> Result<(), KoraError> { + let cap: i128 = env + .storage() + .instance() + .get(&DataKey::WithdrawalCap) + .unwrap_or(0); + if cap == 0 { + return Ok(()); + } + + let now = env.ledger().timestamp(); + let epoch_start: u64 = env + .storage() + .instance() + .get(&DataKey::EpochStart) + .unwrap_or(now); + + let epoch_withdrawn: i128 = if now.saturating_sub(epoch_start) >= EPOCH_DURATION { + // New epoch: reset counters. + env.storage().instance().set(&DataKey::EpochStart, &now); + env.storage().instance().set(&DataKey::EpochWithdrawn, &0i128); + 0 + } else { + env.storage() + .instance() + .get(&DataKey::EpochWithdrawn) + .unwrap_or(0) + }; + + let new_total = epoch_withdrawn + .checked_add(amount) + .ok_or(KoraError::ArithmeticOverflow)?; + if new_total > cap { + return Err(KoraError::WithdrawalRateLimitExceeded); + } + Ok(()) + } + + /// Record a successful withdrawal against the current epoch's running total. + fn record_withdrawal(env: &Env, amount: i128) { + let cap: i128 = env + .storage() + .instance() + .get(&DataKey::WithdrawalCap) + .unwrap_or(0); + if cap == 0 { + return; + } + let current: i128 = env + .storage() + .instance() + .get(&DataKey::EpochWithdrawn) + .unwrap_or(0); + env.storage() + .instance() + .set(&DataKey::EpochWithdrawn, ¤t.saturating_add(amount)); + } + fn bump_persistent(env: &Env, key: &DataKey) { env.storage().persistent().extend_ttl( key, @@ -470,4 +607,39 @@ mod tests { let result = client.try_initialize(&contract_id, &50u32); assert!(result.is_err()); } + + // ── withdrawal rate limit ───────────────────────────────────────────────── + + #[test] + fn test_withdrawal_cap_default_is_uncapped() { + let (_env, _admin, client) = setup(); + assert_eq!(client.get_withdrawal_cap(), 0); + } + + #[test] + fn test_propose_withdrawal_cap_requires_admin() { + let (env, _admin, client) = setup(); + let non_admin = Address::generate(&env); + assert!(client.try_propose_withdrawal_cap(&non_admin, &1_000_000i128).is_err()); + } + + #[test] + fn test_propose_negative_cap_rejected() { + let (_env, admin, client) = setup(); + assert!(client.try_propose_withdrawal_cap(&admin, &-1i128).is_err()); + } + + #[test] + fn test_execute_cap_before_timelock_fails() { + let (_env, admin, client) = setup(); + client.propose_withdrawal_cap(&admin, &1_000_000i128); + // Timelock hasn't elapsed + assert!(client.try_execute_withdrawal_cap(&admin).is_err()); + } + + #[test] + fn test_execute_cap_without_proposal_fails() { + let (_env, admin, client) = setup(); + assert!(client.try_execute_withdrawal_cap(&admin).is_err()); + } }