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 + ); +}