fix(fee_collector): add lifetime_withdrawn tracking and reconciliation view

fee_collector/src/lib.rs

@@ -10,10 +10,39 @@
 //!   KEY_TREASURY → Address
 //!
 //! Persistent storage (keyed by token address):
-//!   (KEY_TOTAL, token) → i128   — lifetime total collected
+//!   (KEY_TOTAL,     token) → i128   — lifetime total collected (reported)
+//!   (KEY_WITHDRAWN, token) → i128   — lifetime total withdrawn
 //!
 //! The actual token balances are tracked by the token contracts themselves;
 //! `get_balance` queries the token contract directly.
+//!
+//! ## Trust-the-caller caveat
+//!
+//! `get_total_collected` is a **reported / claimed** figure, not a value
+//! independently verified against real token movement.  When `collect_fee` is
+//! called, the contract blindly trusts the `amount` parameter supplied by the
+//! caller (by design — the corresponding token transfer has already occurred
+//! before `collect_fee` is called, and re-querying the balance would be
+//! racy).  Consequently `get_total_collected` and `get_balance` can legitimately
+//! diverge for entirely benign reasons:
+//!
+//! * Every `withdraw` call reduces `get_balance` without reducing
+//!   `get_total_collected` (because `KEY_TOTAL` is a *lifetime* counter, not a
+//!   live balance).
+//! * A caller that passes an `amount` inconsistent with what was actually
+//!   transferred (e.g. due to a rounding bug or a fee-on-transfer token) will
+//!   silently skew the counter.
+//!
+//! To aid treasury auditing, the contract also tracks a parallel
+//! `(KEY_WITHDRAWN, token)` lifetime counter so that the invariant
+//!
+//! ```text
+//! get_expected_balance(token) == get_total_collected(token) - get_total_withdrawn(token)
+//! ```
+//!
+//! can be computed on-chain and diffed against `get_balance(token)`.  Any
+//! non-zero difference signals drift that warrants investigation, but the
+//! contract itself cannot auto-correct it.
 
 #![no_std]
 
@@ -32,6 +61,9 @@ const KEY_INIT: Symbol = symbol_short!("INIT");
 /// Persistent key prefix for lifetime-total-collected per token.
 const KEY_TOTAL: Symbol = symbol_short!("TOTAL");
 
+/// Persistent key prefix for lifetime-total-withdrawn per token.
+const KEY_WITHDRAWN: Symbol = symbol_short!("WDRAWN");
+
 // ---------------------------------------------------------------------------
 // Error type
 // ---------------------------------------------------------------------------
