From 23abb8cea4259d383d319d0375674ec30cca9b46 Mon Sep 17 00:00:00 2001 From: Georgechisom Date: Thu, 26 Feb 2026 14:28:45 +0100 Subject: [PATCH] refactor: cleanup and deduplicate stream contract code - Extract validation logic into reusable helper functions (validate_stream_ownership, validate_stream_active) - Add transfer_and_update_stream helper to consolidate token transfer logic - Improve calculate_claimable with better documentation and use .min() for clarity - Add new get_claimable_amount query function for read-only balance checks - Enhance function documentation with detailed parameter and error descriptions - Refactor top_up_stream, withdraw, and cancel_stream to use helper functions - Fix test issues with undefined variables and incorrect contract instances - All 40 tests passing successfully This refactoring reduces code duplication, improves maintainability, and makes the contract more modular while preserving all existing functionality. --- contracts/stream_contract/src/lib.rs | 142 +++++++++++++----- contracts/stream_contract/src/test.rs | 12 -- ...cel_stream_after_partial_withdrawal.1.json | 7 +- .../test_cancel_stream_refunds_sender.1.json | 78 +++++++++- 4 files changed, 185 insertions(+), 54 deletions(-) diff --git a/contracts/stream_contract/src/lib.rs b/contracts/stream_contract/src/lib.rs index a103140..3d5a2c6 100644 --- a/contracts/stream_contract/src/lib.rs +++ b/contracts/stream_contract/src/lib.rs @@ -176,7 +176,8 @@ impl StreamContract { /// Top up an active stream with additional tokens. /// - /// Only the original sender may top up their own stream. + /// Only the original sender may top up their own stream. The top-up amount + /// is subject to protocol fees (if configured) before being added to the stream. /// /// # Errors /// - `InvalidAmount` — `amount` ≤ 0. @@ -197,24 +198,25 @@ impl StreamContract { let mut stream = load_stream(&env, stream_id)?; - if stream.sender != sender { - return Err(StreamError::Unauthorized); - } - if !stream.is_active { - return Err(StreamError::StreamInactive); - } + // Validate ownership and active status using helper functions + Self::validate_stream_ownership(&stream, &sender)?; + Self::validate_stream_active(&stream)?; + // Transfer tokens from sender to contract let token_client = token::Client::new(&env, &stream.token_address); let contract_address = env.current_contract_address(); token_client.transfer(&sender, &contract_address, &amount); + // Collect protocol fee and get net amount let net_amount = Self::collect_fee(&env, &stream.token_address, amount, stream_id); + // Update stream state stream.deposited_amount += net_amount; stream.last_update_time = env.ledger().timestamp(); save_stream(&env, stream_id, &stream); + // Emit top-up event env.events().publish( (Symbol::new(&env, "stream_topped_up"), stream_id), StreamToppedUpEvent { @@ -230,6 +232,17 @@ impl StreamContract { // ─── Internal Helpers ───────────────────────────────────────────────────── + /// Calculate the claimable amount for a stream at a given timestamp. + /// + /// This helper computes how many tokens have been streamed since the last + /// update, capped at the remaining balance to prevent over-withdrawal. + /// + /// # Arguments + /// * `stream` - The stream to calculate claimable amount for + /// * `now` - Current ledger timestamp + /// + /// # Returns + /// The amount of tokens that can be claimed, never exceeding remaining balance fn calculate_claimable(stream: &Stream, now: u64) -> i128 { let elapsed = now.saturating_sub(stream.last_update_time); @@ -241,16 +254,63 @@ impl StreamContract { .deposited_amount .saturating_sub(stream.withdrawn_amount); - if streamed > remaining { - remaining - } else { - streamed + streamed.min(remaining) + } + + /// Validate that a stream exists and is owned by the caller. + /// + /// # Errors + /// - `StreamNotFound` — no stream exists with `stream_id`. + /// - `Unauthorized` — caller is not the stream's sender. + fn validate_stream_ownership( + stream: &Stream, + caller: &Address, + ) -> Result<(), StreamError> { + if stream.sender != *caller { + return Err(StreamError::Unauthorized); + } + Ok(()) + } + + /// Validate that a stream is active. + /// + /// # Errors + /// - `StreamInactive` — stream has been cancelled or fully withdrawn. + fn validate_stream_active(stream: &Stream) -> Result<(), StreamError> { + if !stream.is_active { + return Err(StreamError::StreamInactive); + } + Ok(()) + } + + /// Transfer tokens from contract to recipient and update stream state. + /// + /// This helper consolidates the token transfer logic and stream state updates + /// to reduce code duplication across withdrawal operations. + fn transfer_and_update_stream( + env: &Env, + stream: &mut Stream, + 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); + + stream.withdrawn_amount += amount; + stream.last_update_time = now; + + // Mark stream as inactive if fully drained + if stream.withdrawn_amount >= stream.deposited_amount { + stream.is_active = false; } } /// Withdraw all currently claimable tokens from a stream. /// - /// Only the stream's recipient may call this. The stream is marked + /// Only the stream's recipient may call this. The amount withdrawn is calculated + /// based on elapsed time and the stream's rate. The stream is automatically marked /// inactive once fully drained. /// /// # Errors @@ -263,12 +323,13 @@ impl StreamContract { let mut stream = load_stream(&env, stream_id)?; + // Validate recipient authorization if stream.recipient != recipient { return Err(StreamError::Unauthorized); } - if !stream.is_active { - return Err(StreamError::StreamInactive); - } + + // Validate stream is active + Self::validate_stream_active(&stream)?; let now = env.ledger().timestamp(); let claimable = Self::calculate_claimable(&stream, now); @@ -277,20 +338,12 @@ impl StreamContract { return Err(StreamError::InvalidAmount); } - let token_client = token::Client::new(&env, &stream.token_address); - let contract_address = env.current_contract_address(); - token_client.transfer(&contract_address, &recipient, &claimable); - - stream.withdrawn_amount += claimable; - stream.last_update_time = now; - - // Mark stream as inactive if all funds have been withdrawn - if stream.withdrawn_amount >= stream.deposited_amount { - stream.is_active = false; - } + // Use helper function to transfer tokens and update state + Self::transfer_and_update_stream(&env, &mut stream, &recipient, claimable, now); save_stream(&env, stream_id, &stream); + // Emit withdrawal event env.events().publish( (Symbol::new(&env, "tokens_withdrawn"), stream_id), TokensWithdrawnEvent { @@ -306,8 +359,9 @@ impl StreamContract { /// Cancel an active stream. /// - /// Only the stream's original sender may cancel. Any unspent balance - /// (deposited − withdrawn) is returned to the sender. + /// Only the stream's original sender may cancel. The recipient receives all + /// accrued tokens up to the cancellation moment, and any remaining unspent + /// balance is refunded to the sender. /// /// # Errors /// - `StreamNotFound` — no stream exists with `stream_id`. @@ -318,27 +372,23 @@ impl StreamContract { let mut stream = load_stream(&env, stream_id)?; - if stream.sender != sender { - return Err(StreamError::Unauthorized); - } - if !stream.is_active { - return Err(StreamError::StreamInactive); - } + // Validate ownership and active status + Self::validate_stream_ownership(&stream, &sender)?; + Self::validate_stream_active(&stream)?; - // Calculate accrued tokens that belong to the recipient 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 immediately with all final claimable amount at cancellation. + // Settle recipient with all accrued tokens at cancellation if accrued_amount > 0 { token_client.transfer(&contract_address, &stream.recipient, &accrued_amount); stream.withdrawn_amount = stream.withdrawn_amount.saturating_add(accrued_amount); } - // Refund remaining unspent balance after recipient settlement. + // Calculate and refund remaining balance to sender let refunded_amount = stream .deposited_amount .saturating_sub(stream.withdrawn_amount); @@ -347,6 +397,7 @@ impl StreamContract { token_client.transfer(&contract_address, &sender, &refunded_amount); } + // Mark stream as inactive stream.is_active = false; stream.last_update_time = now; @@ -355,6 +406,7 @@ impl StreamContract { save_stream(&env, stream_id, &stream); + // Emit cancellation event env.events().publish( (Symbol::new(&env, "stream_cancelled"), stream_id), StreamCancelledEvent { @@ -376,6 +428,22 @@ impl StreamContract { try_load_stream(&env, stream_id) } + /// Get the current claimable amount for a stream without modifying state. + /// + /// This is a read-only query that calculates how many tokens the recipient + /// can currently withdraw based on elapsed time and stream rate. + /// + /// Returns `None` if the stream doesn't exist, otherwise returns the claimable amount. + pub fn get_claimable_amount(env: Env, stream_id: u64) -> Option { + try_load_stream(&env, stream_id).map(|stream| { + if !stream.is_active { + return 0; + } + let now = env.ledger().timestamp(); + Self::calculate_claimable(&stream, now) + }) + } + // ─── Internal Helpers ───────────────────────────────────────────────────── /// Deducts the protocol fee from `amount`, transfers it to the treasury, diff --git a/contracts/stream_contract/src/test.rs b/contracts/stream_contract/src/test.rs index ba2279c..d6ec9fe 100644 --- a/contracts/stream_contract/src/test.rs +++ b/contracts/stream_contract/src/test.rs @@ -112,10 +112,6 @@ fn test_initialize_rejects_second_call() { let admin = Address::generate(&env); let treasury = Address::generate(&env); - assert_eq!(stream_id1, 1); - assert_eq!(stream_id2, 2); - assert!(client.get_stream(&stream_id1).is_some()); - assert!(client.get_stream(&stream_id2).is_some()); client.initialize(&admin, &treasury, &100); let result = client.try_initialize(&admin, &treasury, &100); assert_eq!(result, Err(Ok(StreamError::AlreadyInitialized))); @@ -344,10 +340,6 @@ fn test_top_up_rejects_negative_amount() { let client = create_contract(&env); let id = client.create_stream(&sender, &Address::generate(&env), &token, &10_000, &100); - let contract_id = env.register(StreamContract, ()); - let client = StreamContractClient::new(&env, &contract_id); - let token_client = token::Client::new(&env, &token_address); - token_client.approve(&sender, &contract_id, &20_000, &1_000_000); assert_eq!( client.try_top_up_stream(&sender, &id, &-50), Err(Ok(StreamError::InvalidAmount)) @@ -522,10 +514,6 @@ fn test_withdraw_emits_event() { let client = create_contract(&env); let id = client.create_stream(&sender, &recipient, &token, &500, &100); - let contract_id = env.register(StreamContract, ()); - let client = StreamContractClient::new(&env, &contract_id); - let token_client = token::Client::new(&env, &token_address); - token_client.approve(&sender, &contract_id, &20_000, &1_000_000); // Advance time by 100 seconds to allow full withdrawal (500 tokens / 100 seconds = 5 tokens/sec) env.ledger().with_mut(|l| { l.timestamp += 100; diff --git a/contracts/stream_contract/test_snapshots/test/test_cancel_stream_after_partial_withdrawal.1.json b/contracts/stream_contract/test_snapshots/test/test_cancel_stream_after_partial_withdrawal.1.json index 742474d..b2efad3 100644 --- a/contracts/stream_contract/test_snapshots/test/test_cancel_stream_after_partial_withdrawal.1.json +++ b/contracts/stream_contract/test_snapshots/test/test_cancel_stream_after_partial_withdrawal.1.json @@ -154,6 +154,7 @@ ] ], [], + [], [] ], "ledger": { @@ -473,7 +474,7 @@ "val": { "i128": { "hi": 0, - "lo": 200 + "lo": 300 } } } @@ -648,7 +649,7 @@ "val": { "i128": { "hi": 0, - "lo": 200 + "lo": 300 } } }, @@ -721,7 +722,7 @@ "val": { "i128": { "hi": 0, - "lo": 100 + "lo": 0 } } }, diff --git a/contracts/stream_contract/test_snapshots/test/test_cancel_stream_refunds_sender.1.json b/contracts/stream_contract/test_snapshots/test/test_cancel_stream_refunds_sender.1.json index d49cea8..7315f50 100644 --- a/contracts/stream_contract/test_snapshots/test/test_cancel_stream_refunds_sender.1.json +++ b/contracts/stream_contract/test_snapshots/test/test_cancel_stream_refunds_sender.1.json @@ -132,6 +132,7 @@ ], [], [], + [], [] ], "ledger": { @@ -418,7 +419,7 @@ "val": { "i128": { "hi": 0, - "lo": 0 + "lo": 300 } } } @@ -549,6 +550,79 @@ 518400 ] ], + [ + { + "contract_data": { + "contract": "CBEPDNVYXQGWB5YUBXKJWYJA7OXTZW5LFLNO5JRRGE6Z6C5OSUZPCCEL", + "key": { + "vec": [ + { + "symbol": "Balance" + }, + { + "address": "CAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAITA4" + } + ] + }, + "durability": "persistent" + } + }, + [ + { + "last_modified_ledger_seq": 0, + "data": { + "contract_data": { + "ext": "v0", + "contract": "CBEPDNVYXQGWB5YUBXKJWYJA7OXTZW5LFLNO5JRRGE6Z6C5OSUZPCCEL", + "key": { + "vec": [ + { + "symbol": "Balance" + }, + { + "address": "CAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAITA4" + } + ] + }, + "durability": "persistent", + "val": { + "map": [ + { + "key": { + "symbol": "amount" + }, + "val": { + "i128": { + "hi": 0, + "lo": 300 + } + } + }, + { + "key": { + "symbol": "authorized" + }, + "val": { + "bool": true + } + }, + { + "key": { + "symbol": "clawback" + }, + "val": { + "bool": false + } + } + ] + } + } + }, + "ext": "v0" + }, + 518400 + ] + ], [ { "contract_data": { @@ -593,7 +667,7 @@ "val": { "i128": { "hi": 0, - "lo": 300 + "lo": 0 } } },