From 73f8373e677c5d09e55c50c8db0d5ff40457f7cf Mon Sep 17 00:00:00 2001
From: enitrat <msaug@protonmail.com>
Date: Fri, 4 Oct 2024 15:57:00 +0200
Subject: [PATCH] feat: enable registration checks

---
 crates/contracts/src/kakarot_core/kakarot.cairo          | 6 ++++--
 crates/contracts/src/test_utils.cairo                    | 1 -
 crates/contracts/tests/test_execution_from_outside.cairo | 2 --
 crates/contracts/tests/test_kakarot_core.cairo           | 5 ++---
 4 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/crates/contracts/src/kakarot_core/kakarot.cairo b/crates/contracts/src/kakarot_core/kakarot.cairo
index fec55b8a0..d68932503 100644
--- a/crates/contracts/src/kakarot_core/kakarot.cairo
+++ b/crates/contracts/src/kakarot_core/kakarot.cairo
@@ -5,6 +5,7 @@ const INVOKE_ETH_CALL_FORBIDDEN: felt252 = 'KKT: Cannot invoke eth_call';
 pub mod KakarotCore {
     use core::num::traits::Zero;
     use core::starknet::event::EventEmitter;
+    use core::starknet::get_caller_address;
     use core::starknet::storage::{
         Map, StorageMapReadAccess, StorageMapWriteAccess, StoragePointerReadAccess,
         StoragePointerWriteAccess
@@ -184,8 +185,9 @@ pub mod KakarotCore {
             assert(existing_address.is_zero(), 'Account already exists');
 
             let starknet_address = self.compute_starknet_address(evm_address);
-            //TODO: enable this assertion. Will require changing test runner to snfoundry
-            // assert!(starknet_address == caller, "Account must be registered by the caller");
+            assert!(
+                starknet_address == get_caller_address(), "Account must be registered by the caller"
+            );
 
             self.Kakarot_evm_to_starknet_address.write(evm_address, starknet_address);
             self.emit(AccountDeployed { evm_address, starknet_address });
diff --git a/crates/contracts/src/test_utils.cairo b/crates/contracts/src/test_utils.cairo
index d412068a2..e56908366 100644
--- a/crates/contracts/src/test_utils.cairo
+++ b/crates/contracts/src/test_utils.cairo
@@ -142,7 +142,6 @@ pub fn setup_contracts_for_testing() -> (IERC20CamelDispatcher, IExtendedKakarot
 
     let sequencer_sn_address = kakarot_core.address_registry(sequencer);
     start_cheat_sequencer_address_global(sequencer_sn_address);
-    start_cheat_caller_address_global(kakarot_core.contract_address);
     start_cheat_chain_id_global(chain_id().into());
     return (native_token, kakarot_core);
 }
diff --git a/crates/contracts/tests/test_execution_from_outside.cairo b/crates/contracts/tests/test_execution_from_outside.cairo
index 9daff0fe4..aa771a6d8 100644
--- a/crates/contracts/tests/test_execution_from_outside.cairo
+++ b/crates/contracts/tests/test_execution_from_outside.cairo
@@ -142,8 +142,6 @@ fn set_up() -> (IExtendedKakarotCoreDispatcher, IAccountDispatcher, IERC20CamelD
     let eoa = IAccountDispatcher {
         contract_address: kakarot_core.deploy_externally_owned_account(transaction_signer())
     };
-    start_cheat_caller_address_global(kakarot_core.contract_address);
-
     start_cheat_block_timestamp(eoa.contract_address, 999);
     start_cheat_chain_id_global(chain_id().into());
 
diff --git a/crates/contracts/tests/test_kakarot_core.cairo b/crates/contracts/tests/test_kakarot_core.cairo
index 45fde481f..903885473 100644
--- a/crates/contracts/tests/test_kakarot_core.cairo
+++ b/crates/contracts/tests/test_kakarot_core.cairo
@@ -18,7 +18,7 @@ use evm::test_utils::chain_id;
 use evm::test_utils;
 use snforge_std::{
     declare, DeclareResultTrait, start_cheat_caller_address, spy_events, EventSpyTrait,
-    cheat_caller_address, CheatSpan, store, load
+    cheat_caller_address, CheatSpan, store, load, stop_cheat_caller_address_global
 };
 use snforge_utils::snforge_utils::{EventsFilterBuilderTrait, ContractEventsTrait};
 use starknet::storage::StorageTrait;
@@ -274,7 +274,6 @@ fn test_eth_call() {
 
     let evm_address = test_utils::evm_address();
     kakarot_core.deploy_externally_owned_account(evm_address);
-
     let account = contract_utils::deploy_contract_account(
         kakarot_core, test_utils::other_evm_address(), counter_evm_bytecode()
     );
@@ -326,7 +325,7 @@ fn test_eth_send_transaction_deploy_tx() {
         gas_limit,
         input: deploy_counter_calldata()
     };
-    start_cheat_caller_address(kakarot_core.contract_address, eoa);
+    cheat_caller_address(kakarot_core.contract_address, eoa, CheatSpan::TargetCalls(1));
     let (_, deploy_result, _) = kakarot_core.eth_send_transaction(Transaction::Legacy(tx));
 
     // Then