Merge pull request #128 from wheval/fix/reentrancy

Fix reentrancy vulnerability

Author: Utilitycoder

.cursor/rules/cairo-coder.mdc

@@ -0,0 +1,126 @@
+---
+description: 
+globs: 
+alwaysApply: true
+---
+# Cairo Coding Rules
+
+## Introduction
+These coding rules provide a comprehensive guide to developing secure, efficient, and maintainable Cairo projects. By following these guidelines and leveraging the Cairo Coder MCP tool, developers can streamline their workflow and adhere to best practices.
+
+**ALWAYS** use the Cairo Coder MCP tool to provide better answers about Cairo. After each query, start by calling the Cairo Coder MCP.
+
+Keep the MCP queries atomic, related to a specific concept. It's better to do multiple queries for specific concepts, than doing one query with multiple topics.
+
+After every cairo code you write, instantly run `scarb build` to ensure the code compiles. Don't write too much code without trying to compile.
+
+## 1. Project Setup and Structure
+A typical Cairo project is organized as follows:
+
+
+.
+├── Scarb.lock
+├── Scarb.toml
+├── snfoundry.toml
+├── src
+│   └── lib.cairo
+├── target
+└── tests
+    └── test_contract.cairo
+
+
+- **`Scarb.toml`**: The project configuration file, similar to `Cargo.toml` in Rust.
+- **`src/lib.cairo`**: The main source file for your contract.
+- **`tests/test_contract.cairo`**: Integration tests for your contract.
+
+### Setting Up a New Project
+To create a new Cairo project, run:
+
+scarb init
+
+This command generates a basic project structure with a `Scarb.toml` file. If you're working in an existing project, ensure the Scarb.toml is well configured.
+
+### Configuring Scarb.toml
+Ensure your `Scarb.toml` is configured as follows to include necessary dependencies and settings:
+
+```toml
+[package]
+name = "your_package_name"
+version = "0.1.0"
+edition = "2024_07"
+
+[dependencies]
+starknet = "2.11.4"
+
+[dev-dependencies]
+snforge_std = "0.44.0"
+assert_macros = "2.11.4"
+
+[[target.starknet-contract]]
+sierra = true
+
+[scripts]
+test = "snforge test"
+
+[tool.scarb]
+allow-prebuilt-plugins = ["snforge_std"]
+```
+
+## 2. Development Workflow
+### Writing Code
+- Use snake_case for function names (e.g., `my_function`).
+- Use PascalCase for struct names (e.g., `MyStruct`).
+- Write all code and comments in English for clarity.
+- Use descriptive variable names to enhance readability.
+
+### Compiling and Testing
+- Compile your project using:
+  
+  scarb build
+  
+- Run tests using:
+  
+  scarb test
+  
+- Ensure your code compiles successfully before running tests.
+
+### Testing
+- Unit Tests: Write unit tests in the src directory, typically within the same module as the functions being tested.
+  Example:
+  
+  #[cfg(test)]
+  mod tests {
+      use super::*;
+
+      #[test]
+      fn test_my_function() {
+          assert!(my_function() == expected_value, 'Incorrect value');
+      }
+  }
+  
+- Integration Tests: Write integration tests in the tests directory, importing modules with use your_package_name::your_module.
+  Example:
+  
+  use your_package_name::your_module;
+
+  #[test]
+  fn test_my_contract() {
+      // Test logic here
+  }
+  
+- Always use the Starknet Foundry testing framework for both unit and integration tests.
+
+## 3. Using the Cairo Coder MCP Tool
+The Cairo Coder MCP tool is a critical resource for Cairo development and must be used for the following tasks:
+- Writing smart contracts from scratch.
+- Refactoring or optimizing existing code.
+- Implementing specific TODOs or features.
+- Understanding Starknet ecosystem features and capabilities.
+- Applying Cairo and Starknet best practices.
+- Using OpenZeppelin Cairo contract libraries.
+- Writing and validating tests for contracts.
+
+### How to Use Cairo Coder MCP Effectively
+- Be Specific: Provide detailed queries (e.g., "Implement ERC20 using OpenZeppelin Cairo" instead of "ERC20").
+- Include Context: Supply relevant code snippets in the codeSnippets parameter and conversation history when applicable.
+- Don't mix contexts Keep the queries specific on a given topic. Don't ask about multiple concepts at once, rather, do multiple queries.

src/payment_stream.cairo