@@ -150,6 +182,18 @@ impl FeeCollectorContract {
         let token_client = token::Client::new(&env, &token);
         token_client.transfer(&env.current_contract_address(), &recipient, &amount);
 
+        // Update lifetime-withdrawn counter for this token.
+        let withdrawn_key = (KEY_WITHDRAWN, token.clone());
+        let current_withdrawn: i128 = env
+            .storage()
+            .persistent()
+            .get(&withdrawn_key)
+            .unwrap_or(0i128);
+        let new_withdrawn = current_withdrawn
+            .checked_add(amount)
+            .ok_or(FeeCollectorError::ArithmeticOverflow)?;
+        env.storage().persistent().set(&withdrawn_key, &new_withdrawn);
+
         // Emit event.
         env.events().publish(
             (symbol_short!("fee_wdrw"), token, recipient),
@@ -183,6 +227,10 @@ impl FeeCollectorContract {
     }
 
     /// Return the lifetime total amount of `token` ever collected as fees.
+    ///
+    /// **Note:** this is a reported/claimed figure supplied by the caller of
+    /// `collect_fee`.  It is not independently verified against the contract's
+    /// real token balance.  See the module-level documentation for details.
     pub fn get_total_collected(env: Env, token: Address) -> i128 {
         let total_key = (KEY_TOTAL, token);
         env.storage()
@@ -191,6 +239,47 @@ impl FeeCollectorContract {
             .unwrap_or(0i128)
     }
 
+    /// Return the lifetime total amount of `token` ever withdrawn by the admin.
+    pub fn get_total_withdrawn(env: Env, token: Address) -> i128 {
+        let withdrawn_key = (KEY_WITHDRAWN, token);
+        env.storage()
+            .persistent()
+            .get(&withdrawn_key)
+            .unwrap_or(0i128)
+    }
+
+    /// Return the expected current balance derived from on-chain counters:
+    ///
+    /// ```text
+    /// expected = get_total_collected(token) - get_total_withdrawn(token)
+    /// ```
+    ///
+    /// Comparing this value against `get_balance(token)` surfaces any drift
+    /// between what the accounting counters claim and what the contract
+    /// actually holds.  A non-zero difference signals that the `amount`
+    /// passed to one or more `collect_fee` calls did not match the tokens
+    /// that were actually transferred (e.g. a rounding bug or a
+    /// fee-on-transfer token).
+    pub fn get_expected_balance(env: Env, token: Address) -> i128 {
+        let total_key = (KEY_TOTAL, token.clone());
+        let total_collected: i128 = env
+            .storage()
+            .persistent()
+            .get(&total_key)
+            .unwrap_or(0i128);
+
+        let withdrawn_key = (KEY_WITHDRAWN, token);
+        let total_withdrawn: i128 = env
+            .storage()
+            .persistent()
+            .get(&withdrawn_key)
+            .unwrap_or(0i128);
+
+        // Saturating subtraction: the result should never be negative in a
+        // well-behaved deployment, but we guard against it defensively.
+        total_collected.saturating_sub(total_withdrawn)
+    }
+
     /// Return the admin address.
     pub fn get_admin(env: Env) -> Result<Address, FeeCollectorError> {
         env.storage()

fee_collector/src/test.rs

@@ -109,6 +109,68 @@ fn test_withdraw_sends_tokens_to_recipient() {
     assert_eq!(token_client.balance(&contract_id), 100);
 }
 
+/// Demonstrates the reconciliation mechanism introduced by issue #41.
+///
+/// Scenario:
+///   1. 300 tokens are minted to the contract and `collect_fee` is called.
+///   2. 200 tokens are withdrawn to a recipient.
+///
+/// Before the fix there was no on-chain way to reconstruct "300 collected,
+/// 200 withdrawn, 100 should remain" without replaying off-chain events.
+/// After the fix, `get_total_withdrawn` and `get_expected_balance` surface the
+/// full picture on-chain, and `get_expected_balance == get_balance` confirms no
+/// drift for this benign scenario.
+///
+/// A second phase deliberately introduces drift by calling `collect_fee` with
+/// an `amount` (50) that is larger than the tokens actually transferred (30),
+/// simulating a caller-side rounding bug.  In that case
+/// `get_expected_balance > get_balance`, exposing the discrepancy.
+#[test]
+fn test_reconciliation_detects_divergence() {
+    let (env, client, admin, treasury, token, token_admin) = setup();
+    client.initialize(&admin, &treasury);
+
+    let contract_id = client.address.clone();
+
+    // ── Phase 1: normal collect + withdraw ──────────────────────────────────
+    mint(&env, &token, &token_admin, &contract_id, 300);
+    client.collect_fee(&token, &300i128);
+
+    let recipient = Address::generate(&env);
+    client.withdraw(&token, &200i128, &recipient);
+
+    // Counters:
+    //   total_collected  = 300
+    //   total_withdrawn  = 200
+    //   expected_balance = 100   (== actual balance → no drift)
+    assert_eq!(client.get_total_collected(&token), 300);
+    assert_eq!(client.get_total_withdrawn(&token), 200);
+    assert_eq!(client.get_expected_balance(&token), 100);
+    assert_eq!(client.get_balance(&token), 100);
+
+    // No drift yet: expected == actual.
+    assert_eq!(client.get_expected_balance(&token), client.get_balance(&token));
+
+    // ── Phase 2: caller over-reports the collected amount (simulates a ───────
+    //            rounding bug: only 30 tokens arrive but 50 are reported)
+    mint(&env, &token, &token_admin, &contract_id, 30); // only 30 actually transferred
+    client.collect_fee(&token, &50i128); // caller reports 50
+
+    // Counters:
+    //   total_collected  = 350  (300 + 50 reported)
+    //   total_withdrawn  = 200  (unchanged)
+    //   expected_balance = 150  (350 - 200)
+    //   actual balance   = 130  (100 + 30 actually received)
+    assert_eq!(client.get_total_collected(&token), 350);
+    assert_eq!(client.get_total_withdrawn(&token), 200);
+    assert_eq!(client.get_expected_balance(&token), 150);
+    assert_eq!(client.get_balance(&token), 130);
+
+    // Drift is now detectable: expected != actual.
+    let drift = client.get_expected_balance(&token) - client.get_balance(&token);
+    assert_eq!(drift, 20, "expected 20-token drift from over-reported collect_fee");
+}
+
 #[test]
 fn test_withdraw_invalid_amount() {
     let (env, client, admin, treasury, token, _token_admin) = setup();