test and feature added for multi call buy prevention (#81)

1. function `_check_and_update_tx_hash` added in order to check the tx_hash of the recipient in both transfer and transfer_from function. 2. previous test updated in order to give them a non-zero tx_hash (by default it 0x0) 3. new test added which panics when same tx_hash is having two different calls of same recipient. --------- Co-authored-by: enitrat <[email protected]>
keep-starknet-strange · Dec 29, 2023 · 659f01f · 659f01f
1 parent 11a3752
commit 659f01f
Show file tree

Hide file tree

Showing 5 changed files with 149 additions and 14 deletions.
diff --git a/contracts/src/tests.cairo b/contracts/src/tests.cairo
@@ -1,5 +1,6 @@
 mod test_factory;
 
+mod test_memecoin_erc20;
 mod test_token_locker;
 
 mod test_unruggable_memecoin;

diff --git a/contracts/src/tests/test_memecoin_erc20.cairo b/contracts/src/tests/test_memecoin_erc20.cairo
@@ -2,11 +2,12 @@ use openzeppelin::token::erc20::interface::{IERC20, IERC20Dispatcher, IERC20Disp
 use openzeppelin::utils::serde::SerializedAppend;
 
 use snforge_std::{
-    declare, ContractClassTrait, start_prank, stop_prank, RevertedTransaction, CheatTarget
+    declare, ContractClassTrait, start_prank, stop_prank, RevertedTransaction, CheatTarget,
+    TxInfoMock
 };
 use starknet::{ContractAddress, contract_address_const};
 use unruggable::amm::amm::{AMM, AMMV2};
-
+use unruggable::tests::utils::{TxInfoMockTrait};
 use unruggable::tokens::interface::{
     IUnruggableMemecoinDispatcher, IUnruggableMemecoinDispatcherTrait
 };
@@ -183,13 +184,14 @@ mod erc20_entrypoints {
     use core::traits::Into;
     use openzeppelin::token::erc20::interface::IERC20;
     use snforge_std::{
-        declare, ContractClassTrait, start_prank, stop_prank, start_warp, CheatTarget
+        declare, ContractClassTrait, start_prank, stop_prank, start_warp, CheatTarget, TxInfoMock
     };
     use starknet::{ContractAddress, contract_address_const};
     use super::{deploy_contract, instantiate_params};
-    use unruggable::tests_utils::deployer_helper::DeployerHelper::{
+    use unruggable::tests::utils::DeployerHelper::{
         deploy_contracts, deploy_unruggable_memecoin_contract, deploy_memecoin_factory, create_eth
     };
+    use unruggable::tests::utils::{DefaultTxInfoMock};
     use unruggable::tokens::interface::{
         IUnruggableMemecoinDispatcher, IUnruggableMemecoinDispatcherTrait
     };
@@ -319,6 +321,11 @@ mod erc20_entrypoints {
 
         let memecoin = IUnruggableMemecoinDispatcher { contract_address };
 
+        // setting tx_hash here 
+        let mut tx_info: TxInfoMock = Default::default();
+        tx_info.transaction_hash = Option::Some(1234);
+        snforge_std::start_spoof(CheatTarget::One(memecoin.contract_address), tx_info);
+
         // Transfer 20 tokens to recipient.
         start_prank(CheatTarget::One(memecoin.contract_address), initial_holder_1);
         memecoin.transfer(recipient, 20);
@@ -365,6 +372,11 @@ mod erc20_entrypoints {
         start_prank(CheatTarget::One(memecoin.contract_address), initial_holder_1);
         memecoin.approve(spender, initial_supply);
 
+        // setting tx_hash here 
+        let mut tx_info: TxInfoMock = Default::default();
+        tx_info.transaction_hash = Option::Some(1234);
+        snforge_std::start_spoof(CheatTarget::One(memecoin.contract_address), tx_info);
+
         // Transfer 20 tokens to recipient.
         start_prank(CheatTarget::One(memecoin.contract_address), spender);
         memecoin.transfer_from(initial_holder_1, recipient, 20);
@@ -479,6 +491,11 @@ mod erc20_entrypoints {
         start_prank(CheatTarget::One(memecoin.contract_address), initial_holder_1);
         memecoin.approve(spender, initial_supply);
 
+        // setting tx_hash here 
+        let mut tx_info: TxInfoMock = Default::default();
+        tx_info.transaction_hash = Option::Some(1234);
+        snforge_std::start_spoof(CheatTarget::One(memecoin.contract_address), tx_info);
+
         // Transfer 20 tokens to recipient.
         start_prank(CheatTarget::One(memecoin.contract_address), spender);
         memecoin.transferFrom(initial_holder_1, recipient, 20);

diff --git a/contracts/src/tests/test_unruggable_memecoin.cairo b/contracts/src/tests/test_unruggable_memecoin.cairo
@@ -2,11 +2,12 @@ use openzeppelin::token::erc20::interface::{IERC20, ERC20ABIDispatcher, ERC20ABI
 use openzeppelin::utils::serde::SerializedAppend;
 
 use snforge_std::{
-    declare, ContractClassTrait, start_prank, stop_prank, RevertedTransaction, CheatTarget
+    declare, ContractClassTrait, start_prank, stop_prank, RevertedTransaction, CheatTarget,
+    TxInfoMock,
 };
 use starknet::{ContractAddress, contract_address_const};
 use unruggable::amm::amm::{AMM, AMMV2, AMMTrait};
-
+use unruggable::tests::utils::DefaultTxInfoMock;
 use unruggable::tokens::interface::{
     IUnruggableMemecoinDispatcher, IUnruggableMemecoinDispatcherTrait
 };
@@ -90,7 +91,7 @@ mod memecoin_entrypoints {
         IERC20, ERC20ABIDispatcher, ERC20ABIDispatcherTrait
     };
     use snforge_std::{
-        declare, ContractClassTrait, start_prank, stop_prank, CheatTarget, start_warp
+        declare, ContractClassTrait, start_prank, stop_prank, CheatTarget, start_warp, TxInfoMock
     };
     use starknet::{ContractAddress, contract_address_const};
     use super::{deploy_contract, instantiate_params, ETH_UNIT_DECIMALS};
@@ -102,12 +103,12 @@ mod memecoin_entrypoints {
 
     use unruggable::factory::{IFactory, IFactoryDispatcher, IFactoryDispatcherTrait};
     use unruggable::tests::utils::DeployerHelper::{
-        deploy_contracts, deploy_unruggable_memecoin_contract, deploy_memecoin_factory, create_eth
+        deploy_contracts, deploy_unruggable_memecoin_contract, deploy_memecoin_factory, create_eth,
     };
     use unruggable::tests::utils::{
         deploy_amm_factory_and_router, deploy_meme_factory_with_owner, deploy_locker,
         deploy_eth_with_owner, OWNER, NAME, SYMBOL, ETH_INITIAL_SUPPLY, INITIAL_HOLDERS,
-        INITIAL_HOLDERS_AMOUNTS, SALT
+        INITIAL_HOLDERS_AMOUNTS, SALT, DefaultTxInfoMock
     };
     use unruggable::tokens::interface::{
         IUnruggableMemecoin, IUnruggableMemecoinDispatcher, IUnruggableMemecoinDispatcherTrait
@@ -479,6 +480,11 @@ mod memecoin_entrypoints {
 
         let memecoin = IUnruggableMemecoinDispatcher { contract_address };
 
+        // setting tx_hash here 
+        let mut tx_info: TxInfoMock = Default::default();
+        tx_info.transaction_hash = Option::Some(1234);
+        snforge_std::start_spoof(CheatTarget::One(memecoin.contract_address), tx_info);
+
         // Transfer 21 token from owner to alice.
         start_prank(CheatTarget::One(memecoin.contract_address), initial_holder_1);
         let send_amount = memecoin.transfer(alice, 21);
@@ -510,11 +516,53 @@ mod memecoin_entrypoints {
 
         let memecoin = IUnruggableMemecoinDispatcher { contract_address };
 
+        // setting tx_hash here 
+        let mut tx_info: TxInfoMock = Default::default();
+        tx_info.transaction_hash = Option::Some(1234);
+        snforge_std::start_spoof(CheatTarget::One(memecoin.contract_address), tx_info);
+
         // Transfer 21 token from owner to alice.
         start_prank(CheatTarget::One(memecoin.contract_address), initial_holder_1);
         let send_amount = memecoin.transfer_from(initial_holder_1, alice, 500);
     }
 
+    #[test]
+    #[should_panic(expected: ('Multi calls not allowed',))]
+    fn test_transfer_from_multi_call() {
+        let (
+            owner,
+            name,
+            symbol,
+            initial_supply,
+            initial_holder_1,
+            initial_holder_2,
+            initial_holders,
+            initial_holders_amounts
+        ) =
+            instantiate_params();
+        let alice = contract_address_const::<53>();
+        let bob = contract_address_const::<54>();
+
+        let contract_address =
+            match deploy_contract(
+                owner, name, symbol, initial_supply, initial_holders, initial_holders_amounts
+            ) {
+            Result::Ok(address) => address,
+            Result::Err(msg) => panic(msg.panic_data),
+        };
+
+        let memecoin = IUnruggableMemecoinDispatcher { contract_address };
+
+        // setting tx_hash here 
+        let mut tx_info: TxInfoMock = Default::default();
+        tx_info.transaction_hash = Option::Some(1234);
+
+        // Transfer token from owner to alice twice - should fail
+        start_prank(CheatTarget::One(memecoin.contract_address), initial_holder_1);
+        let send_amount = memecoin.transfer_from(initial_holder_1, alice, 0);
+        let send_amount = memecoin.transfer_from(initial_holder_1, alice, 0);
+    }
+
     #[test]
     fn test_classic_max_percentage() {
         let (
@@ -540,6 +588,11 @@ mod memecoin_entrypoints {
 
         let memecoin = IUnruggableMemecoinDispatcher { contract_address };
 
+        // setting tx_hash here 
+        let mut tx_info: TxInfoMock = Default::default();
+        tx_info.transaction_hash = Option::Some(1234);
+        snforge_std::start_spoof(CheatTarget::One(memecoin.contract_address), tx_info);
+
         // Transfer 1 token from owner to alice.
         start_prank(CheatTarget::One(memecoin.contract_address), initial_holder_1);
         let send_amount = memecoin.transfer(alice, 20);
@@ -783,6 +836,7 @@ mod memecoin_internals {
     use openzeppelin::token::erc20::interface::IERC20;
     use snforge_std::{declare, ContractClassTrait, start_prank, stop_prank, CheatTarget};
     use starknet::{ContractAddress, contract_address_const};
+    use super::{TxInfoMock, DefaultTxInfoMock};
     use super::{deploy_contract, instantiate_params};
     use unruggable::tokens::interface::{
         IUnruggableMemecoinDispatcher, IUnruggableMemecoinDispatcherTrait
@@ -826,6 +880,11 @@ mod memecoin_internals {
             // create a unique address
             let unique_recipient: ContractAddress = (index.into() + 9999).try_into().unwrap();
 
+            // creating and setting unique tx_hash here 
+            let mut tx_info: TxInfoMock = Default::default();
+            tx_info.transaction_hash = Option::Some(index.into() + 9999);
+            snforge_std::start_spoof(CheatTarget::One(memecoin.contract_address), tx_info);
+
             // Transfer 1 token to the unique recipient
             memecoin.transfer(unique_recipient, 1);
 
@@ -874,6 +933,11 @@ mod memecoin_internals {
             // create a unique address
             let unique_recipient: ContractAddress = (index.into() + 9999).try_into().unwrap();
 
+            // creating and setting unique tx_hash here 
+            let mut tx_info: TxInfoMock = Default::default();
+            tx_info.transaction_hash = Option::Some(index.into() + 9999);
+            snforge_std::start_spoof(CheatTarget::One(memecoin.contract_address), tx_info);
+
             // Transfer 1 token to the unique recipient
             memecoin.transfer(unique_recipient, 1);
 
@@ -933,6 +997,11 @@ mod memecoin_internals {
                 break;
             }
 
+            // creating and setting unique tx_hash here 
+            let mut tx_info: TxInfoMock = Default::default();
+            tx_info.transaction_hash = Option::Some(index.into() + 9999);
+            snforge_std::start_spoof(CheatTarget::One(memecoin.contract_address), tx_info);
+
             // Self transfer tokens
             memecoin.transfer(initial_holder_2, 1);
 
@@ -978,6 +1047,10 @@ mod memecoin_internals {
             // create a unique address
             let unique_recipient: ContractAddress = (index.into() + 9999).try_into().unwrap();
 
+            // creating and setting unique tx_hash here 
+            let mut tx_info: TxInfoMock = Default::default();
+            tx_info.transaction_hash = Option::Some(index.into() + 9999);
+            snforge_std::start_spoof(CheatTarget::One(memecoin.contract_address), tx_info);
             // Transfer 1 token to the unique recipient
             memecoin.transfer(unique_recipient, 1);
 
@@ -1043,6 +1116,11 @@ mod memecoin_internals {
             // create a unique address
             let unique_recipient: ContractAddress = (index.into() + 9999).try_into().unwrap();
 
+            // creating and setting unique tx_hash here 
+            let mut tx_info: TxInfoMock = Default::default();
+            tx_info.transaction_hash = Option::Some(index.into() + 9999);
+            snforge_std::start_spoof(CheatTarget::One(memecoin.contract_address), tx_info);
+
             // Transfer 1 token to the unique recipient
             memecoin.transfer(unique_recipient, 1);
 

diff --git a/contracts/src/tests/utils.cairo b/contracts/src/tests/utils.cairo
@@ -1,7 +1,9 @@
 use core::traits::TryInto;
 use openzeppelin::token::erc20::interface::ERC20ABIDispatcher;
 use openzeppelin::token::erc20::interface::ERC20ABIDispatcherTrait;
-use snforge_std::{ContractClass, ContractClassTrait, CheatTarget, declare, start_prank, stop_prank};
+use snforge_std::{
+    ContractClass, ContractClassTrait, CheatTarget, declare, start_prank, stop_prank, TxInfoMock
+};
 use starknet::ContractAddress;
 use unruggable::amm::amm::{AMM, AMMV2, AMMTrait};
 use unruggable::factory::{IFactoryDispatcher, IFactoryDispatcherTrait};
@@ -44,6 +46,26 @@ fn SALT() -> felt252 {
     'salty'.try_into().unwrap()
 }
 
+
+trait TxInfoMockTrait {
+    fn default() -> TxInfoMock;
+}
+
+impl DefaultTxInfoMock of Default<TxInfoMock> {
+    fn default() -> TxInfoMock {
+        TxInfoMock {
+            version: Option::None,
+            account_contract_address: Option::None,
+            max_fee: Option::None,
+            signature: Option::None,
+            transaction_hash: Option::None,
+            chain_id: Option::None,
+            nonce: Option::None,
+        }
+    }
+}
+
+
 // NOT THE ACTUAL ETH ADDRESS
 // It's set to a the maximum possible value for a contract address
 // This ensures that in Jediswap pairs, the ETH side is always token1

diff --git a/contracts/src/tokens/memecoin.cairo b/contracts/src/tokens/memecoin.cairo
@@ -3,10 +3,8 @@ use starknet::ContractAddress;
 
 #[starknet::contract]
 mod UnruggableMemecoin {
-    use array::ArrayTrait;
-    use core::box::BoxTrait;
-
     use debug::PrintTrait;
+    use integer::BoundedInt;
     use openzeppelin::access::ownable::OwnableComponent;
     use openzeppelin::access::ownable::ownable::OwnableComponent::InternalTrait as OwnableInternalTrait;
     use openzeppelin::token::erc20::ERC20Component;
@@ -66,6 +64,7 @@ mod UnruggableMemecoin {
         launched: bool,
         pre_launch_holders_count: u8,
         team_allocation: u256,
+        tx_hash_tracker: LegacyMap<ContractAddress, felt252>,
         locker_contract: ContractAddress,
         transfer_delay: u64,
         launch_time: u64,
@@ -265,6 +264,7 @@ mod UnruggableMemecoin {
 
         fn transfer(ref self: ContractState, recipient: ContractAddress, amount: u256) -> bool {
             let sender = get_caller_address();
+            self.ensure_not_multicall(recipient);
             self._check_max_buy_percentage(sender, recipient, amount);
             self._transfer(sender, recipient, amount);
             true
@@ -281,6 +281,7 @@ mod UnruggableMemecoin {
             // which performs a transfer_from() to send the tokens to the pool.
             // Therefore, we need to bypass this validation if the sender is the memecoin contract.
             if sender != get_contract_address() {
+                self.ensure_not_multicall(sender);
                 self._check_max_buy_percentage(sender, recipient, amount);
             }
             self.erc20._spend_allowance(sender, caller, amount);
@@ -435,7 +436,23 @@ mod UnruggableMemecoin {
             }
         }
 
-        /// Constructor logic.
+        /// Ensures that the current call is not a part of a multicall.
+        ///
+        /// By keeping track of the last transaction hash each address has received tokens at,
+        /// we can ensure that the current call is not part of a transaction already performed.
+        ///
+        /// # Arguments
+        /// * `recipient` - The contract address of the recipient.
+        //TODO(audit): Verify whether this can cause a problem for trading through aggregators, that can
+        // do multiple transfers when using complex routes.
+        #[inline(always)]
+        fn ensure_not_multicall(ref self: ContractState, recipient: ContractAddress) {
+            let tx_hash: felt252 = get_tx_info().unbox().transaction_hash;
+            assert(self.tx_hash_tracker.read(recipient) != tx_hash, 'Multi calls not allowed');
+            self.tx_hash_tracker.write(recipient, tx_hash);
+        }
+
+        /// Cons\tructor logic.
         /// # Arguments
         /// * `locker_address` - Token locker contract address.
         /// * `limit_delay` - Delay timestamp to release transfer amount check.