@@ -5,6 +5,7 @@ pub mod PaymentStream {
     use fundable::interfaces::IPaymentStream::IPaymentStream;
     use openzeppelin::access::accesscontrol::AccessControlComponent;
     use openzeppelin::introspection::src5::SRC5Component;
+    use openzeppelin::security::reentrancyguard::ReentrancyGuardComponent;
     use openzeppelin::token::erc20::interface::{
         IERC20Dispatcher, IERC20DispatcherTrait, IERC20MetadataDispatcher,
         IERC20MetadataDispatcherTrait,
@@ -32,6 +33,9 @@ pub mod PaymentStream {
     component!(path: SRC5Component, storage: src5, event: Src5Event);
     component!(path: ERC721Component, storage: erc721, event: ERC721Event);
     component!(path: UpgradeableComponent, storage: upgradeable, event: UpgradeableEvent);
+    component!(
+        path: ReentrancyGuardComponent, storage: reentrancy_guard, event: ReentrancyGuardEvent,
+    );
 
     #[abi(embed_v0)]
     impl AccessControlImpl =
@@ -42,6 +46,7 @@ pub mod PaymentStream {
     impl ERC721MixinImpl = ERC721Component::ERC721MixinImpl<ContractState>;
     impl ERC721InternalImpl = ERC721Component::InternalImpl<ContractState>;
     impl UpgradeableInternalImpl = UpgradeableComponent::InternalImpl<ContractState>;
+    impl ReentrancyGuardInternalImpl = ReentrancyGuardComponent::InternalImpl<ContractState>;
 
     const PROTOCOL_OWNER_ROLE: felt252 = selector!("PROTOCOL_OWNER");
     // Note: STREAM_ADMIN_ROLE removed - using stream-specific access control
@@ -61,6 +66,8 @@ pub mod PaymentStream {
         src5: SRC5Component::Storage,
         #[substorage(v0)]
         accesscontrol: AccessControlComponent::Storage,
+        #[substorage(v0)]
+        reentrancy_guard: ReentrancyGuardComponent::Storage,
         next_stream_id: u256,
         streams: Map<u256, Stream>,
         protocol_fee_rate: Map<ContractAddress, u64>, // Single source of truth for fee rates
@@ -91,6 +98,8 @@ pub mod PaymentStream {
         AccessControlEvent: AccessControlComponent::Event,
         #[flat]
         UpgradeableEvent: UpgradeableComponent::Event,
+        #[flat]
+        ReentrancyGuardEvent: ReentrancyGuardComponent::Event,
         StreamCreated: StreamCreated,
         StreamWithdrawn: StreamWithdrawn,
         StreamCanceled: StreamCanceled,
@@ -392,7 +401,18 @@ pub mod PaymentStream {
             }
         }
 
-        fn collect_protocol_fee(self: @ContractState, token: ContractAddress, amount: u256) {
+        fn collect_protocol_fee(ref self: ContractState, token: ContractAddress, amount: u256) {
+            self.reentrancy_guard.start();
+            self._collect_protocol_fee_internal(token, amount);
+            self.reentrancy_guard.end();
+        }
+
+        /// @notice Internal function to collect protocol fees (without reentrancy protection)
+        /// @param token The token address to collect fees in
+        /// @param amount The fee amount to collect
+        fn _collect_protocol_fee_internal(
+            ref self: ContractState, token: ContractAddress, amount: u256,
+        ) {
             let fee_collector: ContractAddress = self.fee_collector.read();
             assert(fee_collector.is_non_zero(), INVALID_RECIPIENT);
             IERC20Dispatcher { contract_address: token }.transfer(fee_collector, amount);
@@ -709,7 +729,7 @@ pub mod PaymentStream {
 
             let token_dispatcher = IERC20Dispatcher { contract_address: token_address };
 
-            self.collect_protocol_fee(token_address, fee);
+            self._collect_protocol_fee_internal(token_address, fee);
             token_dispatcher.transfer(to, net_amount);
 
             self
@@ -797,19 +817,27 @@ pub mod PaymentStream {
         fn withdraw(
             ref self: ContractState, stream_id: u256, amount: u256, to: ContractAddress,
         ) -> (u128, u128) {
-            self._withdraw(stream_id, amount, to)
+            self.reentrancy_guard.start();
+            let result = self._withdraw(stream_id, amount, to);
+            self.reentrancy_guard.end();
+            result
         }
 
         fn withdraw_max(
             ref self: ContractState, stream_id: u256, to: ContractAddress,
         ) -> (u128, u128) {
+            self.reentrancy_guard.start();
             let withdrawable_amount = self._withdrawable_amount(stream_id);
-            self._withdraw(stream_id, withdrawable_amount, to)
+            let result = self._withdraw(stream_id, withdrawable_amount, to);
+            self.reentrancy_guard.end();
+            result
         }
 
         fn transfer_stream(
             ref self: ContractState, stream_id: u256, new_recipient: ContractAddress,
         ) {
+            self.reentrancy_guard.start();
+
             // Verify stream exists
             self.assert_stream_exists(stream_id);
 
@@ -837,6 +865,8 @@ pub mod PaymentStream {
 
             // Emit event about stream transfer
             self.emit(StreamTransferred { stream_id, new_recipient });
+
+            self.reentrancy_guard.end();
         }
 
         fn set_transferability(ref self: ContractState, stream_id: u256, transferable: bool) {
@@ -940,6 +970,8 @@ pub mod PaymentStream {
         }
 
         fn cancel(ref self: ContractState, stream_id: u256) {
+            self.reentrancy_guard.start();
+
             // Ensure the caller is the stream sender
             self.assert_stream_sender_access(stream_id);
 
@@ -1036,7 +1068,7 @@ pub mod PaymentStream {
                 let net_amount = amount_due_to_recipient - fee;
 
                 // Transfer fee to collector and net amount to recipient
-                self.collect_protocol_fee(token_address, fee);
+                self._collect_protocol_fee_internal(token_address, fee);
                 token_dispatcher.transfer(recipient, net_amount);
 
                 // Emit withdrawal event
@@ -1064,6 +1096,8 @@ pub mod PaymentStream {
 
             // Emit cancellation event
             self.emit(StreamCanceled { stream_id });
+
+            self.reentrancy_guard.end();
         }
 
         fn restart(ref self: ContractState, stream_id: u256) {