From 69aa52c6d4896e203505a08083d238b4e7786aec Mon Sep 17 00:00:00 2001 From: manlikeHB Date: Thu, 4 Sep 2025 23:15:02 +0100 Subject: [PATCH] fix: Critical Security Issues in Token Management Component --- src/component/token_management/mock.cairo | 43 ++++ src/component/token_management/test.cairo | 199 ++++++++++++++++++ .../token_management.cairo | 123 ++++++++--- src/lib.cairo | 6 +- src/starkremit/StarkRemit.cairo | 2 +- 5 files changed, 336 insertions(+), 37 deletions(-) create mode 100644 src/component/token_management/mock.cairo create mode 100644 src/component/token_management/test.cairo rename src/component/{ => token_management}/token_management.cairo (81%) diff --git a/src/component/token_management/mock.cairo b/src/component/token_management/mock.cairo new file mode 100644 index 0000000..96c49da --- /dev/null +++ b/src/component/token_management/mock.cairo @@ -0,0 +1,43 @@ +use starknet::ContractAddress; + +#[starknet::contract] +pub mod MockTokenMGTContract { + use starkremit_contract::component::token_management::token_management::token_management_component; + use super::*; + + component!( + path: token_management_component, + storage: token_management_component, + event: TokenManagementEvent, + ); + + + #[abi(embed_v0)] + impl TokenManagementImpl = + token_management_component::TokenManagement; + impl TokenManagementInternalImpl = token_management_component::InternalImpl; + + + // Event definitions + #[event] + #[derive(Drop, starknet::Event)] + pub enum Event { + #[flat] + TokenManagementEvent: token_management_component::Event, + } + + // Contract storage definition + #[storage] + #[allow(starknet::colliding_storage_paths)] + struct Storage { + #[substorage(v0)] + token_management_component: token_management_component::Storage, + } + + #[constructor] + fn constructor( + ref self: ContractState, name: felt252, symbol: felt252, owner: ContractAddress, + ) { + self.token_management_component.initializer(name, symbol, owner); + } +} diff --git a/src/component/token_management/test.cairo b/src/component/token_management/test.cairo new file mode 100644 index 0000000..937eb8d --- /dev/null +++ b/src/component/token_management/test.cairo @@ -0,0 +1,199 @@ +#[cfg(test)] +mod token_management_tests { + use snforge_std::{ + ContractClassTrait, DeclareResultTrait, EventSpyAssertionsTrait, declare, spy_events, + start_cheat_caller_address, stop_cheat_caller_address, + }; + use starknet::{ContractAddress, contract_address_const}; + use starkremit_contract::component::token_management::token_management::{ + ITokenManagementDispatcher, ITokenManagementDispatcherTrait, + }; + + pub fn OWNER() -> ContractAddress { + contract_address_const::<'OWNER'>() + } + + pub fn MINTER() -> ContractAddress { + contract_address_const::<'MINTER'>() + } + + pub fn USER() -> ContractAddress { + contract_address_const::<'USER'>() + } + + pub fn NEW_OWNER() -> ContractAddress { + contract_address_const::<'NEW_OWNER'>() + } + + fn deploy_mock_contract() -> ContractAddress { + let contract_class = declare("MockTokenMGTContract").unwrap().contract_class(); + let mut constructor_calldata = array![]; + 'token'.serialize(ref constructor_calldata); + 'TK'.serialize(ref constructor_calldata); + OWNER().serialize(ref constructor_calldata); + let (contract_address, _) = contract_class.deploy(@constructor_calldata).unwrap(); + contract_address + } + + // Only owner can add minters + #[test] + fn test_only_owner_can_add_minter() { + let contract_address = deploy_mock_contract(); + let dispatcher = ITokenManagementDispatcher { contract_address }; + + // Owner should be able to add minter + start_cheat_caller_address(contract_address, OWNER()); + let result = dispatcher.add_minter(MINTER()); + assert!(result, "Owner should be able to add minter"); + assert!(dispatcher.is_minter(MINTER()), "Minter should be added"); + stop_cheat_caller_address(contract_address); + } + + #[test] + #[should_panic(expected: ('Token: caller is not owner',))] + fn test_non_owner_cannot_add_minter() { + let contract_address = deploy_mock_contract(); + let dispatcher = ITokenManagementDispatcher { contract_address }; + + // Non-owner should not be able to add minter + start_cheat_caller_address(contract_address, USER()); + dispatcher.add_minter(MINTER()); + stop_cheat_caller_address(contract_address); + } + + #[test] + fn test_only_owner_can_remove_minter() { + let contract_address = deploy_mock_contract(); + let dispatcher = ITokenManagementDispatcher { contract_address }; + + // First add a minter + start_cheat_caller_address(contract_address, OWNER()); + dispatcher.add_minter(MINTER()); + + // Owner should be able to remove minter + let result = dispatcher.remove_minter(MINTER()); + assert!(result, "Owner should be able to remove minter"); + assert!(!dispatcher.is_minter(MINTER()), "Minter should be removed"); + stop_cheat_caller_address(contract_address); + } + + #[test] + fn test_only_owner_can_set_max_supply() { + let contract_address = deploy_mock_contract(); + let dispatcher = ITokenManagementDispatcher { contract_address }; + + start_cheat_caller_address(contract_address, OWNER()); + let result = dispatcher.set_max_supply(2000000); + assert!(result, "Owner should be able to set max supply"); + assert!(dispatcher.get_max_supply() == 2000000, "Max supply should be updated"); + stop_cheat_caller_address(contract_address); + } + + #[test] + fn test_unlimited_allowance_support() { + let contract_address = deploy_mock_contract(); + let dispatcher = ITokenManagementDispatcher { contract_address }; + + // Setup: Add minter and mint tokens to USER + start_cheat_caller_address(contract_address, OWNER()); + dispatcher.add_minter(OWNER()); + stop_cheat_caller_address(contract_address); + + start_cheat_caller_address(contract_address, OWNER()); + dispatcher.mint(USER(), 1000); + stop_cheat_caller_address(contract_address); + + // Set unlimited allowance (max u256) + let unlimited_allowance = + 0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff; + start_cheat_caller_address(contract_address, USER()); + dispatcher.approve(MINTER(), unlimited_allowance); + stop_cheat_caller_address(contract_address); + + // Transfer should work without reducing allowance + start_cheat_caller_address(contract_address, MINTER()); + let result = dispatcher.transfer_from(USER(), OWNER(), 100); + assert!(result, "Transfer from should work with unlimited allowance"); + + // Allowance should still be unlimited + let remaining_allowance = dispatcher.allowance(USER(), MINTER()); + assert!(remaining_allowance == unlimited_allowance, "Allowance should remain unlimited"); + stop_cheat_caller_address(contract_address); + } + + #[test] + fn test_unlimited_max_supply() { + let contract_address = deploy_mock_contract(); + let dispatcher = ITokenManagementDispatcher { contract_address }; + + // Set max supply to 0 (unlimited) + start_cheat_caller_address(contract_address, OWNER()); + dispatcher.set_max_supply(0); + dispatcher.add_minter(OWNER()); + + // Should be able to mint beyond original max supply + let result = dispatcher.mint(USER(), 2000000); // More than original 1M max + assert!(result, "Should be able to mint unlimited when max_supply is 0"); + stop_cheat_caller_address(contract_address); + } + + #[test] + fn test_correct_ownership_transfer() { + let contract_address = deploy_mock_contract(); + let dispatcher = ITokenManagementDispatcher { contract_address }; + let mut spy = spy_events(); + + start_cheat_caller_address(contract_address, OWNER()); + let result = dispatcher.transfer_ownership(NEW_OWNER()); + assert!(result, "Ownership transfer should succeed"); + assert!(dispatcher.get_owner() == NEW_OWNER(), "New owner should be set"); + stop_cheat_caller_address(contract_address); + } + + #[test] + fn test_no_duplicate_transfer_logic() { + let contract_address = deploy_mock_contract(); + let dispatcher = ITokenManagementDispatcher { contract_address }; + + // Setup: Mint tokens to OWNER + start_cheat_caller_address(contract_address, OWNER()); + dispatcher.add_minter(OWNER()); + dispatcher.mint(OWNER(), 1000); + + let initial_balance = dispatcher.balance_of(OWNER()); + let initial_user_balance = dispatcher.balance_of(USER()); + + // Transfer tokens + let result = dispatcher.transfer(USER(), 100); + assert!(result, "Transfer should succeed"); + + // Verify balances are correctly updated + let final_balance = dispatcher.balance_of(OWNER()); + let final_user_balance = dispatcher.balance_of(USER()); + + assert!( + final_balance == initial_balance - 100, + "Sender balance should decrease by exact amount", + ); + assert!( + final_user_balance == initial_user_balance + 100, + "Recipient balance should increase by exact amount", + ); + stop_cheat_caller_address(contract_address); + } + + #[test] + #[should_panic(expected: ('SetMaxSupply: too low',))] + fn test_max_supply_cannot_be_below_current_supply() { + let contract_address = deploy_mock_contract(); + let dispatcher = ITokenManagementDispatcher { contract_address }; + + start_cheat_caller_address(contract_address, OWNER()); + dispatcher.add_minter(OWNER()); + dispatcher.mint(USER(), 1000); + + // Try to set max supply below current supply + dispatcher.set_max_supply(500); // Should fail since current supply is 1000 + stop_cheat_caller_address(contract_address); + } +} diff --git a/src/component/token_management.cairo b/src/component/token_management/token_management.cairo similarity index 81% rename from src/component/token_management.cairo rename to src/component/token_management/token_management.cairo index 1c91ab8..08aebbd 100644 --- a/src/component/token_management.cairo +++ b/src/component/token_management/token_management.cairo @@ -45,6 +45,10 @@ pub mod token_management_component { use starkremit_contract::utils::helpers::{assert_non_zero_amount, assert_not_zero_address}; use super::*; + // Constant for unlimited allowance + const UNLIMITED_ALLOWANCE: u256 = + 0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff; + #[storage] pub struct Storage { // ERC20 Standard Fields @@ -55,7 +59,7 @@ pub mod token_management_component { balances: Map, allowances: Map<(ContractAddress, ContractAddress), u256>, // Token Management Fields - max_supply: u256, + max_supply: u256, // 0 = unlimited supply minters: Map, // Access Control Fields owner: ContractAddress, @@ -79,66 +83,95 @@ pub mod token_management_component { #[derive(Drop, starknet::Event)] pub struct Transfer { - from: ContractAddress, - to: ContractAddress, - value: u256, + #[key] + pub from: ContractAddress, + #[key] + pub to: ContractAddress, + pub value: u256, } #[derive(Drop, starknet::Event)] pub struct Approval { - owner: ContractAddress, - spender: ContractAddress, - value: u256, + #[key] + pub owner: ContractAddress, + #[key] + pub spender: ContractAddress, + pub value: u256, } #[derive(Drop, starknet::Event)] pub struct Minted { - to: ContractAddress, - amount: u256, + #[key] + pub to: ContractAddress, + pub amount: u256, } #[derive(Drop, starknet::Event)] pub struct Burned { - from: ContractAddress, - amount: u256, + #[key] + pub from: ContractAddress, + pub amount: u256, } #[derive(Drop, starknet::Event)] pub struct MinterAdded { - minter: ContractAddress, + #[key] + pub minter: ContractAddress, + #[key] + pub added_by: ContractAddress, } #[derive(Drop, starknet::Event)] pub struct MinterRemoved { - minter: ContractAddress, + #[key] + pub minter: ContractAddress, + #[key] + pub removed_by: ContractAddress, } #[derive(Drop, starknet::Event)] pub struct MaxSupplyUpdated { - old_max_supply: u256, - new_max_supply: u256, + #[key] + pub old_max_supply: u256, + #[key] + pub new_max_supply: u256, } #[derive(Drop, starknet::Event)] pub struct OwnershipTransferred { - previous_owner: ContractAddress, - new_owner: ContractAddress, + #[key] + pub previous_owner: ContractAddress, + #[key] + pub new_owner: ContractAddress, } #[derive(Drop, starknet::Event)] pub struct Paused { - account: ContractAddress, + #[key] + pub account: ContractAddress, } #[derive(Drop, starknet::Event)] pub struct Unpaused { - account: ContractAddress, + #[key] + pub account: ContractAddress, } #[generate_trait] pub impl InternalImpl< TContractState, +HasComponent, > of InternalTrait { + fn initializer( + ref self: ComponentState, + name: felt252, + symbol: felt252, + owner: ContractAddress, + ) { + self.name.write(name); + self.symbol.write(symbol); + self.owner.write(owner); + } + fn _transfer( ref self: ComponentState, from: ContractAddress, @@ -165,6 +198,11 @@ pub mod token_management_component { fn _assert_not_paused(self: @ComponentState) { assert(!self.paused.read(), MintBurnErrors::CONTRACT_PAUSED); } + + fn _assert_only_minter(self: @ComponentState) { + let caller = get_caller_address(); + assert(self.minters.read(caller), MintBurnErrors::NOT_MINTER); + } } #[embeddable_as(TokenManagement)] @@ -204,6 +242,8 @@ pub mod token_management_component { assert_non_zero_amount(value); let caller = get_caller_address(); + + // Use internal _transfer function (no duplicate logic) self._transfer(caller, to, value); true } @@ -239,12 +279,15 @@ pub mod token_management_component { let caller = get_caller_address(); let allowance = self.allowances.read((from, caller)); - assert(allowance >= value, MintBurnErrors::INSUFFICIENT_ALLOWANCE); - // Update allowance - self.allowances.write((from, caller), allowance - value); + // Check allowance (skip for unlimited allowance) + if allowance != UNLIMITED_ALLOWANCE { + assert(allowance >= value, MintBurnErrors::INSUFFICIENT_ALLOWANCE); + // Update allowance + self.allowances.write((from, caller), allowance - value); + } - // Perform transfer + // Perform transfer using internal function self._transfer(from, to, value); true } @@ -267,12 +310,16 @@ pub mod token_management_component { assert_not_zero_address(to); assert_non_zero_amount(amount); - let caller = get_caller_address(); - assert(self.minters.read(caller), MintBurnErrors::NOT_MINTER); + // Check minter permission + self._assert_only_minter(); let supply = self.total_supply.read(); let max_supply = self.max_supply.read(); - assert(supply + amount <= max_supply, MintBurnErrors::MAX_SUPPLY_EXCEEDED); + + // Check max supply (0 means unlimited) + if max_supply != 0 { + assert(supply + amount <= max_supply, MintBurnErrors::MAX_SUPPLY_EXCEEDED); + } // Update total supply self.total_supply.write(supply + amount); @@ -305,8 +352,8 @@ pub mod token_management_component { assert_not_zero_address(from); assert_non_zero_amount(amount); - let caller = get_caller_address(); - assert(self.minters.read(caller), MintBurnErrors::NOT_MINTER); + // Check minter permission + self._assert_only_minter(); let from_balance = self.balances.read(from); assert(from_balance >= amount, MintBurnErrors::INSUFFICIENT_BALANCE_BURN); @@ -338,8 +385,9 @@ pub mod token_management_component { self._assert_only_owner(); assert_not_zero_address(minter); + let caller = get_caller_address(); self.minters.write(minter, true); - self.emit(Event::MinterAdded(MinterAdded { minter })); + self.emit(Event::MinterAdded(MinterAdded { minter, added_by: caller })); true } @@ -349,19 +397,24 @@ pub mod token_management_component { self._assert_only_owner(); assert_not_zero_address(minter); + let caller = get_caller_address(); self.minters.write(minter, false); - self.emit(Event::MinterRemoved(MinterRemoved { minter })); + self.emit(Event::MinterRemoved(MinterRemoved { minter, removed_by: caller })); true } fn set_max_supply(ref self: ComponentState, new_max_supply: u256) -> bool { self._assert_only_owner(); - assert_non_zero_amount(new_max_supply); let current_supply = self.total_supply.read(); - assert(new_max_supply >= current_supply, MintBurnErrors::MAX_SUPPLY_TOO_LOW); + + // If setting non-zero max supply, ensure it's >= current supply + if new_max_supply != 0 { + assert(new_max_supply >= current_supply, MintBurnErrors::MAX_SUPPLY_TOO_LOW); + } let old_max_supply = self.max_supply.read(); + let caller = get_caller_address(); self.max_supply.write(new_max_supply); self.emit(Event::MaxSupplyUpdated(MaxSupplyUpdated { old_max_supply, new_max_supply })); @@ -392,10 +445,9 @@ pub mod token_management_component { self._assert_only_owner(); assert_not_zero_address(new_owner); - let caller = get_caller_address(); - assert(caller != new_owner, MintBurnErrors::OWNERSHIP_TO_SELF); - let previous_owner = self.owner.read(); + assert(previous_owner != new_owner, MintBurnErrors::OWNERSHIP_TO_SELF); + self.owner.write(new_owner); self @@ -424,3 +476,4 @@ pub mod token_management_component { } } } + diff --git a/src/lib.cairo b/src/lib.cairo index d9ebf7b..6d5dc33 100644 --- a/src/lib.cairo +++ b/src/lib.cairo @@ -31,7 +31,11 @@ pub mod component { pub mod kyc; pub mod loan; pub mod savings_group; - pub mod token_management; + pub mod token_management { + pub mod mock; + pub mod test; + pub mod token_management; + } pub mod transfer; pub mod user_management { pub mod mock; diff --git a/src/starkremit/StarkRemit.cairo b/src/starkremit/StarkRemit.cairo index 0a3156b..50121c5 100644 --- a/src/starkremit/StarkRemit.cairo +++ b/src/starkremit/StarkRemit.cairo @@ -33,7 +33,7 @@ pub mod StarkRemit { use starkremit_contract::component::kyc::kyc_component; use starkremit_contract::component::loan::loan_component; use starkremit_contract::component::savings_group::savings_group_component; - use starkremit_contract::component::token_management::token_management_component; + use starkremit_contract::component::token_management::token_management::token_management_component; use starkremit_contract::component::transfer::transfer_component; use starkremit_contract::component::user_management::user_management::user_management_component; use super::*;