diff --git a/fee_collector/src/lib.rs b/fee_collector/src/lib.rs index 2044199..0d719a8 100644 --- a/fee_collector/src/lib.rs +++ b/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 { env.storage() diff --git a/fee_collector/src/test.rs b/fee_collector/src/test.rs index 14e6388..0a70a31 100644 --- a/fee_collector/src/test.rs +++ b/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();