Add operator withdrawal caps to token vault

contracts/token-vault/src/events.rs

@@ -60,6 +60,15 @@ pub fn limit_set(env: &Env, caller: &Address, old_limit: i128, new_limit: i128)
     );
 }
 
+/// Emitted by `set_operator_withdraw_limit` when the owner configures how much
+/// a delegated operator may withdraw in a single call.
+pub fn operator_withdraw_limit_set(env: &Env, caller: &Address, old_limit: i128, new_limit: i128) {
+    env.events().publish(
+        (symbol_short!("op_limit"), caller.clone()),
+        (old_limit, new_limit),
+    );
+}
+
 /// Emitted by `set_operator` when the owner delegates a new operator.
 ///
 /// Topics: `("set_op", caller)` — the owner.

contracts/token-vault/src/lib.rs

@@ -10,8 +10,10 @@ use drip_common::{is_zero_address, TTL_EXTEND_TO, TTL_THRESHOLD};
 use errors::Error;
 use soroban_sdk::{contract, contractimpl, token, Address, Env};
 use storage::{
-    get_max_limit, get_operator, get_owner, get_token, is_paused, remove_operator, set_max_limit,
-    set_operator, set_owner, set_paused, set_token,
+    get_max_limit, get_operator, get_operator_withdraw_limit, get_owner, get_pending_owner,
+    get_pending_owner_proposer, get_token, is_paused, remove_operator, remove_pending_owner,
+    remove_pending_owner_proposer, set_max_limit, set_operator, set_operator_withdraw_limit,
+    set_owner, set_paused, set_pending_owner, set_pending_owner_proposer, set_token,
 };
 
 #[contract]
@@ -144,6 +146,13 @@ impl TokenVault {
             return Err(Error::InvalidAmount);
         }
 
+        if caller != owner {
+            let limit = get_operator_withdraw_limit(&env).ok_or(Error::LimitExceeded)?;
+            if amount > limit {
+                return Err(Error::LimitExceeded);
+            }
+        }
+
         let balance = vault_balance(&env)?;
         let new_balance = balance
             .checked_sub(amount)
@@ -157,6 +166,29 @@ impl TokenVault {
         Ok(())
     }
 
+    pub fn set_operator_withdraw_limit(
+        env: Env,
+        caller: Address,
+        new_limit: i128,
+    ) -> Result<(), Error> {
+        assert_not_paused(&env)?;
+        let owner = get_owner(&env).ok_or(Error::NotInitialized)?;
+        if caller != owner {
+            return Err(Error::NotAuthorized);
+        }
+        caller.require_auth();
+
+        if new_limit <= 0 {
+            return Err(Error::InvalidAmount);
+        }
+
+        let old_limit = get_operator_withdraw_limit(&env).unwrap_or(0);
+        bump_instance(&env);
+        set_operator_withdraw_limit(&env, &new_limit);
+        events::operator_withdraw_limit_set(&env, &caller, old_limit, new_limit);
+        Ok(())
+    }
+
     pub fn set_limit(env: Env, caller: Address, new_limit: i128) -> Result<(), Error> {
         assert_not_paused(&env)?;
         let owner = get_owner(&env).ok_or(Error::NotInitialized)?;
@@ -272,6 +304,12 @@ impl TokenVault {
         get_operator(&env)
     }
 
+    /// Read-only: the maximum single-call withdrawal a delegated operator may
+    /// execute before the owner raises or removes the cap.
+    pub fn operator_withdraw_limit(env: Env) -> Option<i128> {
+        get_operator_withdraw_limit(&env)
+    }
+
     /// Read-only: the current owner address, if any.
     pub fn owner(env: Env) -> Option<Address> {
         get_owner(&env)

contracts/token-vault/src/storage.rs

@@ -14,6 +14,13 @@ pub enum DataKey {
     /// pattern matching `DripStream`'s `set_operator` design.
     /// Absent key means no operator has been delegated.
     Operator,
+    /// Maximum amount a delegated operator may withdraw in a single call.
+    ///
+    /// When set, `withdraw` enforces this cap for operator-authenticated
+    /// withdrawals while the owner remains unbounded. If this key is absent,
+    /// the operator is effectively unable to withdraw until the owner sets a
+    /// positive cap.
+    OperatorWithdrawLimit,
     /// Emergency-pause flag. When `true`, all state-mutating entry points
     /// (`deposit`, `withdraw`, `set_limit`) revert before touching state.
     Paused,
@@ -71,6 +78,18 @@ pub fn remove_operator(env: &Env) {
     env.storage().instance().remove(&DataKey::Operator);
 }
 
+pub fn set_operator_withdraw_limit(env: &Env, v: &i128) {
+    env.storage().instance().set(&DataKey::OperatorWithdrawLimit, v);
+}
+
+pub fn get_operator_withdraw_limit(env: &Env) -> Option<i128> {
+    env.storage().instance().get(&DataKey::OperatorWithdrawLimit)
+}
+
+pub fn remove_operator_withdraw_limit(env: &Env) {
+    env.storage().instance().remove(&DataKey::OperatorWithdrawLimit);
+}
+
 pub fn set_paused(env: &Env, paused: bool) {
     env.storage().instance().set(&DataKey::Paused, &paused);
 }

contracts/token-vault/src/tests.rs

@@ -310,6 +310,37 @@ fn operator_can_set_limit() {
     s.client.set_limit(&op, &2_000_000);
 }
 
+#[test]
+fn operator_withdraw_limit_is_enforced() {
+    let s = Setup::new(1_000_000);
+    s.client.deposit(&s.owner, &500);
+
+    let op = Address::generate(&s.env);
+    s.client.set_operator(&s.owner, &op);
+    s.client.set_operator_withdraw_limit(&s.owner, &100);
+
+    let recipient = Address::generate(&s.env);
+    let result = s.client.try_withdraw(&op, &recipient, &200);
+    assert_eq!(result, Err(Ok(Error::LimitExceeded)));
+
+    s.client.withdraw(&op, &recipient, &100);
+    assert_eq!(s.token.balance(&recipient), 100);
+}
+
+#[test]
+fn owner_withdraw_is_not_limited_by_operator_cap() {
+    let s = Setup::new(1_000_000);
+    s.client.deposit(&s.owner, &500);
+
+    let op = Address::generate(&s.env);
+    s.client.set_operator(&s.owner, &op);
+    s.client.set_operator_withdraw_limit(&s.owner, &100);
+
+    let recipient = Address::generate(&s.env);
+    s.client.withdraw(&s.owner, &recipient, &500);
+    assert_eq!(s.token.balance(&recipient), 500);
+}
+
 #[test]
 fn stranger_cannot_withdraw_even_with_operator_set() {
     let s = Setup::new(1_000_000);