From cd95784c65fb80707de38f1ee5a6a87ee763823d Mon Sep 17 00:00:00 2001
From: Juan Rigada <62958725+Jrigada@users.noreply.github.com>
Date: Fri, 31 May 2024 14:08:27 -0300
Subject: [PATCH] fix: Override blockhash in zk context (#400)

---
 crates/zksync/core/src/vm/tracer.rs | 19 +++++++++++++++++++
 zk-tests/src/Basic.t.sol            | 27 ++++++++++++++++++++++++++-
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/crates/zksync/core/src/vm/tracer.rs b/crates/zksync/core/src/vm/tracer.rs
index fee2d0cc9..b628fb94b 100644
--- a/crates/zksync/core/src/vm/tracer.rs
+++ b/crates/zksync/core/src/vm/tracer.rs
@@ -62,6 +62,12 @@ const SELECTOR_SYSTEM_CONTEXT_BLOCK_TIMESTAMP: [u8; 4] = hex!("796b89b9");
 /// Selector for `baseFee()`
 const SELECTOR_BASE_FEE: [u8; 4] = hex!("6ef25c3a");
 
+/// Selector for retrieving the blockhash of a given block.
+/// This is used to override the current `blockhash()` to foundry test's context.
+///
+/// Selector for `getBlockHashEVM(uint256)`
+const SELECTOR_BLOCK_HASH: [u8; 4] = hex!("80b41246");
+
 /// Represents the context for [CheatcodeContext]
 #[derive(Debug, Default)]
 pub struct CheatcodeTracerContext<'a> {
@@ -286,6 +292,19 @@ impl<S: Send, H: HistoryMode> DynTracer<S, SimpleMemory<H>> for CheatcodeTracer
             }
         }
 
+        // Override blockhash
+        if let Opcode::FarCall(_call) = data.opcode.variant.opcode {
+            let calldata = get_calldata(&state, memory);
+            let current = state.vm_local_state.callstack.current;
+
+            if current.code_address == SYSTEM_CONTEXT_ADDRESS &&
+                calldata.starts_with(&SELECTOR_BLOCK_HASH)
+            {
+                self.farcall_handler.set_immediate_return(rU256::ZERO.to_be_bytes_vec());
+                return
+            }
+        }
+
         if let Some(delegate_as) = self.call_context.delegate_as {
             if let Opcode::FarCall(_call) = data.opcode.variant.opcode {
                 let current = state.vm_local_state.callstack.current;
diff --git a/zk-tests/src/Basic.t.sol b/zk-tests/src/Basic.t.sol
index 8d02e629a..14bbb7cae 100644
--- a/zk-tests/src/Basic.t.sol
+++ b/zk-tests/src/Basic.t.sol
@@ -15,6 +15,10 @@ contract BlockEnv {
         basefee = block.basefee;
         chainid = block.chainid;
     }
+
+    function ZkBlockhash(uint256 _blockNumber) public view returns (bytes32) {
+        return blockhash(_blockNumber);
+    }
 }
 
 contract ZkBasicTest is Test {
@@ -28,13 +32,17 @@ contract ZkBasicTest is Test {
 
     uint256 forkEra;
     uint256 forkEth;
+    uint256 latestForkEth;
 
     function setUp() public {
         forkEra = vm.createFork("mainnet", ERA_FORK_BLOCK);
         forkEth = vm.createFork(
-            "https://eth-mainnet.alchemyapi.io/v2/Lc7oIGYeL_QvInzI0Wiu_pOZZDEKBrdf",
+            "https://eth-mainnet.alchemyapi.io/v2/Lc7oIGYeL_QvInzI0Wiu_pOZZDEKBrdf", //trufflehog:ignore
             ETH_FORK_BLOCK
         );
+        latestForkEth = vm.createFork(
+            "https://eth-mainnet.alchemyapi.io/v2/Lc7oIGYeL_QvInzI0Wiu_pOZZDEKBrdf" //trufflehog:ignore
+        );
     }
 
     function testZkBasicBlockNumber() public {
@@ -88,6 +96,14 @@ contract ZkBasicTest is Test {
             be.chainid() == block.chainid,
             "propagated block chainid is the same as current"
         );
+        require(
+            be.ZkBlockhash(block.number) == blockhash(block.number),
+            "propagated blockhash is the same as current"
+        );
+        require(
+            be.ZkBlockhash(block.number) == bytes32(0),
+            "blockhash mismatch"
+        );
 
         be = new BlockEnv();
         require(
@@ -140,4 +156,13 @@ contract ZkBasicTest is Test {
             "propagated block basefee is the same as before"
         );
     }
+
+    function testZkBlockhashWithNewerBlocks() public {
+        vm.selectFork(latestForkEth);
+        BlockEnv be = new BlockEnv();
+        require(
+            be.ZkBlockhash(block.number) == blockhash(block.number),
+            "blockhash mismatch"
+        );
+    }
 }