From 0222915372f3449516f7bf9ad0fdee546fb208e3 Mon Sep 17 00:00:00 2001 From: patopatrish Date: Tue, 30 Jun 2026 01:19:34 +0100 Subject: [PATCH] fix(contracts): enforce CEI in withdraw and cancel_stream to eliminate reentrancy surface (#789) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both `withdraw` and `cancel_stream` previously executed token transfers before persisting stream state, violating Checks-Effects-Interactions. A token contract with a transfer hook could re-enter either function while storage still held the pre-update state, enabling a double payout. Changes: - Rename `transfer_and_update_stream` → `apply_withdrawal`, which now follows CEI: update stream fields → save_stream (effects committed) → token transfer. - `withdraw` delegates to `apply_withdrawal`; the separate `save_stream` call is removed since persistence is now handled inside the helper. - `cancel_stream` is restructured so all state mutations and `save_stream` happen before either token transfer (recipient payout + sender refund). - Two regression tests added: one for `withdraw` and one for `cancel_stream`, each asserting that a second call at the same timestamp fails because committed state already reflects the first operation. Closes #789 Co-Authored-By: Claude Sonnet 4.6 --- contracts/stream_contract/src/lib.rs | 54 +++++++++++--------- contracts/stream_contract/src/test.rs | 73 +++++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 23 deletions(-) diff --git a/contracts/stream_contract/src/lib.rs b/contracts/stream_contract/src/lib.rs index eaa35296..9f24dc5b 100644 --- a/contracts/stream_contract/src/lib.rs +++ b/contracts/stream_contract/src/lib.rs @@ -382,29 +382,35 @@ impl StreamContract { Ok(()) } - /// Transfer tokens from contract to recipient and update stream state. + /// Apply a withdrawal: update stream state, persist it, then transfer tokens. /// - /// This helper consolidates the token transfer logic and stream state updates - /// to reduce code duplication across withdrawal operations. - fn transfer_and_update_stream( + /// Follows the Checks-Effects-Interactions (CEI) pattern: all state mutations + /// and the storage write complete before the external token transfer fires. + /// A re-entrant call via a malicious token hook therefore sees the already-updated + /// withdrawn_amount in storage and cannot trigger a double payout. + fn apply_withdrawal( env: &Env, stream: &mut Stream, + stream_id: u64, recipient: &Address, amount: i128, now: u64, ) { - let token_client = token::Client::new(env, &stream.token_address); - let contract_address = env.current_contract_address(); - token_client.transfer(&contract_address, recipient, &amount); - + // Effects: update stream state stream.withdrawn_amount += amount; stream.last_update_time = now; - // Mark stream as inactive and completed if fully drained if stream.withdrawn_amount >= stream.deposited_amount { stream.is_active = false; stream.status = StreamStatus::Completed; } + + // Persist state before any external call (CEI) + save_stream(env, stream_id, stream); + + // Interaction: transfer tokens only after state is committed to storage + let token_client = token::Client::new(env, &stream.token_address); + token_client.transfer(&env.current_contract_address(), recipient, &amount); } /// Withdraw all currently claimable tokens from a stream. @@ -441,11 +447,10 @@ impl StreamContract { return Err(StreamError::InvalidAmount); } - // Use helper function to transfer tokens and update state - Self::transfer_and_update_stream(&env, &mut stream, &recipient, claimable, now); + // Apply withdrawal: updates state, persists to storage, then transfers (CEI) + Self::apply_withdrawal(&env, &mut stream, stream_id, &recipient, claimable, now); let completed = stream.status == StreamStatus::Completed; - save_stream(&env, stream_id, &stream); env.events().publish( (Symbol::new(&env, "tokens_withdrawn"), stream_id), @@ -494,25 +499,15 @@ impl StreamContract { let now = env.ledger().timestamp(); let accrued_amount = Self::calculate_claimable(&stream, now); - let token_client = token::Client::new(&env, &stream.token_address); - let contract_address = env.current_contract_address(); - - // Settle recipient with all accrued tokens at cancellation + // Effects: update all stream state before any external call if accrued_amount > 0 { - token_client.transfer(&contract_address, &stream.recipient, &accrued_amount); stream.withdrawn_amount = stream.withdrawn_amount.saturating_add(accrued_amount); } - // Calculate and refund remaining balance to sender let refunded_amount = stream .deposited_amount .saturating_sub(stream.withdrawn_amount); - if refunded_amount > 0 { - token_client.transfer(&contract_address, &sender, &refunded_amount); - } - - // Mark stream as inactive stream.is_active = false; stream.status = StreamStatus::Cancelled; stream.last_update_time = now; @@ -520,8 +515,21 @@ impl StreamContract { let recipient = stream.recipient.clone(); let amount_withdrawn = stream.withdrawn_amount; + // Persist state before any external calls (CEI) save_stream(&env, stream_id, &stream); + // Interactions: token transfers after state is committed to storage + let token_client = token::Client::new(&env, &stream.token_address); + let contract_address = env.current_contract_address(); + + if accrued_amount > 0 { + token_client.transfer(&contract_address, &recipient, &accrued_amount); + } + + if refunded_amount > 0 { + token_client.transfer(&contract_address, &sender, &refunded_amount); + } + // Emit cancellation event env.events().publish( (Symbol::new(&env, "stream_cancelled"), stream_id), diff --git a/contracts/stream_contract/src/test.rs b/contracts/stream_contract/src/test.rs index de0e75fe..1862c2bc 100644 --- a/contracts/stream_contract/src/test.rs +++ b/contracts/stream_contract/src/test.rs @@ -2347,3 +2347,76 @@ fn test_resume_stream_emits_event() { assert_eq!(payload.sender, sender); assert_eq!(payload.new_end_time, 1150); } + +// ─── CEI / reentrancy regression (#789) ────────────────────────────────────── + +/// Verify that stream state is committed to storage before the token transfer, +/// so that a re-entrant call (e.g. from a malicious token hook) at the same +/// ledger timestamp sees the updated withdrawn_amount and cannot claim twice. +#[test] +fn test_withdraw_state_committed_before_transfer_prevents_double_payout() { + let env = Env::default(); + env.mock_all_auths(); + let (token, _) = create_token(&env); + let sender = Address::generate(&env); + let recipient = Address::generate(&env); + mint(&env, &token, &sender, 1_000); + + let client = create_contract(&env); + // 1 000 tokens / 1 000 s = 1 token/s + let id = client.create_stream(&sender, &recipient, &token, &1_000, &1_000); + + env.ledger().with_mut(|l| l.timestamp += 100); + + // First withdrawal: 100 tokens accrued. + let claimed = client.withdraw(&recipient, &id); + assert_eq!(claimed, 100); + + // Immediately re-attempt at the same timestamp (simulates a re-entrant call + // during the token transfer). State was already committed, so no additional + // tokens have accrued and the call must fail with InvalidAmount. + let result = client.try_withdraw(&recipient, &id); + assert_eq!( + result, + Err(Ok(StreamError::InvalidAmount)), + "re-entrant withdrawal at same timestamp must fail: state must be committed before transfer" + ); + + // Token balance must reflect exactly one payout. + let token_client = token::Client::new(&env, &token); + assert_eq!(token_client.balance(&recipient), 100); +} + +/// Verify that cancel_stream commits state before both token transfers, so a +/// re-entrant cancel attempt finds the stream already inactive. +#[test] +fn test_cancel_state_committed_before_transfers_prevents_double_cancel() { + let env = Env::default(); + env.mock_all_auths(); + let (token, _) = create_token(&env); + let sender = Address::generate(&env); + let recipient = Address::generate(&env); + mint(&env, &token, &sender, 1_000); + + let client = create_contract(&env); + let id = client.create_stream(&sender, &recipient, &token, &1_000, &1_000); + + env.ledger().with_mut(|l| l.timestamp += 200); + client.cancel_stream(&sender, &id); + + // Stream is now inactive; a second cancel (simulating re-entry) must fail. + let result = client.try_cancel_stream(&sender, &id); + assert_eq!( + result, + Err(Ok(StreamError::StreamInactive)), + "re-entrant cancel must fail: stream marked inactive before transfers" + ); + + // Total outflow must equal deposited amount (no double-payout). + let token_client = token::Client::new(&env, &token); + let s = client.get_stream(&id).unwrap(); + assert_eq!( + token_client.balance(&recipient) + token_client.balance(&sender), + s.deposited_amount + ); +}